From 562664f117abaa7a98b6b6f85877cf5acf8d9ffd Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 31 Aug 2026 16:23:36 -0700
Subject: [PATCH 01/38] fix(anthropic_endpoints): return Anthropic type:error
 envelope for /v1/messages errors

---
 .../exceptions/exceptions.py                  |   4 +-
 .../proxy/anthropic_endpoints/endpoints.py    |  42 +++++--
 .../anthropic_endpoints/test_endpoints.py     | 111 ++++++++++++------
 3 files changed, 113 insertions(+), 44 deletions(-)

diff --git a/litellm/anthropic_interface/exceptions/exceptions.py b/litellm/anthropic_interface/exceptions/exceptions.py
index ae333d1f4ad..91bcf82f455 100644
--- a/litellm/anthropic_interface/exceptions/exceptions.py
+++ b/litellm/anthropic_interface/exceptions/exceptions.py
@@ -1,8 +1,9 @@
 """Anthropic error format type definitions."""
 
+from collections.abc import Mapping
 from typing import Literal
 
-from typing_extensions import Required, TypedDict
+from typing_extensions import NotRequired, ReadOnly, Required, TypedDict
 
 # Known Anthropic error types
 # Source: https://docs.anthropic.com/en/api/errors
@@ -23,6 +24,7 @@ class AnthropicErrorDetail(TypedDict):
 
     type: AnthropicErrorType
     message: str
+    provider_specific_fields: NotRequired[ReadOnly[Mapping[str, object]]]
 
 
 class AnthropicErrorResponse(TypedDict, total=False):
diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py
index 7f0045c1d93..b243b737b0a 100644
--- a/litellm/proxy/anthropic_endpoints/endpoints.py
+++ b/litellm/proxy/anthropic_endpoints/endpoints.py
@@ -9,7 +9,7 @@ from fastapi.responses import JSONResponse
 
 import litellm
 from litellm._logging import verbose_proxy_logger
-from litellm.anthropic_interface.exceptions import AnthropicExceptionMapping
+from litellm.anthropic_interface.exceptions import AnthropicErrorResponse, AnthropicExceptionMapping
 from litellm.integrations.custom_guardrail import ModifyResponseException
 from litellm.llms.anthropic.experimental_pass_through.context_management import (
     AnthropicContextManagementError,
@@ -30,6 +30,27 @@ from litellm.types.utils import TokenCountResponse
 router: Final = APIRouter()
 
 
+def _anthropic_error_json_response(exc: ProxyException, request: Request) -> JSONResponse:
+    from litellm.proxy.proxy_server import (
+        _close_dangling_otel_server_span,  # pyright: ignore[reportPrivateUsage]  # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does
+    )
+
+    status_code: Final = int(exc.code) if exc.code is not None and exc.code.isdigit() else 500
+    _close_dangling_otel_server_span(request, status_code, exc=exc)
+    envelope: Final = AnthropicExceptionMapping.transform_to_anthropic_error(
+        status_code=status_code,
+        raw_message=exc.message,
+        request_id=request.headers.get("x-request-id"),
+    )
+    if not exc.provider_specific_fields:
+        return JSONResponse(status_code=status_code, content=envelope, headers=exc.headers)
+    content: Final[AnthropicErrorResponse] = {
+        **envelope,
+        "error": {**envelope["error"], "provider_specific_fields": exc.provider_specific_fields},
+    }
+    return JSONResponse(status_code=status_code, content=content, headers=exc.headers)
+
+
 def _strip_total_tokens_from_anthropic_response(response: Any) -> None:
     """Remove the OpenAI-flavored `usage.total_tokens` field that LiteLLM
     injects into Anthropic /v1/messages responses.
@@ -195,7 +216,7 @@ async def anthropic_response(
         verbose_proxy_logger.exception("litellm.proxy.proxy_server.anthropic_response(): Exception occured - %s", e)
 
         if isinstance(e, ProxyException):
-            raise
+            return _anthropic_error_json_response(e, request)
 
         # Extract model_id from request metadata (same as success path)
         litellm_metadata: Final = data.get("litellm_metadata", {}) or {}
@@ -216,15 +237,18 @@ async def anthropic_response(
         )
 
         if isinstance(e, HTTPException):
-            raise proxy_exception_from_http_exception(e, headers)
+            return _anthropic_error_json_response(proxy_exception_from_http_exception(e, headers), request)
 
         error_msg: Final = f"{e}"
-        raise ProxyException(
-            message=getattr(e, "message", error_msg),
-            type=getattr(e, "type", "None"),
-            param=getattr(e, "param", "None"),
-            code=getattr(e, "status_code", 500),
-            headers=headers,
+        return _anthropic_error_json_response(
+            ProxyException(
+                message=getattr(e, "message", error_msg),
+                type=getattr(e, "type", "None"),
+                param=getattr(e, "param", "None"),
+                code=getattr(e, "status_code", 500),
+                headers=headers,
+            ),
+            request,
         )
 
 
diff --git a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py
index c83ba142011..f809fadc879 100644
--- a/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/anthropic_endpoints/test_endpoints.py
@@ -125,11 +125,12 @@ class TestBlockedResponseUsage:
         mock_logging.post_call_failure_hook.assert_awaited_once()
 
 
-class TestProxyExceptionPassthrough:
+class TestProxyExceptionAnthropicEnvelope:
     @pytest.mark.asyncio
-    async def test_anthropic_response_reraises_proxy_exception_unwrapped(self):
-        """A 400 ProxyException from request validation must surface as-is,
-        not be re-wrapped into a code-500 ProxyException."""
+    async def test_anthropic_response_maps_proxy_exception_to_anthropic_envelope(self):
+        """LIT-6468: a 400 ProxyException from request validation must surface as
+        Anthropic's documented {"type": "error", "error": {...}} envelope with the
+        original status and message, not the OpenAI {"error": {...}} envelope."""
         import litellm.proxy.anthropic_endpoints.endpoints as ep
         import litellm.proxy.proxy_server as proxy_server
         from litellm.proxy._types import ProxyErrorTypes, ProxyException
@@ -140,6 +141,8 @@ class TestProxyExceptionPassthrough:
             param="metadata",
             code=400,
         )
+        request = MagicMock()
+        request.headers = {"x-request-id": "req_test_6468"}
 
         with (
             patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})),
@@ -151,30 +154,61 @@ class TestProxyExceptionPassthrough:
             patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
         ):
             mock_logging.post_call_failure_hook = AsyncMock()
-            with pytest.raises(ProxyException) as exc_info:
-                await ep.anthropic_response(
-                    fastapi_response=MagicMock(),
-                    request=MagicMock(),
-                    user_api_key_dict=MagicMock(),
-                )
+            response = await ep.anthropic_response(
+                fastapi_response=MagicMock(),
+                request=request,
+                user_api_key_dict=MagicMock(),
+            )
 
-        assert exc_info.value is exc
-        assert exc_info.value.code == "400"
-        assert exc_info.value.param == "metadata"
+        assert response.status_code == 400
+        body = json.loads(response.body)
+        assert body == {
+            "type": "error",
+            "error": {
+                "type": "invalid_request_error",
+                "message": "Invalid type for 'metadata': expected an object, but got a string instead.",
+            },
+            "request_id": "req_test_6468",
+        }
         mock_logging.post_call_failure_hook.assert_awaited_once()
 
+    @pytest.mark.asyncio
+    async def test_anthropic_response_maps_429_to_rate_limit_error(self):
+        """The Anthropic error type follows the status code (429 -> rate_limit_error),
+        and a code-less exception falls back to 500 api_error."""
+        import litellm.proxy.anthropic_endpoints.endpoints as ep
+        from litellm.proxy._types import ProxyException
+
+        request = MagicMock()
+        request.headers = {}
+
+        response = ep._anthropic_error_json_response(
+            ProxyException(message="Rate limit exceeded", type="rate_limit_error", param=None, code=429),
+            request,
+        )
+        assert response.status_code == 429
+        assert json.loads(response.body)["error"]["type"] == "rate_limit_error"
+
+        fallback = ep._anthropic_error_json_response(
+            ProxyException(message="boom", type="None", param=None, code=None),
+            request,
+        )
+        assert fallback.status_code == 500
+        assert json.loads(fallback.body)["error"]["type"] == "api_error"
+
 
 class TestHttpExceptionDictDetail:
     @pytest.mark.asyncio
     async def test_anthropic_response_serializes_dict_detail_http_exception(self):
-        """LIT-6466: a post_call guardrail's HTTPException(detail=<dict>) must
-        surface with a clean message plus provider_specific_fields, matching
-        /v1/chat/completions and /v1/responses, not the str() of the exception."""
+        """LIT-6466 + LIT-6468: a post_call guardrail's HTTPException(detail=<dict>)
+        must surface as Anthropic's {"type": "error", "error": {...}} envelope with
+        the guardrail's clean message plus provider_specific_fields, not the str()
+        of the exception and not the OpenAI envelope."""
         from fastapi import HTTPException
 
         import litellm.proxy.anthropic_endpoints.endpoints as ep
         import litellm.proxy.proxy_server as proxy_server
-        from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+        from litellm.proxy._types import UserAPIKeyAuth
 
         detail = {
             "error": "Content blocked: keyword 'kumquat' detected",
@@ -182,6 +216,8 @@ class TestHttpExceptionDictDetail:
             "guardrail": "keyword-block",
         }
         exc = HTTPException(status_code=400, detail=detail)
+        request = MagicMock()
+        request.headers = {}
 
         with (
             patch.object(ep, "_read_request_body", new=AsyncMock(return_value={})),  # test-quality-ok: endpoint reads the body via a module function; no injection seam
@@ -193,17 +229,19 @@ class TestHttpExceptionDictDetail:
             patch.object(proxy_server, "proxy_logging_obj") as mock_logging,  # test-quality-ok: module global imported at call time; no injection seam
         ):
             mock_logging.post_call_failure_hook = AsyncMock()
-            with pytest.raises(ProxyException) as exc_info:
-                await ep.anthropic_response(
-                    fastapi_response=MagicMock(),
-                    request=MagicMock(),
-                    user_api_key_dict=UserAPIKeyAuth(),
-                )
+            response = await ep.anthropic_response(
+                fastapi_response=MagicMock(),
+                request=request,
+                user_api_key_dict=UserAPIKeyAuth(),
+            )
 
-        assert exc_info.value.message == "Content blocked: keyword 'kumquat' detected"
-        assert "{'error'" not in exc_info.value.message
-        assert exc_info.value.provider_specific_fields == detail
-        assert exc_info.value.code == "400"
+        assert response.status_code == 400
+        body = json.loads(response.body)
+        assert body["type"] == "error"
+        assert body["error"]["type"] == "invalid_request_error"
+        assert body["error"]["message"] == "Content blocked: keyword 'kumquat' detected"
+        assert "{'error'" not in body["error"]["message"]
+        assert body["error"]["provider_specific_fields"] == detail
         mock_logging.post_call_failure_hook.assert_awaited_once()
 
 
@@ -215,7 +253,7 @@ class TestFailureHookRequestData:
         handler must pass that replaced dict, not the raw request body dict."""
         import litellm.proxy.anthropic_endpoints.endpoints as ep
         import litellm.proxy.proxy_server as proxy_server
-        from litellm.proxy._types import ProxyException, UserAPIKeyAuth
+        from litellm.proxy._types import UserAPIKeyAuth
 
         captured = {}
 
@@ -224,18 +262,23 @@ class TestFailureHookRequestData:
             captured["processor_data"] = self.data
             raise RuntimeError("provider timeout")
 
+        request = MagicMock()
+        request.headers = {}
+
         with (
             patch.object(ep, "_read_request_body", new=AsyncMock(return_value={"model": "claude-sonnet"})),
             patch.object(ep.ProxyBaseLLMRequestProcessing, "base_process_llm_request", new=fake_process),
             patch.object(proxy_server, "proxy_logging_obj") as mock_logging,
         ):
             mock_logging.post_call_failure_hook = AsyncMock()
-            with pytest.raises(ProxyException):
-                await ep.anthropic_response(
-                    fastapi_response=MagicMock(),
-                    request=MagicMock(),
-                    user_api_key_dict=UserAPIKeyAuth(),
-                )
+            response = await ep.anthropic_response(
+                fastapi_response=MagicMock(),
+                request=request,
+                user_api_key_dict=UserAPIKeyAuth(),
+            )
+
+        assert response.status_code == 500
+        assert json.loads(response.body)["error"]["message"] == "provider timeout"
 
         hook_request_data = mock_logging.post_call_failure_hook.await_args.kwargs["request_data"]
         assert hook_request_data is captured["processor_data"]

From 6a469c2159bb0086aa752788a2c0e6300eff0d8b Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 1 Sep 2026 15:46:09 -0700
Subject: [PATCH 02/38] fix(access_groups): derive attached teams from the team
 table and reject unknown team ids

GET /v1/access_group and GET /v1/access_group/{id} (and the /v1/unified_access_group aliases) used to
return the assigned_team_ids column verbatim. That column is a denormalized mirror of
LiteLLM_TeamTable.access_group_ids and can be stale or hold ids of teams that no longer exist, so the
Attached Teams view drifted from reality.

The read path now runs one team find_many per request, unioning teams whose access_group_ids carry any
group in the response with teams listed in the stored columns. Only real team rows come back, so ghost
ids drop out and teams the mirror missed are added. The stored order is kept for ids that survive and
newly discovered teams are appended.

Create and update now resolve the requested assigned_team_ids inside the transaction and answer 400
with the missing ids before anything is written, instead of silently storing ids that point nowhere.

Refs LIT-6593

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL
---
 .../access_group_endpoints.py                 |  62 ++++++-
 .../test_access_group_endpoints.py            | 161 +++++++++++++++---
 2 files changed, 195 insertions(+), 28 deletions(-)

diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py
index 0357bc7dbc6..363f312336c 100644
--- a/litellm/proxy/management_endpoints/access_group_endpoints.py
+++ b/litellm/proxy/management_endpoints/access_group_endpoints.py
@@ -1,4 +1,5 @@
 from collections.abc import Mapping, Sequence
+from types import MappingProxyType
 from typing import Final, Protocol
 
 from fastapi import APIRouter, Depends, HTTPException, status
@@ -119,8 +120,57 @@ def _require_admin_view(user_api_key_dict: UserAPIKeyAuth) -> None:
         )
 
 
-def _record_to_response(record: _AccessGroupRecord) -> AccessGroupResponse:
-    return AccessGroupResponse.model_validate(record.dict())
+def _record_to_response(
+    record: _AccessGroupRecord, *, assigned_team_ids: Sequence[str] | None = None
+) -> AccessGroupResponse:
+    stored: Final = record.dict()
+    payload: Final = (
+        stored if assigned_team_ids is None else MappingProxyType({**stored, "assigned_team_ids": assigned_team_ids})
+    )
+    return AccessGroupResponse.model_validate(payload)
+
+
+def _attached_team_ids_by_group(
+    records: Sequence[_AccessGroupRecord], teams: Sequence[_TeamRecord]
+) -> Mapping[str, tuple[str, ...]]:
+    """Teams really attached to each group: the stored column minus ghosts, plus teams the mirror missed."""
+    real_team_ids: Final = frozenset(team.team_id for team in teams)
+
+    def attached(record: _AccessGroupRecord) -> tuple[str, ...]:
+        stored: Final = (team_id for team_id in (record.assigned_team_ids or ()) if team_id in real_team_ids)
+        carrying: Final = (team.team_id for team in teams if record.access_group_id in (team.access_group_ids or ()))
+        return tuple(dict.fromkeys((*stored, *carrying)))
+
+    return MappingProxyType({record.access_group_id: attached(record) for record in records})
+
+
+async def _attached_team_ids_for(
+    team_table: _TeamTable, records: Sequence[_AccessGroupRecord]
+) -> Mapping[str, tuple[str, ...]]:
+    if not records:
+        return MappingProxyType({})
+    group_ids: Final = tuple(record.access_group_id for record in records)
+    stored_team_ids: Final = tuple(
+        dict.fromkeys(team_id for record in records for team_id in (record.assigned_team_ids or ()))
+    )
+    carrying: Final = {"access_group_ids": {"hasSome": group_ids}}  # mutable-ok: prisma where must be a dict
+    listed: Final = {"team_id": {"in": stored_team_ids}}  # mutable-ok: prisma where must be a dict
+    clauses: Final = (carrying, listed) if stored_team_ids else (carrying,)
+    where: Final = {"OR": clauses}  # mutable-ok: prisma where must be a dict
+    return _attached_team_ids_by_group(records, await team_table.find_many(where=where))
+
+
+async def _require_teams_exist(tx: _AccessGroupTx, team_ids: Sequence[str]) -> None:
+    if not team_ids:
+        return
+    where: Final = {"team_id": {"in": team_ids}}  # mutable-ok: prisma where must be a dict
+    found: Final = await tx.litellm_teamtable.find_many(where=where)
+    missing: Final = frozenset(team_ids) - frozenset(team.team_id for team in found)
+    if missing:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"Unknown team ids: {', '.join(sorted(missing))}",
+        )
 
 
 def _record_to_access_group_table(record: _AccessGroupRecord) -> LiteLLM_AccessGroupTable:
@@ -330,6 +380,7 @@ async def create_access_group(
                     status_code=status.HTTP_409_CONFLICT,
                     detail=f"Access group '{data.access_group_name}' already exists",
                 )
+            await _require_teams_exist(tx, data.assigned_team_ids or ())
 
             record: Final = await tx.litellm_accessgrouptable.create(
                 data={
@@ -390,7 +441,8 @@ async def list_access_groups(
 
     table: Final = AccessGroupRepository(prisma_client).table
     records: Final = await table.find_many(order={"created_at": "desc"})
-    return [_record_to_response(r) for r in records]
+    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, records)
+    return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records]
 
 
 @router.get(
@@ -411,7 +463,8 @@ async def get_access_group(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=f"Access group '{access_group_id}' not found",
         )
-    return _record_to_response(record)
+    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, (record,))
+    return _record_to_response(record, assigned_team_ids=attached[record.access_group_id])
 
 
 @router.put(
@@ -461,6 +514,7 @@ async def update_access_group(
                     status_code=status.HTTP_404_NOT_FOUND,
                     detail=f"Access group '{access_group_id}' not found",
                 )
+            await _require_teams_exist(tx, data.assigned_team_ids or ())
 
             old_team_ids: Final[set[str]] = set(existing.assigned_team_ids or [])
             old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or [])
diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
index e8f768c14ef..39a6f78d14d 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
@@ -12,13 +12,12 @@ from fastapi.testclient import TestClient
 from prisma.errors import PrismaError
 
 import litellm.proxy.proxy_server as ps
-from litellm.proxy.proxy_server import app
 from litellm.proxy._types import (
     CommonProxyErrors,
     LitellmUserRoles,
     UserAPIKeyAuth,
 )
-
+from litellm.proxy.proxy_server import app
 
 
 def _make_access_group_record(
@@ -58,6 +57,10 @@ def _make_access_group_record(
     return record
 
 
+def _make_team_record(team_id: str, access_group_ids: list[str] | None = None):
+    return types.SimpleNamespace(team_id=team_id, access_group_ids=access_group_ids or [])
+
+
 @pytest.fixture
 def client_and_mocks(monkeypatch):
     """Setup mock prisma and admin auth for access group endpoints."""
@@ -185,7 +188,8 @@ ACCESS_GROUP_PATHS = ["/v1/access_group", "/v1/unified_access_group"]
 )
 def test_create_access_group_success(client_and_mocks, base_path, payload):
     """Create access group with various payloads returns 201."""
-    client, _, mock_table, *_ = client_and_mocks
+    client, mock_prisma, mock_table, *_ = client_and_mocks
+    mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[_make_team_record("team-1")])
 
     resp = client.post(base_path, json=payload)
     assert resp.status_code == 201
@@ -277,13 +281,44 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks
 
 @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
 def test_list_access_groups_success_empty(client_and_mocks, base_path):
-    """List access groups returns empty list when none exist."""
-    client, _, mock_table, *_ = client_and_mocks
+    """List access groups returns empty list when none exist, without querying teams."""
+    client, mock_prisma, mock_table, *_ = client_and_mocks
 
     resp = client.get(base_path)
     assert resp.status_code == 200
     assert resp.json() == []
     mock_table.find_many.assert_awaited_once()
+    mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited()
+
+
+@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
+def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and_mocks, base_path):
+    """List derives each group's teams from the team table in a single query, attributed per group."""
+    client, mock_prisma, mock_table, *_ = client_and_mocks
+    mock_team_table = mock_prisma.db.litellm_teamtable
+
+    records = [
+        _make_access_group_record(access_group_id="ag-1", access_group_name="group-1"),
+        _make_access_group_record(access_group_id="ag-2", access_group_name="group-2"),
+    ]
+    mock_table.find_many = AsyncMock(return_value=records)
+    mock_team_table.find_many = AsyncMock(
+        return_value=[
+            _make_team_record("team-x", ["ag-1"]),
+            _make_team_record("team-y", ["ag-2"]),
+            _make_team_record("team-z", ["ag-1", "ag-2"]),
+        ]
+    )
+
+    resp = client.get(base_path)
+    assert resp.status_code == 200
+    body = resp.json()
+    assert body[0]["assigned_team_ids"] == ["team-x", "team-z"]
+    assert body[1]["assigned_team_ids"] == ["team-y", "team-z"]
+
+    mock_team_table.find_many.assert_awaited_once()
+    (carrying,) = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
+    assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-1", "ag-2"]
 
 
 @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@@ -373,6 +408,43 @@ def test_get_access_group_success(client_and_mocks, base_path, access_group_id):
     assert resp.json()["access_group_id"] == access_group_id
 
 
+@pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
+def test_get_access_group_derives_assigned_teams_from_team_table(client_and_mocks, base_path):
+    """Get drops ghost ids from the stored column and adds teams that carry the group but were never mirrored."""
+    client, mock_prisma, mock_table, *_ = client_and_mocks
+    mock_team_table = mock_prisma.db.litellm_teamtable
+
+    record = _make_access_group_record(access_group_id="ag-123", assigned_team_ids=["team-a", "ghost-team"])
+    mock_table.find_unique = AsyncMock(return_value=record)
+    mock_team_table.find_many = AsyncMock(
+        return_value=[
+            _make_team_record("team-a", ["ag-123"]),
+            _make_team_record("team-b", ["ag-123"]),
+            _make_team_record("team-c", ["ag-123"]),
+        ]
+    )
+
+    resp = client.get(f"{base_path}/ag-123")
+    assert resp.status_code == 200
+    assert resp.json()["assigned_team_ids"] == ["team-a", "team-b", "team-c"]
+
+    carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
+    assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-123"]
+    assert list(listed["team_id"]["in"]) == ["team-a", "ghost-team"]
+
+
+def test_get_access_group_empty_column_and_no_teams_returns_empty(client_and_mocks):
+    """Get returns [] when the column is empty and no team carries the group."""
+    client, mock_prisma, mock_table, *_ = client_and_mocks
+
+    mock_table.find_unique = AsyncMock(return_value=_make_access_group_record(access_group_id="ag-123"))
+    mock_prisma.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
+
+    resp = client.get("/v1/access_group/ag-123")
+    assert resp.status_code == 200
+    assert resp.json()["assigned_team_ids"] == []
+
+
 def test_get_access_group_not_found(client_and_mocks):
     """Get access group returns 404 when not found."""
     client, _, mock_table, *_ = client_and_mocks
@@ -985,6 +1057,28 @@ def test_record_to_access_group_table():
     assert result.access_agent_ids == ["agent-1"]
 
 
+def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_teams():
+    """Stored ids that resolve keep their order, ghosts drop, carriers the mirror missed append once, per group."""
+    from litellm.proxy.management_endpoints.access_group_endpoints import (
+        _attached_team_ids_by_group,
+    )
+
+    records = [
+        _make_access_group_record(access_group_id="ag-1", assigned_team_ids=["team-b", "ghost", "team-a"]),
+        _make_access_group_record(access_group_id="ag-2", assigned_team_ids=[]),
+    ]
+    teams = [
+        _make_team_record("team-a", ["ag-1"]),
+        _make_team_record("team-b", []),
+        _make_team_record("team-c", ["ag-1"]),
+        _make_team_record("team-d", ["ag-2"]),
+    ]
+
+    result = _attached_team_ids_by_group(records, teams)
+
+    assert dict(result) == {"ag-1": ("team-b", "team-a", "team-c"), "ag-2": ("team-d",)}
+
+
 # ---------------------------------------------------------------------------
 # Sync tests: CREATE
 # ---------------------------------------------------------------------------
@@ -997,9 +1091,8 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks):
     )
     mock_team_table = mock_prisma.db.litellm_teamtable
 
-    team_record = MagicMock()
-    team_record.team_id = "team-1"
-    team_record.access_group_ids = []
+    team_record = _make_team_record("team-1")
+    mock_team_table.find_many = AsyncMock(return_value=[team_record])
     mock_team_table.find_unique = AsyncMock(return_value=team_record)
 
     resp = client.post(
@@ -1043,20 +1136,22 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks):
     assert "ag-new" in call_kwargs["data"]["access_group_ids"]
 
 
-def test_create_access_group_skips_sync_for_nonexistent_team(client_and_mocks):
-    """Create skips updating a team that doesn't exist in DB."""
-    client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks
+def test_create_access_group_rejects_nonexistent_team(client_and_mocks):
+    """Create refuses to store a team id that does not resolve to a team row."""
+    client, mock_prisma, mock_access_group_table, *_ = client_and_mocks
     mock_team_table = mock_prisma.db.litellm_teamtable
-    mock_team_table.find_unique = AsyncMock(return_value=None)
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-real")])
 
     resp = client.post(
         "/v1/access_group",
         json={
             "access_group_name": "new-group",
-            "assigned_team_ids": ["nonexistent-team"],
+            "assigned_team_ids": ["team-real", "nonexistent-team", "also-missing"],
         },
     )
-    assert resp.status_code == 201
+    assert resp.status_code == 400
+    assert resp.json()["detail"] == "Unknown team ids: also-missing, nonexistent-team"
+    mock_access_group_table.create.assert_not_awaited()
     mock_team_table.update.assert_not_awaited()
 
 
@@ -1065,9 +1160,8 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks):
     client, mock_prisma, _, mock_cache, mock_proxy_logging = client_and_mocks
     mock_team_table = mock_prisma.db.litellm_teamtable
 
-    team_record = MagicMock()
-    team_record.team_id = "team-1"
-    team_record.access_group_ids = ["ag-new"]  # already synced
+    team_record = _make_team_record("team-1", ["ag-new"])
+    mock_team_table.find_many = AsyncMock(return_value=[team_record])
     mock_team_table.find_unique = AsyncMock(return_value=team_record)
 
     resp = client.post(
@@ -1095,9 +1189,8 @@ def test_update_access_group_syncs_added_teams(client_and_mocks):
     )
     mock_access_group_table.find_unique = AsyncMock(return_value=existing)
 
-    team_record = MagicMock()
-    team_record.team_id = "team-new"
-    team_record.access_group_ids = []
+    team_record = _make_team_record("team-new")
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"]), team_record])
     mock_team_table.find_unique = AsyncMock(return_value=team_record)
 
     resp = client.put(
@@ -1113,6 +1206,25 @@ def test_update_access_group_syncs_added_teams(client_and_mocks):
     assert "ag-update" in call_kwargs["data"]["access_group_ids"]
 
 
+def test_update_access_group_rejects_nonexistent_team(client_and_mocks):
+    """Update refuses to store a team id that does not resolve to a team row and leaves the group untouched."""
+    client, mock_prisma, mock_access_group_table, *_ = client_and_mocks
+    mock_team_table = mock_prisma.db.litellm_teamtable
+
+    existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-existing"])
+    mock_access_group_table.find_unique = AsyncMock(return_value=existing)
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-existing", ["ag-update"])])
+
+    resp = client.put(
+        "/v1/access_group/ag-update",
+        json={"assigned_team_ids": ["team-existing", "team-ghost"]},
+    )
+    assert resp.status_code == 400
+    assert resp.json()["detail"] == "Unknown team ids: team-ghost"
+    mock_access_group_table.update.assert_not_awaited()
+    mock_team_table.update.assert_not_awaited()
+
+
 def test_update_access_group_syncs_removed_teams(client_and_mocks):
     """Update removes access_group_id from de-assigned teams."""
     client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
@@ -1125,9 +1237,8 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
     )
     mock_access_group_table.find_unique = AsyncMock(return_value=existing)
 
-    team_to_remove = MagicMock()
-    team_to_remove.team_id = "team-remove"
-    team_to_remove.access_group_ids = ["ag-update"]
+    team_to_remove = _make_team_record("team-remove", ["ag-update"])
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"])])
     mock_team_table.find_unique = AsyncMock(return_value=team_to_remove)
 
     resp = client.put(
@@ -1160,6 +1271,7 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc
     resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"})
     assert resp.status_code == 200
 
+    mock_team_table.find_many.assert_not_awaited()
     mock_team_table.find_unique.assert_not_awaited()
     mock_team_table.update.assert_not_awaited()
 
@@ -1293,7 +1405,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks)
 
 def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks):
     """Update with explicit null for assigned_*_ids clears the list and writes [] to DB."""
-    client, _, mock_table, *_ = client_and_mocks
+    client, mock_prisma, mock_table, *_ = client_and_mocks
 
     existing = _make_access_group_record(
         access_group_id="ag-update",
@@ -1313,3 +1425,4 @@ def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks
     update_call_kwargs = mock_table.update.call_args.kwargs
     assert update_call_kwargs["data"]["assigned_team_ids"] == []
     assert update_call_kwargs["data"]["assigned_key_ids"] == []
+    mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited()

From c14cf9d173b3af9d1e3dea64ede016b72ee762b7 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Tue, 1 Sep 2026 16:03:22 -0700
Subject: [PATCH 03/38] fix(access_groups): reconcile update deltas against the
 derived team set

Team membership deltas on PUT now start from the teams that really carry
the group, so a team the mirror column missed can be detached. Read
endpoints go through a typed TeamRepository instead of the untyped db
handle, and the where clause always carries both OR arms.

Claude-Session: https://claude.ai/code/session_01QvQzYztinxj8ZuD5YxbVdL
---
 .../access_group_endpoints.py                 | 26 ++++++++--------
 litellm/repositories/table_repositories.py    |  4 +++
 .../test_access_group_endpoints.py            | 31 ++++++++++++++++---
 3 files changed, 43 insertions(+), 18 deletions(-)

diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py
index 363f312336c..1f91eeedf64 100644
--- a/litellm/proxy/management_endpoints/access_group_endpoints.py
+++ b/litellm/proxy/management_endpoints/access_group_endpoints.py
@@ -21,7 +21,7 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
 from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache
 from litellm.proxy.utils import get_prisma_client_or_throw
-from litellm.repositories.table_repositories import AccessGroupRepository
+from litellm.repositories.table_repositories import AccessGroupRepository, TeamRepository
 from litellm.types.access_group import (
     AccessGroupCreateRequest,
     AccessGroupResponse,
@@ -75,11 +75,11 @@ class _AccessGroupTable(Protocol):
 
 
 class _TeamTable(Protocol):
-    async def find_unique(self, where: Mapping[str, object]) -> _TeamRecord | None: ...
+    async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ...
 
-    async def find_many(self, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ...
+    async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_TeamRecord]: ...
 
-    async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
+    async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
 
 
 class _KeyTable(Protocol):
@@ -153,17 +153,16 @@ async def _attached_team_ids_for(
     stored_team_ids: Final = tuple(
         dict.fromkeys(team_id for record in records for team_id in (record.assigned_team_ids or ()))
     )
-    carrying: Final = {"access_group_ids": {"hasSome": group_ids}}  # mutable-ok: prisma where must be a dict
-    listed: Final = {"team_id": {"in": stored_team_ids}}  # mutable-ok: prisma where must be a dict
-    clauses: Final = (carrying, listed) if stored_team_ids else (carrying,)
-    where: Final = {"OR": clauses}  # mutable-ok: prisma where must be a dict
-    return _attached_team_ids_by_group(records, await team_table.find_many(where=where))
+    carrying: Final = {"access_group_ids": {"hasSome": group_ids}}  # mutable-ok: prisma where is a dict
+    listed: Final = {"team_id": {"in": stored_team_ids}}  # mutable-ok: prisma where is a dict
+    teams: Final = await team_table.find_many(where={"OR": (carrying, listed)})  # mutable-ok: prisma where is a dict
+    return _attached_team_ids_by_group(records, teams)
 
 
 async def _require_teams_exist(tx: _AccessGroupTx, team_ids: Sequence[str]) -> None:
     if not team_ids:
         return
-    where: Final = {"team_id": {"in": team_ids}}  # mutable-ok: prisma where must be a dict
+    where: Final = {"team_id": {"in": team_ids}}  # mutable-ok: prisma where is a dict
     found: Final = await tx.litellm_teamtable.find_many(where=where)
     missing: Final = frozenset(team_ids) - frozenset(team.team_id for team in found)
     if missing:
@@ -441,7 +440,7 @@ async def list_access_groups(
 
     table: Final = AccessGroupRepository(prisma_client).table
     records: Final = await table.find_many(order={"created_at": "desc"})
-    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, records)
+    attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, records)
     return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records]
 
 
@@ -463,7 +462,7 @@ async def get_access_group(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=f"Access group '{access_group_id}' not found",
         )
-    attached: Final = await _attached_team_ids_for(prisma_client.db.litellm_teamtable, (record,))
+    attached: Final = await _attached_team_ids_for(TeamRepository(prisma_client).table, (record,))
     return _record_to_response(record, assigned_team_ids=attached[record.access_group_id])
 
 
@@ -516,7 +515,8 @@ async def update_access_group(
                 )
             await _require_teams_exist(tx, data.assigned_team_ids or ())
 
-            old_team_ids: Final[set[str]] = set(existing.assigned_team_ids or [])
+            attached: Final = await _attached_team_ids_for(tx.litellm_teamtable, (existing,))
+            old_team_ids: Final[set[str]] = set(attached[access_group_id])
             old_key_ids: Final[set[str]] = set(existing.assigned_key_ids or [])
             new_team_ids: Final[set[str]] = (
                 set(update_fields["assigned_team_ids"] or []) if "assigned_team_ids" in update_fields else old_team_ids
diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py
index 18cf884f267..739d6ada71c 100644
--- a/litellm/repositories/table_repositories.py
+++ b/litellm/repositories/table_repositories.py
@@ -176,6 +176,10 @@ class PolicyAttachmentRepository(PrismaTableRepository["prisma_models.LiteLLM_Po
     table_name = "litellm_policyattachmenttable"
 
 
+class TeamRepository(PrismaTableRepository["prisma_models.LiteLLM_TeamTable"]):
+    table_name = "litellm_teamtable"
+
+
 class DeletedTeamRepository(PrismaTableRepository["prisma_models.LiteLLM_DeletedTeamTable"]):
     table_name = "litellm_deletedteamtable"
 
diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
index 39a6f78d14d..81816e21c10 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py
@@ -317,8 +317,9 @@ def test_list_access_groups_attributes_teams_per_group_with_one_query(client_and
     assert body[1]["assigned_team_ids"] == ["team-y", "team-z"]
 
     mock_team_table.find_many.assert_awaited_once()
-    (carrying,) = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
+    carrying, listed = mock_team_table.find_many.call_args.kwargs["where"]["OR"]
     assert list(carrying["access_group_ids"]["hasSome"]) == ["ag-1", "ag-2"]
+    assert list(listed["team_id"]["in"]) == []
 
 
 @pytest.mark.parametrize("base_path", ACCESS_GROUP_PATHS)
@@ -1238,7 +1239,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
     mock_access_group_table.find_unique = AsyncMock(return_value=existing)
 
     team_to_remove = _make_team_record("team-remove", ["ag-update"])
-    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"])])
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), team_to_remove])
     mock_team_table.find_unique = AsyncMock(return_value=team_to_remove)
 
     resp = client.put(
@@ -1256,6 +1257,28 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks):
     assert "ag-update" not in call_kwargs["data"]["access_group_ids"]
 
 
+def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks):
+    """Update removes the group from a team that carries it but was never written to the stored column."""
+    client, mock_prisma, mock_access_group_table, *_ = client_and_mocks
+    mock_team_table = mock_prisma.db.litellm_teamtable
+
+    existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep"])
+    mock_access_group_table.find_unique = AsyncMock(return_value=existing)
+
+    unmirrored = _make_team_record("team-unmirrored", ["ag-update", "ag-other"])
+    mock_team_table.find_many = AsyncMock(return_value=[_make_team_record("team-keep", ["ag-update"]), unmirrored])
+    mock_team_table.find_unique = AsyncMock(return_value=unmirrored)
+
+    resp = client.put("/v1/access_group/ag-update", json={"assigned_team_ids": ["team-keep"]})
+    assert resp.status_code == 200
+
+    mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-unmirrored"})
+    mock_team_table.update.assert_awaited_once()
+    call_kwargs = mock_team_table.update.call_args.kwargs
+    assert call_kwargs["where"] == {"team_id": "team-unmirrored"}
+    assert call_kwargs["data"]["access_group_ids"] == ["ag-other"]
+
+
 def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks):
     """Update does not sync teams when assigned_team_ids is absent from the payload."""
     client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = (
@@ -1271,7 +1294,6 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc
     resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"})
     assert resp.status_code == 200
 
-    mock_team_table.find_many.assert_not_awaited()
     mock_team_table.find_unique.assert_not_awaited()
     mock_team_table.update.assert_not_awaited()
 
@@ -1405,7 +1427,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks)
 
 def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks):
     """Update with explicit null for assigned_*_ids clears the list and writes [] to DB."""
-    client, mock_prisma, mock_table, *_ = client_and_mocks
+    client, _, mock_table, *_ = client_and_mocks
 
     existing = _make_access_group_record(
         access_group_id="ag-update",
@@ -1425,4 +1447,3 @@ def test_update_access_group_null_assigned_ids_treated_as_empty(client_and_mocks
     update_call_kwargs = mock_table.update.call_args.kwargs
     assert update_call_kwargs["data"]["assigned_team_ids"] == []
     assert update_call_kwargs["data"]["assigned_key_ids"] == []
-    mock_prisma.db.litellm_teamtable.find_many.assert_not_awaited()

From e29796882668a7359a05acab53fb2810f2b12009 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 13:45:06 -0700
Subject: [PATCH 04/38] feat(cost): honor off_peak_pricing reasoning and
 cache-creation rates

The block accepts output_cost_per_reasoning_token and cache_creation_input_token_cost. The generic
cost path and the DashScope calculator swap them in while a window is open, and unset keys keep the
standard rate. One shared TokenRates value replaces the DashScope-local copy, and
apply_off_peak_pricing takes and returns it.
---
 .../litellm_core_utils/llm_cost_calc/utils.py | 155 +++++++---
 litellm/llms/dashscope/cost_calculator.py     |  25 +-
 litellm/types/utils.py                        |   2 +
 .../llm_cost_calc/test_llm_cost_calc_utils.py | 267 ++++++++++++++++++
 .../test_dashscope_cost_calculator.py         |  89 ++++++
 5 files changed, 470 insertions(+), 68 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index 21587af73aa..e03c2c93c26 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -415,40 +415,69 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None =
     return False
 
 
-def _coerce_off_peak_rate(value: object, default: float) -> float:
+@dataclass(frozen=True, slots=True)
+class TokenRates:
+    """The per-token rates one request bills at. reasoning_rate is None when reasoning bills at
+    output_rate: the model has no dedicated reasoning rate, or the caller resolves reasoning on
+    its own.
+    """
+
+    input_rate: float
+    output_rate: float
+    cache_read_rate: float
+    cache_creation_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 _parse_off_peak_rate(value: object) -> float | None:
     if isinstance(value, bool):
-        return default
+        return None
     if isinstance(value, (int, float)):
         return float(value)
     if isinstance(value, str):
         try:
             return float(value)
         except ValueError:
-            return default
-    return default
+            return None
+    return None
 
 
-def apply_off_peak_pricing(
-    model_info: ModelInfo,
-    current_time: datetime | None,
-    prompt_base_cost: float,
-    completion_base_cost: float,
-    cache_read_cost: float,
-) -> tuple[float, float, float]:
+def _off_peak_rate(off_peak: Mapping[str, object], key: str, standard_rate: float) -> float:
+    parsed: Final = _parse_off_peak_rate(off_peak.get(key))
+    return standard_rate if parsed is None else parsed
+
+
+def _open_off_peak_block(model_info: ModelInfo, current_time: datetime | None) -> Mapping[str, object] | None:
+    off_peak: Final = model_info.get("off_peak_pricing")
+    if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time):
+        return None
+    return off_peak
+
+
+def apply_off_peak_pricing(model_info: ModelInfo, current_time: datetime | None, rates: TokenRates) -> TokenRates:
     """Swap in off-peak per-token rates when the current UTC time is inside one of the model's
     off_peak_pricing rules, the every-day hours_utc windows or a day-of-week-qualified entry in
     windows. An off-peak rate replaces the rate that would otherwise apply rather than
     discounting it, so a model that also has tiered or above-threshold pricing bills the flat
     off-peak rate for the whole request while the window is open. Any rate left unset in
-    off_peak_pricing falls back to the standard rate.
+    off_peak_pricing falls back to the standard rate, so a block without
+    output_cost_per_reasoning_token keeps the model's own reasoning rate, or its off-peak output
+    rate when reasoning has no dedicated rate at all.
     """
-    off_peak: Final = model_info.get("off_peak_pricing")
-    if not isinstance(off_peak, Mapping) or not _is_off_peak(off_peak, current_time):
-        return prompt_base_cost, completion_base_cost, cache_read_cost
-    return (
-        _coerce_off_peak_rate(off_peak.get("input_cost_per_token"), prompt_base_cost),
-        _coerce_off_peak_rate(off_peak.get("output_cost_per_token"), completion_base_cost),
-        _coerce_off_peak_rate(off_peak.get("cache_read_input_token_cost"), cache_read_cost),
+    off_peak: Final = _open_off_peak_block(model_info, current_time)
+    if off_peak is None:
+        return rates
+    off_peak_reasoning_rate: Final = _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token"))
+    return TokenRates(
+        input_rate=_off_peak_rate(off_peak, "input_cost_per_token", rates.input_rate),
+        output_rate=_off_peak_rate(off_peak, "output_cost_per_token", rates.output_rate),
+        cache_read_rate=_off_peak_rate(off_peak, "cache_read_input_token_cost", rates.cache_read_rate),
+        cache_creation_rate=_off_peak_rate(off_peak, "cache_creation_input_token_cost", rates.cache_creation_rate),
+        reasoning_rate=rates.reasoning_rate if off_peak_reasoning_rate is None else off_peak_reasoning_rate,
     )
 
 
@@ -458,14 +487,28 @@ def _apply_off_peak_to_base_costs(
     base_costs: tuple[float, float, float, float, float],
 ) -> tuple[float, float, float, float, float]:
     """Apply off-peak rates to an already-resolved set of base costs, whichever pricing path
-    produced them. Cache-creation rates are passed through untouched, since off_peak_pricing
-    has no field for them.
+    produced them. The one-hour cache-creation rate passes through untouched, since
+    off_peak_pricing has no field for it, and reasoning is left to _resolve_billed_reasoning_rate.
     """
     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(
-        model_info, current_time, prompt, completion, cache_read
+    rates: Final = apply_off_peak_pricing(
+        model_info,
+        current_time,
+        TokenRates(
+            input_rate=prompt,
+            output_rate=completion,
+            cache_read_rate=cache_read,
+            cache_creation_rate=cache_creation,
+            reasoning_rate=None,
+        ),
+    )
+    return (
+        rates.input_rate,
+        rates.output_rate,
+        rates.cache_creation_rate,
+        cache_creation_above_1hr,
+        rates.cache_read_rate,
     )
-    return (off_peak_prompt, off_peak_completion, cache_creation, cache_creation_above_1hr, off_peak_cache_read)
 
 
 def _get_token_base_cost(
@@ -1029,6 +1072,29 @@ def _resolve_reasoning_token_cost(
     return standard_reasoning_cost if standard_reasoning_cost is not None else completion_base_cost
 
 
+def _resolve_billed_reasoning_rate(
+    model_info: ModelInfo,
+    usage: Usage,
+    service_tier: str | None,
+    completion_base_cost: float,
+    current_time: datetime | None,
+) -> float:
+    off_peak: Final = _open_off_peak_block(model_info, current_time)
+    off_peak_reasoning_rate: Final = (
+        None if off_peak is None else _parse_off_peak_rate(off_peak.get("output_cost_per_reasoning_token"))
+    )
+    if off_peak_reasoning_rate is not None:
+        return off_peak_reasoning_rate
+    tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage)
+    if tiered_reasoning_rate is not None:
+        return tiered_reasoning_rate
+    return _resolve_reasoning_token_cost(
+        model_info=model_info,
+        service_tier=service_tier,
+        completion_base_cost=completion_base_cost,
+    )
+
+
 def generic_cost_per_token(
     model: str,
     usage: Usage,
@@ -1037,6 +1103,7 @@ def generic_cost_per_token(
     data_residency: str | None = None,
     model_info: ModelInfo | None = None,
     vertex_location: str | None = None,
+    current_time: datetime | None = None,
 ) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens.
@@ -1051,6 +1118,7 @@ def generic_cost_per_token(
         - vertex_location: optional Vertex AI location the request was served from
           (e.g. "us-east5", "global"), used to apply the per-model
           regional-endpoint uplift multiplier when non-global.
+        - current_time: the moment the request is billed at, for off_peak_pricing; defaults to now, UTC
 
     Returns:
         Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
@@ -1117,6 +1185,7 @@ def generic_cost_per_token(
             usage.prompt_tokens - cache_hit - audio_tokens - cache_creation - image_tokens - video_tokens, 0
         )
 
+    billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
     (
         prompt_base_cost,
         completion_base_cost,
@@ -1127,6 +1196,7 @@ def generic_cost_per_token(
         model_info=model_info,
         usage=usage,
         service_tier=service_tier,
+        current_time=billing_time,
         threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
     )
 
@@ -1185,17 +1255,13 @@ def generic_cost_per_token(
 
     ## REASONING COST
     if not is_text_tokens_total and reasoning_tokens and reasoning_tokens > 0:
-        tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage)
-        _output_cost_per_reasoning_token = (
-            tiered_reasoning_rate
-            if tiered_reasoning_rate is not None
-            else _resolve_reasoning_token_cost(
-                model_info=model_info,
-                service_tier=service_tier,
-                completion_base_cost=completion_base_cost,
-            )
+        completion_cost += float(reasoning_tokens) * _resolve_billed_reasoning_rate(
+            model_info=model_info,
+            usage=usage,
+            service_tier=service_tier,
+            completion_base_cost=completion_base_cost,
+            current_time=billing_time,
         )
-        completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token
 
     ## IMAGE COST
     if not is_text_tokens_total and image_tokens and image_tokens > 0:
@@ -1247,6 +1313,7 @@ def get_token_type_cost_breakdown(
     service_tier: str | None = None,
     data_residency: str | None = None,
     vertex_location: str | None = None,
+    current_time: datetime | None = None,
 ) -> TokenTypeCostBreakdown:
     """
     Provider-agnostic cost of reasoning and cache tokens, derived from the usage
@@ -1265,6 +1332,7 @@ def get_token_type_cost_breakdown(
     except Exception:
         return TokenTypeCostBreakdown(0.0, 0.0, 0.0)
 
+    billing_time: Final = current_time if current_time is not None else datetime.now(timezone.utc)
     (
         _prompt_base_cost,
         completion_base_cost,
@@ -1275,6 +1343,7 @@ def get_token_type_cost_breakdown(
         model_info=model_info,
         usage=usage,
         service_tier=service_tier,
+        current_time=billing_time,
         threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
     )
 
@@ -1284,18 +1353,12 @@ def get_token_type_cost_breakdown(
     if not reasoning_tokens:
         reasoning_tokens = _coerce_token_count(getattr(usage, "reasoning_tokens", 0))
 
-    # Reasoning is billed at the selected tier's reasoning rate for tiered models,
-    # else at the service-tier-aware per-reasoning-token rate - this mirrors how the
-    # total completion cost is computed, so the breakdown can never diverge from it.
-    tiered_reasoning_rate: Final = _get_tiered_reasoning_rate(model_info=model_info, usage=usage)
-    reasoning_rate: Final = (
-        tiered_reasoning_rate
-        if tiered_reasoning_rate is not None
-        else _resolve_reasoning_token_cost(
-            model_info=model_info,
-            service_tier=service_tier,
-            completion_base_cost=completion_base_cost,
-        )
+    reasoning_rate: Final = _resolve_billed_reasoning_rate(
+        model_info=model_info,
+        usage=usage,
+        service_tier=service_tier,
+        completion_base_cost=completion_base_cost,
+        current_time=billing_time,
     )
     reasoning_cost = float(reasoning_tokens) * reasoning_rate
 
diff --git a/litellm/llms/dashscope/cost_calculator.py b/litellm/llms/dashscope/cost_calculator.py
index d8eb1f9f8d7..17f70ec5db7 100644
--- a/litellm/llms/dashscope/cost_calculator.py
+++ b/litellm/llms/dashscope/cost_calculator.py
@@ -7,12 +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, replace
+from dataclasses import dataclass
 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 (
+    TokenRates,
     apply_off_peak_pricing,
     parse_completion_tokens_details,
     parse_prompt_tokens_details,
@@ -34,19 +35,6 @@ 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"]
@@ -105,13 +93,6 @@ def _tier_rates(model_info: ModelInfo, tier: dict) -> TokenRates:
     )
 
 
-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 _bill(breakdown: TokenBreakdown, rates: TokenRates) -> tuple[float, float]:
     prompt_cost: Final = (
         (breakdown.text_tokens * rates.input_rate)
@@ -155,6 +136,6 @@ def cost_per_token(
         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)
+    rates: Final = apply_off_peak_pricing(model_info, current_time, standard_rates)
 
     return _bill(breakdown, rates)
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 569fce4f7b8..a8e1f1b7ee8 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -222,7 +222,9 @@ class OffPeakPricing(TypedDict, total=False):
     weekday_timezone: ReadOnly[str]
     input_cost_per_token: ReadOnly[float]
     output_cost_per_token: ReadOnly[float]
+    output_cost_per_reasoning_token: ReadOnly[float]
     cache_read_input_token_cost: ReadOnly[float]
+    cache_creation_input_token_cost: ReadOnly[float]
 
 
 class ModelInfoBase(ProviderSpecificModelInfo, total=False):
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..7bc02145841 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
@@ -29,11 +29,13 @@ from litellm.types.utils import (
 from litellm.litellm_core_utils.llm_cost_calc.utils import (
     CostCalculatorUtils,
     PromptTokensDetailsResult,
+    TokenRates,
     TokenTypeCostBreakdown,
     _calculate_input_cost,
     _get_token_base_cost,
     _is_off_peak,
     _is_within_off_peak_window,
+    apply_off_peak_pricing,
     calculate_cache_writing_cost,
     generic_cost_per_token,
     get_token_type_cost_breakdown,
@@ -782,6 +784,271 @@ def test_get_token_base_cost_off_peak_wins_over_tiered_pricing():
     assert outside[:2] == (3e-6, 6e-6)
 
 
+def _register_off_peak_reasoning_model(
+    model_name: str, off_peak_pricing: dict, reasoning_rate: float | None = 4e-6, **service_tier_rates: float
+) -> None:
+    reasoning_entry = {} if reasoning_rate is None else {"output_cost_per_reasoning_token": reasoning_rate}
+    litellm.register_model(
+        {
+            model_name: {
+                "litellm_provider": "openai",
+                "mode": "chat",
+                "input_cost_per_token": 1e-6,
+                "output_cost_per_token": 2e-6,
+                "cache_read_input_token_cost": 1e-7,
+                "cache_creation_input_token_cost": 1.25e-6,
+                "off_peak_pricing": off_peak_pricing,
+                **reasoning_entry,
+                **service_tier_rates,
+            }
+        }
+    )
+
+
+def _off_peak_reasoning_usage() -> Usage:
+    return Usage(
+        prompt_tokens=100,
+        completion_tokens=80,
+        total_tokens=180,
+        completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50),
+    )
+
+
+def test_generic_cost_per_token_off_peak_reasoning_rate():
+    """Regression (LIT-6887): the block's output_cost_per_reasoning_token used to be ignored, so
+    reasoning tokens billed at the model's standard reasoning rate all through the window."""
+    from datetime import datetime, timezone
+
+    model_name = "litellm-test-off-peak-reasoning"
+    _register_off_peak_reasoning_model(
+        model_name,
+        {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7},
+    )
+
+    _, inside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc),
+    )
+    assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7)
+
+    _, outside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc),
+    )
+    assert outside == pytest.approx(50 * 2e-6 + 30 * 4e-6)
+
+
+def test_generic_cost_per_token_off_peak_block_without_reasoning_rate():
+    """A block that leaves output_cost_per_reasoning_token unset keeps the model's own reasoning
+    rate, and a model with no reasoning rate at all follows the off-peak output rate."""
+    from datetime import datetime, timezone
+
+    inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)
+    block = {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6}
+
+    _register_off_peak_reasoning_model("litellm-test-off-peak-model-reasoning-rate", block)
+    _, with_model_rate = generic_cost_per_token(
+        model="litellm-test-off-peak-model-reasoning-rate",
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=inside_window,
+    )
+    assert with_model_rate == pytest.approx(50 * 1e-6 + 30 * 4e-6)
+
+    _register_off_peak_reasoning_model("litellm-test-off-peak-no-reasoning-rate", block, reasoning_rate=None)
+    _, without_model_rate = generic_cost_per_token(
+        model="litellm-test-off-peak-no-reasoning-rate",
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=inside_window,
+    )
+    assert without_model_rate == pytest.approx(80 * 1e-6)
+
+
+def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier():
+    """Tiered models resolve reasoning on their own path, so the block has to win there too."""
+    from datetime import datetime, timezone
+
+    model_name = "litellm-test-off-peak-tiered-reasoning"
+    litellm.register_model(
+        {
+            model_name: {
+                "litellm_provider": "openai",
+                "mode": "chat",
+                "tiered_pricing": [
+                    {
+                        "range": [0, 128000],
+                        "input_cost_per_token": 3e-6,
+                        "output_cost_per_token": 6e-6,
+                        "output_cost_per_reasoning_token": 8e-6,
+                    },
+                ],
+                "off_peak_pricing": {
+                    "hours_utc": "16:30-00:30",
+                    "output_cost_per_token": 1e-6,
+                    "output_cost_per_reasoning_token": 5e-7,
+                },
+            }
+        }
+    )
+
+    _, inside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc),
+    )
+    assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7)
+
+    _, outside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc),
+    )
+    assert outside == pytest.approx(50 * 6e-6 + 30 * 8e-6)
+
+
+def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_tier():
+    """A priority request bills its service-tier reasoning rate outside the window and the block's
+    rate inside it."""
+    from datetime import datetime, timezone
+
+    model_name = "litellm-test-off-peak-reasoning-service-tier"
+    _register_off_peak_reasoning_model(
+        model_name,
+        {"hours_utc": "16:30-00:30", "output_cost_per_token": 1e-6, "output_cost_per_reasoning_token": 5e-7},
+        output_cost_per_token_priority=3e-6,
+        output_cost_per_reasoning_token_priority=6e-6,
+    )
+
+    _, inside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        service_tier="priority",
+        current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc),
+    )
+    assert inside == pytest.approx(50 * 1e-6 + 30 * 5e-7)
+
+    _, outside = generic_cost_per_token(
+        model=model_name,
+        usage=_off_peak_reasoning_usage(),
+        custom_llm_provider="openai",
+        service_tier="priority",
+        current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc),
+    )
+    assert outside == pytest.approx(50 * 3e-6 + 30 * 6e-6)
+
+
+def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings():
+    """A YAML true never turns into a rate of 1.0, and a quoted number still counts."""
+    from datetime import datetime, timezone
+
+    model_name = "litellm-test-off-peak-odd-values"
+    _register_off_peak_reasoning_model(
+        model_name,
+        {
+            "hours_utc": "16:30-00:30",
+            "cache_creation_input_token_cost": True,
+            "output_cost_per_reasoning_token": "5e-7",
+        },
+    )
+    standard = TokenRates(
+        input_rate=1e-6, output_rate=2e-6, cache_read_rate=1e-7, cache_creation_rate=1.25e-6, reasoning_rate=4e-6
+    )
+
+    rates = apply_off_peak_pricing(
+        litellm.get_model_info(model_name, custom_llm_provider="openai"),
+        datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc),
+        standard,
+    )
+    assert rates.cache_creation_rate == 1.25e-6
+    assert rates.reasoning_rate == 5e-7
+
+
+def test_get_token_base_cost_off_peak_cache_creation_rate():
+    """Regression (LIT-6887): the block's cache_creation_input_token_cost used to be ignored. It
+    replaces the five-minute cache-creation rate inside the window; the one-hour rate, and a
+    block without the key, keep the standard rate."""
+    from datetime import datetime, timezone
+    from typing import cast
+
+    from litellm.types.utils import ModelInfo
+
+    model_info = cast(
+        ModelInfo,
+        {
+            "input_cost_per_token": 1e-6,
+            "output_cost_per_token": 2e-6,
+            "cache_creation_input_token_cost": 1.25e-6,
+            "cache_creation_input_token_cost_above_1hr": 2e-6,
+            "off_peak_pricing": {"hours_utc": "16:30-00:30", "cache_creation_input_token_cost": 5e-7},
+        },
+    )
+    usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
+    inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)
+
+    inside = _get_token_base_cost(model_info, usage, current_time=inside_window)
+    assert inside[2] == 5e-7
+    assert inside[3] == 2e-6
+
+    outside = _get_token_base_cost(model_info, usage, current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc))
+    assert outside[2] == 1.25e-6
+
+    without_key = cast(
+        ModelInfo,
+        {**model_info, "off_peak_pricing": {"hours_utc": "16:30-00:30", "input_cost_per_token": 5e-7}},
+    )
+    assert _get_token_base_cost(without_key, usage, current_time=inside_window)[2] == 1.25e-6
+
+
+def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates():
+    """The per-token-type breakdown feeds the spend logs, so it has to bill the new keys the same
+    way the total does."""
+    from datetime import datetime, timezone
+
+    model_name = "litellm-test-off-peak-breakdown"
+    _register_off_peak_reasoning_model(
+        model_name,
+        {
+            "hours_utc": "16:30-00:30",
+            "output_cost_per_reasoning_token": 5e-7,
+            "cache_creation_input_token_cost": 5e-7,
+        },
+    )
+    usage = Usage(
+        prompt_tokens=1000,
+        completion_tokens=80,
+        total_tokens=1080,
+        completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=30, text_tokens=50),
+        prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=100, cache_creation_tokens=400, text_tokens=500),
+    )
+
+    inside = get_token_type_cost_breakdown(
+        model=model_name,
+        custom_llm_provider="openai",
+        usage=usage,
+        current_time=datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc),
+    )
+    assert inside.reasoning_cost == pytest.approx(30 * 5e-7)
+    assert inside.cache_creation_cost == pytest.approx(400 * 5e-7)
+    assert inside.cache_read_cost == pytest.approx(100 * 1e-7)
+
+    outside = get_token_type_cost_breakdown(
+        model=model_name,
+        custom_llm_provider="openai",
+        usage=usage,
+        current_time=datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc),
+    )
+    assert outside.reasoning_cost == pytest.approx(30 * 4e-6)
+    assert outside.cache_creation_cost == pytest.approx(400 * 1.25e-6)
+
+
 def test_generic_cost_per_token_gpt54_above_272k_tokens(_local_model_cost_map):
     """GPT-5.4/5.4-pro: prompts >272K input tokens priced at 2x input, 1.5x output."""
     model = "gpt-5.4"
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 b6281834f24..f0949ce041a 100644
--- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
+++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
@@ -649,6 +649,95 @@ class TestDashscopeCostCalculator:
 
         assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
 
+    def test_dashscope_off_peak_reasoning_rate_replaces_the_dedicated_reasoning_rate(self):
+        """Regression (LIT-6887): a block carrying output_cost_per_reasoning_token bills reasoning
+        tokens at it inside the window, over the model's own reasoning rate, which returns outside."""
+        self._register_off_peak_flat_model(
+            "dashscope/qwen-reasoning-rate-off-peak-test",
+            {
+                "hours_utc": self.OFF_PEAK_WINDOW,
+                "output_cost_per_token": 2.4e-06,
+                "output_cost_per_reasoning_token": 4.5e-06,
+            },
+        )
+        litellm.model_cost["dashscope/qwen-reasoning-rate-off-peak-test"]["output_cost_per_reasoning_token"] = 9e-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-rate-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+        assert math.isclose(completion_cost, (150 * 2.4e-06) + (50 * 4.5e-06), rel_tol=1e-10)
+
+        _, peak_completion_cost = dashscope_cost_per_token(
+            model="qwen-reasoning-rate-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
+        )
+        assert math.isclose(peak_completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10)
+
+    def test_dashscope_off_peak_cache_creation_rate_replaces_the_standard_rate(self):
+        """Regression (LIT-6887): a block carrying cache_creation_input_token_cost bills cache-creation
+        tokens at it inside the window, while the cache-read rate it leaves unset stays standard."""
+        self._register_off_peak_flat_model(
+            "dashscope/qwen-cache-creation-off-peak-test",
+            {"hours_utc": self.OFF_PEAK_WINDOW, "cache_creation_input_token_cost": 1.5e-06},
+        )
+        usage = Usage(
+            prompt_tokens=1000,
+            completion_tokens=10,
+            prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=300, cache_creation_tokens=100),
+        )
+
+        prompt_cost, _ = dashscope_cost_per_token(
+            model="qwen-cache-creation-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+        assert math.isclose(prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 1.5e-06), rel_tol=1e-10)
+
+        peak_prompt_cost, _ = dashscope_cost_per_token(
+            model="qwen-cache-creation-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)
+
+    def test_dashscope_off_peak_reasoning_and_cache_creation_rates_override_the_selected_tier(self):
+        """The new keys override the selected tier the way the input and output rates already do."""
+        self._register_tiered_model(
+            "dashscope/qwen-tiered-reasoning-off-peak-test",
+            [
+                {
+                    "range": [0, 1000],
+                    "input_cost_per_token": 4e-07,
+                    "cache_creation_input_token_cost": 3e-07,
+                    "output_cost_per_token": 1.6e-06,
+                    "output_cost_per_reasoning_token": 3.2e-06,
+                },
+            ],
+        )
+        litellm.model_cost["dashscope/qwen-tiered-reasoning-off-peak-test"]["off_peak_pricing"] = {
+            "hours_utc": self.OFF_PEAK_WINDOW,
+            "cache_creation_input_token_cost": 1e-07,
+            "output_cost_per_reasoning_token": 8e-07,
+        }
+        usage = Usage(
+            prompt_tokens=500,
+            completion_tokens=100,
+            prompt_tokens_details=PromptTokensDetailsWrapper(cache_creation_tokens=200),
+            completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=40),
+        )
+
+        prompt_cost, completion_cost = dashscope_cost_per_token(
+            model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.INSIDE_WINDOW
+        )
+        assert math.isclose(prompt_cost, (300 * 4e-07) + (200 * 1e-07), rel_tol=1e-10)
+        assert math.isclose(completion_cost, (60 * 1.6e-06) + (40 * 8e-07), rel_tol=1e-10)
+
+        peak_prompt_cost, peak_completion_cost = dashscope_cost_per_token(
+            model="qwen-tiered-reasoning-off-peak-test", usage=usage, current_time=self.OUTSIDE_WINDOW
+        )
+        assert math.isclose(peak_prompt_cost, (300 * 4e-07) + (200 * 3e-07), rel_tol=1e-10)
+        assert math.isclose(peak_completion_cost, (60 * 1.6e-06) + (40 * 3.2e-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."""

From 19da217167766a0dbb7d1e4ef6ea4a8685f0f16e Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 14:38:47 -0700
Subject: [PATCH 05/38] fix(openai): mint workload identity tokens for
 PrivateLink and regional api.openai.com hosts

---
 litellm/llms/openai/workload_identity.py      |  5 +--
 .../openai/test_openai_workload_identity.py   | 43 +++++++++++++++++++
 2 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py
index ecec161ed46..283fdfb92c2 100644
--- a/litellm/llms/openai/workload_identity.py
+++ b/litellm/llms/openai/workload_identity.py
@@ -8,7 +8,7 @@ from urllib.parse import urlparse
 import litellm
 from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
 
-from .common_utils import OpenAIError
+from .common_utils import OpenAIError, is_openai_backed_api_base
 
 if TYPE_CHECKING:
     from collections.abc import Callable
@@ -16,7 +16,6 @@ if TYPE_CHECKING:
     from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth
 
 OPENAI_WIF_CLIENT_ID: Final = "litellm"
-_OPENAI_API_HOST: Final = "api.openai.com"
 _SDK_UPGRADE_MESSAGE: Final = (
     "OpenAI workload identity federation requires openai>=2.32.0. "
     "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / "
@@ -75,7 +74,7 @@ def _targets_openai_api(api_base: str | None) -> bool:
     if api_base is None:
         return True
     parsed: Final = urlparse(api_base)
-    return parsed.scheme == "https" and parsed.hostname == _OPENAI_API_HOST
+    return parsed.scheme == "https" and is_openai_backed_api_base(api_base)
 
 
 @lru_cache(maxsize=16)
diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py
index d8d9936e9a1..db107e00df0 100644
--- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py
+++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py
@@ -85,6 +85,31 @@ class TestResolveConfig:
     def test_plaintext_http_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
         assert resolve_openai_workload_identity_config(api_key=None, api_base="http://api.openai.com/v1") is None
 
+    @pytest.mark.parametrize(
+        "api_base",
+        (
+            "https://southcentralus.privatelink.api.openai.com/v1",
+            "https://eu.api.openai.com/v1",
+            "https://us.api.openai.com/v1",
+        ),
+    )
+    def test_openai_backed_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) == wif_env
+
+    @pytest.mark.parametrize(
+        "api_base",
+        (
+            "https://api.openai.com.evil.example/v1",
+            "https://openai.com/v1",
+            "https://euapi.openai.com/v1",
+            "http://southcentralus.privatelink.api.openai.com/v1",
+        ),
+    )
+    def test_lookalike_or_plaintext_api_base_disables(
+        self, wif_env: OpenAIWorkloadIdentityConfig, api_base: str
+    ) -> None:
+        assert resolve_openai_workload_identity_config(api_key=None, api_base=api_base) is None
+
     def test_foreign_env_base_url_disables(
         self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch
     ) -> None:
@@ -158,6 +183,14 @@ class TestClientConstruction:
         assert client.api_key == "workload-identity-auth"
         assert client._workload_identity_auth is not None
 
+    def test_privatelink_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        client: Final = OpenAIChatCompletion()._get_openai_client(
+            is_async=False, api_key=None, api_base="https://southcentralus.privatelink.api.openai.com/v1"
+        )
+        assert isinstance(client, OpenAI)
+        assert client.api_key == "workload-identity-auth"
+        assert client._workload_identity_auth is not None
+
     def test_static_key_client_unaffected(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
         client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key="sk-static", api_base=None)
         assert isinstance(client, OpenAI)
@@ -231,6 +264,16 @@ class TestResponsesValidateEnvironment:
         )
         assert headers["Authorization"] == "Bearer None"
 
+    @respx.mock
+    def test_privatelink_api_base_mints_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
+        mock_token_exchange()
+        headers: Final = OpenAIResponsesAPIConfig().validate_environment(
+            headers={},
+            model="gpt-4o-mini",
+            litellm_params=GenericLiteLLMParams(api_base="https://southcentralus.privatelink.api.openai.com/v1"),
+        )
+        assert headers["Authorization"] == "Bearer exchanged-bearer-token"
+
     def test_litellm_proxy_subclass_never_mints_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None:
         headers: Final = LiteLLMProxyResponsesAPIConfig().validate_environment(
             headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams()

From 942a46ffd7793ca1c535bf73ae1c8f3366976617 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 22:18:20 +0000
Subject: [PATCH 06/38] fix(caching): don't trip redis circuit breaker on short
 timeout bursts (#38999)

* fix(caching): don't trip redis circuit breaker on short timeout bursts

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

* fix(caching): scope timeout duration gate to timeout failures and count breaker states per label

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

* fix(caching): reset the timeout streak on hard failures so stale timeouts cannot pre-age the duration gate

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

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/caching/redis_cache.py                | 142 +++++++++++++++--
 litellm/constants.py                          |   3 +
 .../test_litellm/caching/test_redis_cache.py  | 150 +++++++++++++++++-
 3 files changed, 284 insertions(+), 11 deletions(-)

diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py
index 58733b384b9..aaee7188d86 100644
--- a/litellm/caching/redis_cache.py
+++ b/litellm/caching/redis_cache.py
@@ -27,6 +27,7 @@ from litellm.constants import (
     REDIS_CIRCUIT_BREAKER_ENABLED,
     REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
     REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT,
+    REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
 )
 from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
 from litellm.litellm_core_utils.coroutine_checker import coroutine_checker
@@ -41,6 +42,8 @@ from .base_cache import BaseCache
 
 if TYPE_CHECKING:
     from opentelemetry.trace import Span as _Span
+    from prometheus_client import Counter as _PromCounter
+    from prometheus_client import Gauge as _PromGauge
     from redis.asyncio import Redis, RedisCluster
     from redis.asyncio.client import Pipeline
     from redis.asyncio.cluster import ClusterPipeline
@@ -135,10 +138,20 @@ class RedisCircuitBreaker:
       HALF_OPEN - recovery probe: allow one request through
 
     Transitions:
-      CLOSED    -> OPEN      after failure_threshold consecutive failures
+      CLOSED    -> OPEN      after failure_threshold consecutive hard connectivity
+                             failures, or after an unbroken run of timeout failures
+                             (no success or hard failure in between) that reaches
+                             failure_threshold and spans timeout_min_duration seconds
       OPEN      -> HALF_OPEN after recovery_timeout seconds
       HALF_OPEN -> CLOSED    on success
       HALF_OPEN -> OPEN      on failure (resets timer)
+
+    Timeouts are accounted separately from hard connectivity failures because the async
+    Redis timeout includes time waiting for the worker event loop to resume: one loop
+    stall makes every in-flight operation time out together, which satisfies a purely
+    consecutive threshold instantly even though Redis is healthy. Requiring a
+    timeout-only streak to also span timeout_min_duration filters such bursts while a
+    real outage that surfaces as timeouts still opens the breaker after that duration.
     """
 
     CLOSED = "closed"
@@ -150,13 +163,19 @@ class RedisCircuitBreaker:
         failure_threshold: int,
         recovery_timeout: int,
         enabled: bool = True,
+        timeout_min_duration: float = REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION,
     ) -> None:
         self.failure_threshold = failure_threshold
         self.recovery_timeout = recovery_timeout
         self.enabled = enabled
+        self.timeout_min_duration = timeout_min_duration
         self._failure_count = 0
+        self._hard_failure_count = 0
+        self._timeout_count = 0
+        self._timeout_streak_started_at: float | None = None
         self._opened_at: float | None = None
         self._state = self.CLOSED
+        _breaker_metrics().record_state_change(None, self._state)
 
     def is_open(self) -> bool:
         """Returns True if Redis calls should be skipped."""
@@ -169,24 +188,45 @@ class RedisCircuitBreaker:
             return True
         if self._state == self.OPEN:
             if time.time() - (self._opened_at or 0) > self.recovery_timeout:
-                self._state = self.HALF_OPEN
+                self._set_state(self.HALF_OPEN)
                 return False  # this caller is the designated probe
             return True
         return False
 
-    def record_failure(self) -> None:
+    def _should_open(self, now: float) -> bool:
+        if self._state == self.HALF_OPEN:
+            return True
+        if self._hard_failure_count >= self.failure_threshold:
+            return True
+        if self._timeout_count < self.failure_threshold:
+            return False
+        return now - (self._timeout_streak_started_at or now) >= self.timeout_min_duration
+
+    def record_failure(self, is_timeout: bool = False) -> None:
         if not self.enabled:
             return
+        now: Final = time.time()
         self._failure_count += 1
-        self._opened_at = time.time()
-        if self._failure_count >= self.failure_threshold:
+        if is_timeout:
+            self._timeout_count += 1
+            if self._timeout_streak_started_at is None:
+                self._timeout_streak_started_at = now
+        else:
+            self._hard_failure_count += 1
+            self._timeout_count = 0
+            self._timeout_streak_started_at = None
+        self._opened_at = now
+        _breaker_metrics().record_failure("timeout" if is_timeout else "connectivity")
+        if self._should_open(now):
             if self._state != self.OPEN:
                 verbose_logger.warning(
-                    "Redis circuit breaker OPENED after %d consecutive failures — fast-failing Redis calls for %ds",
+                    "Redis circuit breaker OPENED after %d consecutive failures"
+                    " (%d hard connectivity) — fast-failing Redis calls for %ds",
                     self._failure_count,
+                    self._hard_failure_count,
                     self.recovery_timeout,
                 )
-            self._state = self.OPEN
+            self._set_state(self.OPEN)
 
     def record_success(self) -> None:
         if not self.enabled:
@@ -194,7 +234,17 @@ class RedisCircuitBreaker:
         if self._state == self.HALF_OPEN:
             verbose_logger.info("Redis circuit breaker CLOSED — Redis recovered")
         self._failure_count = 0
-        self._state = self.CLOSED
+        self._hard_failure_count = 0
+        self._timeout_count = 0
+        self._timeout_streak_started_at = None
+        self._set_state(self.CLOSED)
+
+    def _set_state(self, state: str) -> None:
+        if state == self._state:
+            return
+        _breaker_metrics().record_transition(state)
+        _breaker_metrics().record_state_change(self._state, state)
+        self._state = state
 
 
 _RedisCallResult = TypeVar("_RedisCallResult")
@@ -234,6 +284,78 @@ def _is_redis_health_failure(exc: BaseException) -> bool:
         return True
 
 
+@functools.lru_cache(maxsize=1)
+def _redis_timeout_error_types() -> tuple[type, ...]:
+    """Health failures that are timeouts rather than unambiguous connectivity errors.
+
+    ``builtins.TimeoutError`` covers ``asyncio.TimeoutError`` and ``socket.timeout``
+    (aliases since py3.11 / py3.10). ``redis.exceptions.TimeoutError`` does not subclass
+    either, so it is listed explicitly.
+    """
+    try:
+        from redis.exceptions import TimeoutError as RedisTimeoutError
+    except ImportError:
+        return (TimeoutError,)
+    return (RedisTimeoutError, TimeoutError)
+
+
+def _is_redis_timeout_failure(exc: BaseException) -> bool:
+    return isinstance(exc, _redis_timeout_error_types())
+
+
+class _BreakerMetrics:
+    """Prometheus metrics for the Redis circuit breaker; no-ops when the client is absent.
+
+    Registered lazily on the default registry (which /metrics serves) via the module-level
+    ``_breaker_metrics`` singleton so repeated RedisCache construction never re-registers.
+    """
+
+    def __init__(self) -> None:
+        self._state_gauge: _PromGauge | None = None
+        self._transitions: _PromCounter | None = None
+        self._failures: _PromCounter | None = None
+        try:
+            from prometheus_client import Counter as PromCounter
+            from prometheus_client import Gauge
+        except ImportError:
+            return
+        self._state_gauge = Gauge(
+            "litellm_redis_circuit_breaker_state",
+            "Number of Redis circuit breakers currently in each state",
+            labelnames=("state",),
+        )
+        self._transitions = PromCounter(
+            "litellm_redis_circuit_breaker_transitions",
+            "Redis circuit breaker state transitions",
+            labelnames=("state",),
+        )
+        self._failures = PromCounter(
+            "litellm_redis_circuit_breaker_failures",
+            "Redis health failures counted by the circuit breaker",
+            labelnames=("failure_class",),
+        )
+
+    def record_state_change(self, old_state: str | None, new_state: str) -> None:
+        if self._state_gauge is None:
+            return
+        if old_state is not None:
+            self._state_gauge.labels(old_state).dec()
+        self._state_gauge.labels(new_state).inc()
+
+    def record_transition(self, state: str) -> None:
+        if self._transitions is not None:
+            self._transitions.labels(state).inc()
+
+    def record_failure(self, failure_class: str) -> None:
+        if self._failures is not None:
+            self._failures.labels(failure_class).inc()
+
+
+@functools.lru_cache(maxsize=1)
+def _breaker_metrics() -> _BreakerMetrics:
+    return _BreakerMetrics()
+
+
 def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseException) -> None:
     """Record a Redis failure that the calling method is about to swallow.
 
@@ -245,7 +367,7 @@ def _record_swallowed_redis_failure(breaker: RedisCircuitBreaker, exc: BaseExcep
     """
     if not _is_redis_health_failure(exc):
         return
-    breaker.record_failure()
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(exc))
     _swallowed_redis_failures.set(_swallowed_redis_failures.get() + 1)
 
 
@@ -281,7 +403,7 @@ async def _run_under_circuit_breaker(
         result: Final = await call()
     except Exception as e:
         if _is_redis_health_failure(e):
-            breaker.record_failure()
+            breaker.record_failure(is_timeout=_is_redis_timeout_failure(e))
         raise
     _exit_circuit_breaker(breaker, swallowed_before)
     return result
diff --git a/litellm/constants.py b/litellm/constants.py
index f5acadc32ab..cd72adc3db5 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -432,6 +432,9 @@ REDIS_CONNECTION_POOL_TIMEOUT: Final = int(os.getenv("REDIS_CONNECTION_POOL_TIME
 REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD", 5))
 REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT: Final = int(os.getenv("REDIS_CIRCUIT_BREAKER_RECOVERY_TIMEOUT", 60))
 REDIS_CIRCUIT_BREAKER_ENABLED: Final = os.getenv("REDIS_CIRCUIT_BREAKER_ENABLED", "true").lower() == "true"
+# minimum seconds a timeout-only failure streak must span before it can open the breaker,
+# so one event-loop stall timing out many queued calls at once does not trip it
+REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION: Final = float(os.getenv("REDIS_CIRCUIT_BREAKER_TIMEOUT_MIN_DURATION", 5.0))
 # Seconds of idle before a Redis cluster connection is validated with a PING and
 # reconnected if dead, so a connection silently dropped by a cluster restart
 # (e.g. ElastiCache Serverless maintenance) is not reused while broken
diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py
index 71be8730df1..e4724ff8705 100644
--- a/tests/test_litellm/caching/test_redis_cache.py
+++ b/tests/test_litellm/caching/test_redis_cache.py
@@ -779,7 +779,7 @@ async def test_concurrent_success_is_not_cancelled_by_another_calls_failure():
     "error, opens_breaker",
     [
         pytest.param("ConnectionError", True, id="connection_refused_is_unhealthy"),
-        pytest.param("TimeoutError", True, id="timeout_is_unhealthy"),
+        pytest.param("TimeoutError", False, id="timeout_burst_is_ambiguous"),
         pytest.param("BusyLoadingError", True, id="loading_is_unhealthy"),
         pytest.param("ResponseError", False, id="wrong_type_command_is_not"),
         pytest.param("DataError", False, id="bad_data_is_not"),
@@ -791,6 +791,10 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker)
     They say nothing about connectivity, and a caller able to provoke them (an INCR against
     a non-numeric value, say) could otherwise trip the shared breaker on demand and drop
     rate limiting to per-process counters, which spreading traffic across replicas outruns.
+
+    A rapid burst of timeouts is ambiguous too: the async timeout includes event-loop
+    scheduling delay, so a loop stall times out every queued call at once against a
+    healthy Redis. It must not open the breaker until the streak spans a minimum duration.
     """
     import redis.exceptions
 
@@ -810,3 +814,147 @@ async def test_only_connectivity_failures_open_the_breaker(error, opens_breaker)
         await _run_under_circuit_breaker(breaker, "op", failing_call)
 
     assert breaker.is_open() is opens_breaker
+
+
+@pytest.mark.asyncio
+async def test_event_loop_stall_timeout_burst_keeps_breaker_closed():
+    """One blocking stall of the worker event loop must not trip the breaker.
+
+    Every operation already waiting on the loop times out together when the loop resumes,
+    so a purely consecutive threshold is satisfied instantly even though the Redis on the
+    other end (here an in-process fake that answers immediately) is healthy.
+    """
+    import time as time_mod
+
+    from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker
+
+    breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0)
+
+    async def healthy_redis_call_with_client_timeout():
+        return await asyncio.wait_for(asyncio.sleep(0.001, result="ok"), timeout=0.05)
+
+    async def stall_the_loop():
+        await asyncio.sleep(0)
+        time_mod.sleep(0.2)
+
+    results = await asyncio.gather(
+        *(_run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) for _ in range(8)),
+        stall_the_loop(),
+        return_exceptions=True,
+    )
+    timeouts = [r for r in results if isinstance(r, asyncio.TimeoutError)]
+    assert len(timeouts) >= breaker.failure_threshold, "the stall must time out a full burst"
+
+    assert breaker.is_open() is False, "a healthy Redis behind one loop stall must stay in the pool"
+    assert await _run_under_circuit_breaker(breaker, "op", healthy_redis_call_with_client_timeout) == "ok"
+
+
+@pytest.mark.asyncio
+async def test_persistent_timeouts_still_open_the_breaker():
+    """A real outage that surfaces only as timeouts must still open the breaker.
+
+    Once the timeout-only streak spans the minimum duration with no success in between,
+    Redis is genuinely unusable from this worker and protection has to kick in.
+    """
+    from redis.exceptions import TimeoutError as RedisTimeoutError
+
+    from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker
+
+    breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.1)
+
+    async def timing_out_call():
+        raise RedisTimeoutError("read timed out")
+
+    for _ in range(breaker.failure_threshold):
+        with pytest.raises(RedisTimeoutError):
+            await _run_under_circuit_breaker(breaker, "op", timing_out_call)
+    assert breaker.is_open() is False, "the burst has not spanned the minimum duration yet"
+
+    await asyncio.sleep(0.12)
+    with pytest.raises(RedisTimeoutError):
+        await _run_under_circuit_breaker(breaker, "op", timing_out_call)
+
+    assert breaker.is_open() is True
+
+
+@pytest.mark.asyncio
+async def test_stale_timeout_does_not_let_sub_threshold_hard_failures_open_the_breaker():
+    """Hard connectivity failures below the threshold must not open the breaker just
+    because an old timeout already started the streak and the duration has elapsed.
+
+    Each class has to earn the open on its own terms: hard failures by reaching the
+    threshold, timeouts by reaching the threshold and spanning the minimum duration.
+    """
+    from redis.exceptions import ConnectionError as RedisConnectionError
+    from redis.exceptions import TimeoutError as RedisTimeoutError
+
+    from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure
+
+    breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05)
+
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out")))
+    await asyncio.sleep(0.06)
+    for _ in range(breaker.failure_threshold - 1):
+        breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused")))
+    assert breaker.is_open() is False, "2 hard failures and 1 stale timeout are below both thresholds"
+
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused")))
+    assert breaker.is_open() is True, "the threshold-th hard failure must still open it"
+
+
+@pytest.mark.asyncio
+async def test_hard_failure_resets_timeout_streak_so_a_later_burst_must_earn_its_own_duration():
+    """A stale timeout followed by hard failures must not pre-age the duration gate:
+    a later short timeout burst has to span timeout_min_duration on its own.
+    """
+    from redis.exceptions import ConnectionError as RedisConnectionError
+    from redis.exceptions import TimeoutError as RedisTimeoutError
+
+    from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure
+
+    breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=0.05)
+
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out")))
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused")))
+    await asyncio.sleep(0.06)
+    for _ in range(breaker.failure_threshold):
+        breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out")))
+    assert breaker.is_open() is False, "the burst is instantaneous, so the duration gate must hold it closed"
+
+    await asyncio.sleep(0.06)
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("read timed out")))
+    assert breaker.is_open() is True, "the same run of timeouts persisting past the duration must open it"
+
+
+@pytest.mark.asyncio
+async def test_breaker_metrics_track_state_and_failure_class():
+    """Breaker accounting must be observable: failure class, transitions, and state."""
+    from prometheus_client import REGISTRY
+    from redis.exceptions import ConnectionError as RedisConnectionError
+    from redis.exceptions import TimeoutError as RedisTimeoutError
+
+    from litellm.caching.redis_cache import RedisCircuitBreaker, _is_redis_timeout_failure
+
+    def sample(name, labels=None):
+        return REGISTRY.get_sample_value(name, labels) or 0.0
+
+    timeout_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"})
+    hard_before = sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"})
+    opened_before = sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"})
+    open_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "open"})
+    closed_gauge_before = sample("litellm_redis_circuit_breaker_state", {"state": "closed"})
+
+    breaker = RedisCircuitBreaker(failure_threshold=2, recovery_timeout=60, timeout_min_duration=5.0)
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisTimeoutError("t")))
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused")))
+    breaker.record_failure(is_timeout=_is_redis_timeout_failure(RedisConnectionError("refused")))
+
+    assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "timeout"}) == timeout_before + 1
+    assert sample("litellm_redis_circuit_breaker_failures_total", {"failure_class": "connectivity"}) == hard_before + 2
+    assert sample("litellm_redis_circuit_breaker_transitions_total", {"state": "open"}) == opened_before + 1
+    assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before + 1
+    assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before
+
+    breaker.record_success()
+    assert sample("litellm_redis_circuit_breaker_state", {"state": "open"}) == open_gauge_before
+    assert sample("litellm_redis_circuit_breaker_state", {"state": "closed"}) == closed_gauge_before + 1

From 8e4397d5f1358d8465ce80c28a56cb60f3faa997 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:19:49 -0700
Subject: [PATCH 07/38] chore(ci): rebuild the PR merge ref against staging's
 router coverage fix


From 4811041048fb8bca24bf00938d0111e060a8b686 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 15:20:13 -0700
Subject: [PATCH 08/38] feat(proxy): add a search param to key, memory, audit,
 and spend log listings

GET /key/list?search= matches the key hash (a raw sk- key is hashed
first) or a case-insensitive alias substring, and key_hash= now hashes a
raw sk- value too. GET /v1/memory?search= matches a key prefix or an
exact memory_id. GET /audit?search= matches id, object_id, changed_by,
or changed_by_api_key. GET /spend/logs/ui?search= matches request_id
across all time and api_key, team_id, user, end_user, session_id, or
model_id inside the date window; session grouping is skipped while a
search is active.

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
---
 .../proxy/audit_logging_endpoints.py          |  28 +-
 .../key_management_endpoints.py               |  30 +-
 litellm/proxy/memory/memory_endpoints.py      |  56 +++-
 .../spend_management_endpoints.py             |  62 +++-
 .../key_management_endpoints.py               |  17 ++
 .../proxy/test_audit_logging_endpoints.py     |  66 +++-
 .../test_key_management_endpoints.py          | 152 +++++++++
 .../proxy/memory/test_memory_endpoints.py     |  90 ++++++
 .../test_spend_management_endpoints.py        | 288 +++++++++++++++++-
 9 files changed, 767 insertions(+), 22 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
index b6f8bf2dc5b..72f7a66a420 100644
--- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
+++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
@@ -18,6 +18,7 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
 
 from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.utils import _hash_token_if_needed
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import AuditLogRepository
 
@@ -48,6 +49,19 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec
     }
 
 
+def _build_search_condition(search: str) -> dict[str, object]:
+    """Match any id column; a raw sk- key is hashed for the two columns that store key hashes."""
+    hashed: Final = _hash_token_if_needed(search)
+    return {
+        "OR": (
+            {"id": search},
+            {"changed_by": search},
+            {"object_id": hashed},
+            {"changed_by_api_key": hashed},
+        )
+    }
+
+
 @router.get(
     "/audit",
     tags=["Audit Logging"],
@@ -83,6 +97,13 @@ async def get_audit_logs(
         None,
         description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
     ),
+    search: str | None = Query(
+        None,
+        description=(
+            "Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value "
+            "(a raw sk- virtual key is hashed first)"
+        ),
+    ),
     # Sorting parameters
     sort_by: str | None = Query(
         None,
@@ -118,6 +139,11 @@ async def get_audit_logs(
         *([_build_json_field_or_condition("token", object_key_hash)] if object_key_hash else []),
     ]
 
+    and_conditions: Final[tuple[dict[str, object], ...]] = (
+        *json_field_conditions,
+        *((_build_search_condition(search),) if search else ()),
+    )
+
     # Build filter conditions
     where_conditions: Final[dict[str, object]] = {
         **({"changed_by": changed_by} if changed_by else {}),
@@ -126,7 +152,7 @@ async def get_audit_logs(
         **({"table_name": table_name} if table_name else {}),
         **({"object_id": object_id} if object_id else {}),
         **({"updated_at": date_filter} if start_date or end_date else {}),
-        **({"AND": json_field_conditions} if json_field_conditions else {}),
+        **({"AND": and_conditions} if and_conditions else {}),
     }
 
     order_by: Final[dict[str, str]] = (
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index d7d20d168b5..758644ff01b 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -147,6 +147,7 @@ from litellm.types.proxy.management_endpoints.key_management_endpoints import (
     BulkUpdateKeyResponse,
     BulkUpdateTeamKeysRequest,
     FailedKeyUpdate,
+    KeySearchWhere,
     SuccessfulKeyUpdate,
 )
 from litellm.types.router import Deployment
@@ -5800,6 +5801,10 @@ async def list_keys(
         None,
         description="Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching.",
     ),
+    search: str | None = Query(
+        None,
+        description="Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive).",
+    ),
     return_full_object: bool = Query(False, description="Return full key object"),
     include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."),
     include_created_by_keys: bool = Query(False, description="Include keys created by the user"),
@@ -5862,13 +5867,17 @@ async def list_keys(
                 detail={"error": "Invalid expires value. Supported: 'active', 'expired'."},
             )
 
+        hashed_key_hash: Final[str | None] = (
+            _hash_token_if_needed(token=key_hash) if isinstance(key_hash, str) else None
+        )
+
         complete_user_info: Final = await validate_key_list_check(
             user_api_key_dict=user_api_key_dict,
             user_id=user_id,
             team_id=team_id,
             organization_id=organization_id,
             key_alias=key_alias,
-            key_hash=key_hash,
+            key_hash=hashed_key_hash,
             prisma_client=prisma_client,
         )
 
@@ -5928,7 +5937,7 @@ async def list_keys(
             user_id=user_id,
             team_id=team_id,
             key_alias=key_alias,
-            key_hash=key_hash,
+            key_hash=hashed_key_hash,
             return_full_object=return_full_object,
             organization_id=organization_id,
             admin_team_ids=admin_team_ids,
@@ -5943,6 +5952,7 @@ async def list_keys(
             agent_id=agent_id,
             use_substring_matching=use_substring_matching,
             expires_filter=expires if isinstance(expires, str) else None,
+            search=search,
         )
 
         verbose_proxy_logger.debug("Successfully prepared response")
@@ -6162,6 +6172,16 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str,
     return {"OR": [{"expires": None}, {"expires": {"gte": now}}]}
 
 
+def _build_key_search_where(search: str) -> KeySearchWhere:
+    search_where: Final[KeySearchWhere] = {
+        "OR": (
+            {"token": _hash_token_if_needed(token=search)},
+            {"key_alias": {"contains": search, "mode": "insensitive"}},
+        )
+    }
+    return search_where
+
+
 def _build_key_filter_conditions(
     user_id: str | None,
     team_id: str | None,
@@ -6177,6 +6197,7 @@ def _build_key_filter_conditions(
     agent_id: str | None = None,
     use_substring_matching: bool = False,
     expires_filter: str | None = None,
+    search: str | None = None,
 ) -> Mapping[str, object]:
     """Build filter conditions for key listing.
 
@@ -6266,7 +6287,7 @@ def _build_key_filter_conditions(
 
     # Apply team_id, project_id and access_group_id as global AND filters so they
     # narrow results across all visibility conditions (own keys, team keys, etc.)
-    global_filters: Final[tuple[dict[str, object], ...]] = (
+    global_filters: Final[tuple[Mapping[str, object], ...]] = (
         *(
             (
                 {"key_alias": {"contains": key_alias, "mode": "insensitive"}}
@@ -6277,6 +6298,7 @@ def _build_key_filter_conditions(
             else ()
         ),
         *(({"token": key_hash},) if key_hash and isinstance(key_hash, str) else ()),
+        *((_build_key_search_where(search),) if isinstance(search, str) and search else ()),
         *(({"team_id": team_id},) if team_id and isinstance(team_id, str) else ()),
         *(({"project_id": project_id},) if project_id else ()),
         *(({"access_group_ids": {"hasSome": [access_group_id]}},) if access_group_id else ()),
@@ -6316,6 +6338,7 @@ async def _list_key_helper(
     agent_id: str | None = None,
     use_substring_matching: bool = False,
     expires_filter: str | None = None,
+    search: str | None = None,
 ) -> KeyListResponseObject:
     """
     Helper function to list keys
@@ -6354,6 +6377,7 @@ async def _list_key_helper(
         agent_id=agent_id,
         use_substring_matching=use_substring_matching,
         expires_filter=expires_filter,
+        search=search,
     )
 
     # Calculate skip for pagination
diff --git a/litellm/proxy/memory/memory_endpoints.py b/litellm/proxy/memory/memory_endpoints.py
index 98c5fdd198c..d8f72d200c7 100644
--- a/litellm/proxy/memory/memory_endpoints.py
+++ b/litellm/proxy/memory/memory_endpoints.py
@@ -22,6 +22,7 @@ from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final
 
 from fastapi import APIRouter, Depends, HTTPException, Query
+from typing_extensions import ReadOnly, TypedDict
 
 from litellm._logging import verbose_proxy_logger
 from litellm.proxy._types import (
@@ -91,6 +92,36 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object
     return {"OR": ors}
 
 
+class _StartsWith(TypedDict):
+    startsWith: ReadOnly[str]
+
+
+class _MemoryKeyWhere(TypedDict):
+    key: ReadOnly[str | _StartsWith]
+
+
+class _MemoryIdWhere(TypedDict):
+    memory_id: ReadOnly[str]
+
+
+class _MemorySearchWhere(TypedDict):
+    OR: ReadOnly[tuple[_MemoryKeyWhere, _MemoryIdWhere]]
+
+
+def _key_filter(search: str | None, key_prefix: str | None, key: str | None) -> Mapping[str, object] | None:
+    """`search` matches a key prefix or an exact memory_id; otherwise `key_prefix` wins over `key`."""
+    if search is not None:
+        search_where: Final[_MemorySearchWhere] = {"OR": ({"key": {"startsWith": search}}, {"memory_id": search})}
+        return search_where
+    if key_prefix is not None:
+        prefix_where: Final[_MemoryKeyWhere] = {"key": {"startsWith": key_prefix}}
+        return prefix_where
+    if key is not None:
+        exact_where: Final[_MemoryKeyWhere] = {"key": key}
+        return exact_where
+    return None
+
+
 def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow:
     return LiteLLM_MemoryRow(
         memory_id=row.memory_id,
@@ -326,6 +357,13 @@ async def list_memory(
             "Mutually exclusive with `key`; if both are provided, `key_prefix` wins."
         ),
     ),
+    search: str | None = Query(
+        None,
+        description=(
+            "Match entries whose key starts with this value or whose memory_id equals it. "
+            "Takes precedence over `key_prefix` and `key` when provided."
+        ),
+    ),
     page: int = Query(1, ge=1),
     page_size: int = Query(50, ge=1, le=500),
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
@@ -333,22 +371,16 @@ async def list_memory(
     """List memory entries visible to the caller."""
     prisma_client: Final = _require_prisma()
 
-    # Build the key filter first (prefix wins if both `key` and `key_prefix`
-    # are passed). Then AND it with the visibility filter via an explicit
-    # top-level "AND" — safer than `dict.update` since future visibility
-    # filters could grow an "OR" key that would clobber this one if merged
-    # by key.
-    key_filter: Final[dict[str, object]] = {}
-    if key_prefix is not None:
-        key_filter["key"] = {"startsWith": key_prefix}
-    elif key is not None:
-        key_filter["key"] = key
+    # AND the key filter with the visibility filter via an explicit top-level
+    # "AND": both sides can carry an "OR" key (`search`, non-admin visibility),
+    # so merging them by key would let one clobber the other and leak rows.
+    key_filter: Final = _key_filter(search=search, key_prefix=key_prefix, key=key)
 
     vis: Final = _visibility_filter(user_api_key_dict)
-    where: Mapping[str, object]
+    where: Mapping[str, object] | None
     if vis is None:
         where = key_filter
-    elif not key_filter:
+    elif key_filter is None:
         where = vis
     else:
         where = {"AND": [key_filter, vis]}
diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 2a50d5170f0..9ec8dd205a6 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -2229,6 +2229,33 @@ async def calculate_spend(request: SpendCalculateRequest):
         )
 
 
+class _SpendLogSearchCondition(NamedTuple):
+    sql: str
+    params: tuple[object, ...]
+
+
+def _build_spend_log_search_condition(
+    search: str,
+    start_date: datetime,
+    end_date: datetime,
+    next_param_index: int,
+) -> _SpendLogSearchCondition:
+    """request_id (indexed) matches across all time; the unindexed id columns only inside the window (sk- keys hashed)."""
+    raw: Final = f"${next_param_index}"
+    hashed: Final = f"${next_param_index + 1}"
+    window_start: Final = f"${next_param_index + 2}"
+    window_end: Final = f"${next_param_index + 3}"
+    sql: Final = (
+        f"(request_id = {raw} OR ("
+        f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') "
+        f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') "
+        f'AND (api_key = {hashed} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} '
+        f"OR session_id = {raw} OR model_id = {raw})))"
+    )
+    hashed_search: Final = hash_token(token=search) if search.startswith("sk-") else search
+    return _SpendLogSearchCondition(sql=sql, params=(search, hashed_search, start_date, end_date))
+
+
 @router.get(
     "/spend/logs/v2",
     tags=["Budget & Spend Tracking"],
@@ -2329,6 +2356,14 @@ async def ui_view_spend_logs(
             "UI route only, honored when sorting by startTime"
         ),
     ),
+    search: str | None = fastapi.Query(
+        default=None,
+        description=(
+            "Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, "
+            "session_id, or model_id equals this value. request_id matches across all time; the other columns "
+            "match inside start_date/end_date, which stay required"
+        ),
+    ),
 ):
     """
     View spend logs with pagination support.
@@ -2392,8 +2427,10 @@ async def ui_view_spend_logs(
     try:
         is_admin_view: Final = _is_admin_view_safe(user_api_key_dict=user_api_key_dict)
         is_request_id_lookup: Final = request_id is not None and not is_v2
+        is_search_lookup: Final = search is not None
+        search_owns_window: Final = is_search_lookup and not is_v2
 
-        if is_request_id_lookup:
+        if is_request_id_lookup and not is_search_lookup:
             # request_id is the @id primary key: it identifies a single row, so a
             # time window is meaningless. The dashboard always sends a default 24h
             # window, which hid ids copied from an older page (LIT-3981). Drop the
@@ -2576,7 +2613,7 @@ async def ui_view_spend_logs(
         # Date range. Wrap the param side with `AT TIME ZONE 'UTC'` so comparison
         # against the plain `timestamp` column does not depend on the DB session
         # timezone (see #22529). Absent for a request_id-only lookup (see above).
-        if start_date_obj is not None and end_date_obj is not None:
+        if start_date_obj is not None and end_date_obj is not None and not search_owns_window:
             sql_conditions.append(f"\"startTime\" >= (${p}::timestamptz AT TIME ZONE 'UTC')")
             sql_params.append(start_date_obj)
             p += 1
@@ -2584,6 +2621,17 @@ async def ui_view_spend_logs(
             sql_params.append(end_date_obj)
             p += 1
 
+        if search is not None and start_date_obj is not None and end_date_obj is not None:
+            search_condition: Final = _build_spend_log_search_condition(
+                search=search,
+                start_date=start_date_obj,
+                end_date=end_date_obj,
+                next_param_index=p,
+            )
+            sql_conditions.append(search_condition.sql)
+            sql_params.extend(search_condition.params)
+            p += len(search_condition.params)  # rebind-ok: advances the file's shared $N placeholder counter
+
         # Equality filters - read effective values from where_conditions (post-authorization)
         for sql_col, wc_key in [
             ("team_id", "team_id"),
@@ -2662,7 +2710,13 @@ async def ui_view_spend_logs(
             sql_params.append(f"%{error_message}%")
             p += 1
 
-        if group_by_session is True and not is_v2 and not is_request_id_lookup and sort_by == "startTime":
+        if (
+            group_by_session is True
+            and not is_v2
+            and not is_request_id_lookup
+            and not is_search_lookup
+            and sort_by == "startTime"
+        ):
             return await _ui_session_grouped_spend_logs(
                 prisma_client=prisma_client,
                 sql_conditions=sql_conditions,
@@ -2696,7 +2750,7 @@ async def ui_view_spend_logs(
             _order_expr = order_column
 
         joined_conditions: Final = " AND ".join(sql_conditions)
-        session_grouping: Final = group_by_session is True
+        session_grouping: Final = group_by_session is True and not is_search_lookup
         count_group_clause: Final = f"GROUP BY {_SESSION_GROUP_KEY_SQL}" if session_grouping else ""
         count_query: Final = f"""
             SELECT COUNT(*) AS total_count
diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py
index 0f17f2f23ab..5d410d7b55b 100644
--- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py
@@ -2,6 +2,23 @@ from datetime import datetime
 from typing import Any, Final, Literal
 
 from pydantic import BaseModel, ConfigDict, model_validator
+from typing_extensions import ReadOnly, TypedDict
+
+from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains
+
+
+class KeyTokenWhere(TypedDict):
+    token: ReadOnly[str]
+
+
+class KeyAliasContainsWhere(TypedDict):
+    key_alias: ReadOnly[InsensitiveContains]
+
+
+class KeySearchWhere(TypedDict):
+    """Prisma filter behind `/key/list?search=`: exact token (sk- keys hashed) or alias substring, case-insensitive."""
+
+    OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
 
 
 class BulkUpdateKeyRequestItem(BaseModel):
diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
index a0a26c089eb..cd2c8b0b904 100644
--- a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
+++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
@@ -1,5 +1,7 @@
+import hashlib
 from datetime import datetime, timedelta
-from unittest.mock import AsyncMock, patch
+from typing import Final
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastapi import FastAPI
@@ -8,10 +10,12 @@ from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_rou
 from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse
 
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 
 # Create an app with just the audit router for testing
 app = FastAPI()
 app.include_router(audit_router)
+app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin")
 client = TestClient(app)
 
 # Mock data for testing
@@ -130,3 +134,63 @@ async def test_get_audit_log_by_id_not_found(mock_prisma_client):
         data = response.json()
         assert "message" in data["detail"]
         assert "not found" in data["detail"]["message"].lower()
+
+
+def _list_audit_logs_where(mock_prisma_client: MagicMock, query: str) -> dict[str, object]:
+    mock_prisma_client.db.litellm_auditlog.find_many.return_value = []
+    mock_prisma_client.db.litellm_auditlog.count.return_value = 0
+
+    response: Final = client.get(f"/audit?{query}")
+
+    assert response.status_code == 200, response.text
+    find_many_where: Final = mock_prisma_client.db.litellm_auditlog.find_many.call_args.kwargs["where"]
+    assert mock_prisma_client.db.litellm_auditlog.count.call_args.kwargs["where"] == find_many_where
+    return find_many_where
+
+
+def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_client):
+    where: Final = _list_audit_logs_where(mock_prisma_client, "search=abc-123&action=create&object_team_id=team-1")
+
+    assert where == {
+        "action": "create",
+        "AND": (
+            {
+                "OR": [
+                    {"before_value": {"path": ["team_id"], "string_contains": "team-1"}},
+                    {"updated_values": {"path": ["team_id"], "string_contains": "team-1"}},
+                ]
+            },
+            {
+                "OR": (
+                    {"id": "abc-123"},
+                    {"changed_by": "abc-123"},
+                    {"object_id": "abc-123"},
+                    {"changed_by_api_key": "abc-123"},
+                )
+            },
+        ),
+    }
+
+
+def test_search_hashes_a_raw_virtual_key_for_the_hashed_columns(mock_prisma_client):
+    where: Final = _list_audit_logs_where(mock_prisma_client, "search=sk-raw")
+
+    hashed: Final = hashlib.sha256(b"sk-raw").hexdigest()
+    assert where == {
+        "AND": (
+            {
+                "OR": (
+                    {"id": "sk-raw"},
+                    {"changed_by": "sk-raw"},
+                    {"object_id": hashed},
+                    {"changed_by_api_key": hashed},
+                )
+            },
+        )
+    }
+
+
+def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client):
+    where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=")
+
+    assert where == {"action": "create"}
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 7954a4693cc..1e399e4fb58 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
@@ -6347,6 +6347,114 @@ def test_build_key_filter_conditions_key_hash_narrows_team_admin_visibility():
     assert {"token": "hashed-token-123"} in where["AND"], f"key_hash not ANDed: {where}"
 
 
+def _search_clause(search: str, token: str) -> dict:
+    return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]}
+
+
+def test_build_key_filter_conditions_search_hashes_raw_key_and_ors_alias_contains():
+    """
+    LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by
+    its ID. A pasted raw sk- key is hashed to its token first; an already-hashed
+    value is used verbatim.
+    """
+    from litellm.proxy._types import hash_token
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _build_key_filter_conditions,
+    )
+
+    raw_where = json.loads(
+        json.dumps(
+            _build_key_filter_conditions(
+                user_id=None,
+                team_id=None,
+                organization_id=None,
+                key_alias=None,
+                key_hash=None,
+                exclude_team_id=None,
+                admin_team_ids=None,
+                search="sk-raw",
+            )
+        )
+    )
+    assert _search_clause("sk-raw", hash_token("sk-raw")) in raw_where["AND"], f"raw search not ANDed: {raw_where}"
+
+    hashed_where = json.loads(
+        json.dumps(
+            _build_key_filter_conditions(
+                user_id=None,
+                team_id=None,
+                organization_id=None,
+                key_alias=None,
+                key_hash=None,
+                exclude_team_id=None,
+                admin_team_ids=None,
+                search="already-hashed-token",
+            )
+        )
+    )
+    assert _search_clause("already-hashed-token", "already-hashed-token") in hashed_where["AND"], (
+        f"hashed search not used verbatim: {hashed_where}"
+    )
+
+
+def test_build_key_filter_conditions_search_narrows_team_admin_visibility():
+    """
+    LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it
+    narrows a team admin's admin-team branch instead of being bypassed by it.
+    """
+    from litellm.proxy._types import hash_token
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _build_key_filter_conditions,
+    )
+
+    where = json.loads(
+        json.dumps(
+            _build_key_filter_conditions(
+                user_id="team-admin-user",
+                team_id=None,
+                organization_id=None,
+                key_alias=None,
+                key_hash=None,
+                exclude_team_id=None,
+                admin_team_ids=["team-a"],
+                member_team_ids=["team-a"],
+                include_created_by_keys=False,
+                search="sk-member",
+            )
+        )
+    )
+
+    assert where.get("AND"), f"expected top-level AND, got: {where}"
+    assert _search_clause("sk-member", hash_token("sk-member")) in where["AND"], f"search not ANDed: {where}"
+    assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where)
+
+
+@pytest.mark.asyncio
+async def test_list_key_helper_applies_search_to_prisma_where():
+    """LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause."""
+    from litellm.proxy._types import hash_token
+
+    mock_prisma_client = AsyncMock()
+    mock_find_many = AsyncMock(return_value=[])
+    mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
+    mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0)
+
+    await _list_key_helper(
+        prisma_client=mock_prisma_client,
+        page=1,
+        size=50,
+        user_id=None,
+        team_id=None,
+        organization_id=None,
+        key_alias=None,
+        key_hash=None,
+        search="sk-raw",
+    )
+
+    where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"]))
+    assert _search_clause("sk-raw", hash_token("sk-raw")) in where["AND"], f"search not in Prisma where: {where}"
+
+
 @pytest.mark.asyncio
 async def test_generate_key_negative_max_budget():
     """
@@ -14870,6 +14978,50 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
     assert kwargs["user_id"] == "alice"
 
 
+@pytest.mark.asyncio
+async def test_list_keys_hashes_raw_key_hash_before_validation():
+    """LIT-4741: a raw sk- key pasted as key_hash is hashed before the ownership
+    check and the query, so a non-admin filtering by their own raw key gets the
+    row instead of the 'Key Hash not found.' 403."""
+    from litellm.proxy._types import hash_token
+
+    user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
+    validate = AsyncMock(
+        return_value=LiteLLM_UserTable(
+            user_id="alice", user_email="alice@example.com", teams=[], organization_memberships=[]
+        )
+    )
+    helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0})
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
+            validate,
+        ),
+        patch("litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", helper),
+    ):
+        await list_keys(
+            request=MagicMock(),
+            user_api_key_dict=user,
+            status=None,
+            user_id=None,
+            key_hash="sk-raw",
+        )
+
+    assert validate.call_args.kwargs["key_hash"] == hash_token("sk-raw")
+    assert helper.call_args.kwargs["key_hash"] == hash_token("sk-raw")
+
+
+@pytest.mark.asyncio
+async def test_list_keys_search_is_honored_for_non_admin():
+    """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's
+    search reaches the helper while their own-user scoping stays in place."""
+    user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
+    kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="sk-raw")
+    assert kwargs["search"] == "sk-raw"
+    assert kwargs["user_id"] == "alice"
+
+
 @pytest.mark.asyncio
 async def test_cli_session_token_delegation_ceiling_blocked_by_team_budget():
     team = LiteLLM_TeamTableCachedObj(team_id="team-1", max_budget=50.0)
diff --git a/tests/test_litellm/proxy/memory/test_memory_endpoints.py b/tests/test_litellm/proxy/memory/test_memory_endpoints.py
index be75d980d9d..dff0e80fa77 100644
--- a/tests/test_litellm/proxy/memory/test_memory_endpoints.py
+++ b/tests/test_litellm/proxy/memory/test_memory_endpoints.py
@@ -615,6 +615,96 @@ class TestMemoryEndpoints:
         assert keys == {"user:profile"}
         assert body["total"] == 1
 
+    def test_list_memory_search_matches_key_prefix_or_memory_id_within_scope(self):
+        """
+        `search` matches a key prefix OR an exact memory_id, and stays ANDed
+        with the visibility filter so a pasted foreign id cannot leak a row.
+        """
+        table = self.prisma.db.litellm_memorytable
+        table.rows.extend(
+            [
+                _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None),
+                _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None),
+                _make_row(memory_id="mem-foreign", key="user:secret", user_id="user-b", team_id=None),
+            ]
+        )
+        client = _make_client(_user_auth("user-a", "team-a"))
+        with _patch_prisma(self.prisma):
+            by_id = client.get("/v1/memory?search=mem-target")
+            by_prefix = client.get("/v1/memory?search=user:")
+            foreign_id = client.get("/v1/memory?search=mem-foreign")
+
+        assert by_id.status_code == 200, by_id.text
+        assert [m["memory_id"] for m in by_id.json()["memories"]] == ["mem-target"]
+        assert by_id.json()["total"] == 1
+
+        assert by_prefix.status_code == 200, by_prefix.text
+        assert {m["key"] for m in by_prefix.json()["memories"]} == {"user:profile"}
+        assert by_prefix.json()["total"] == 1
+
+        assert foreign_id.status_code == 200, foreign_id.text
+        assert foreign_id.json()["memories"] == []
+        assert foreign_id.json()["total"] == 0
+
+    def test_list_memory_search_by_memory_id_for_admin_sees_any_scope(self):
+        """Admins have no visibility filter, so an id search returns the row whoever owns it."""
+        table = self.prisma.db.litellm_memorytable
+        table.rows.extend(
+            [
+                _make_row(memory_id="mem-a", key="a", user_id="user-a", team_id=None),
+                _make_row(memory_id="mem-b", key="b", user_id="user-b", team_id=None),
+            ]
+        )
+        client = _make_client(_admin_auth())
+        with _patch_prisma(self.prisma):
+            resp = client.get("/v1/memory?search=mem-b")
+        assert resp.status_code == 200, resp.text
+        assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-b"]
+        assert resp.json()["total"] == 1
+
+    def test_list_memory_search_wins_over_key_prefix(self):
+        """When both are sent, `search` decides the match and `key_prefix` is ignored."""
+        table = self.prisma.db.litellm_memorytable
+        table.rows.extend(
+            [
+                _make_row(memory_id="mem-own", key="user:profile", user_id="user-a", team_id=None),
+                _make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None),
+            ]
+        )
+        client = _make_client(_user_auth("user-a", "team-a"))
+        with _patch_prisma(self.prisma):
+            resp = client.get("/v1/memory?search=mem-target&key_prefix=user:")
+        assert resp.status_code == 200, resp.text
+        assert [m["memory_id"] for m in resp.json()["memories"]] == ["mem-target"]
+        assert resp.json()["total"] == 1
+
+    def test_list_memory_key_prefix_never_matches_memory_id(self):
+        """`key_prefix` stays a pure key-prefix match; only `search` consults memory_id."""
+        table = self.prisma.db.litellm_memorytable
+        table.rows.append(_make_row(memory_id="mem-target", key="project:context", user_id="user-a", team_id=None))
+        client = _make_client(_user_auth("user-a", "team-a"))
+        with _patch_prisma(self.prisma):
+            resp = client.get("/v1/memory?key_prefix=mem-target")
+        assert resp.status_code == 200, resp.text
+        assert resp.json()["memories"] == []
+        assert resp.json()["total"] == 0
+
+    def test_list_memory_key_exact_filter(self):
+        """`key` is an exact match, never a prefix."""
+        table = self.prisma.db.litellm_memorytable
+        table.rows.extend(
+            [
+                _make_row(memory_id="m1", key="user:profile", user_id="user-a", team_id=None),
+                _make_row(memory_id="m2", key="user:profile:archived", user_id="user-a", team_id=None),
+            ]
+        )
+        client = _make_client(_user_auth("user-a", "team-a"))
+        with _patch_prisma(self.prisma):
+            resp = client.get("/v1/memory?key=user:profile")
+        assert resp.status_code == 200, resp.text
+        assert [m["memory_id"] for m in resp.json()["memories"]] == ["m1"]
+        assert resp.json()["total"] == 1
+
     def test_list_memory_admin_sees_all(self):
         table = self.prisma.db.litellm_memorytable
         table.rows.extend(
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 f4cd8814bc1..f869f3ffba2 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
@@ -58,6 +58,25 @@ def _filter_logs_by_date_range(logs, where):
     return filtered
 
 
+_SEARCH_CLAUSE_RE = re.compile(
+    r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
+    r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
+    r'AND \(api_key = \$(\d+) OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 '
+    r"OR session_id = \$\1 OR model_id = \$\1\)\)\)"
+)
+
+
+def _matches_spend_log_search(log, search):
+    """Mirror the search clause: request_id across all time, the other id columns inside the window."""
+    if log.get("request_id") == search["value"]:
+        return True
+    if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}):
+        return False
+    if log.get("api_key") == search["api_key"]:
+        return True
+    return any(log.get(col) == search["value"] for col in ("team_id", "user", "end_user", "session_id", "model_id"))
+
+
 def _reconstruct_ui_where_from_sql(sql_query, params):
     """
     Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the
@@ -77,6 +96,17 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
     def _iso(value):
         return value.isoformat() if hasattr(value, "isoformat") else str(value)
 
+    search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1))
+    if search_clause:
+        raw_index, start_index, end_index, hashed_index = (int(g) for g in search_clause.groups())
+        where["search"] = {
+            "value": params[raw_index - 1],
+            "api_key": params[hashed_index - 1],
+            "gte": _iso(params[start_index - 1]),
+            "lte": _iso(params[end_index - 1]),
+        }
+    remaining = clause.group(1) if search_clause is None else clause.group(1).replace(search_clause.group(0), "")
+
     eq_cols = {
         "team_id": "team_id",
         '"user"': "user",
@@ -89,7 +119,7 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
     }
     date_bounds: dict = {}
     metadata_conds: list = []
-    for cond in (c.strip() for c in clause.group(1).split(" AND ")):
+    for cond in (c.strip() for c in remaining.split(" AND ")):
         gte = re.search(r'"startTime" >= \(\$(\d+)', cond)
         lte = re.search(r'"startTime" <= \(\$(\d+)', cond)
         alias = re.search(r"user_api_key_alias' LIKE \$(\d+)", cond)
@@ -2352,6 +2382,219 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
         app.dependency_overrides.pop(ps.user_api_key_auth, None)
 
 
+def test_build_spend_log_search_condition_windows_every_branch_except_request_id():
+    """LIT-4741: request_id matches across all time; the six other id columns only inside the window,
+    and a raw sk- key is hashed for the api_key branch alone."""
+    start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
+    end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
+
+    condition = spend_management_endpoints._build_spend_log_search_condition(
+        search="sk-raw-key", start_date=start, end_date=end, next_param_index=3
+    )
+
+    assert condition.sql == (
+        "(request_id = $3 OR (\"startTime\" >= ($5::timestamptz AT TIME ZONE 'UTC') "
+        "AND \"startTime\" <= ($6::timestamptz AT TIME ZONE 'UTC') "
+        'AND (api_key = $4 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))'
+    )
+    assert condition.params == ("sk-raw-key", hashlib.sha256(b"sk-raw-key").hexdigest(), start, end)
+
+
+def test_build_spend_log_search_condition_leaves_non_key_values_unhashed():
+    start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
+    end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
+
+    condition = spend_management_endpoints._build_spend_log_search_condition(
+        search="sess-42", start_date=start, end_date=end, next_param_index=1
+    )
+
+    assert condition.params == ("sess-42", "sess-42", start, end)
+
+
+def _search_fixture_logs(today):
+    recent = (today - datetime.timedelta(days=1)).isoformat()
+    old = (today - datetime.timedelta(days=90)).isoformat()
+    base = {
+        "api_key": "hashed-other",
+        "user": "user-x",
+        "team_id": "team-x",
+        "end_user": "cust-x",
+        "session_id": "sess-x",
+        "model_id": "mdl-x",
+        "spend": 0.01,
+        "model": "gpt-4",
+    }
+    return [
+        {**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent},
+        {**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old},
+        {**base, "request_id": "req-key", "api_key": hashlib.sha256(b"sk-raw-key").hexdigest(), "startTime": recent},
+        {**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent},
+        {**base, "request_id": "req-user", "user": "user-7", "startTime": recent},
+        {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent},
+        {**base, "request_id": "req-model", "model_id": "mdl-7", "startTime": recent},
+    ]
+
+
+def _search_filter_fn(logs, captured):
+    def filter_fn(where):
+        captured["where"] = where
+        rows = _filter_logs_by_date_range(logs, where)
+        if "user" in where:
+            rows = [row for row in rows if row["user"] == where["user"]]
+        if "search" in where:
+            rows = [row for row in rows if _matches_spend_log_search(row, where["search"])]
+        return rows
+
+    return filter_fn
+
+
+def _five_day_window(today):
+    return {
+        "start_date": (today - datetime.timedelta(days=5)).strftime("%Y-%m-%d %H:%M:%S"),
+        "end_date": today.strftime("%Y-%m-%d %H:%M:%S"),
+    }
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "search,expected_request_ids",
+    [
+        ("req-session-old", {"req-session-old"}),
+        ("sess-42", {"req-session"}),
+        ("sk-raw-key", {"req-key"}),
+        ("team-7", {"req-team"}),
+        ("user-7", {"req-user"}),
+        ("cust-7", {"req-end-user"}),
+        ("mdl-7", {"req-model"}),
+        ("no-such-id", set()),
+    ],
+)
+async def test_ui_view_spend_logs_search_matches_any_id(client, monkeypatch, search, expected_request_ids):
+    """LIT-4741: one box matches any id column. A request_id is found across all time (the 5-day
+    window excludes the 90-day-old row), every other column only inside the window, and a raw
+    sk- key is hashed before it is compared with api_key. The window is not applied globally."""
+    today = datetime.datetime.now(timezone.utc)
+    logs = _search_fixture_logs(today)
+    captured = {}
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client",
+        make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
+    )
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    try:
+        response = client.get(
+            "/spend/logs/ui",
+            params={"search": search, **_five_day_window(today)},
+            headers={"Authorization": "Bearer sk-test"},
+        )
+        assert response.status_code == 200, response.text
+        data = response.json()
+        assert {row["request_id"] for row in data["data"]} == expected_request_ids
+        assert data["total"] == len(expected_request_ids)
+        assert "startTime" not in captured["where"]
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+async def test_spend_logs_v2_search_keeps_global_window(client, monkeypatch):
+    """The public route keeps the caller's window on the whole query, so a search only finds rows
+    inside it even by request_id; the windowless request_id branch is a dashboard-only relaxation."""
+    today = datetime.datetime.now(timezone.utc)
+    logs = _search_fixture_logs(today)
+    captured = {}
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client",
+        make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
+    )
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    try:
+        response = client.get(
+            "/spend/logs/v2",
+            params={"search": "req-session-old", **_five_day_window(today)},
+            headers={"Authorization": "Bearer sk-test"},
+        )
+        assert response.status_code == 200, response.text
+        data = response.json()
+        assert data["data"] == []
+        assert data["total"] == 0
+        assert "startTime" in captured["where"]
+        assert captured["where"]["search"]["value"] == "req-session-old"
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"search": "req-old"},
+        {"search": "req-old", "request_id": "req-old"},
+    ],
+)
+async def test_ui_view_spend_logs_search_requires_dates(client, monkeypatch, params):
+    """A search needs the window for its non-request_id branches, so it stays required even
+    alongside a request_id, which on its own may drop the window."""
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client",
+        make_ui_spend_logs_mock_prisma([], lambda where: []),
+    )
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+    try:
+        response = client.get("/spend/logs/ui", params=params, headers={"Authorization": "Bearer sk-test"})
+        assert response.status_code == 400
+        assert "date" in response.text.lower()
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "search,expected_request_ids",
+    [("sess-9", {"req-own"}), ("req-foreign", set())],
+)
+async def test_ui_view_spend_logs_search_keeps_non_admin_scope(client, monkeypatch, search, expected_request_ids):
+    """A search is scoped like any other listing: an internal user only sees their own rows even
+    when the id is on someone else's row, and the request_id ownership shortcut is not used."""
+    yesterday = (datetime.datetime.now(timezone.utc) - datetime.timedelta(days=1)).isoformat()
+    base = {"api_key": "hashed-key", "team_id": None, "spend": 0.01, "startTime": yesterday, "model": "gpt-4"}
+    logs = [
+        {**base, "request_id": "req-own", "user": "internal_user_1", "session_id": "sess-9"},
+        {**base, "request_id": "req-own-other", "user": "internal_user_1", "session_id": "sess-other"},
+        {**base, "request_id": "req-foreign", "user": "internal_user_2", "session_id": "sess-9"},
+    ]
+    captured = {}
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client",
+        make_ui_spend_logs_mock_prisma(logs, _search_filter_fn(logs, captured)),
+    )
+    monkeypatch.setattr(
+        "litellm.proxy.spend_tracking.spend_management_endpoints._get_permitted_team_ids_for_spend_logs",
+        AsyncMock(return_value=[]),
+    )
+    ownership_check = AsyncMock()
+    monkeypatch.setattr(
+        "litellm.proxy.spend_tracking.spend_management_endpoints._assert_user_can_view_request_id",
+        ownership_check,
+    )
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.INTERNAL_USER, user_id="internal_user_1"
+    )
+    try:
+        start_date, end_date = _default_date_range()
+        response = client.get(
+            "/spend/logs/ui",
+            params={"search": search, "start_date": start_date, "end_date": end_date},
+            headers={"Authorization": "Bearer sk-test"},
+        )
+        assert response.status_code == 200, response.text
+        assert {row["request_id"] for row in response.json()["data"]} == expected_request_ids
+        assert captured["where"]["user"] == "internal_user_1"
+        ownership_check.assert_not_awaited()
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
 @pytest.mark.asyncio
 async def test_ui_view_spend_logs_unauthorized(client):
     # Test without authorization header
@@ -6351,3 +6594,46 @@ async def test_ui_view_spend_logs_group_by_session_offset_for_non_starttime_sort
         assert "OFFSET" in emitted_sql[1]
     finally:
         app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+async def test_ui_view_spend_logs_search_returns_flat_rows_when_grouping_by_session(client, monkeypatch):
+    """The dashboard lists sessions by default; a search for an id lists every matching row instead,
+    so both calls of a session show up rather than one representative, and no session cursor is returned."""
+    rows = [_session_representative_row("req-1", "sess-1"), _session_representative_row("req-2", "sess-1")]
+
+    async def mock_query_raw(sql_query, *params):
+        if "mcp_tool_call_count" in sql_query:
+            return []
+        grouped = "DISTINCT ON" in sql_query or "GROUP BY" in sql_query
+        visible = rows[:1] if grouped else rows
+        if "COUNT(*)" in sql_query:
+            return [{"total_count": len(visible)}]
+        return visible
+
+    mock_prisma = MagicMock()
+    mock_prisma.db = MagicMock()
+    mock_prisma.db.query_raw = AsyncMock(side_effect=mock_query_raw)
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user"
+    )
+    try:
+        start_date, end_date = _default_date_range()
+        response = client.get(
+            "/spend/logs/ui",
+            params={
+                "search": "sess-1",
+                "group_by_session": "true",
+                "start_date": start_date,
+                "end_date": end_date,
+            },
+            headers={"Authorization": "Bearer sk-test"},
+        )
+        assert response.status_code == 200, response.text
+        data = response.json()
+        assert [row["request_id"] for row in data["data"]] == ["req-1", "req-2"]
+        assert data["total"] == 2
+        assert "next_session_cursor" not in data
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)

From ea12510f1a5799e7d9fd2cf87bd5c94263440f46 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:20:42 -0700
Subject: [PATCH 09/38] test(router): drop duplicate get_configured_mode test
 failing ruff F811

PRs #39630 and #39634 both added test_get_configured_mode_reads_deployment_model_info
to tests/test_litellm/test_router.py, so the staging tip defines it twice and the
required lint check fails with F811 on every PR synced past 321636ef5d.

Keep the four tests from #39634 (mode read, None for unset or unknown, no wildcard
pattern routing, malformed values treated as absent), which subsume the #39630 pair,
and delete that pair. Five hand-applied mutations of Router.get_configured_mode are
all still killed by the surviving tests.
---
 tests/test_litellm/test_router.py | 30 ------------------------------
 1 file changed, 30 deletions(-)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index e4236afc586..417c58b95f6 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -12701,33 +12701,3 @@ 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/some-unmapped-mode-model"},
-                "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/some-unmapped-mode-model"},
-                "model_info": model_info,
-            }
-        ]
-    )
-
-    assert router.get_configured_mode("plain-model") is None
-    assert router.get_configured_mode("unknown-model") is None

From a6b7ef6abed774bb6461f9b09b3621b0533cd8c6 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 15:20:46 -0700
Subject: [PATCH 10/38] feat(ui): find a row by its pasted ID on every list
 page

Virtual Keys and Team Virtual Keys send the search box to the new
key/list search param so a key hash matches. Agents matches agent_id
client-side. Memory sends the box as search so a memory_id matches.
Audit Logs gains a search box. Request Logs sends the box as search so
a session, team, user, key hash, or model id matches without opening
the filter drawer.

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
---
 .../agents/_components/AgentsTable.test.tsx   |  29 +++-
 .../agents/_components/AgentsTable.tsx        |   9 +-
 .../(dashboard)/hooks/keys/useKeys.test.ts    |  18 +++
 .../src/app/(dashboard)/hooks/keys/useKeys.ts |   2 +
 .../memory/_components/MemoryTable.test.tsx   |   2 +
 .../memory/_components/MemoryTable.tsx        |   4 +-
 .../memory/_components/MemoryView.test.tsx    |  39 ++++-
 .../memory/_components/MemoryView.tsx         |   4 +-
 .../VirtualKeysPage/VirtualKeysTable.test.tsx |  17 +-
 .../VirtualKeysPage/VirtualKeysTable.tsx      |   4 +-
 .../src/components/networking.test.ts         |  45 ++++++
 .../src/components/networking.tsx             |   9 +-
 .../team/TeamVirtualKeysTable.test.tsx        |  39 ++++-
 .../components/team/TeamVirtualKeysTable.tsx  |  41 +++--
 .../view_logs/AuditLogsPanel.test.tsx         | 145 ++++++++++++++++++
 .../components/view_logs/AuditLogsPanel.tsx   |  16 +-
 .../view_logs/AuditLogsTable.test.tsx         |  18 +++
 .../components/view_logs/AuditLogsTable.tsx   |  10 +-
 .../view_logs/RequestLogsPanel.test.tsx       |  61 +++++++-
 .../components/view_logs/RequestLogsPanel.tsx |  23 ++-
 .../components/view_logs/RequestLogsTable.tsx |   2 +-
 .../view_logs/log_filter_logic.test.tsx       |   1 +
 .../components/view_logs/log_filter_logic.tsx |   3 +
 23 files changed, 489 insertions(+), 52 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx
index 4d18ec2ef5f..bef938cd31c 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.test.tsx
@@ -90,7 +90,7 @@ describe("AgentsTable", () => {
       />,
     );
 
-    const search = screen.getByPlaceholderText("Search agent names or descriptions...");
+    const search = screen.getByPlaceholderText("Search agents by name, ID, or description...");
     await user.type(search, "billing");
     expect(screen.getByText("Billing Router")).toBeInTheDocument();
     expect(screen.queryByText("Second Agent")).not.toBeInTheDocument();
@@ -101,11 +101,36 @@ describe("AgentsTable", () => {
     expect(screen.queryByText("Billing Router")).not.toBeInTheDocument();
   });
 
+  it("filters agents by a pasted agent_id so only that agent's row survives", async () => {
+    const user = userEvent.setup();
+    render(
+      <AgentsTable
+        agents={[
+          makeAgent({ agent_id: "5f3c2a1b-9d8e-4f7a-b6c5-d4e3f2a1b0c9", agent_name: "Billing Router" }),
+          makeAgent({ agent_id: "0a9b8c7d-6e5f-4a3b-8c2d-1e0f9a8b7c6d", agent_name: "Second Agent" }),
+        ]}
+        {...baseProps}
+      />,
+    );
+
+    const search = screen.getByPlaceholderText("Search agents by name, ID, or description...");
+    await user.click(search);
+    await user.paste("5f3c2a1b-9d8e-4f7a-b6c5-d4e3f2a1b0c9");
+    expect(screen.getByText("Billing Router")).toBeInTheDocument();
+    expect(screen.queryByText("Second Agent")).not.toBeInTheDocument();
+
+    await user.clear(search);
+    await user.paste("ffffffff-0000-4000-8000-000000000000");
+    expect(screen.queryByText("Billing Router")).not.toBeInTheDocument();
+    expect(screen.queryByText("Second Agent")).not.toBeInTheDocument();
+    expect(screen.getByText("No matching agents")).toBeInTheDocument();
+  });
+
   it("shows the no-match empty state when the search matches nothing", async () => {
     const user = userEvent.setup();
     render(<AgentsTable agents={[makeAgent()]} {...baseProps} />);
 
-    await user.type(screen.getByPlaceholderText("Search agent names or descriptions..."), "zzzz");
+    await user.type(screen.getByPlaceholderText("Search agents by name, ID, or description..."), "zzzz");
     expect(screen.queryByText("Test Agent")).not.toBeInTheDocument();
     expect(screen.getByText("No matching agents")).toBeInTheDocument();
   });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx
index 35ed6b66425..aceb07e2e9a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentsTable.tsx
@@ -55,7 +55,12 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
   const [sorting, setSorting] = useState<SortingState>(DEFAULT_SORTING);
   const [searchTerm, setSearchTerm] = useState("");
   const filteredAgents = useMemo(
-    () => filterBySearchTerm(agents, searchTerm, (agent) => [agent.agent_name, agent.agent_card_params?.description]),
+    () =>
+      filterBySearchTerm(agents, searchTerm, (agent) => [
+        agent.agent_name,
+        agent.agent_id,
+        agent.agent_card_params?.description,
+      ]),
     [agents, searchTerm],
   );
 
@@ -83,7 +88,7 @@ const AgentsTable: React.FC<AgentsTableProps> = ({
               <SearchIcon className="size-4 text-muted-foreground" />
             </InputGroupAddon>
             <InputGroupInput
-              placeholder="Search agent names or descriptions..."
+              placeholder="Search agents by name, ID, or description..."
               value={searchTerm}
               onChange={(e) => setSearchTerm(e.target.value)}
             />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
index 8c9b33f2c3e..f23fcf811f2 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
@@ -518,6 +518,24 @@ describe("useKeys", () => {
     const callUrl = mockFetch.mock.calls[0][0];
     expect(callUrl).not.toContain("agent_id");
   });
+
+  it("sends the combined alias-or-ID search as the search param, separate from key_alias and key_hash", async () => {
+    mockFetch.mockResolvedValueOnce({
+      ok: true,
+      json: async () => mockKeysResponse,
+    });
+
+    const { result } = renderHook(() => useKeys(1, 10, { search: "sk-pasted-key" }), { wrapper });
+
+    await waitFor(() => {
+      expect(result.current.isLoading).toBe(false);
+    });
+
+    const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost");
+    expect(callUrl.searchParams.get("search")).toBe("sk-pasted-key");
+    expect(callUrl.searchParams.has("key_alias")).toBe(false);
+    expect(callUrl.searchParams.has("key_hash")).toBe(false);
+  });
 });
 
 describe("useDeletedKeys", () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts
index 94ded01679d..7e7089e685f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts
@@ -40,6 +40,7 @@ export interface KeyListCallOptions {
   selectedKeyAlias?: string | null;
   userID?: string | null;
   keyHash?: string | null;
+  search?: string | null;
   sortBy?: string | null;
   sortOrder?: string | null;
   expand?: string | null;
@@ -61,6 +62,7 @@ const keyListCall = async (accessToken: string, page: number, pageSize: number,
         organization_id: options.organizationID,
         key_alias: options.selectedKeyAlias,
         key_hash: options.keyHash,
+        search: options.search,
         user_id: options.userID,
         page,
         size: pageSize,
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx
index 5100b998b80..984b8135466 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.test.tsx
@@ -95,6 +95,7 @@ describe("MemoryTable", () => {
   it("shows the filtered-empty copy when a search is active", () => {
     render(<MemoryTable {...baseProps} data={[]} rowCount={0} hasActiveSearch={true} />);
     expect(screen.getByText("No matching memories")).toBeInTheDocument();
+    expect(screen.getByText("No memories match your search.")).toBeInTheDocument();
     expect(screen.queryByText("No memories stored yet")).not.toBeInTheDocument();
   });
 
@@ -128,6 +129,7 @@ describe("MemoryTable", () => {
     const onRefresh = vi.fn();
     render(<MemoryTable {...baseProps} onSearchChange={onSearchChange} onRefresh={onRefresh} />);
 
+    expect(screen.getByPlaceholderText("Search by key prefix or memory ID…")).toBeInTheDocument();
     fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "u" } });
     expect(onSearchChange).toHaveBeenCalledWith("u");
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx
index 50dd04ee14c..3e37faafe15 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryTable.tsx
@@ -36,7 +36,7 @@ function MemoryEmptyState({ hasActiveSearch }: { hasActiveSearch: boolean }) {
       </div>
       <div className="text-sm text-muted-foreground">
         {hasActiveSearch
-          ? "No memories have keys starting with your search."
+          ? "No memories match your search."
           : "Memories your agents store under /v1/memory will appear here."}
       </div>
     </div>
@@ -81,7 +81,7 @@ export function MemoryTable({
           table={table}
           searchValue={searchValue}
           onSearchChange={onSearchChange}
-          searchPlaceholder='Filter by key prefix, e.g. "user:"'
+          searchPlaceholder="Search by key prefix or memory ID…"
           onRefresh={onRefresh}
           isRefreshing={isRefreshing}
           showViewOptions={false}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx
index 9ccef5357b9..b703df652c2 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.test.tsx
@@ -1,8 +1,9 @@
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
-import { act, render, screen } from "@testing-library/react";
+import type { PaginationState } from "@tanstack/react-table";
+import { act, render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import React from "react";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
 import { MemoryRow } from "@/components/networking";
 
@@ -13,10 +14,13 @@ interface CapturedTableProps {
   rowCount: number;
   data: MemoryRow[];
   hasActiveSearch: boolean;
+  onSearchChange: (value: string) => void;
+  onPaginationChange: (state: PaginationState) => void;
   onViewClick: (row: MemoryRow) => void;
 }
 
 const captured = vi.hoisted(() => ({ current: null as CapturedTableProps | null }));
+const fetchMemoryListMock = vi.hoisted(() => vi.fn());
 
 vi.mock("./MemoryTable", () => ({
   MemoryTable: function MemoryTableMock(props: CapturedTableProps) {
@@ -25,6 +29,15 @@ vi.mock("./MemoryTable", () => ({
   },
 }));
 
+vi.mock("@/components/networking", async (importOriginal) => ({
+  ...(await importOriginal<typeof import("@/components/networking")>()),
+  fetchMemoryList: fetchMemoryListMock,
+}));
+
+vi.mock("@tanstack/react-pacer/debouncer", () => ({
+  useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }],
+}));
+
 const renderView = (accessToken: string | null) => {
   const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
   return render(
@@ -35,6 +48,28 @@ const renderView = (accessToken: string | null) => {
 };
 
 describe("MemoryView", () => {
+  beforeEach(() => {
+    fetchMemoryListMock.mockReset();
+    fetchMemoryListMock.mockResolvedValue({ memories: [], total: 0 });
+  });
+
+  it("queries the server with the search box value as `search` and resets to page 1", async () => {
+    renderView("token");
+    await waitFor(() => expect(fetchMemoryListMock).toHaveBeenCalled());
+
+    act(() => captured.current?.onPaginationChange({ pageIndex: 2, pageSize: 50 }));
+    await waitFor(() =>
+      expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", expect.objectContaining({ page: 3 })),
+    );
+
+    act(() => captured.current?.onSearchChange("mem-abc123"));
+
+    await waitFor(() =>
+      expect(fetchMemoryListMock).toHaveBeenLastCalledWith("token", { search: "mem-abc123", page: 1, pageSize: 50 }),
+    );
+    expect(captured.current?.hasActiveSearch).toBe(true);
+  });
+
   it("keeps the table out of the skeleton state when the token is null (disabled query)", () => {
     renderView(null);
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx
index 1d2e5150a62..58d4e42aa96 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/memory/_components/MemoryView.tsx
@@ -43,10 +43,8 @@ export const MemoryView: React.FC<MemoryViewProps> = ({ accessToken }) => {
     queryKey: [MEMORY_LIST_KEY, debouncedSearch, pagination.pageIndex, pagination.pageSize],
     queryFn: () => {
       if (!accessToken) throw new Error("Access token required");
-      // Prefix search matches the Redis-style mental model (namespace scan):
-      // typing "user:" finds "user:profile", "user:prefs", etc.
       return fetchMemoryList(accessToken, {
-        keyPrefix: debouncedSearch || undefined,
+        search: debouncedSearch || undefined,
         page: pagination.pageIndex + 1,
         pageSize: pagination.pageSize,
       });
diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
index c69662b69dc..881b0b93ff9 100644
--- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx
@@ -542,6 +542,19 @@ describe("server-side filtering – the LIT-4080 regression guard", () => {
       expect((lastCall[2] ?? {}).userID).toBeUndefined();
     });
   });
+
+  it("sends the search box as the combined alias-or-ID search rather than the key-alias filter", async () => {
+    renderWithProviders(<VirtualKeysTable />);
+
+    fireEvent.change(screen.getByPlaceholderText(/Search by key alias or ID/), { target: { value: mockKey.token } });
+
+    await waitFor(() => {
+      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: mockKey.token }));
+    });
+    const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2];
+    expect(lastOptions?.selectedKeyAlias).toBeUndefined();
+    expect(lastOptions?.keyHash).toBeUndefined();
+  });
 });
 
 describe("pagination display – total count comes from useKeys", () => {
@@ -663,7 +676,7 @@ describe("table state lives in the URL so it survives leaving and returning to t
       expect(mockUseKeys).toHaveBeenLastCalledWith(
         3,
         25,
-        expect.objectContaining({ selectedKeyAlias: "prod", sortBy: "spend", sortOrder: "asc" }),
+        expect.objectContaining({ search: "prod", sortBy: "spend", sortOrder: "asc" }),
       );
     });
     expect(screen.getByPlaceholderText(/Search by key alias/)).toHaveValue("prod");
@@ -736,7 +749,7 @@ describe("table state lives in the URL so it survives leaving and returning to t
     fireEvent.change(screen.getByPlaceholderText(/Search by key alias/), { target: { value: "prod" } });
 
     await waitFor(() => {
-      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "prod" }));
+      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: "prod" }));
     });
     await waitFor(() => {
       expect(lastSearchParam(onUrlUpdate, "page")).toBeNull();
diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
index c424966a0a3..ebedf57af45 100644
--- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
@@ -118,7 +118,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
   const keyListOptions = {
     teamID: appliedFilters.team_id || undefined,
     organizationID: appliedFilters.org_id || undefined,
-    selectedKeyAlias: searchQuery.trim() || undefined,
+    search: searchQuery.trim() || undefined,
     userID: appliedFilters.user_id || undefined,
     keyHash: appliedFilters.key_hash || undefined,
     sortBy,
@@ -291,7 +291,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
               table={table}
               searchValue={searchInput}
               onSearchChange={handleSearchChange}
-              searchPlaceholder="Search by key alias…"
+              searchPlaceholder="Search by key alias or ID…"
               onRefresh={() => refetch?.()}
               isRefreshing={isFetching}
               onOpenFilters={() => setFiltersOpen(true)}
diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts
index 7a220dd4711..9df9e9a9209 100644
--- a/ui/litellm-dashboard/src/components/networking.test.ts
+++ b/ui/litellm-dashboard/src/components/networking.test.ts
@@ -854,3 +854,48 @@ describe("userListCall search serialization", () => {
     expect(lastParams(mockFetch).get("user_email")).toBe("ada@example.com");
   });
 });
+
+describe("fetchMemoryList search serialization", () => {
+  const originalFetch = global.fetch;
+
+  afterEach(() => {
+    global.fetch = originalFetch;
+  });
+
+  const mockOkFetch = () => {
+    const emptyPage = { memories: [], total: 0 };
+    const mockFetch = vi.fn().mockResolvedValue({ ok: true, json: vi.fn().mockResolvedValue(emptyPage) } as any);
+    global.fetch = mockFetch as any;
+    return mockFetch;
+  };
+
+  const lastParams = (mockFetch: ReturnType<typeof vi.fn>) => {
+    const [url] = mockFetch.mock.calls.at(-1) ?? [];
+    return new URL(url as string, "http://example.com").searchParams;
+  };
+
+  it("sends the search box value as search and omits key_prefix and key", async () => {
+    const mockFetch = mockOkFetch();
+
+    await Networking.fetchMemoryList("token", { search: "mem-abc123", page: 1, pageSize: 50 });
+
+    const params = lastParams(mockFetch);
+    expect(params.get("search")).toBe("mem-abc123");
+    expect(params.has("key_prefix")).toBe(false);
+    expect(params.has("key")).toBe(false);
+    expect(params.get("page")).toBe("1");
+    expect(params.get("page_size")).toBe("50");
+  });
+
+  it("keeps key_prefix and key working when no search is given", async () => {
+    const mockFetch = mockOkFetch();
+
+    await Networking.fetchMemoryList("token", { keyPrefix: "user:" });
+    expect(lastParams(mockFetch).get("key_prefix")).toBe("user:");
+    expect(lastParams(mockFetch).has("search")).toBe(false);
+
+    await Networking.fetchMemoryList("token", { key: "user:profile" });
+    expect(lastParams(mockFetch).get("key")).toBe("user:profile");
+    expect(lastParams(mockFetch).has("search")).toBe(false);
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx
index d8762565e08..1384679a88a 100644
--- a/ui/litellm-dashboard/src/components/networking.tsx
+++ b/ui/litellm-dashboard/src/components/networking.tsx
@@ -2056,6 +2056,7 @@ interface UiSpendLogsParams {
   exclude_internal_health_checks?: boolean;
   group_by_session?: boolean;
   session_cursor?: string;
+  search?: string;
 }
 
 interface UiSpendLogsCallOptions {
@@ -6563,6 +6564,7 @@ interface UiAuditLogsParams {
   changed_by_api_key?: string;
   object_team_id?: string;
   object_key_hash?: string;
+  search?: string | null;
   sort_by?: string;
   sort_order?: "asc" | "desc";
 }
@@ -8061,15 +8063,18 @@ export const fetchMemoryList = async (
   options: {
     key?: string;
     keyPrefix?: string;
+    search?: string;
     page?: number;
     pageSize?: number;
   } = {},
 ): Promise<MemoryListResponse> => {
   const base = proxyBaseUrl ? `${proxyBaseUrl}/v1/memory` : `/v1/memory`;
   const params = new URLSearchParams();
-  // keyPrefix takes precedence — backend also does, but we omit `key`
+  // Backend precedence is search > key_prefix > key; only the winner is sent
   // to keep the URL clean and intent obvious.
-  if (options.keyPrefix) {
+  if (options.search) {
+    params.append("search", options.search);
+  } else if (options.keyPrefix) {
     params.append("key_prefix", options.keyPrefix);
   } else if (options.key) {
     params.append("key", options.key);
diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx
index 8bf5d639d6c..0d9d0988aa1 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.test.tsx
@@ -30,6 +30,8 @@ vi.mock("@tanstack/react-pacer/debouncer", () => ({
 
 const mockUseKeys = useKeys as MockedFunction<typeof useKeys>;
 
+const KEY_HASH = "88a145505dd6e87e2ea166fcef1e4b53948dbdb32af6431dfd05ec06b571ee52";
+
 const createMockKey = (overrides: Partial<KeyResponse> = {}): KeyResponse =>
   ({
     token: "sk-test123",
@@ -277,7 +279,7 @@ describe("TeamVirtualKeysTable", () => {
     );
   });
 
-  it("maps the search box to a server-side key-alias query", async () => {
+  it("maps the Key ID drawer filter to a server-side useKeys query and clears it", async () => {
     const user = userEvent.setup();
     mockUseKeys.mockReturnValue({
       data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 },
@@ -288,11 +290,42 @@ describe("TeamVirtualKeysTable", () => {
 
     renderWithProviders(<TeamVirtualKeysTable {...defaultProps} />);
 
-    fireEvent.change(await screen.findByTestId("datatable-search"), { target: { value: "check-002" } });
+    await user.click(await screen.findByTestId("datatable-filters-trigger"));
+    const drawerBody = await screen.findByTestId("filter-drawer-body");
+    fireEvent.change(within(drawerBody).getByPlaceholderText("Enter Key ID…"), { target: { value: KEY_HASH } });
+    await user.click(screen.getByTestId("filter-drawer-apply"));
 
     await waitFor(() =>
-      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ selectedKeyAlias: "check-002" })),
+      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: KEY_HASH })),
     );
+    expect(screen.getByTestId("filter-chip-key_hash")).toHaveTextContent("Key ID");
+
+    await user.click(screen.getByTestId("datatable-clear-filters"));
+    await waitFor(() =>
+      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ keyHash: undefined })),
+    );
+  });
+
+  it("maps the search box to the combined alias-or-ID search rather than the key-alias filter", async () => {
+    mockUseKeys.mockReturnValue({
+      data: { keys: [createMockKey()], total_count: 1, current_page: 1, total_pages: 1 },
+      isPending: false,
+      isFetching: false,
+      refetch: vi.fn(),
+    } as unknown as ReturnType<typeof useKeys>);
+
+    renderWithProviders(<TeamVirtualKeysTable {...defaultProps} />);
+
+    const searchBox = await screen.findByTestId("datatable-search");
+    expect(searchBox).toHaveAttribute("placeholder", "Search by key alias or ID…");
+    fireEvent.change(searchBox, { target: { value: KEY_HASH } });
+
+    await waitFor(() =>
+      expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ search: KEY_HASH })),
+    );
+    const lastOptions = mockUseKeys.mock.calls.at(-1)?.[2];
+    expect(lastOptions?.selectedKeyAlias).toBeUndefined();
+    expect(lastOptions?.keyHash).toBeUndefined();
   });
 
   it("should show Loading keys when isPending", async () => {
diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
index bd7faa41ee9..aa9df4a0319 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
@@ -68,19 +68,17 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
   const pageIndex = tablePagination.pageIndex;
   const pageSize = tablePagination.pageSize;
 
-  const {
-    data: keys,
-    isPending: isLoading,
-    isFetching,
-    refetch,
-  } = useKeys(pageIndex + 1, pageSize, {
+  const keyListOptions = {
     teamID: teamId,
-    selectedKeyAlias: searchQuery.trim() || undefined,
+    search: searchQuery.trim() || undefined,
     userID: getFilterValue("user_id"),
+    keyHash: getFilterValue("key_hash"),
     sortBy: sortBy || undefined,
     sortOrder: sortOrder || undefined,
     expand: "user",
-  });
+  };
+
+  const { data: keys, isPending: isLoading, isFetching, refetch } = useKeys(pageIndex + 1, pageSize, keyListOptions);
 
   const displayKeys = useMemo(() => {
     const kList = keys?.keys || [];
@@ -481,11 +479,11 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
                   table={table}
                   searchValue={searchInput}
                   onSearchChange={handleSearchChange}
-                  searchPlaceholder="Search by key alias…"
+                  searchPlaceholder="Search by key alias or ID…"
                   onRefresh={() => refetch?.()}
                   isRefreshing={isFetching}
                   onOpenFilters={() => setFiltersOpen(true)}
-                  filterLabels={{ user_id: "User ID" }}
+                  filterLabels={{ user_id: "User ID", key_hash: "Key ID" }}
                 />
                 <DataTableFilterDrawer
                   table={table}
@@ -495,13 +493,22 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
                   description={`Narrow down keys for ${teamAlias ?? "this team"}`}
                 >
                   {({ get, set }) => (
-                    <DataTableFilterField label="User ID">
-                      <Input
-                        value={(get("user_id") as string) ?? ""}
-                        onChange={(event) => set("user_id", event.target.value)}
-                        placeholder="Filter by user ID…"
-                      />
-                    </DataTableFilterField>
+                    <>
+                      <DataTableFilterField label="User ID">
+                        <Input
+                          value={(get("user_id") as string) ?? ""}
+                          onChange={(event) => set("user_id", event.target.value)}
+                          placeholder="Filter by user ID…"
+                        />
+                      </DataTableFilterField>
+                      <DataTableFilterField label="Key ID">
+                        <Input
+                          value={(get("key_hash") as string) ?? ""}
+                          onChange={(event) => set("key_hash", event.target.value)}
+                          placeholder="Enter Key ID…"
+                        />
+                      </DataTableFilterField>
+                    </>
                   )}
                 </DataTableFilterDrawer>
               </>
diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx
new file mode 100644
index 00000000000..3b27f663b8e
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.test.tsx
@@ -0,0 +1,145 @@
+import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
+import { fireEvent, render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+import { chooseSelectOption } from "../../../tests/test-utils";
+import AuditLogsPanel from "./AuditLogsPanel";
+
+vi.mock("../networking", async (importOriginal) => {
+  const actual = await importOriginal<typeof import("../networking")>();
+  return { ...actual, uiAuditLogsCall: vi.fn() };
+});
+
+// Resolve the debounced search synchronously so typed input reaches the query within the test tick.
+vi.mock("@tanstack/react-pacer/debouncer", () => ({
+  useDebouncedValue: (value: unknown) => [value, { cancel: vi.fn(), flush: vi.fn() }],
+}));
+
+import { uiAuditLogsCall } from "../networking";
+
+type AuditLogsParams = NonNullable<Parameters<typeof uiAuditLogsCall>[0]["params"]>;
+
+const PAGE_SIZE = 50;
+
+const ID_PARAM_KEYS = [
+  "search",
+  "object_id",
+  "changed_by",
+  "object_team_id",
+  "object_key_hash",
+  "action",
+  "table_name",
+] as const satisfies readonly (keyof AuditLogsParams)[];
+
+const respondWith = (total: number) => {
+  const response = { audit_logs: [], total, page: 1, page_size: PAGE_SIZE, total_pages: Math.ceil(total / PAGE_SIZE) };
+  return vi.mocked(uiAuditLogsCall).mockResolvedValue(response);
+};
+
+const lastCall = () => vi.mocked(uiAuditLogsCall).mock.calls.at(-1)?.[0];
+const sentIdParams = () => ID_PARAM_KEYS.filter((key) => lastCall()?.params?.[key] !== undefined);
+
+const defaultProps = {
+  accessToken: "sk-test",
+  token: "jwt-test",
+  userRole: "Admin",
+  userID: "user-1",
+  isActive: true,
+  premiumUser: true,
+};
+
+const renderPanel = () => {
+  const queryClient = new QueryClient({ defaultOptions: { queries: { retry: false } } });
+  return render(
+    <QueryClientProvider client={queryClient}>
+      <AuditLogsPanel {...defaultProps} />
+    </QueryClientProvider>,
+  );
+};
+
+const TEXT_FILTERS: { filterId: string; placeholder: string; paramKey: keyof AuditLogsParams }[] = [
+  { filterId: "object_id", placeholder: "Enter object ID…", paramKey: "object_id" },
+  { filterId: "changed_by", placeholder: "Enter user ID…", paramKey: "changed_by" },
+  { filterId: "team_id", placeholder: "Enter team ID…", paramKey: "object_team_id" },
+  { filterId: "key_hash", placeholder: "Enter key hash…", paramKey: "object_key_hash" },
+];
+
+const SELECT_FILTERS: {
+  label: string;
+  comboboxIndex: number;
+  option: string;
+  paramKey: keyof AuditLogsParams;
+  value: string;
+}[] = [
+  { label: "Action", comboboxIndex: 0, option: "Created", paramKey: "action", value: "created" },
+  { label: "Table", comboboxIndex: 1, option: "Teams", paramKey: "table_name", value: "LiteLLM_TeamTable" },
+];
+
+describe("AuditLogsPanel", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    respondWith(0);
+  });
+
+  it("sends the typed search as params.search and returns to the first page", async () => {
+    const user = userEvent.setup();
+    respondWith(120);
+    renderPanel();
+    await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled());
+    expect(lastCall()?.params?.search).toBeUndefined();
+
+    await user.click(screen.getByTestId("pagination-next"));
+    await waitFor(() => expect(lastCall()?.page).toBe(2));
+
+    await user.type(screen.getByTestId("datatable-search"), "team-abc");
+
+    await waitFor(() => expect(lastCall()?.params?.search).toBe("team-abc"));
+    expect(lastCall()?.page).toBe(1);
+    expect(sentIdParams()).toEqual(["search"]);
+  });
+
+  it("trims the search and drops params.search once the box is cleared", async () => {
+    const user = userEvent.setup();
+    renderPanel();
+    const input = await screen.findByTestId("datatable-search");
+
+    await user.type(input, "  abc");
+    await waitFor(() => expect(lastCall()?.params?.search).toBe("abc"));
+
+    await user.clear(input);
+
+    await waitFor(() => expect(lastCall()?.params?.search).toBeUndefined());
+    expect(sentIdParams()).toEqual([]);
+  });
+
+  it.each(TEXT_FILTERS)("maps the $filterId drawer filter to params.$paramKey", async ({ placeholder, paramKey }) => {
+    const user = userEvent.setup();
+    renderPanel();
+    await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled());
+
+    await user.click(screen.getByTestId("datatable-filters-trigger"));
+    fireEvent.change(await screen.findByPlaceholderText(placeholder), { target: { value: "val-1" } });
+    await user.click(screen.getByTestId("filter-drawer-apply"));
+
+    await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe("val-1"));
+    expect(sentIdParams()).toEqual([paramKey]);
+  });
+
+  it.each(SELECT_FILTERS)(
+    "maps the $label drawer select to params.$paramKey",
+    async ({ comboboxIndex, option, paramKey, value }) => {
+      const user = userEvent.setup();
+      renderPanel();
+      await waitFor(() => expect(uiAuditLogsCall).toHaveBeenCalled());
+
+      await user.click(screen.getByTestId("datatable-filters-trigger"));
+      const triggers = await screen.findAllByRole("combobox");
+      await chooseSelectOption(user, triggers[comboboxIndex], option);
+      await user.click(screen.getByTestId("filter-drawer-apply"));
+
+      await waitFor(() => expect(lastCall()?.params?.[paramKey]).toBe(value));
+      expect(sentIdParams()).toEqual([paramKey]);
+    },
+  );
+});
diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx
index 81bd4a19f76..5ce4ca6053f 100644
--- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsPanel.tsx
@@ -1,7 +1,9 @@
 import { useCallback, useState } from "react";
+import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
 import { useQuery, keepPreviousData } from "@tanstack/react-query";
 import { ColumnFiltersState, OnChangeFn, PaginationState } from "@tanstack/react-table";
 import { resolveLogoSrc } from "@/lib/assetPaths";
+import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
 import { uiAuditLogsCall } from "../networking";
 import { AuditLogEntry } from "./AuditLogsTableColumns";
 import { AuditLogsTable } from "./AuditLogsTable";
@@ -39,9 +41,13 @@ export default function AuditLogsPanel({
 }: AuditLogsProps) {
   const [pagination, setPagination] = useState<PaginationState>({ pageIndex: 0, pageSize: PAGE_SIZE });
   const [columnFilters, setColumnFilters] = useState<ColumnFiltersState>([]);
+  const [searchInput, setSearchInput] = useState("");
+  const [debouncedSearch] = useDebouncedValue(searchInput, { wait: DEBOUNCE_WAIT_MS });
   const [selectedLog, setSelectedLog] = useState<AuditLogEntry | null>(null);
   const [drawerOpen, setDrawerOpen] = useState(false);
 
+  const searchTerm = debouncedSearch.trim();
+
   const getFilterValue = (columnId: string): string | undefined => {
     const entry = columnFilters.find((filter) => filter.id === columnId);
     return typeof entry?.value === "string" && entry.value.trim() ? entry.value.trim() : undefined;
@@ -50,7 +56,7 @@ export default function AuditLogsPanel({
   const canQueryAuditLogs = !!accessToken && !!token && !!userRole && !!userID && isActive && premiumUser;
 
   const query = useQuery<AuditLogsResponse>({
-    queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters],
+    queryKey: ["audit_logs", pagination.pageIndex, pagination.pageSize, columnFilters, searchTerm],
     queryFn: async () => {
       if (!accessToken) {
         return { audit_logs: [], total: 0, page: 1, page_size: pagination.pageSize, total_pages: 0 };
@@ -60,6 +66,7 @@ export default function AuditLogsPanel({
         page: pagination.pageIndex + 1,
         page_size: pagination.pageSize,
         params: {
+          search: searchTerm || undefined,
           object_id: getFilterValue("object_id"),
           changed_by: getFilterValue("changed_by"),
           object_key_hash: getFilterValue("key_hash"),
@@ -80,6 +87,11 @@ export default function AuditLogsPanel({
     setPagination((prev) => ({ ...prev, pageIndex: 0 }));
   }, []);
 
+  const handleSearchChange = useCallback((value: string) => {
+    setSearchInput(value);
+    setPagination((prev) => ({ ...prev, pageIndex: 0 }));
+  }, []);
+
   const handleViewLog = useCallback((log: AuditLogEntry) => {
     setSelectedLog(log);
     setDrawerOpen(true);
@@ -128,6 +140,8 @@ export default function AuditLogsPanel({
         onPaginationChange={setPagination}
         columnFilters={columnFilters}
         onColumnFiltersChange={handleColumnFiltersChange}
+        searchValue={searchInput}
+        onSearchChange={handleSearchChange}
         onRefresh={() => query.refetch()}
         onViewLog={handleViewLog}
       />
diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx
index 7349c3019ae..d8e549715af 100644
--- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.test.tsx
@@ -120,6 +120,24 @@ describe("AuditLogsTable", () => {
     expect(screen.getByText("No matching audit logs")).toBeInTheDocument();
   });
 
+  it("renders the toolbar search box from the search props and forwards typed input", () => {
+    const onSearchChange = vi.fn();
+    renderTable({ searchValue: "team-", onSearchChange });
+
+    const input = screen.getByPlaceholderText("Search audit logs by ID…");
+    expect(input).toHaveValue("team-");
+
+    fireEvent.change(input, { target: { value: "team-7" } });
+    expect(onSearchChange).toHaveBeenCalledWith("team-7");
+  });
+
+  it("treats an active search as a filter for the empty state", () => {
+    const emptySearchResult = { data: [], rowCount: 0, searchValue: "zzz", onSearchChange: vi.fn() };
+    renderTable(emptySearchResult);
+
+    expect(screen.getByText("No matching audit logs")).toBeInTheDocument();
+  });
+
   it("renders active filter chips with human-readable labels", () => {
     const filters: ColumnFiltersState = [{ id: "action", value: "created" }];
     renderTable({ columnFilters: filters });
diff --git a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx
index 799505ed07c..cef828838e2 100644
--- a/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/AuditLogsTable.tsx
@@ -24,6 +24,8 @@ interface AuditLogsTableProps {
   onPaginationChange: OnChangeFn<PaginationState>;
   columnFilters: ColumnFiltersState;
   onColumnFiltersChange: OnChangeFn<ColumnFiltersState>;
+  searchValue?: string;
+  onSearchChange?: (value: string) => void;
   onRefresh: () => void;
   onViewLog: (log: AuditLogEntry) => void;
 }
@@ -102,11 +104,14 @@ export function AuditLogsTable({
   onPaginationChange,
   columnFilters,
   onColumnFiltersChange,
+  searchValue,
+  onSearchChange,
   onRefresh,
   onViewLog,
 }: AuditLogsTableProps) {
   const [filtersOpen, setFiltersOpen] = useState(false);
   const columns = useMemo(() => getAuditLogsTableColumns({ onViewLog }), [onViewLog]);
+  const hasActiveSearch = Boolean(searchValue?.trim());
 
   return (
     <DataTable
@@ -122,12 +127,15 @@ export function AuditLogsTable({
       onColumnFiltersChange={onColumnFiltersChange}
       isLoading={isLoading}
       loadingMessage="Loading audit logs…"
-      noDataMessage={<AuditLogsEmptyState filtered={columnFilters.length > 0} />}
+      noDataMessage={<AuditLogsEmptyState filtered={columnFilters.length > 0 || hasActiveSearch} />}
       size="compact"
       toolbar={(table) => (
         <>
           <DataTableToolbar
             table={table}
+            searchValue={searchValue}
+            onSearchChange={onSearchChange}
+            searchPlaceholder="Search audit logs by ID…"
             onRefresh={onRefresh}
             isRefreshing={isRefreshing}
             onOpenFilters={() => setFiltersOpen(true)}
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx
index 93066c106ee..be0d0049c13 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.test.tsx
@@ -54,6 +54,14 @@ vi.mock("./LogDetailsDrawer", () => ({
   },
 }));
 
+const debounce = vi.hoisted(() => ({ settled: null as string | null }));
+
+vi.mock("@tanstack/react-pacer/debouncer", () => ({
+  useDebouncedValue: vi.fn((value: unknown) => [debounce.settled ?? value, { cancel: vi.fn(), flush: vi.fn() }]),
+}));
+
+import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
+import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
 import { uiSpendLogsCall } from "../networking";
 
 const logEntry = (overrides: Partial<LogEntry>): LogEntry => ({
@@ -136,6 +144,7 @@ describe("RequestLogsPanel", () => {
     sessionStorage.clear();
     testQueryClient.clear();
     respondWith([]);
+    debounce.settled = null;
   });
 
   describe("server-grouped session pagination (#38060)", () => {
@@ -322,9 +331,8 @@ describe("RequestLogsPanel", () => {
     });
   });
 
-  describe("search by request id (LIT-3981)", () => {
-    it("sends the typed request id to the server on the first page instead of filtering the loaded rows", async () => {
-      const user = userEvent.setup();
+  describe("search by any id (LIT-3981, LIT-4741)", () => {
+    it("sends the typed id to the server as search on the first page instead of filtering the loaded rows", async () => {
       renderPanel();
 
       await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
@@ -334,9 +342,54 @@ describe("RequestLogsPanel", () => {
       await waitFor(() => {
         const call = lastCall();
         if (!call) throw new Error("uiSpendLogsCall was not called");
-        expect(call.params?.request_id).toBe("req-on-another-page");
+        expect(call.params?.search).toBe("req-on-another-page");
         expect(call.page).toBe(1);
       });
+      expect(lastCall()?.params?.request_id).toBeUndefined();
+      expect(lastCall()?.params?.session_cursor).toBeUndefined();
+    });
+
+    it("sends the debounced value to the server while the box shows what is being typed", async () => {
+      debounce.settled = "settled-id";
+      renderPanel();
+
+      await waitFor(() => expect(uiSpendLogsCall).toHaveBeenCalled());
+
+      fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "still-typing" } });
+
+      expect(screen.getByTestId("datatable-search")).toHaveValue("still-typing");
+      await waitFor(() =>
+        expect(useDebouncedValue).toHaveBeenLastCalledWith("still-typing", { wait: DEBOUNCE_WAIT_MS }),
+      );
+      await waitFor(() => expect(lastCall()?.params?.search).toBe("settled-id"));
+      const sentLiveValue = vi
+        .mocked(uiSpendLogsCall)
+        .mock.calls.some(([options]) => options.params?.search === "still-typing");
+      expect(sentLiveValue).toBe(false);
+    });
+
+    it("shows a Search chip whose remove button clears the box and restores the unsearched listing", async () => {
+      const user = userEvent.setup();
+      vi.mocked(uiSpendLogsCall).mockImplementation(async ({ params }) => {
+        const data =
+          params?.search === "sess-42"
+            ? [logEntry({ request_id: "req-sess", session_id: "sess-42" })]
+            : [logEntry({ request_id: "req-initial" })];
+        return { data, total: data.length, page: 1, page_size: 50, total_pages: 1 };
+      });
+      renderPanel();
+
+      await waitFor(() => expect(row("req-initial")).not.toBeNull());
+      fireEvent.change(screen.getByTestId("datatable-search"), { target: { value: "sess-42" } });
+      await waitFor(() => expect(row("req-sess")).not.toBeNull());
+      expect(row("req-initial")).toBeNull();
+      expect(screen.getByTestId("filter-chip-search")).toHaveTextContent("Search:sess-42");
+
+      await user.click(screen.getByRole("button", { name: "Remove Search filter" }));
+
+      expect(screen.getByTestId("datatable-search")).toHaveValue("");
+      await waitFor(() => expect(row("req-initial")).not.toBeNull());
+      expect(row("req-sess")).toBeNull();
     });
   });
 
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx
index 96cf2bd5dde..9b99c6af923 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsPanel.tsx
@@ -1,11 +1,13 @@
 "use client";
 
+import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
 import { useQuery, type UseQueryOptions } from "@tanstack/react-query";
 import type { ColumnFiltersState, OnChangeFn, PaginationState, SortingState } from "@tanstack/react-table";
 import moment from "moment";
 import { useCallback, useEffect, useMemo, useState } from "react";
 
 import { AutoRouterModelGroupsProvider } from "@/components/shared/table_cells";
+import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
 import type { KeyResponse } from "../key_team_helpers/key_list";
 import { keyInfoV1Call, uiSpendLogsCall } from "../networking";
 import KeyInfoView from "../templates/key_info_view";
@@ -75,12 +77,22 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
     sessionStorage.setItem("excludeInternalHealthChecks", JSON.stringify(excludeInternalHealthChecks));
   }, [excludeInternalHealthChecks]);
 
+  const searchTerm = useMemo(() => {
+    const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.SEARCH);
+    return typeof entry?.value === "string" ? entry.value : "";
+  }, [columnFilters]);
+  const [debouncedSearch] = useDebouncedValue(searchTerm, { wait: DEBOUNCE_WAIT_MS });
+  const queryColumnFilters = useMemo<ColumnFiltersState>(() => {
+    const others = columnFilters.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH);
+    return debouncedSearch === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value: debouncedSearch }];
+  }, [columnFilters, debouncedSearch]);
+
   const { logsQuery, filteredLogs, allTeams, usesSessionCursor } = useLogFilterLogic({
     accessToken,
     token,
     userRole,
     userID,
-    columnFilters,
+    columnFilters: queryColumnFilters,
     activeTab: isActive ? "request logs" : "inactive",
     isLiveTail,
     excludeInternalHealthChecks,
@@ -155,15 +167,10 @@ export default function RequestLogsPanel({ accessToken, token, userRole, userID,
 
   const rows: LogEntry[] = filteredLogs.data;
 
-  const searchTerm = useMemo(() => {
-    const entry = columnFilters.find((filter) => filter.id === LOG_FILTER_IDS.REQUEST_ID);
-    return typeof entry?.value === "string" ? entry.value : "";
-  }, [columnFilters]);
-
   const handleSearchChange = useCallback((value: string) => {
     setColumnFilters((previous) => {
-      const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.REQUEST_ID);
-      return value === "" ? others : [...others, { id: LOG_FILTER_IDS.REQUEST_ID, value }];
+      const others = previous.filter((filter) => filter.id !== LOG_FILTER_IDS.SEARCH);
+      return value === "" ? others : [...others, { id: LOG_FILTER_IDS.SEARCH, value }];
     });
     setSessionCursors({});
     setPagination((previous) => ({ ...previous, pageIndex: 0 }));
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
index 4159b3b699b..17caa4466fa 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
@@ -108,7 +108,7 @@ export function RequestLogsTable({
             table={table}
             searchValue={searchValue}
             onSearchChange={onSearchChange}
-            searchPlaceholder="Search by Request ID"
+            searchPlaceholder="Search logs by ID…"
             onRefresh={onRefresh}
             isRefreshing={isRefreshing}
             onOpenFilters={() => setFiltersOpen(true)}
diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx
index 1b738db097d..6af791cc50e 100644
--- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.test.tsx
@@ -91,6 +91,7 @@ describe("useLogFilterLogic", () => {
       { id: LOG_FILTER_IDS.ERROR_CODE, value: "429", param: "error_code" },
       { id: LOG_FILTER_IDS.ERROR_MESSAGE, value: "rate limited", param: "error_message" },
       { id: LOG_FILTER_IDS.USER_ID, value: "user-9", param: "user_id" },
+      { id: LOG_FILTER_IDS.SEARCH, value: "any-id", param: "search" },
     ];
 
     it.each(cases)("sends $id as $param", async ({ id, value, param }) => {
diff --git a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx
index 90f0f0a60f1..3d368527ad9 100644
--- a/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/log_filter_logic.tsx
@@ -33,6 +33,7 @@ export const LOG_FILTER_IDS = {
   PUBLIC_MODEL_OR_SEARCH_TOOL: "model",
   REQUEST_ID: "request_id",
   USER_ID: "user_id",
+  SEARCH: "search",
 } as const;
 
 export const LOG_FILTER_LABELS: Record<string, string> = {
@@ -48,6 +49,7 @@ export const LOG_FILTER_LABELS: Record<string, string> = {
   [LOG_FILTER_IDS.SESSION_ID]: "Session ID",
   [LOG_FILTER_IDS.MODEL_ID]: "Model",
   [LOG_FILTER_IDS.PUBLIC_MODEL_OR_SEARCH_TOOL]: "Public model / search tool",
+  [LOG_FILTER_IDS.SEARCH]: "Search",
 };
 
 export interface LogsWindow {
@@ -175,6 +177,7 @@ export function useLogFilterLogic({
           api_key: getFilterValue(columnFilters, LOG_FILTER_IDS.KEY_HASH),
           team_id: getFilterValue(columnFilters, LOG_FILTER_IDS.TEAM_ID),
           request_id: getFilterValue(columnFilters, LOG_FILTER_IDS.REQUEST_ID),
+          search: getFilterValue(columnFilters, LOG_FILTER_IDS.SEARCH),
           session_id: getFilterValue(columnFilters, LOG_FILTER_IDS.SESSION_ID),
           user_id: userIdFilter,
           end_user: getFilterValue(columnFilters, LOG_FILTER_IDS.END_USER),

From e504477a696ca0b8c82c2083510e4ae0c9a391f6 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 15:20:55 -0700
Subject: [PATCH 11/38] chore(ui): regenerate schema.d.ts for the new search
 params

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
---
 ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 6742076fa78..48a8cfd54d6 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -40904,6 +40904,8 @@ export interface operations {
                 object_team_id?: string | null;
                 /** @description Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only) */
                 object_key_hash?: string | null;
+                /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value (a raw sk- virtual key is hashed first) */
+                search?: string | null;
                 /** @description Column to sort by (e.g. 'updated_at', 'action', 'table_name') */
                 sort_by?: string | null;
                 /** @description Sort order ('asc' or 'desc') */
@@ -49609,6 +49611,8 @@ export interface operations {
                 key_hash?: string | null;
                 /** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */
                 key_alias?: string | null;
+                /** @description Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive). */
+                search?: string | null;
                 /** @description Return full key object */
                 return_full_object?: boolean;
                 /** @description Include all keys for teams that user is an admin of. */
@@ -56867,6 +56871,8 @@ export interface operations {
                 group_by_session?: boolean;
                 /** @description Keyset cursor '<last_activity>|<api_key>|<session_key>' from a previous group_by_session page. UI route only, honored when sorting by startTime */
                 session_cursor?: string | null;
+                /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
+                search?: string | null;
             };
             header?: never;
             path?: never;
@@ -56983,6 +56989,8 @@ export interface operations {
                 group_by_session?: boolean;
                 /** @description Keyset cursor '<last_activity>|<api_key>|<session_key>' from a previous group_by_session page. UI route only, honored when sorting by startTime */
                 session_cursor?: string | null;
+                /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
+                search?: string | null;
             };
             header?: never;
             path?: never;
@@ -63282,6 +63290,8 @@ export interface operations {
                 key?: string | null;
                 /** @description Filter by key prefix (Redis-style namespace scan). Mutually exclusive with `key`; if both are provided, `key_prefix` wins. */
                 key_prefix?: string | null;
+                /** @description Match entries whose key starts with this value or whose memory_id equals it. Takes precedence over `key_prefix` and `key` when provided. */
+                search?: string | null;
                 page?: number;
                 page_size?: number;
             };

From 57da95a77cbbc27a01372786798acae2a7987489 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:27:47 -0700
Subject: [PATCH 12/38] fix(proxy): apply default_vertex_config location before
 building the Vertex passthrough base URL

Routes without /projects/<project>/locations/<location>/ built the upstream host from the URL's
still-empty location and 500ed even with default_vertex_config set. Build the base URL once after
the configured project and location are applied, drop the hook that re-derived it afterwards, and
answer 400 with a fix-it message when no location is available at all.

Resolves LIT-6905
---
 .../llm_passthrough_endpoints.py              |  42 ++---
 .../test_llm_pass_through_endpoints.py        | 147 ++++++++++++++++--
 .../test_vertex_passthrough_load_balancing.py |  25 +--
 3 files changed, 147 insertions(+), 67 deletions(-)

diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 29f216fd450..688123c9d41 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -1659,6 +1659,12 @@ async def azure_proxy_route(
 
 from abc import ABC, abstractmethod
 
+_VERTEX_LOCATION_REQUIRED_DETAIL: Final = (
+    "No Vertex AI location for this request. Include /projects/<project>/locations/<location>/ in the "
+    "route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the "
+    "model to model_list with use_in_pass_through: true."
+)
+
 
 class BaseVertexAIPassThroughHandler(ABC):
     @staticmethod
@@ -1666,29 +1672,18 @@ class BaseVertexAIPassThroughHandler(ABC):
     def get_default_base_target_url(vertex_location: str | None) -> str:
         pass
 
-    @staticmethod
-    @abstractmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
-        pass
-
 
 class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
     @staticmethod
     def get_default_base_target_url(vertex_location: str | None) -> str:
         return "https://discoveryengine.googleapis.com/"
 
-    @staticmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
-        return base_target_url
-
 
 class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
     @staticmethod
     def get_default_base_target_url(vertex_location: str | None) -> str:
-        return get_vertex_base_url(vertex_location)
-
-    @staticmethod
-    def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
+        if vertex_location is None:
+            raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL)
         return get_vertex_base_url(vertex_location)
 
 
@@ -1911,10 +1906,8 @@ async def _prepare_vertex_auth_headers(
     router_credentials: LiteLLM_ManagedVectorStore | None,
     vertex_project: str | None,
     vertex_location: str | None,
-    base_target_url: str | None,
-    get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
     user_api_key_dict: UserAPIKeyAuth,
-) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]:
+) -> tuple[Mapping[str, str], bool, str | None, str | None]:
     """
     Prepare authentication headers for Vertex AI pass-through requests.
 
@@ -1924,15 +1917,12 @@ async def _prepare_vertex_auth_headers(
         router_credentials: Optional vector store credentials from registry
         vertex_project: Vertex project ID
         vertex_location: Vertex location
-        base_target_url: Base URL for the Vertex AI service
-        get_vertex_pass_through_handler: Handler for the specific Vertex AI service
         user_api_key_dict: The caller's resolved authentication, so only the secret that
             authenticated them is stripped on the credential-less branch
 
     Returns:
         tuple containing:
             - headers: dict - Authentication headers to use
-            - base_target_url: str | None - Updated base target URL
             - headers_passed_through: bool - Whether headers were passed through from request
             - vertex_project: str | None - Updated vertex project ID
             - vertex_location: str | None - Updated vertex location
@@ -1985,14 +1975,8 @@ async def _prepare_vertex_auth_headers(
         # Add the Authorization header with vendor credentials
         headers["Authorization"] = f"Bearer {auth_header}"
 
-        if base_target_url is not None:
-            base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
-                base_target_url, vertex_location
-            )
-
     return (
         headers,
-        base_target_url,
         headers_passed_through,
         vertex_project,
         vertex_location,
@@ -2085,12 +2069,9 @@ async def _base_vertex_proxy_route(
         location=vertex_location,
     )
 
-    base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
-
     # Prepare authentication headers
     (
         headers,
-        base_target_url,
         headers_passed_through,
         vertex_project,
         vertex_location,
@@ -2100,13 +2081,10 @@ async def _base_vertex_proxy_route(
         router_credentials=router_credentials,
         vertex_project=vertex_project,
         vertex_location=vertex_location,
-        base_target_url=base_target_url,
-        get_vertex_pass_through_handler=get_vertex_pass_through_handler,
         user_api_key_dict=user_api_key_dict,
     )
 
-    if base_target_url is None:
-        base_target_url = get_vertex_base_url(vertex_location)
+    base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
 
     request_route: Final = encoded_endpoint
     verbose_proxy_logger.debug("request_route %s", request_route)
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
index 1d4b0264879..e37060493f7 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
@@ -319,9 +319,6 @@ class TestVertexAIPassThroughHandler:
             mock_handler.get_default_base_target_url.return_value = (
                 f"https://{test_location}-aiplatform.googleapis.com/"
             )
-            mock_handler.update_base_target_url_with_credential_location = Mock(
-                return_value=f"https://{test_location}-aiplatform.googleapis.com/"
-            )
             mock_get_handler.return_value = mock_handler
 
             # Mock create_pass_through_route to return a function that returns a mock response
@@ -427,9 +424,6 @@ class TestVertexAIPassThroughHandler:
             mock_handler.get_default_base_target_url.return_value = (
                 "https://aiplatform.googleapis.com/"
             )
-            mock_handler.update_base_target_url_with_credential_location = Mock(
-                return_value="https://aiplatform.googleapis.com/"
-            )
             mock_get_handler.return_value = mock_handler
 
             # Mock create_pass_through_route to return a function that returns a mock response
@@ -530,9 +524,6 @@ class TestVertexAIPassThroughHandler:
             mock_handler.get_default_base_target_url.return_value = (
                 f"https://{default_location}-aiplatform.googleapis.com/"
             )
-            mock_handler.update_base_target_url_with_credential_location = Mock(
-                return_value=f"https://{default_location}-aiplatform.googleapis.com/"
-            )
             mock_get_handler.return_value = mock_handler
 
             # Mock create_pass_through_route to return a function that returns a mock response
@@ -1308,9 +1299,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
             mock_handler.get_default_base_target_url.return_value = (
                 "https://discoveryengine.googleapis.com"
             )
-            mock_handler.update_base_target_url_with_credential_location = Mock(
-                return_value="https://discoveryengine.googleapis.com"
-            )
             mock_get_handler.return_value = mock_handler
 
             # Mock create_pass_through_route to return a function that returns a mock response
@@ -3650,7 +3638,6 @@ class TestVertexRawPredictStreamingClassification:
         base_url = "https://us-east5-aiplatform.googleapis.com/"
         mock_handler = Mock()
         mock_handler.get_default_base_target_url.return_value = base_url
-        mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url)
 
         module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
         with (
@@ -4234,6 +4221,140 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
         assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
 
 
+class TestVertexPassthroughDefaultLocationOnShortRoutes:
+    """Regression coverage for LIT-6905.
+
+    ``default_vertex_config`` carries the project and location, yet a route that
+    omits ``/projects/<project>/locations/<location>/`` built the upstream base URL
+    from the still-unresolved URL location and 500ed with ``vertex_location is
+    required``. The base URL must be built after the configured location is
+    applied, and a request with no location anywhere must fail with a clean 400
+    that says where a location can come from, never a 500.
+    """
+
+    PROJECT = "test-project"
+    SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent"
+
+    async def _forward(
+        self,
+        monkeypatch,
+        endpoint: str,
+        default_config: dict | None,
+        headers: list[tuple[bytes, bytes]],
+    ) -> tuple[HTTPException | None, dict]:
+        from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
+            PassthroughEndpointRouter,
+        )
+
+        async def receive():
+            return {"type": "http.request", "body": b"{}", "more_body": False}
+
+        request = Request(
+            {
+                "type": "http",
+                "method": "POST",
+                "path": f"/vertex_ai/{endpoint}",
+                "headers": headers,
+                "query_string": b"",
+            },
+            receive=receive,
+        )
+
+        captured: dict = {}
+
+        def fake_create_pass_through_route(**kwargs):
+            captured.update(kwargs)
+            return AsyncMock(return_value={"status": "success"})
+
+        router = PassthroughEndpointRouter()
+        if default_config is not None:
+            router.set_default_vertex_config(dict(default_config))
+        module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
+        monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router)
+        monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
+        mock_credentials = Mock()
+        mock_credentials.token = "test-token"
+        caller: Final = UserAPIKeyAuth(api_key="test-key")
+        raised: HTTPException | None = None
+        with (
+            mock.patch(  # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
+                "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
+                return_value=(mock_credentials, self.PROJECT),
+            ),
+            mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
+            mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
+        ):
+            try:
+                await vertex_proxy_route(
+                    endpoint=endpoint,
+                    request=request,
+                    fastapi_response=Response(),
+                    user_api_key_dict=caller,
+                )
+            except HTTPException as exc:
+                raised = exc
+        return raised, captured
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("endpoint", "location", "expected_target"),
+        [
+            (
+                SHORT_ROUTE,
+                "global",
+                "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
+            ),
+            (
+                f"v1/{SHORT_ROUTE}",
+                "global",
+                "https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
+            ),
+            (
+                f"v1beta1/{SHORT_ROUTE}",
+                "global",
+                "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/" + SHORT_ROUTE,
+            ),
+            (
+                SHORT_ROUTE,
+                "us-central1",
+                "https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/"
+                + SHORT_ROUTE,
+            ),
+        ],
+    )
+    async def test_default_vertex_config_location_fills_routes_without_project_and_location(
+        self, monkeypatch, endpoint, location, expected_target
+    ):
+        raised, captured = await self._forward(
+            monkeypatch,
+            endpoint,
+            {"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"},
+            [(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
+        )
+        assert raised is None
+        assert str(captured["target"]) == expected_target
+        assert captured["custom_headers"]["Authorization"] == "Bearer test-token"
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("default_config", "headers"),
+        [
+            (None, [(b"content-type", b"application/json"), (b"authorization", b"Bearer ya29.byo-google-oauth")]),
+            (
+                {"vertex_project": PROJECT, "vertex_credentials": "test-creds"},
+                [(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
+            ),
+        ],
+    )
+    async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers):
+        raised, captured = await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers)
+        assert not captured, "a request with no location must never reach the upstream forwarder"
+        assert raised is not None
+        assert raised.status_code == 400
+        assert "/projects/<project>/locations/<location>/" in str(raised.detail)
+        assert "default_vertex_config" in str(raised.detail)
+
+
 class TestGetAzureAISearchIndexFromEndpoint:
     """The operable index is only the segment right after ``indexes``.
 
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 6735f2a3780..e8fd5579631 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
@@ -4,6 +4,7 @@ import pytest
 
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
+    VertexAIPassThroughHandler,
     _base_vertex_proxy_route,
     _upstream_headers_for_vertex_route,
 )
@@ -20,6 +21,7 @@ async def test_vertex_passthrough_load_balancing():
     mock_request = MagicMock()
     mock_response = MagicMock()
     mock_handler = MagicMock()
+    mock_handler.get_default_base_target_url.return_value = "https://test.url"
 
     # Mock the router
     mock_router = MagicMock()
@@ -68,7 +70,6 @@ async def test_vertex_passthrough_load_balancing():
         mock_pt_router.get_vertex_credentials.return_value = MagicMock()
         mock_prep_headers.return_value = (
             {},
-            "https://test.url",
             False,
             "test-project-lb",
             "us-central1-lb",
@@ -290,12 +291,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
     mock_vertex_credentials.vertex_location = "us-central1"
     mock_vertex_credentials.vertex_credentials = "test-credentials"
 
-    # Create mock handler
-    mock_handler = MagicMock()
-    mock_handler.update_base_target_url_with_credential_location.return_value = (
-        "https://us-central1-aiplatform.googleapis.com"
-    )
-
     with (
         patch.object(
             VertexBase,
@@ -313,7 +308,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
         # Call the function
         (
             headers,
-            base_target_url,
             headers_passed_through,
             vertex_project,
             vertex_location,
@@ -323,8 +317,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
             router_credentials=None,
             vertex_project="test-project",
             vertex_location="us-central1",
-            base_target_url="https://us-central1-aiplatform.googleapis.com",
-            get_vertex_pass_through_handler=mock_handler,
             user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
         )
 
@@ -394,7 +386,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
                 "content-type": "application/json",
                 "Authorization": "Bearer vertex-access-token",
             },
-            "https://aiplatform.googleapis.com",
             False,
             "test-project",
             "global",
@@ -406,7 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
             endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
             request=MagicMock(),
             fastapi_response=MagicMock(),
-            get_vertex_pass_through_handler=MagicMock(),
+            get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
         )
 
         upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
@@ -473,12 +464,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
     mock_vertex_credentials.vertex_location = "us-central1"
     mock_vertex_credentials.vertex_credentials = "test-credentials"
 
-    # Create mock handler
-    mock_handler = MagicMock()
-    mock_handler.update_base_target_url_with_credential_location.return_value = (
-        "https://us-central1-aiplatform.googleapis.com"
-    )
-
     with (
         patch.object(
             VertexBase,
@@ -495,7 +480,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
 
         (
             headers,
-            _base_target_url,
             _headers_passed_through,
             _vertex_project,
             _vertex_location,
@@ -505,8 +489,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
             router_credentials=None,
             vertex_project="test-project",
             vertex_location="us-central1",
-            base_target_url="https://us-central1-aiplatform.googleapis.com",
-            get_vertex_pass_through_handler=mock_handler,
             user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
         )
 
@@ -742,7 +724,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
         mock_pt_router.get_vertex_credentials.return_value = MagicMock()
         mock_prep_headers.return_value = (
             {},
-            "https://global-aiplatform.googleapis.com",
             False,
             "nv-gcpllmgwit-20250411173346",
             "global",

From 24531ee576b4a5c535f46e840d27e2d89e990878 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:37:56 -0700
Subject: [PATCH 13/38] refactor(cost): drop the docstrings that restate
 TokenRates and the new tests

---
 litellm/litellm_core_utils/llm_cost_calc/utils.py   |  5 -----
 .../llm_cost_calc/test_llm_cost_calc_utils.py       | 13 -------------
 .../dashscope/test_dashscope_cost_calculator.py     |  5 -----
 3 files changed, 23 deletions(-)

diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py
index e03c2c93c26..8e24302b440 100644
--- a/litellm/litellm_core_utils/llm_cost_calc/utils.py
+++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py
@@ -417,11 +417,6 @@ def _is_off_peak(off_peak: Mapping[str, object], current_time: datetime | None =
 
 @dataclass(frozen=True, slots=True)
 class TokenRates:
-    """The per-token rates one request bills at. reasoning_rate is None when reasoning bills at
-    output_rate: the model has no dedicated reasoning rate, or the caller resolves reasoning on
-    its own.
-    """
-
     input_rate: float
     output_rate: float
     cache_read_rate: float
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 7bc02145841..d5d7b6a47f4 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
@@ -815,8 +815,6 @@ def _off_peak_reasoning_usage() -> Usage:
 
 
 def test_generic_cost_per_token_off_peak_reasoning_rate():
-    """Regression (LIT-6887): the block's output_cost_per_reasoning_token used to be ignored, so
-    reasoning tokens billed at the model's standard reasoning rate all through the window."""
     from datetime import datetime, timezone
 
     model_name = "litellm-test-off-peak-reasoning"
@@ -843,8 +841,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate():
 
 
 def test_generic_cost_per_token_off_peak_block_without_reasoning_rate():
-    """A block that leaves output_cost_per_reasoning_token unset keeps the model's own reasoning
-    rate, and a model with no reasoning rate at all follows the off-peak output rate."""
     from datetime import datetime, timezone
 
     inside_window = datetime(2026, 1, 1, 18, 0, tzinfo=timezone.utc)
@@ -870,7 +866,6 @@ def test_generic_cost_per_token_off_peak_block_without_reasoning_rate():
 
 
 def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier():
-    """Tiered models resolve reasoning on their own path, so the block has to win there too."""
     from datetime import datetime, timezone
 
     model_name = "litellm-test-off-peak-tiered-reasoning"
@@ -914,8 +909,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_tier():
 
 
 def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_tier():
-    """A priority request bills its service-tier reasoning rate outside the window and the block's
-    rate inside it."""
     from datetime import datetime, timezone
 
     model_name = "litellm-test-off-peak-reasoning-service-tier"
@@ -946,7 +939,6 @@ def test_generic_cost_per_token_off_peak_reasoning_rate_wins_over_the_service_ti
 
 
 def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings():
-    """A YAML true never turns into a rate of 1.0, and a quoted number still counts."""
     from datetime import datetime, timezone
 
     model_name = "litellm-test-off-peak-odd-values"
@@ -972,9 +964,6 @@ def test_apply_off_peak_pricing_treats_bool_as_unset_and_parses_strings():
 
 
 def test_get_token_base_cost_off_peak_cache_creation_rate():
-    """Regression (LIT-6887): the block's cache_creation_input_token_cost used to be ignored. It
-    replaces the five-minute cache-creation rate inside the window; the one-hour rate, and a
-    block without the key, keep the standard rate."""
     from datetime import datetime, timezone
     from typing import cast
 
@@ -1008,8 +997,6 @@ def test_get_token_base_cost_off_peak_cache_creation_rate():
 
 
 def test_get_token_type_cost_breakdown_reflects_off_peak_reasoning_and_cache_creation_rates():
-    """The per-token-type breakdown feeds the spend logs, so it has to bill the new keys the same
-    way the total does."""
     from datetime import datetime, timezone
 
     model_name = "litellm-test-off-peak-breakdown"
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 f0949ce041a..a30d35d46f2 100644
--- a/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
+++ b/tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py
@@ -650,8 +650,6 @@ class TestDashscopeCostCalculator:
         assert math.isclose(completion_cost, 200 * 2.4e-06, rel_tol=1e-10)
 
     def test_dashscope_off_peak_reasoning_rate_replaces_the_dedicated_reasoning_rate(self):
-        """Regression (LIT-6887): a block carrying output_cost_per_reasoning_token bills reasoning
-        tokens at it inside the window, over the model's own reasoning rate, which returns outside."""
         self._register_off_peak_flat_model(
             "dashscope/qwen-reasoning-rate-off-peak-test",
             {
@@ -678,8 +676,6 @@ class TestDashscopeCostCalculator:
         assert math.isclose(peak_completion_cost, (150 * 4.8e-06) + (50 * 9e-06), rel_tol=1e-10)
 
     def test_dashscope_off_peak_cache_creation_rate_replaces_the_standard_rate(self):
-        """Regression (LIT-6887): a block carrying cache_creation_input_token_cost bills cache-creation
-        tokens at it inside the window, while the cache-read rate it leaves unset stays standard."""
         self._register_off_peak_flat_model(
             "dashscope/qwen-cache-creation-off-peak-test",
             {"hours_utc": self.OFF_PEAK_WINDOW, "cache_creation_input_token_cost": 1.5e-06},
@@ -701,7 +697,6 @@ class TestDashscopeCostCalculator:
         assert math.isclose(peak_prompt_cost, (600 * 2.4e-06) + (300 * 2e-07) + (100 * 3e-06), rel_tol=1e-10)
 
     def test_dashscope_off_peak_reasoning_and_cache_creation_rates_override_the_selected_tier(self):
-        """The new keys override the selected tier the way the input and output rates already do."""
         self._register_tiered_model(
             "dashscope/qwen-tiered-reasoning-off-peak-test",
             [

From 1d71e306cc13008f551d8964623ca813ca44bccb Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 15:38:15 -0700
Subject: [PATCH 14/38] test(router): assert a non-chat configured mode
 round-trips

With every kept get_configured_mode test using mode "chat", a Router that
answered "chat" for any non-blank configured mode passed all four of them
(the deleted #39630 pair's audio_speech case was the only test catching it).
Read the mode back as audio_speech on an unmapped model so the configured
value itself is what the test checks. Six hand-applied mutations of
Router.get_configured_mode, including that hardcoded-chat one, are now all
killed by the four surviving tests.
---
 tests/test_litellm/test_router.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 417c58b95f6..228588d974f 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -7362,14 +7362,14 @@ def test_get_configured_mode_reads_deployment_model_info():
     router = litellm.Router(
         model_list=[
             {
-                "model_name": "chat-model",
-                "litellm_params": {"model": "openai/some-unmapped-model"},
-                "model_info": {"mode": "chat"},
+                "model_name": "tts-model",
+                "litellm_params": {"model": "openai/some-unmapped-tts-model"},
+                "model_info": {"mode": "audio_speech"},
             }
         ]
     )
 
-    assert router.get_configured_mode("chat-model") == "chat"
+    assert router.get_configured_mode("tts-model") == "audio_speech"
 
 
 def test_get_configured_mode_returns_none_for_unset_or_unknown():

From 9baa19c7d13c4d5937d7c6f337659b94b514f773 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 15:54:42 -0700
Subject: [PATCH 15/38] fix(proxy): stop hashing raw sk- values in list
 searches

The search= param on /key/list, /audit, and /spend/logs/ui, plus key_hash= on /key/list, now compare the pasted value verbatim. Only a copied key ID (the hash) matches, so a raw virtual key never needs to travel in a GET query string

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
---
 .../proxy/audit_logging_endpoints.py          | 17 ++---
 .../key_management_endpoints.py               | 12 ++--
 .../spend_management_endpoints.py             | 14 ++--
 .../key_management_endpoints.py               |  2 +-
 .../proxy/test_audit_logging_endpoints.py     | 19 -----
 .../test_key_management_endpoints.py          | 71 +++----------------
 .../test_spend_management_endpoints.py        | 37 ++++------
 .../(dashboard)/hooks/keys/useKeys.test.ts    |  4 +-
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  8 +--
 9 files changed, 43 insertions(+), 141 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
index 72f7a66a420..7df14565c3f 100644
--- a/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
+++ b/enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py
@@ -18,7 +18,6 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
 
 from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
-from litellm.proxy.utils import _hash_token_if_needed
 from litellm.repositories.prisma_protocols import TableActions
 from litellm.repositories.table_repositories import AuditLogRepository
 
@@ -50,14 +49,13 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec
 
 
 def _build_search_condition(search: str) -> dict[str, object]:
-    """Match any id column; a raw sk- key is hashed for the two columns that store key hashes."""
-    hashed: Final = _hash_token_if_needed(search)
+    """Match a row whose id, changed_by, object_id, or changed_by_api_key equals the search value."""
     return {
         "OR": (
             {"id": search},
             {"changed_by": search},
-            {"object_id": hashed},
-            {"changed_by_api_key": hashed},
+            {"object_id": search},
+            {"changed_by_api_key": search},
         )
     }
 
@@ -99,10 +97,7 @@ async def get_audit_logs(
     ),
     search: str | None = Query(
         None,
-        description=(
-            "Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value "
-            "(a raw sk- virtual key is hashed first)"
-        ),
+        description="Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value",
     ),
     # Sorting parameters
     sort_by: str | None = Query(
@@ -159,7 +154,7 @@ async def get_audit_logs(
         {sort_by: sort_order} if sort_by and isinstance(sort_by, str) else {"updated_at": sort_order}
     )
 
-    audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
+    audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table
 
     # Get paginated results
     audit_logs: Final = await audit_log_table.find_many(
@@ -221,7 +216,7 @@ async def get_audit_log_by_id(
             detail={"message": CommonProxyErrors.db_not_connected_error.value},
         )
 
-    audit_log_table: Final[TableActions["prisma_models.LiteLLM_AuditLog"]] = AuditLogRepository(prisma_client).table
+    audit_log_table: Final[TableActions[prisma_models.LiteLLM_AuditLog]] = AuditLogRepository(prisma_client).table
 
     # Get the audit log by ID
     audit_log: Final = await audit_log_table.find_unique(where={"id": id})
diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 758644ff01b..324d380b85b 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -5803,7 +5803,7 @@ async def list_keys(
     ),
     search: str | None = Query(
         None,
-        description="Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive).",
+        description="Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive).",
     ),
     return_full_object: bool = Query(False, description="Return full key object"),
     include_team_keys: bool = Query(False, description="Include all keys for teams that user is an admin of."),
@@ -5867,17 +5867,13 @@ async def list_keys(
                 detail={"error": "Invalid expires value. Supported: 'active', 'expired'."},
             )
 
-        hashed_key_hash: Final[str | None] = (
-            _hash_token_if_needed(token=key_hash) if isinstance(key_hash, str) else None
-        )
-
         complete_user_info: Final = await validate_key_list_check(
             user_api_key_dict=user_api_key_dict,
             user_id=user_id,
             team_id=team_id,
             organization_id=organization_id,
             key_alias=key_alias,
-            key_hash=hashed_key_hash,
+            key_hash=key_hash,
             prisma_client=prisma_client,
         )
 
@@ -5937,7 +5933,7 @@ async def list_keys(
             user_id=user_id,
             team_id=team_id,
             key_alias=key_alias,
-            key_hash=hashed_key_hash,
+            key_hash=key_hash,
             return_full_object=return_full_object,
             organization_id=organization_id,
             admin_team_ids=admin_team_ids,
@@ -6175,7 +6171,7 @@ def _build_expires_where_clause(expires_filter: str, now: datetime) -> dict[str,
 def _build_key_search_where(search: str) -> KeySearchWhere:
     search_where: Final[KeySearchWhere] = {
         "OR": (
-            {"token": _hash_token_if_needed(token=search)},
+            {"token": search},
             {"key_alias": {"contains": search, "mode": "insensitive"}},
         )
     }
diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index 9ec8dd205a6..b86a877e8f9 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -2240,20 +2240,18 @@ def _build_spend_log_search_condition(
     end_date: datetime,
     next_param_index: int,
 ) -> _SpendLogSearchCondition:
-    """request_id (indexed) matches across all time; the unindexed id columns only inside the window (sk- keys hashed)."""
+    """request_id (indexed) matches across all time; the unindexed id columns only inside the window."""
     raw: Final = f"${next_param_index}"
-    hashed: Final = f"${next_param_index + 1}"
-    window_start: Final = f"${next_param_index + 2}"
-    window_end: Final = f"${next_param_index + 3}"
+    window_start: Final = f"${next_param_index + 1}"
+    window_end: Final = f"${next_param_index + 2}"
     sql: Final = (
         f"(request_id = {raw} OR ("
         f"\"startTime\" >= ({window_start}::timestamptz AT TIME ZONE 'UTC') "
         f"AND \"startTime\" <= ({window_end}::timestamptz AT TIME ZONE 'UTC') "
-        f'AND (api_key = {hashed} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} '
+        f'AND (api_key = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} '
         f"OR session_id = {raw} OR model_id = {raw})))"
     )
-    hashed_search: Final = hash_token(token=search) if search.startswith("sk-") else search
-    return _SpendLogSearchCondition(sql=sql, params=(search, hashed_search, start_date, end_date))
+    return _SpendLogSearchCondition(sql=sql, params=(search, start_date, end_date))
 
 
 @router.get(
@@ -2359,7 +2357,7 @@ async def ui_view_spend_logs(
     search: str | None = fastapi.Query(
         default=None,
         description=(
-            "Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, "
+            "Match a log whose request_id, api_key (hash), team_id, user, end_user, "
             "session_id, or model_id equals this value. request_id matches across all time; the other columns "
             "match inside start_date/end_date, which stay required"
         ),
diff --git a/litellm/types/proxy/management_endpoints/key_management_endpoints.py b/litellm/types/proxy/management_endpoints/key_management_endpoints.py
index 5d410d7b55b..9fb5bea81e3 100644
--- a/litellm/types/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/types/proxy/management_endpoints/key_management_endpoints.py
@@ -16,7 +16,7 @@ class KeyAliasContainsWhere(TypedDict):
 
 
 class KeySearchWhere(TypedDict):
-    """Prisma filter behind `/key/list?search=`: exact token (sk- keys hashed) or alias substring, case-insensitive."""
+    """Prisma filter behind `/key/list?search=`: exact token or case-insensitive alias substring."""
 
     OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
 
diff --git a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
index cd2c8b0b904..fd1b05ff060 100644
--- a/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
+++ b/tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py
@@ -1,4 +1,3 @@
-import hashlib
 from datetime import datetime, timedelta
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -172,24 +171,6 @@ def test_search_matches_any_id_column_alongside_the_other_filters(mock_prisma_cl
     }
 
 
-def test_search_hashes_a_raw_virtual_key_for_the_hashed_columns(mock_prisma_client):
-    where: Final = _list_audit_logs_where(mock_prisma_client, "search=sk-raw")
-
-    hashed: Final = hashlib.sha256(b"sk-raw").hexdigest()
-    assert where == {
-        "AND": (
-            {
-                "OR": (
-                    {"id": "sk-raw"},
-                    {"changed_by": "sk-raw"},
-                    {"object_id": hashed},
-                    {"changed_by_api_key": hashed},
-                )
-            },
-        )
-    }
-
-
 def test_an_empty_search_leaves_the_where_clause_unchanged(mock_prisma_client):
     where: Final = _list_audit_logs_where(mock_prisma_client, "action=create&search=")
 
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 1e399e4fb58..0e4af9f75a5 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
@@ -6351,33 +6351,15 @@ def _search_clause(search: str, token: str) -> dict:
     return {"OR": [{"token": token}, {"key_alias": {"contains": search, "mode": "insensitive"}}]}
 
 
-def test_build_key_filter_conditions_search_hashes_raw_key_and_ors_alias_contains():
+def test_build_key_filter_conditions_search_ors_token_and_alias_contains():
     """
     LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by
-    its ID. A pasted raw sk- key is hashed to its token first; an already-hashed
-    value is used verbatim.
+    its ID (the token column), with the pasted value used verbatim.
     """
-    from litellm.proxy._types import hash_token
     from litellm.proxy.management_endpoints.key_management_endpoints import (
         _build_key_filter_conditions,
     )
 
-    raw_where = json.loads(
-        json.dumps(
-            _build_key_filter_conditions(
-                user_id=None,
-                team_id=None,
-                organization_id=None,
-                key_alias=None,
-                key_hash=None,
-                exclude_team_id=None,
-                admin_team_ids=None,
-                search="sk-raw",
-            )
-        )
-    )
-    assert _search_clause("sk-raw", hash_token("sk-raw")) in raw_where["AND"], f"raw search not ANDed: {raw_where}"
-
     hashed_where = json.loads(
         json.dumps(
             _build_key_filter_conditions(
@@ -6402,7 +6384,6 @@ def test_build_key_filter_conditions_search_narrows_team_admin_visibility():
     LIT-4741, same class as LIT-3243: `search` must be a top-level AND so it
     narrows a team admin's admin-team branch instead of being bypassed by it.
     """
-    from litellm.proxy._types import hash_token
     from litellm.proxy.management_endpoints.key_management_endpoints import (
         _build_key_filter_conditions,
     )
@@ -6419,21 +6400,19 @@ def test_build_key_filter_conditions_search_narrows_team_admin_visibility():
                 admin_team_ids=["team-a"],
                 member_team_ids=["team-a"],
                 include_created_by_keys=False,
-                search="sk-member",
+                search="member-key-id",
             )
         )
     )
 
     assert where.get("AND"), f"expected top-level AND, got: {where}"
-    assert _search_clause("sk-member", hash_token("sk-member")) in where["AND"], f"search not ANDed: {where}"
+    assert _search_clause("member-key-id", "member-key-id") in where["AND"], f"search not ANDed: {where}"
     assert json.dumps({"team_id": {"in": ["team-a"]}}) in json.dumps(where)
 
 
 @pytest.mark.asyncio
 async def test_list_key_helper_applies_search_to_prisma_where():
     """LIT-4741: `search` given to _list_key_helper must reach the Prisma where clause."""
-    from litellm.proxy._types import hash_token
-
     mock_prisma_client = AsyncMock()
     mock_find_many = AsyncMock(return_value=[])
     mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many
@@ -6448,11 +6427,11 @@ async def test_list_key_helper_applies_search_to_prisma_where():
         organization_id=None,
         key_alias=None,
         key_hash=None,
-        search="sk-raw",
+        search="key-id-123",
     )
 
     where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"]))
-    assert _search_clause("sk-raw", hash_token("sk-raw")) in where["AND"], f"search not in Prisma where: {where}"
+    assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}"
 
 
 @pytest.mark.asyncio
@@ -14978,47 +14957,13 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
     assert kwargs["user_id"] == "alice"
 
 
-@pytest.mark.asyncio
-async def test_list_keys_hashes_raw_key_hash_before_validation():
-    """LIT-4741: a raw sk- key pasted as key_hash is hashed before the ownership
-    check and the query, so a non-admin filtering by their own raw key gets the
-    row instead of the 'Key Hash not found.' 403."""
-    from litellm.proxy._types import hash_token
-
-    user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
-    validate = AsyncMock(
-        return_value=LiteLLM_UserTable(
-            user_id="alice", user_email="alice@example.com", teams=[], organization_memberships=[]
-        )
-    )
-    helper = AsyncMock(return_value={"keys": [], "total_count": 0, "current_page": 1, "total_pages": 0})
-    with (
-        patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()),
-        patch(
-            "litellm.proxy.management_endpoints.key_management_endpoints.validate_key_list_check",
-            validate,
-        ),
-        patch("litellm.proxy.management_endpoints.key_management_endpoints._list_key_helper", helper),
-    ):
-        await list_keys(
-            request=MagicMock(),
-            user_api_key_dict=user,
-            status=None,
-            user_id=None,
-            key_hash="sk-raw",
-        )
-
-    assert validate.call_args.kwargs["key_hash"] == hash_token("sk-raw")
-    assert helper.call_args.kwargs["key_hash"] == hash_token("sk-raw")
-
-
 @pytest.mark.asyncio
 async def test_list_keys_search_is_honored_for_non_admin():
     """LIT-4741: unlike substring_matching, `search` is not admin-gated. A non-admin's
     search reaches the helper while their own-user scoping stays in place."""
     user = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, user_id="alice")
-    kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="sk-raw")
-    assert kwargs["search"] == "sk-raw"
+    kwargs = await _list_keys_capture_helper_kwargs(user, user_id=None, search="key-id-123")
+    assert kwargs["search"] == "key-id-123"
     assert kwargs["user_id"] == "alice"
 
 
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 f869f3ffba2..73a29afd9b9 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
@@ -61,7 +61,7 @@ def _filter_logs_by_date_range(logs, where):
 _SEARCH_CLAUSE_RE = re.compile(
     r'\(request_id = \$(\d+) OR \("startTime" >= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
     r'AND "startTime" <= \(\$(\d+)::timestamptz AT TIME ZONE \'UTC\'\) '
-    r'AND \(api_key = \$(\d+) OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 '
+    r'AND \(api_key = \$\1 OR team_id = \$\1 OR "user" = \$\1 OR end_user = \$\1 '
     r"OR session_id = \$\1 OR model_id = \$\1\)\)\)"
 )
 
@@ -72,9 +72,8 @@ def _matches_spend_log_search(log, search):
         return True
     if not _filter_logs_by_date_range([log], {"startTime": {"gte": search["gte"], "lte": search["lte"]}}):
         return False
-    if log.get("api_key") == search["api_key"]:
-        return True
-    return any(log.get(col) == search["value"] for col in ("team_id", "user", "end_user", "session_id", "model_id"))
+    columns = ("api_key", "team_id", "user", "end_user", "session_id", "model_id")
+    return any(log.get(col) == search["value"] for col in columns)
 
 
 def _reconstruct_ui_where_from_sql(sql_query, params):
@@ -98,10 +97,9 @@ def _reconstruct_ui_where_from_sql(sql_query, params):
 
     search_clause = _SEARCH_CLAUSE_RE.search(clause.group(1))
     if search_clause:
-        raw_index, start_index, end_index, hashed_index = (int(g) for g in search_clause.groups())
+        raw_index, start_index, end_index = (int(g) for g in search_clause.groups())
         where["search"] = {
             "value": params[raw_index - 1],
-            "api_key": params[hashed_index - 1],
             "gte": _iso(params[start_index - 1]),
             "lte": _iso(params[end_index - 1]),
         }
@@ -2384,31 +2382,20 @@ async def test_ui_view_spend_logs_request_id_owner_scoped_by_id_only(
 
 def test_build_spend_log_search_condition_windows_every_branch_except_request_id():
     """LIT-4741: request_id matches across all time; the six other id columns only inside the window,
-    and a raw sk- key is hashed for the api_key branch alone."""
+    all comparing the pasted value verbatim."""
     start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
     end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
 
     condition = spend_management_endpoints._build_spend_log_search_condition(
-        search="sk-raw-key", start_date=start, end_date=end, next_param_index=3
+        search="key-hash-7", start_date=start, end_date=end, next_param_index=3
     )
 
     assert condition.sql == (
-        "(request_id = $3 OR (\"startTime\" >= ($5::timestamptz AT TIME ZONE 'UTC') "
-        "AND \"startTime\" <= ($6::timestamptz AT TIME ZONE 'UTC') "
-        'AND (api_key = $4 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))'
+        "(request_id = $3 OR (\"startTime\" >= ($4::timestamptz AT TIME ZONE 'UTC') "
+        "AND \"startTime\" <= ($5::timestamptz AT TIME ZONE 'UTC') "
+        'AND (api_key = $3 OR team_id = $3 OR "user" = $3 OR end_user = $3 OR session_id = $3 OR model_id = $3)))'
     )
-    assert condition.params == ("sk-raw-key", hashlib.sha256(b"sk-raw-key").hexdigest(), start, end)
-
-
-def test_build_spend_log_search_condition_leaves_non_key_values_unhashed():
-    start = datetime.datetime(2026, 8, 1, tzinfo=timezone.utc)
-    end = datetime.datetime(2026, 8, 2, tzinfo=timezone.utc)
-
-    condition = spend_management_endpoints._build_spend_log_search_condition(
-        search="sess-42", start_date=start, end_date=end, next_param_index=1
-    )
-
-    assert condition.params == ("sess-42", "sess-42", start, end)
+    assert condition.params == ("key-hash-7", start, end)
 
 
 def _search_fixture_logs(today):
@@ -2427,7 +2414,7 @@ def _search_fixture_logs(today):
     return [
         {**base, "request_id": "req-session", "session_id": "sess-42", "startTime": recent},
         {**base, "request_id": "req-session-old", "session_id": "sess-42", "startTime": old},
-        {**base, "request_id": "req-key", "api_key": hashlib.sha256(b"sk-raw-key").hexdigest(), "startTime": recent},
+        {**base, "request_id": "req-key", "api_key": "hashed-7", "startTime": recent},
         {**base, "request_id": "req-team", "team_id": "team-7", "startTime": recent},
         {**base, "request_id": "req-user", "user": "user-7", "startTime": recent},
         {**base, "request_id": "req-end-user", "end_user": "cust-7", "startTime": recent},
@@ -2461,7 +2448,7 @@ def _five_day_window(today):
     [
         ("req-session-old", {"req-session-old"}),
         ("sess-42", {"req-session"}),
-        ("sk-raw-key", {"req-key"}),
+        ("hashed-7", {"req-key"}),
         ("team-7", {"req-team"}),
         ("user-7", {"req-user"}),
         ("cust-7", {"req-end-user"}),
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
index f23fcf811f2..84be7e2ef49 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts
@@ -525,14 +525,14 @@ describe("useKeys", () => {
       json: async () => mockKeysResponse,
     });
 
-    const { result } = renderHook(() => useKeys(1, 10, { search: "sk-pasted-key" }), { wrapper });
+    const { result } = renderHook(() => useKeys(1, 10, { search: "pasted-key-id" }), { wrapper });
 
     await waitFor(() => {
       expect(result.current.isLoading).toBe(false);
     });
 
     const callUrl = new URL(mockFetch.mock.calls[0][0], "http://localhost");
-    expect(callUrl.searchParams.get("search")).toBe("sk-pasted-key");
+    expect(callUrl.searchParams.get("search")).toBe("pasted-key-id");
     expect(callUrl.searchParams.has("key_alias")).toBe(false);
     expect(callUrl.searchParams.has("key_hash")).toBe(false);
   });
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 48a8cfd54d6..324085fdae4 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -40904,7 +40904,7 @@ export interface operations {
                 object_team_id?: string | null;
                 /** @description Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only) */
                 object_key_hash?: string | null;
-                /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value (a raw sk- virtual key is hashed first) */
+                /** @description Match a row whose id, object_id, changed_by, or changed_by_api_key equals this value */
                 search?: string | null;
                 /** @description Column to sort by (e.g. 'updated_at', 'action', 'table_name') */
                 sort_by?: string | null;
@@ -49611,7 +49611,7 @@ export interface operations {
                 key_hash?: string | null;
                 /** @description Filter keys by key alias. Exact match by default; set substring_matching=true (admin only) for case-insensitive substring matching. */
                 key_alias?: string | null;
-                /** @description Combined search: matches keys whose token (key hash) equals the value, hashing a raw sk- key first, OR whose key_alias contains it (case-insensitive). */
+                /** @description Combined search: matches keys whose token (key hash) equals the value OR whose key_alias contains it (case-insensitive). */
                 search?: string | null;
                 /** @description Return full key object */
                 return_full_object?: boolean;
@@ -56871,7 +56871,7 @@ export interface operations {
                 group_by_session?: boolean;
                 /** @description Keyset cursor '<last_activity>|<api_key>|<session_key>' from a previous group_by_session page. UI route only, honored when sorting by startTime */
                 session_cursor?: string | null;
-                /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
+                /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
                 search?: string | null;
             };
             header?: never;
@@ -56989,7 +56989,7 @@ export interface operations {
                 group_by_session?: boolean;
                 /** @description Keyset cursor '<last_activity>|<api_key>|<session_key>' from a previous group_by_session page. UI route only, honored when sorting by startTime */
                 session_cursor?: string | null;
-                /** @description Match a log whose request_id, api_key (a raw sk- key is hashed first), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
+                /** @description Match a log whose request_id, api_key (hash), team_id, user, end_user, session_id, or model_id equals this value. request_id matches across all time; the other columns match inside start_date/end_date, which stay required */
                 search?: string | null;
             };
             header?: never;

From 9464888ee9064df4083eee8843424b16eacd7da0 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 16:04:41 -0700
Subject: [PATCH 16/38] test(proxy): pass search=None in direct
 ui_view_spend_logs calls

Calling the endpoint without going through FastAPI leaves the new search param set to its Query default object, which is not None, so the grouped-session and request_id lookup tests started taking the search branch

Claude-Session: https://claude.ai/code/session_01Q5sbiogJzPcCRmYSbaHxZf
---
 .../proxy/spend_tracking/test_spend_query_optimization.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py
index 9ae932ff01f..a7de3f1d8d6 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_spend_query_optimization.py
@@ -184,6 +184,7 @@ async def test_spend_logs_ui_wraps_params_in_at_time_zone_utc(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=1,
@@ -247,6 +248,7 @@ async def test_spend_logs_ui_uses_bounded_count_not_full_scan(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=1,
@@ -314,6 +316,7 @@ async def test_spend_logs_ui_caps_total_for_large_result_sets(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=1,
@@ -359,6 +362,7 @@ async def test_spend_logs_ui_empty_page_reports_zero_total(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=1,
@@ -406,6 +410,7 @@ async def test_spend_logs_ui_out_of_range_page_keeps_total(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=99,
@@ -552,6 +557,7 @@ async def test_spend_logs_ui_group_by_session_paginates_sessions(monkeypatch):
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=1,
@@ -616,6 +622,7 @@ async def test_spend_logs_ui_group_by_session_offset_pages_for_other_sorts(monke
         api_key=None,
         user_id=None,
         request_id=None,
+        search=None,
         start_date="2026-02-16 00:00:00",
         end_date="2026-02-16 23:59:59",
         page=2,
@@ -664,6 +671,7 @@ async def test_spend_logs_ui_request_id_lookup_with_grouping_returns_exact_row(m
         api_key=None,
         user_id=None,
         request_id="req-deep-link",
+        search=None,
         start_date=None,
         end_date=None,
         page=1,

From 7bdd148f38f35e5baed4bced6fd980dd77a83bdd Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 3 Sep 2026 16:10:44 -0700
Subject: [PATCH 17/38] test(proxy-extras): fake run_prisma instead of
 subprocess.run in the migrate deploy harness

---
 tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
index 9ffb57924b6..3fab20a28ad 100644
--- a/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
+++ b/tests/litellm-proxy-extras/test_litellm_proxy_extras_utils.py
@@ -763,7 +763,7 @@ class _MigrateDeployHarness:
             "_resolve_specific_migration",
             staticmethod(self.resolved.append),
         )
-        monkeypatch.setattr(utils_module.subprocess, "run", self._fake_run)
+        monkeypatch.setattr(utils_module.prisma_toolchain, "run_prisma", self._fake_run)
         monkeypatch.setattr(utils_module.time, "sleep", lambda seconds: None)
 
         self.baseline_succeeds = True

From dc98901dc1645391986e3434a72cd256617837cf Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Thu, 3 Sep 2026 19:53:32 +0000
Subject: [PATCH 18/38] fix(scim): apply default_internal_user_params.teams to
 SCIM-created users

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../management_endpoints/scim/scim_v2.py      |  3 +-
 .../scim/test_scim_v2_endpoints.py            | 80 ++++++++++++++++++-
 2 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py
index 069f86c852c..0f0124ee9d9 100644
--- a/litellm/proxy/management_endpoints/scim/scim_v2.py
+++ b/litellm/proxy/management_endpoints/scim/scim_v2.py
@@ -1123,7 +1123,6 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
             user_id=user_id,
             user_email=user_id,  # We don't have email from group membership
             user_alias=None,
-            teams=[],  # Teams will be added separately
             metadata={"created_via": created_via},
             auto_create_key=False,
             user_role=default_role,
@@ -1699,7 +1698,7 @@ async def create_user(
             user_id=user_id,
             user_email=user_data["user_email"],
             user_alias=user_data["user_alias"],
-            teams=user_data["teams"],
+            teams=user_data["teams"] or None,
             metadata=metadata,
             auto_create_key=False,
             user_role=resolved_role if admin_group is not None else default_role,
diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
index 1697b77b99a..f4627f82506 100644
--- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
@@ -24,6 +24,7 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
     SCIMRosterSyncError,
     UserProvisionerHelpers,
     _apply_group_patch_updates,
+    _create_user_if_not_exists,
     _extract_group_member_ids,
     _extract_ids_from_path_filter,
     _handle_group_membership_changes,
@@ -37,8 +38,8 @@ from litellm.proxy.management_endpoints.scim.scim_v2 import (
     delete_group,
     delete_user,
     get_groups,
-    get_users,
     get_service_provider_config,
+    get_users,
     merge_placeholder,
     patch_group,
     patch_team_membership,
@@ -304,6 +305,83 @@ async def test_create_user_uses_default_internal_user_params_role(mocker, monkey
     assert called_args.user_role == LitellmUserRoles.PROXY_ADMIN
 
 
+def _mock_scim_create_user_deps(mocker: MockerFixture, scim_user: SCIMUser) -> AsyncMock:
+    mock_prisma_client = mocker.MagicMock()
+    mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+    mock_prisma_client.db.litellm_usertable.find_many = AsyncMock(return_value=())
+    mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None)
+    mocker.patch(  # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user
+        "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
+        AsyncMock(return_value=mock_prisma_client),
+    )
+    mocker.patch(  # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user
+        "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
+        AsyncMock(return_value=scim_user),
+    )
+    return mocker.patch(  # test-quality-ok: endpoint collaborators are module-level, not injectable into create_user
+        "litellm.proxy.management_endpoints.scim.scim_v2.new_user",
+        AsyncMock(return_value=NewUserRequest(user_id=scim_user.userName)),
+    )
+
+
+@pytest.mark.asyncio
+async def test_create_user_without_groups_defers_to_default_team(mocker: MockerFixture, monkeypatch):
+    """IdPs omit groups on POST /Users; teams must stay unset so new_user applies default_internal_user_params.teams"""
+    scim_user = SCIMUser(
+        schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
+        userName="new-user",
+        emails=[SCIMUserEmail(value="new@example.com")],
+    )
+    monkeypatch.setattr(
+        "litellm.default_internal_user_params",
+        {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]},
+        raising=False,
+    )
+    new_user_mock = _mock_scim_create_user_deps(mocker, scim_user)
+
+    await create_user(user=scim_user)
+
+    assert new_user_mock.call_args.kwargs["data"].teams is None
+
+
+@pytest.mark.asyncio
+async def test_create_user_with_groups_keeps_idp_teams(mocker: MockerFixture, monkeypatch):
+    scim_user = SCIMUser(
+        schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
+        userName="new-user",
+        emails=[SCIMUserEmail(value="new@example.com")],
+        groups=[SCIMUserGroup(value="idp-team")],
+    )
+    monkeypatch.setattr(
+        "litellm.default_internal_user_params",
+        {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]},
+        raising=False,
+    )
+    new_user_mock = _mock_scim_create_user_deps(mocker, scim_user)
+
+    await create_user(user=scim_user)
+
+    assert new_user_mock.call_args.kwargs["data"].teams == ["idp-team"]
+
+
+@pytest.mark.asyncio
+async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFixture, monkeypatch):
+    monkeypatch.setattr(
+        "litellm.default_internal_user_params",
+        {"teams": [{"team_id": "default-team", "max_budget_in_team": 25}]},
+        raising=False,
+    )
+    new_user_mock = mocker.patch(  # test-quality-ok: new_user is imported inside the helper, not injectable
+        "litellm.proxy.management_endpoints.internal_user_endpoints.new_user",
+        AsyncMock(return_value=NewUserResponse(user_id="group-user", key="k")),
+    )
+
+    created = await _create_user_if_not_exists(user_id="group-user")
+
+    assert created is not None
+    assert new_user_mock.call_args.kwargs["data"].teams is None
+
+
 @pytest.mark.asyncio
 async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch):
     """

From 0429339204ac41f2c0420d1693f78155244b6205 Mon Sep 17 00:00:00 2001
From: ryan <ryan@berri.ai>
Date: Thu, 3 Sep 2026 20:16:09 +0000
Subject: [PATCH 19/38] fix(scim): pass proxy admin auth to new_user so default
 team add succeeds

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/proxy/management_endpoints/scim/scim_v2.py          | 6 +++++-
 .../management_endpoints/scim/test_scim_v2_endpoints.py     | 3 +++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py
index 0f0124ee9d9..98770cf9c2b 100644
--- a/litellm/proxy/management_endpoints/scim/scim_v2.py
+++ b/litellm/proxy/management_endpoints/scim/scim_v2.py
@@ -1128,7 +1128,10 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
             user_role=default_role,
         )
 
-        created_user: Final = await new_user(data=new_user_request)
+        created_user: Final = await new_user(
+            data=new_user_request,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
         verbose_proxy_logger.info("Created user %s via %s", user_id, created_via)
         return created_user
 
@@ -1716,6 +1719,7 @@ async def create_user(
 
         created_user: Final = await new_user(
             data=new_user_request,
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
         )
 
         scim_user: Final = await ScimTransformations.transform_litellm_user_to_scim_user(user=created_user)
diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
index f4627f82506..88f67cfe0e4 100644
--- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
@@ -19,6 +19,7 @@ from litellm.proxy._types import (
     NewUserResponse,
     ProxyErrorTypes,
     ProxyException,
+    UserAPIKeyAuth,
 )
 from litellm.proxy.management_endpoints.scim.scim_v2 import (
     SCIMRosterSyncError,
@@ -342,6 +343,7 @@ async def test_create_user_without_groups_defers_to_default_team(mocker: MockerF
     await create_user(user=scim_user)
 
     assert new_user_mock.call_args.kwargs["data"].teams is None
+    assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
 
 
 @pytest.mark.asyncio
@@ -380,6 +382,7 @@ async def test_create_user_if_not_exists_defers_to_default_team(mocker: MockerFi
 
     assert created is not None
     assert new_user_mock.call_args.kwargs["data"].teams is None
+    assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
 
 
 @pytest.mark.asyncio

From 07dd8a7e47957a020841439da35da1d287998e08 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 15:46:44 -0700
Subject: [PATCH 20/38] fix(scim): keep team memberships when PUT /Users
 carries no groups

Okta sends profile updates as full PUTs with no groups or groups: [], since SCIM User.groups is readOnly and membership is synced through /Groups. The PUT handler diffed that empty list against the stored teams, removed the user from every team (which also deletes their team keys) and recomputed the role from an empty group list. Treat an empty groups list on PUT as unspecified: keep the stored teams and leave the role alone. Explicit non-empty groups still replace memberships as before

Claude-Session: https://claude.ai/code/session_01CqwUV4Ywnu5aUjXx1UhJrM
---
 .../management_endpoints/scim/scim_v2.py      | 11 ++--
 .../scim/test_scim_v2_endpoints.py            | 61 +++++++++++++++++++
 type-discipline-budget.json                   |  2 +-
 3 files changed, 69 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py
index 98770cf9c2b..ceb67e3eee8 100644
--- a/litellm/proxy/management_endpoints/scim/scim_v2.py
+++ b/litellm/proxy/management_endpoints/scim/scim_v2.py
@@ -1774,22 +1774,25 @@ async def update_user(
             roles=user_data["roles"],
         )
 
+        # SCIM User.groups is readOnly (RFC 7643 4.1.2): IdPs sync membership via /Groups and send
+        # no groups or `groups: []` on profile PUTs, so empty means unspecified, not "remove from every team"
+        target_teams: Final = user_data["teams"] or existing_user.teams
         await _handle_team_membership_changes(
             user_id=user_id,
-            existing_teams=existing_user.teams or [],
-            new_teams=user_data["teams"],
+            existing_teams=existing_user.teams,
+            new_teams=target_teams,
         )
 
         update_data: Final = {
             "user_email": user_data["user_email"],
             "user_alias": user_data["user_alias"],
             "sso_user_id": user_data["sso_user_id"],
-            "teams": user_data["teams"],
+            "teams": target_teams,
             "metadata": safe_dumps(metadata),
         }
 
         admin_group: Final = await _get_scim_admin_group()
-        if admin_group is not None:
+        if admin_group is not None and user_data["teams"]:
             update_data["user_role"] = _resolve_scim_user_role(
                 user.groups or [], admin_group, _default_scim_user_role()
             )
diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
index 88f67cfe0e4..60f9a1a55e2 100644
--- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py
@@ -1257,6 +1257,67 @@ async def test_update_user_success(mocker):
     assert call_args[1]["data"]["teams"] == ["new-team"]
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("groups", [None, []], ids=["groups-omitted", "groups-empty"])
+async def test_update_user_without_groups_preserves_memberships_and_role(mocker, monkeypatch, groups):
+    """Okta profile PUTs carry no `groups` or `groups: []`; neither may drop teams (and their keys) or recompute role"""
+    from litellm.proxy.proxy_server import proxy_config
+
+    async def mock_get_config():
+        return {"litellm_settings": {"scim_admin_group": "litellm-admins"}}
+
+    monkeypatch.setattr(proxy_config, "get_config", mock_get_config)
+    monkeypatch.setattr("litellm.default_internal_user_params", None, raising=False)
+
+    existing_user = mocker.MagicMock()
+    existing_user.teams = ["litellm-admins", "engineering"]
+    existing_user.metadata = {}
+
+    scim_user = SCIMUser(
+        schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
+        userName="okta-user",
+        name=SCIMUserName(familyName="Renamed", givenName="Okta"),
+        emails=[SCIMUserEmail(value="okta@example.com")],
+        **({} if groups is None else {"groups": groups}),
+    )
+    response_scim_user = SCIMUser(
+        schemas=["urn:ietf:params:scim:schemas:core:2.0:User"],
+        id="okta-user",
+        userName="okta-user",
+        emails=[SCIMUserEmail(value="okta@example.com")],
+    )
+
+    mock_prisma_client = mocker.MagicMock()
+    mock_prisma_client.db = mocker.MagicMock()
+    mock_prisma_client.db.litellm_usertable = mocker.MagicMock()
+    mock_prisma_client.db.litellm_usertable.update = AsyncMock(return_value={"user_id": "okta-user"})
+
+    mocker.patch(  # test-quality-ok: update_user's collaborators are module-level, not injectable
+        "litellm.proxy.management_endpoints.scim.scim_v2._get_prisma_client_or_raise_exception",
+        AsyncMock(return_value=mock_prisma_client),
+    )
+    mocker.patch(  # test-quality-ok: update_user's collaborators are module-level, not injectable
+        "litellm.proxy.management_endpoints.scim.scim_v2._check_user_exists",
+        AsyncMock(return_value=existing_user),
+    )
+    mocker.patch(  # test-quality-ok: update_user's collaborators are module-level, not injectable
+        "litellm.proxy.management_endpoints.scim.scim_v2.ScimTransformations.transform_litellm_user_to_scim_user",
+        AsyncMock(return_value=response_scim_user),
+    )
+    patch_membership = mocker.patch(  # test-quality-ok: roster writes are module-level, not injectable
+        "litellm.proxy.management_endpoints.scim.scim_v2.patch_team_membership",
+        AsyncMock(),
+    )
+
+    result = await update_user(user_id="okta-user", user=scim_user)
+
+    assert result == response_scim_user
+    patch_membership.assert_not_awaited()
+    update_data = mock_prisma_client.db.litellm_usertable.update.call_args.kwargs["data"]
+    assert update_data["teams"] == ["litellm-admins", "engineering"]
+    assert "user_role" not in update_data
+
+
 @pytest.mark.asyncio
 async def test_update_user_not_found(mocker):
     """Should raise 404 when user doesn't exist"""
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 2f85128b4b6..78090779109 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -3,7 +3,7 @@
     "limit": 22328
   },
   "LIT002": {
-    "limit": 26760
+    "limit": 26758
   },
   "LIT003": {
     "limit": 261

From 0b7773dd44aaf5c7da2e591e995ff3a41688ba0b 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 16:38:48 -0700
Subject: [PATCH 21/38] fix(router): count tools and Anthropic system prompt in
 context-window pre-call check (#39663)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../messages/utils.py                         |  11 ++
 litellm/router.py                             |  36 +++-
 tests/test_litellm/test_router.py             | 160 +++++++++++++++++-
 3 files changed, 198 insertions(+), 9 deletions(-)

diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py
index 9deff950724..242300c7b6d 100644
--- a/litellm/llms/anthropic/experimental_pass_through/messages/utils.py
+++ b/litellm/llms/anthropic/experimental_pass_through/messages/utils.py
@@ -6,6 +6,7 @@ from litellm.types.llms.anthropic import AnthropicMessagesRequestOptionalParams
 from litellm.types.llms.anthropic_messages.anthropic_response import (
     AnthropicMessagesResponse,
 )
+from litellm.types.llms.openai import ChatCompletionSystemMessage
 
 if TYPE_CHECKING:
     from litellm.exceptions import ContentPolicyViolationError
@@ -36,6 +37,16 @@ def safeguard_refusal_error(model: str, stop_details: Mapping[str, object]) -> "
     )
 
 
+def anthropic_system_to_openai_message(system: object) -> ChatCompletionSystemMessage | None:
+    """
+    Return the Anthropic Messages top-level ``system`` (a string or a list of text
+    blocks) as an OpenAI-style system message, or None when the request has none.
+    """
+    if not isinstance(system, (str, list)) or not system:
+        return None
+    return ChatCompletionSystemMessage(role="system", content=system)
+
+
 @lru_cache(maxsize=1)
 def _anthropic_messages_optional_param_keys() -> frozenset[str]:
     """
diff --git a/litellm/router.py b/litellm/router.py
index f33dfbba7bf..dea9aa62729 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -197,6 +197,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import (
 from litellm.scheduler import FlowItem, Scheduler
 from litellm.types.llms.openai import (
     AllMessageValues,
+    ChatCompletionToolParam,
     FileTypes,
     OpenAIFileObject,
     OpenAIFilesPurpose,
@@ -11762,7 +11763,7 @@ class Router:
         self,
         messages: list[dict[str, str]] | None,
         input: str | list | None,
-        instructions: str | None = None,
+        request_kwargs: Mapping[str, object] | None = None,
     ) -> int:
         """
         Count input tokens for context-window pre-call checks.
@@ -11772,9 +11773,28 @@ class Router:
         The Responses payload is normalized to chat messages via the shared
         LiteLLMCompletionResponsesConfig transform so the same token_counter path covers
         both API surfaces and `instructions` tokens are included in the count.
+
+        Prompt content the message list never carries is read from `request_kwargs`:
+        `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the
+        Anthropic Messages top-level `system` block.
         """
+        from litellm.llms.anthropic.experimental_pass_through.messages.utils import (
+            anthropic_system_to_openai_message,
+        )
+
+        extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({})
+        raw_instructions: Final = extras.get("instructions")
+        instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
+        raw_tools: Final = extras.get("tools")
+        tools: Final = (
+            cast(list[ChatCompletionToolParam], raw_tools)  # cast-ok: token_counter formats any tool dict shape
+            if isinstance(raw_tools, list) and raw_tools
+            else None
+        )
+        system_message: Final = anthropic_system_to_openai_message(extras.get("system"))
         if messages is not None:
-            return litellm.token_counter(messages=messages)
+            counted_messages: Final = (system_message, *messages) if system_message is not None else messages
+            return litellm.token_counter(messages=counted_messages, tools=tools)
         if input is not None:
             from openai.types.responses.response_create_params import ResponseInputParam
 
@@ -11787,7 +11807,10 @@ class Router:
                 input=typed_input,
                 responses_api_request={"instructions": instructions} if instructions is not None else {},
             )
-            return litellm.token_counter(messages=cast(list, input_messages))  # cast-ok: transformed chat messages
+            return litellm.token_counter(
+                messages=cast(list, input_messages),  # cast-ok: transformed chat messages
+                tools=tools,
+            )
         raise ValueError("Either messages or input must be provided to count tokens")
 
     def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None:
@@ -11833,14 +11856,13 @@ class Router:
         """
         if messages is None and input is None:
             return None
-        raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
         try:
             if not self._pre_call_checks_need_token_count(model, healthy_deployments):
                 return None
             return await asyncify(self._count_pre_call_check_tokens)(
                 messages=cast(list[dict[str, str]] | None, messages),  # cast-ok: forwarded to the sync counter
                 input=cast(str | list | None, input),  # cast-ok: forwarded to the sync counter
-                instructions=raw_instructions if isinstance(raw_instructions, str) else None,
+                request_kwargs=request_kwargs,
             )
         except Exception as e:  # noqa: BLE001  # best-effort: an uncountable prompt must not fail the request
             verbose_router_logger.error(
@@ -11887,8 +11909,6 @@ class Router:
         _rate_limit_error = False
         parent_otel_span: Final = _get_parent_otel_span_from_kwargs(request_kwargs)
 
-        raw_instructions: Final = request_kwargs.get("instructions") if request_kwargs else None
-        instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None
         has_countable_input: Final = messages is not None or input is not None
 
         ## get model group RPM ##
@@ -11919,7 +11939,7 @@ class Router:
                             return _returned_deployments
                         try:
                             input_tokens = self._count_pre_call_check_tokens(
-                                messages=messages, input=input, instructions=instructions
+                                messages=messages, input=input, request_kwargs=request_kwargs
                             )
                         except Exception as e:
                             verbose_router_logger.error(
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 228588d974f..f7f0d79b4fd 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -3855,7 +3855,7 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
 
     input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input)
     with_instructions_tokens = router._count_pre_call_check_tokens(
-        messages=None, input=short_input, instructions=long_instructions
+        messages=None, input=short_input, request_kwargs={"instructions": long_instructions}
     )
     assert with_instructions_tokens > input_only_tokens
 
@@ -3871,6 +3871,164 @@ def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch):
         )
 
 
+_OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40
+
+
+@pytest.mark.parametrize(
+    "prompt_kwargs, tool",
+    [
+        pytest.param(
+            {"messages": [{"role": "user", "content": "hi"}]},
+            {
+                "type": "function",
+                "function": {
+                    "name": "lookup",
+                    "description": _OVERSIZED_TOOL_DESCRIPTION,
+                    "parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
+                },
+            },
+            id="chat_completions_tool",
+        ),
+        pytest.param(
+            {"input": "hi"},
+            {
+                "type": "function",
+                "name": "lookup",
+                "description": _OVERSIZED_TOOL_DESCRIPTION,
+                "parameters": {"type": "object", "properties": {"q": {"type": "string"}}},
+            },
+            id="responses_tool",
+        ),
+        pytest.param(
+            {"messages": [{"role": "user", "content": "hi"}]},
+            {
+                "name": "lookup",
+                "description": _OVERSIZED_TOOL_DESCRIPTION,
+                "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}},
+            },
+            id="anthropic_messages_tool",
+        ),
+    ],
+)
+def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool):
+    """
+    Tool definitions are sent to the model as prompt tokens but never appear in
+    `messages` or `input`. A request whose prompt alone fits the context window but
+    whose prompt plus `tools` exceeds it must be rejected before dispatch, for the
+    Chat Completions, Responses and Anthropic Messages tool shapes alike.
+    """
+    router = litellm.Router(
+        model_list=[
+            {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
+        ],
+        enable_pre_call_checks=True,
+    )
+    deployments = [
+        {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
+    ]
+
+    prompt_only_tokens = router._count_pre_call_check_tokens(
+        messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input")
+    )
+    monkeypatch.setattr(
+        router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens}
+    )
+
+    assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1
+    with pytest.raises(litellm.ContextWindowExceededError):
+        router._pre_call_checks(
+            model="m",
+            healthy_deployments=deployments,
+            request_kwargs={"tools": [tool]},
+            **prompt_kwargs,
+        )
+
+
+@pytest.mark.parametrize(
+    "system",
+    [
+        pytest.param("You are a meticulous assistant. " * 40, id="system_string"),
+        pytest.param(
+            [{"type": "text", "text": "You are a meticulous assistant. " * 40}],
+            id="system_blocks",
+        ),
+    ],
+)
+def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system):
+    """
+    The Anthropic Messages API carries the system prompt as a top-level `system` field,
+    not as a message. Its tokens reach the model, so a request whose `messages` fit but
+    whose `messages` plus `system` exceed the context window must be rejected.
+    """
+    router = litellm.Router(
+        model_list=[
+            {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}},
+        ],
+        enable_pre_call_checks=True,
+    )
+    deployments = [
+        {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}},
+    ]
+    messages = [{"role": "user", "content": "hi"}]
+
+    messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None)
+    monkeypatch.setattr(router, "get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens})
+
+    assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1
+    with pytest.raises(litellm.ContextWindowExceededError):
+        router._pre_call_checks(
+            model="m",
+            healthy_deployments=deployments,
+            messages=messages,
+            request_kwargs={"system": system},
+        )
+
+
+@pytest.mark.asyncio
+async def test_aanthropic_messages_enforces_context_window_with_system_and_tools():
+    """
+    End-to-end router regression for /v1/messages: a request whose only oversized
+    content lives in the top-level `system` field or in `tools` must trip the pre-call
+    context-window check instead of being dispatched (the deployment uses mock_response,
+    so reaching the provider handler would return a response rather than raise).
+    """
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "small-ctx",
+                "litellm_params": {"model": "anthropic/claude-3-5-haiku-20241022", "mock_response": "hi"},
+                "model_info": {"max_input_tokens": 20},
+            }
+        ],
+        enable_pre_call_checks=True,
+    )
+    messages = [{"role": "user", "content": "hi"}]
+
+    response = await router.aanthropic_messages(model="small-ctx", messages=messages, max_tokens=5)
+    assert response is not None
+
+    with pytest.raises(litellm.ContextWindowExceededError):
+        await router.aanthropic_messages(
+            model="small-ctx",
+            messages=messages,
+            max_tokens=5,
+            system="You are a meticulous assistant. " * 40,
+        )
+    with pytest.raises(litellm.ContextWindowExceededError):
+        await router.aanthropic_messages(
+            model="small-ctx",
+            messages=messages,
+            max_tokens=5,
+            tools=[
+                {
+                    "name": "lookup",
+                    "description": _OVERSIZED_TOOL_DESCRIPTION,
+                    "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}},
+                }
+            ],
+        )
+
+
 def test_count_pre_call_check_tokens_across_api_surfaces():
     """
     _count_pre_call_check_tokens must count tokens from chat `messages`, a Responses

From a330bc98a68725d7e5afa078ac5eb05d0cfcf8ff Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 16:39:03 -0700
Subject: [PATCH 22/38] test(vertex-passthrough): inject the forwarder into the
 short-route regression helper

---
 .../test_llm_pass_through_endpoints.py        | 74 ++++++++-----------
 1 file changed, 30 insertions(+), 44 deletions(-)

diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
index e37060493f7..5154f738e9a 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
@@ -4222,26 +4222,21 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
 
 
 class TestVertexPassthroughDefaultLocationOnShortRoutes:
-    """Regression coverage for LIT-6905.
-
-    ``default_vertex_config`` carries the project and location, yet a route that
-    omits ``/projects/<project>/locations/<location>/`` built the upstream base URL
-    from the still-unresolved URL location and 500ed with ``vertex_location is
-    required``. The base URL must be built after the configured location is
-    applied, and a request with no location anywhere must fail with a clean 400
-    that says where a location can come from, never a 500.
-    """
-
     PROJECT = "test-project"
     SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent"
 
+    @staticmethod
+    def _forwarder() -> Mock:
+        return Mock(return_value=AsyncMock(return_value={"status": "success"}))
+
     async def _forward(
         self,
         monkeypatch,
         endpoint: str,
         default_config: dict | None,
         headers: list[tuple[bytes, bytes]],
-    ) -> tuple[HTTPException | None, dict]:
+        forwarder: Mock,
+    ) -> None:
         from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
             PassthroughEndpointRouter,
         )
@@ -4249,7 +4244,7 @@ class TestVertexPassthroughDefaultLocationOnShortRoutes:
         async def receive():
             return {"type": "http.request", "body": b"{}", "more_body": False}
 
-        request = Request(
+        request: Final = Request(
             {
                 "type": "http",
                 "method": "POST",
@@ -4259,41 +4254,29 @@ class TestVertexPassthroughDefaultLocationOnShortRoutes:
             },
             receive=receive,
         )
-
-        captured: dict = {}
-
-        def fake_create_pass_through_route(**kwargs):
-            captured.update(kwargs)
-            return AsyncMock(return_value={"status": "success"})
-
-        router = PassthroughEndpointRouter()
+        router: Final = PassthroughEndpointRouter()
         if default_config is not None:
             router.set_default_vertex_config(dict(default_config))
-        module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
+        module: Final = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
         monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router)
         monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
-        mock_credentials = Mock()
+        mock_credentials: Final = Mock()
         mock_credentials.token = "test-token"
         caller: Final = UserAPIKeyAuth(api_key="test-key")
-        raised: HTTPException | None = None
         with (
             mock.patch(  # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
                 "litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
                 return_value=(mock_credentials, self.PROJECT),
             ),
-            mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
+            mock.patch(f"{module}.create_pass_through_route", new=forwarder),
             mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
         ):
-            try:
-                await vertex_proxy_route(
-                    endpoint=endpoint,
-                    request=request,
-                    fastapi_response=Response(),
-                    user_api_key_dict=caller,
-                )
-            except HTTPException as exc:
-                raised = exc
-        return raised, captured
+            await vertex_proxy_route(
+                endpoint=endpoint,
+                request=request,
+                fastapi_response=Response(),
+                user_api_key_dict=caller,
+            )
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -4325,15 +4308,17 @@ class TestVertexPassthroughDefaultLocationOnShortRoutes:
     async def test_default_vertex_config_location_fills_routes_without_project_and_location(
         self, monkeypatch, endpoint, location, expected_target
     ):
-        raised, captured = await self._forward(
+        forwarder: Final = self._forwarder()
+        await self._forward(
             monkeypatch,
             endpoint,
             {"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"},
             [(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
+            forwarder,
         )
-        assert raised is None
-        assert str(captured["target"]) == expected_target
-        assert captured["custom_headers"]["Authorization"] == "Bearer test-token"
+        forwarded: Final = forwarder.call_args.kwargs
+        assert str(forwarded["target"]) == expected_target
+        assert forwarded["custom_headers"]["Authorization"] == "Bearer test-token"
 
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
@@ -4347,12 +4332,13 @@ class TestVertexPassthroughDefaultLocationOnShortRoutes:
         ],
     )
     async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers):
-        raised, captured = await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers)
-        assert not captured, "a request with no location must never reach the upstream forwarder"
-        assert raised is not None
-        assert raised.status_code == 400
-        assert "/projects/<project>/locations/<location>/" in str(raised.detail)
-        assert "default_vertex_config" in str(raised.detail)
+        forwarder: Final = self._forwarder()
+        with pytest.raises(HTTPException) as raised:
+            await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers, forwarder)
+        forwarder.assert_not_called()
+        assert raised.value.status_code == 400
+        assert "/projects/<project>/locations/<location>/" in str(raised.value.detail)
+        assert "default_vertex_config" in str(raised.value.detail)
 
 
 class TestGetAzureAISearchIndexFromEndpoint:

From 39a17898ffdb55e4b54c0a7fe750b4d5afe49ea6 Mon Sep 17 00:00:00 2001
From: mubashir1osmani <mubashir.osmani777@gmail.com>
Date: Thu, 3 Sep 2026 19:45:09 -0400
Subject: [PATCH 23/38] test(proxy-extras): repoint the migrate-deploy harness
 at the run_prisma seam (#39673)


From cf3af0f486590633ee875f6d120c416509fad5d2 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 23:53:59 +0000
Subject: [PATCH 24/38] perf(spend): group /spend/logs summary by day in
 Postgres instead of per-row Prisma group_by (#39351)

* perf(spend): group /spend/logs summary by day in Postgres instead of per-row Prisma group_by

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

* refactor(spend): simplify /spend/logs daily summary aggregation

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

* fix(spend): preserve spend logs response schema

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

* chore: ratchet lint budgets after merge

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

* chore: ratchet lint budgets after merge

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

* fix(spend): compare spend log range bounds as naive UTC timestamps

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

* chore: ratchet lint budgets after merge

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

* chore: ratchet lint budgets after merge

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

* chore: ratchet lint budgets after merge

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

* chore: ratchet lint budgets after merge

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

* test: cover spend logs summary edge cases

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

* chore: ratchet lint budgets after merge

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

* test: cover spend summary request filters

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

* chore: ratchet lint budgets after merge

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

* chore: ratchet lint budgets after merge

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>
---
 basedpyright-code-budget.json                 |  10 +-
 .../spend_management_endpoints.py             | 164 +++++++-----
 ruff-strict-budget.json                       |   2 +-
 .../test_spend_management_endpoints.py        | 235 +++++++++++++++---
 type-discipline-budget.json                   |   6 +-
 5 files changed, 303 insertions(+), 114 deletions(-)

diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json
index 9a9b1138a7d..cb3756575f7 100644
--- a/basedpyright-code-budget.json
+++ b/basedpyright-code-budget.json
@@ -3,7 +3,7 @@
     "limit": 14074
   },
   "reportArgumentType": {
-    "limit": 2214
+    "limit": 2206
   },
   "reportAssignmentType": {
     "limit": 319
@@ -57,7 +57,7 @@
     "limit": 5601
   },
   "reportMissingTypeArgument": {
-    "limit": 15287
+    "limit": 15285
   },
   "reportMissingTypeStubs": {
     "limit": 40
@@ -99,19 +99,19 @@
     "limit": 0
   },
   "reportUnknownArgumentType": {
-    "limit": 44362
+    "limit": 44360
   },
   "reportUnknownLambdaType": {
     "limit": 109
   },
   "reportUnknownMemberType": {
-    "limit": 38323
+    "limit": 38311
   },
   "reportUnknownParameterType": {
     "limit": 19624
   },
   "reportUnknownVariableType": {
-    "limit": 29861
+    "limit": 29847
   },
   "reportUnnecessaryCast": {
     "limit": 111
diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py
index b86a877e8f9..dff100bdea7 100644
--- a/litellm/proxy/spend_tracking/spend_management_endpoints.py
+++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py
@@ -3,7 +3,8 @@ import collections
 import json
 import os
 from collections.abc import Mapping, Sequence
-from datetime import datetime, timedelta, timezone
+from datetime import date, datetime, timedelta, timezone
+from itertools import groupby
 from types import MappingProxyType
 from typing import (
     TYPE_CHECKING,
@@ -16,7 +17,6 @@ from typing import (
     TypeAlias,
     TypedDict,
     TypeVar,
-    cast,  # noqa: TID251  # prisma group_by returns untyped aggregate mappings
 )
 
 import fastapi
@@ -201,16 +201,12 @@ class _SessionSpendStats(NamedTuple):
 _SessionSpendMap: TypeAlias = Mapping[tuple[str, str], _SessionSpendStats]
 
 
-class _SpendSumAggregate(TypedDict, total=False):
-    spend: ReadOnly[float]
-
-
-class _SpendGroupByRow(TypedDict):
+class _SpendDailySummaryRow(TypedDict):
+    day: ReadOnly[str]
     api_key: ReadOnly[str]
     user: ReadOnly[str | None]
     model: ReadOnly[str]
-    startTime: ReadOnly[object]
-    _sum: ReadOnly[_SpendSumAggregate]
+    spend: ReadOnly[float]
 
 
 async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
@@ -251,6 +247,66 @@ def _verification_token_table(prisma_client: PrismaClient) -> _VerificationToken
     return VerificationTokenRepository(prisma_client).table
 
 
+def _spend_logs_daily_summary_sql(
+    *,
+    start_date_iso: str,
+    end_date_iso: str,
+    api_key: str | None,
+    request_id: str | None,
+    user_id: str | None,
+) -> tuple[str, tuple[object, ...]]:
+    filter_params: Final[tuple[tuple[str, object], ...]] = tuple(
+        (column, value)
+        for column, value in (
+            ("api_key", api_key),
+            ("request_id", request_id),
+            ('"user"', user_id),
+        )
+        if value is not None
+    )
+    filter_clauses: Final[tuple[str, ...]] = tuple(
+        f"AND {column} = ${index}" for index, (column, _) in enumerate(filter_params, start=3)
+    )
+    filter_sql: Final = "\n".join(filter_clauses)
+    sql_query: Final = f"""
+SELECT
+    to_char(date_trunc('day', "startTime"), 'YYYY-MM-DD') AS day,
+    api_key,
+    "user",
+    model,
+    SUM(spend) AS spend
+FROM "LiteLLM_SpendLogs"
+WHERE "startTime" >= ($1::timestamptz AT TIME ZONE 'UTC') AND "startTime" <= ($2::timestamptz AT TIME ZONE 'UTC')
+{filter_sql}
+GROUP BY 1, 2, 3, 4
+ORDER BY 1
+"""
+    params: Final[tuple[object, ...]] = (
+        start_date_iso,
+        end_date_iso,
+        *(value for _, value in filter_params),
+    )
+    return sql_query, params
+
+
+def _sum_spend_by(
+    rows: Sequence[_SpendDailySummaryRow], column: Literal["api_key", "user", "model"]
+) -> Mapping[str | None, float]:
+    keys: Final = frozenset(row[column] for row in rows)
+    return {key: sum(float(row["spend"]) for row in rows if row[column] == key) for key in keys}
+
+
+def _daily_summary_item(summary_date: date, rows: Sequence[_SpendDailySummaryRow]) -> Mapping[str, object]:
+    api_key_spend: Final = {key: value for key, value in _sum_spend_by(rows, "api_key").items() if key is not None}
+    return {
+        **api_key_spend,
+        "startTime": summary_date,
+        "spend": sum(float(row["spend"]) for row in rows),
+        "users": _sum_spend_by(rows, "user"),
+        "models": _sum_spend_by(rows, "model"),
+    }
+
+
 async def _find_spend_logs(
     prisma_client: PrismaClient,
     where: Mapping[str, object],
@@ -3266,18 +3322,22 @@ async def view_spend_logs(
             start_date_iso: Final = start_date_obj.isoformat()
             end_date_iso: Final = end_date_obj.isoformat()
 
-            filter_query: Final = {
+            filter_query: Final[
+                dict[str, object]
+            ] = {  # mutable-ok: legacy filters are extended for optional parameters
                 "startTime": {
                     "gte": start_date_iso,  # Greater than or equal to Start Date
                     "lte": end_date_iso,  # Less than or equal to End Date
                 }
             }
 
+            summary_api_key: Final[str | None] = (
+                prisma_client.hash_token(token=api_key)
+                if api_key is not None and api_key.startswith("sk-")
+                else api_key
+            )
             if api_key is not None and isinstance(api_key, str):
-                if api_key.startswith("sk-"):
-                    filter_query["api_key"] = prisma_client.hash_token(token=api_key)
-                else:
-                    filter_query["api_key"] = api_key
+                filter_query["api_key"] = summary_api_key
             if request_id is not None and isinstance(request_id, str):
                 filter_query["request_id"] = request_id
             if user_id is not None and isinstance(user_id, str):
@@ -3296,58 +3356,34 @@ async def view_spend_logs(
                 return data
 
             # Legacy behavior: return summarized data (when summarize=true)
-            # SQL query
-            response: Final = await SpendLogsRepository(prisma_client).table.group_by(
-                by=["api_key", "user", "model", "startTime"],
-                where=filter_query,
-                sum={
-                    "spend": True,
-                },
+            summary_sql_and_params: Final = _spend_logs_daily_summary_sql(
+                start_date_iso=start_date_iso,
+                end_date_iso=end_date_iso,
+                api_key=summary_api_key,
+                request_id=request_id,
+                user_id=user_id,
             )
+            sql_query, params = summary_sql_and_params
+            rows: Final[Sequence[_SpendDailySummaryRow]] = await _query_raw(prisma_client, sql_query, *params)
+            if len(rows) == 0:
+                return []  # pyright: ignore[reportUnknownVariableType]  # empty summary has no element type
 
-            if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
-                spend_rows: Final = cast(Sequence[_SpendGroupByRow], response)  # cast-ok: by/sum fix the shape
-                result: Final[dict] = {}
-                for record in spend_rows:
-                    dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
-                    date = dt_object.date()
-                    if date not in result:
-                        result[date] = {"users": {}, "models": {}}
-                    api_key = record["api_key"]
-                    user_id = record["user"]
-                    model = record["model"]
-                    result[date]["spend"] = result[date].get("spend", 0) + record.get("_sum", {}).get("spend", 0)
-                    result[date][api_key] = result[date].get(api_key, 0) + record.get("_sum", {}).get("spend", 0)
-                    result[date]["users"][user_id] = result[date]["users"].get(user_id, 0) + record.get("_sum", {}).get(
-                        "spend", 0
-                    )
-                    result[date]["models"][model] = result[date]["models"].get(model, 0) + record.get("_sum", {}).get(
-                        "spend", 0
-                    )
-                return_list: Final = []
-                final_date = None
-                for k, v in sorted(result.items()):
-                    return_list.append({**v, "startTime": k})
-                    final_date = k
-
-                end_date_date: Final = end_date_obj.date()
-                if final_date is not None and final_date < end_date_date:
-                    current_date = final_date + timedelta(days=1)
-                    while current_date <= end_date_date:
-                        # Represent current_date as string because original response has it this way
-                        return_list.append(
-                            {
-                                "startTime": current_date,
-                                "spend": 0,
-                                "users": {},
-                                "models": {},
-                            }
-                        )  # If no data, will stay as zero
-                        current_date += timedelta(days=1)  # Move on to the next day
-
-                return return_list
-
-            return response
+            summary_items: Final = tuple(
+                _daily_summary_item(date.fromisoformat(day), tuple(day_rows))
+                for day, day_rows in groupby(rows, key=lambda row: row["day"])
+            )
+            final_date: Final = date.fromisoformat(rows[-1]["day"])
+            end_date_date: Final = end_date_obj.date()
+            padding: Final[tuple[Mapping[str, object], ...]] = tuple(
+                {
+                    "startTime": final_date + timedelta(days=offset),
+                    "spend": 0,
+                    "users": {},
+                    "models": {},
+                }
+                for offset in range(1, (end_date_date - final_date).days + 1)
+            )
+            return [*summary_items, *padding]
 
         else:
             scoped_filter: Final[dict[str, str]] = {}
diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json
index 24c0ff6b181..8763318b4eb 100644
--- a/ruff-strict-budget.json
+++ b/ruff-strict-budget.json
@@ -96,7 +96,7 @@
     "limit": 10
   },
   "DTZ007": {
-    "limit": 17
+    "limit": 6
   },
   "DTZ011": {
     "limit": 3
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 73a29afd9b9..329a33eb440 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
@@ -5,14 +5,12 @@ import hashlib
 import json
 import re
 from datetime import timezone
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastapi import HTTPException
 from fastapi.testclient import TestClient
 
-
-from unittest.mock import AsyncMock, MagicMock, patch
-
 import litellm
 import litellm.proxy.proxy_server as ps
 
@@ -3325,7 +3323,7 @@ def _compare_nested_dicts(
         return differences
 
     # Check for keys in actual but not in expected
-    for key in actual.keys():
+    for key in actual:
         current_path = f"{path}.{key}" if path else key
         if current_path not in ignore_keys and key not in expected:
             differences.append(f"Extra key in actual: {current_path}")
@@ -3495,24 +3493,22 @@ async def test_view_spend_logs_summarize_parameter(client, monkeypatch):
             # Return individual log entries when summarize=false
             return mock_spend_logs
 
-        async def group_by(self, *args, **kwargs):
-            # Return grouped data when summarize=true
-            # Simplified mock response for grouped data
+        async def query_raw(self, sql_query, *params):
             yesterday = datetime.datetime.now(timezone.utc) - timedelta(days=1)
             return [
                 {
                     "api_key": "sk-test-key",
                     "user": "test_user_1",
                     "model": "gpt-3.5-turbo",
-                    "startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
-                    "_sum": {"spend": 0.05},
+                    "day": yesterday.date().isoformat(),
+                    "spend": 0.05,
                 },
                 {
                     "api_key": "sk-test-key",
                     "user": "test_user_1",
                     "model": "gpt-4",
-                    "startTime": yesterday.strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
-                    "_sum": {"spend": 0.10},
+                    "day": yesterday.date().isoformat(),
+                    "spend": 0.10,
                 },
             ]
 
@@ -3850,47 +3846,30 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
     """
     from datetime import datetime, timedelta, timezone
 
-    # This simulates the summarized data that Prisma's `group_by` would return.
     mock_summarized_response = [
         {
             "api_key": "sk-test-key",
             "user": "test_user_1",
             "model": "gpt-4",
-            "startTime": (datetime.now(timezone.utc) - timedelta(days=1)).strftime(
-                "%Y-%m-%dT%H:%M:%S.%fZ"
-            ),
-            "_sum": {"spend": 0.15},
+            "day": (datetime.now(timezone.utc) - timedelta(days=1)).date().isoformat(),
+            "spend": 0.15,
         }
     ]
 
-    # This mock class will replace the real Prisma client.
     class MockDB:
-        def __init__(self):
-            self.litellm_spendlogs = self
-
-        async def group_by(self, *args, **kwargs):
-            # We assert that the `gte` and `lte` values are strings in ISO format.
-            # If they were datetime objects, this test would fail.
-            where_clause = kwargs.get("where", {})
-            start_time_filter = where_clause.get("startTime", {})
-
-            assert "gte" in start_time_filter
-            assert "lte" in start_time_filter
-            assert isinstance(start_time_filter["gte"], str)
-            assert isinstance(start_time_filter["lte"], str)
-            assert "T" in start_time_filter["gte"]  # Check for ISO format 'T' separator
-
-            # If the assertions pass, return the mock response.
+        async def query_raw(self, sql_query, *params):
+            assert isinstance(params[0], str)
+            assert isinstance(params[1], str)
+            assert "T" in params[0]
+            assert "T" in params[1]
             return mock_summarized_response
 
     class MockPrismaClient:
         def __init__(self):
             self.db = MockDB()
 
-    # Apply the monkeypatch to replace the real prisma_client with our mock.
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
 
-    # Define a date range for the test.
     start_date = (datetime.now(timezone.utc) - timedelta(days=2)).strftime("%Y-%m-%d")
     end_date = datetime.now(timezone.utc).strftime("%Y-%m-%d")
 
@@ -3898,8 +3877,6 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
         user_role=LitellmUserRoles.PROXY_ADMIN
     )
     try:
-        # Call the endpoint with both start and end dates.
-        # We don't need `summarize=true` as it's the default.
         response = client.get(
             "/spend/logs",
             params={
@@ -3909,11 +3886,9 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
             headers={"Authorization": "Bearer sk-test"},
         )
 
-        # ASSERTIONS
         assert response.status_code == 200
         data = response.json()
 
-        # Check that the response is not empty and has the summarized structure.
         assert isinstance(data, list)
         assert len(data) > 0
         assert "startTime" in data[0]
@@ -3924,6 +3899,183 @@ async def test_view_spend_logs_with_date_range_summarized(client, monkeypatch):
         app.dependency_overrides.pop(ps.user_api_key_auth, None)
 
 
+@pytest.mark.asyncio
+async def test_view_spend_logs_summarize_groups_by_day_in_sql(client, monkeypatch):
+    mock_rows = [
+        {
+            "day": "2024-01-01",
+            "api_key": "hashed::sk-abc",
+            "user": "u1",
+            "model": "gpt-4",
+            "spend": 0.1,
+        },
+        {
+            "day": "2024-01-01",
+            "api_key": "hashed::sk-abc",
+            "user": "u1",
+            "model": "gpt-4o",
+            "spend": 0.2,
+        },
+    ]
+
+    class MockDB:
+        def __init__(self):
+            self.captured_sql = None
+            self.captured_params = None
+
+        async def query_raw(self, sql_query, *params):
+            self.captured_sql = sql_query
+            self.captured_params = params
+            return mock_rows
+
+    class MockPrismaClient:
+        def __init__(self):
+            self.db = MockDB()
+
+        def hash_token(self, token):
+            return "hashed::" + token
+
+    mock_prisma_client = MockPrismaClient()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.get(
+            "/spend/logs",
+            params={
+                "start_date": "2024-01-01",
+                "end_date": "2024-01-03",
+                "api_key": "sk-abc",
+                "request_id": "req-123",
+                "user_id": "u1",
+            },
+            headers={"Authorization": "Bearer sk-test"},
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        sql = mock_prisma_client.db.captured_sql
+        assert "date_trunc('day'" in sql
+        assert "GROUP BY" in sql
+        assert "find_many" not in sql
+        assert not hasattr(mock_prisma_client.db, "group_by")
+        assert mock_prisma_client.db.captured_params == (
+            "2024-01-01T00:00:00+00:00",
+            "2024-01-03T00:00:00+00:00",
+            "hashed::sk-abc",
+            "req-123",
+            "u1",
+        )
+        assert len(data) == 3
+        assert data[0]["startTime"] == "2024-01-01"
+        assert data[0]["spend"] == pytest.approx(0.3)
+        assert data[0]["models"] == {"gpt-4": 0.1, "gpt-4o": 0.2}
+        assert data[0]["users"] == {"u1": pytest.approx(0.3)}
+        assert data[0]["hashed::sk-abc"] == pytest.approx(0.3)
+        assert data[1] == {
+            "startTime": "2024-01-02",
+            "spend": 0,
+            "users": {},
+            "models": {},
+        }
+        assert data[2] == {
+            "startTime": "2024-01-03",
+            "spend": 0,
+            "users": {},
+            "models": {},
+        }
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+async def test_view_spend_logs_summarize_empty_rows(client, monkeypatch):
+    class MockDB:
+        async def query_raw(self, sql_query, *params):
+            return []
+
+    class MockPrismaClient:
+        def __init__(self):
+            self.db = MockDB()
+
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MockPrismaClient())
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.get(
+            "/spend/logs",
+            params={"start_date": "2024-01-01", "end_date": "2024-01-01"},
+            headers={"Authorization": "Bearer sk-test"},
+        )
+
+        assert response.status_code == 200
+        assert response.json() == []
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
+@pytest.mark.asyncio
+async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client, monkeypatch):
+    mock_rows = [
+        {
+            "day": "2024-01-01",
+            "api_key": "plain-key",
+            "user": "u1",
+            "model": "gpt-4",
+            "spend": 0.4,
+        }
+    ]
+
+    class MockDB:
+        def __init__(self):
+            self.captured_params = None
+
+        async def query_raw(self, sql_query, *params):
+            self.captured_params = params
+            return mock_rows
+
+    class MockPrismaClient:
+        def __init__(self):
+            self.db = MockDB()
+
+    mock_prisma_client = MockPrismaClient()
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
+    app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(
+        user_role=LitellmUserRoles.PROXY_ADMIN
+    )
+    try:
+        response = client.get(
+            "/spend/logs",
+            params={
+                "start_date": "2024-01-01",
+                "end_date": "2024-01-01",
+                "api_key": "plain-key",
+            },
+            headers={"Authorization": "Bearer sk-test"},
+        )
+
+        assert response.status_code == 200
+        data = response.json()
+        assert mock_prisma_client.db.captured_params == (
+            "2024-01-01T00:00:00+00:00",
+            "2024-01-01T00:00:00+00:00",
+            "plain-key",
+        )
+        assert data == [
+            {
+                "startTime": "2024-01-01",
+                "spend": pytest.approx(0.4),
+                "plain-key": pytest.approx(0.4),
+                "users": {"u1": pytest.approx(0.4)},
+                "models": {"gpt-4": pytest.approx(0.4)},
+            }
+        ]
+    finally:
+        app.dependency_overrides.pop(ps.user_api_key_auth, None)
+
+
 @pytest.mark.asyncio
 async def test_ui_view_spend_logs_with_error_code(client):
     """Test filtering spend logs by error code"""
@@ -4832,13 +4984,14 @@ class _CaptureFilterDB:
     def __init__(self):
         self.litellm_spendlogs = self
         self.captured_where = None
+        self.captured_params = None
 
     async def find_many(self, *args, **kwargs):
         self.captured_where = kwargs.get("where")
         return []
 
-    async def group_by(self, *args, **kwargs):
-        self.captured_where = kwargs.get("where")
+    async def query_raw(self, sql_query, *params):
+        self.captured_params = params
         return []
 
 
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 78090779109..ab1a793e09d 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -3,7 +3,7 @@
     "limit": 22328
   },
   "LIT002": {
-    "limit": 26758
+    "limit": 26750
   },
   "LIT003": {
     "limit": 261
@@ -27,10 +27,10 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16470
+    "limit": 16468
   },
   "LIT011": {
-    "limit": 5516
+    "limit": 5514
   },
   "LIT012": {
     "limit": 4489

From c1a607f90f95f235af8df663199bc94912bae287 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 3 Sep 2026 16:57:10 -0700
Subject: [PATCH 25/38] test(e2e): match the Internal Users search placeholder
 shipped by #39604
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

PR #39604 renamed the Internal Users search box placeholder to "Search by email or ID…" but left searchUsers.spec.ts looking for the old "Search by email…" copy, so e2e_ui_testing has been red on litellm_internal_staging since it merged. Point the locator at the shipped placeholder
---
 tests/e2e/ui/tests/users/searchUsers.spec.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/e2e/ui/tests/users/searchUsers.spec.ts b/tests/e2e/ui/tests/users/searchUsers.spec.ts
index fa8f32764e8..64fce24bf2a 100644
--- a/tests/e2e/ui/tests/users/searchUsers.spec.ts
+++ b/tests/e2e/ui/tests/users/searchUsers.spec.ts
@@ -17,7 +17,7 @@ test.describe("Internal Users Search", () => {
   test("narrows the table to the matching email, and restores it when cleared", async ({ page }) => {
     await goToInternalUsers(page);
 
-    const search = page.getByPlaceholder("Search by email…");
+    const search = page.getByPlaceholder("Search by email or ID…");
     await expect(search).toBeVisible();
 
     await search.fill("noteam@");

From dc9f40c11fce86cb6c473d8235366a049894d583 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 17:13:18 -0700
Subject: [PATCH 26/38] fix(ui): scroll admin table rows inside the table
 instead of the page

Virtual Keys, Teams, Request Logs and Tags now hand DataTable a bounded
flex chain and use fillHeight, so the app shell main stays the only page
scroller, the rows scroll under a pinned header and the pagination footer
sits at the bottom of the page. DataTable keeps the sticky header inside
its own scroller in maxBodyHeight mode too, which is what let the header
scroll away with the rows on Keys, Teams and Models. Model Hub, Vector
Stores and the team detail keys tab drop their 75vh boxes and flow with
the page scroller.

Adds an e2e spec that fails on the merge base for every one of those
pages and passes at this tip.

Refs LIT-4738

Claude-Session: https://claude.ai/code/session_018yW93iDaEMhoQUXcYjus7D
---
 .../ui/tests/tables/tableScrolling.spec.ts    | 239 ++++++++++++++++++
 .../tag-management/_components/TagTable.tsx   |   1 +
 .../tag-management/_components/index.tsx      |  30 +--
 .../vector-stores/_components/index.tsx       |   4 +-
 .../src/components/AIHub/ModelHubTable.tsx    |   2 +-
 ui/litellm-dashboard/src/components/Teams.tsx |   9 +-
 .../src/components/TeamsPage/TeamsTable.tsx   |   2 +-
 .../VirtualKeysPage/VirtualKeysTable.tsx      |   4 +-
 .../shared/DataTable/DataTable.test.tsx       |   8 +-
 .../components/shared/DataTable/DataTable.tsx |  25 +-
 .../components/team/TeamVirtualKeysTable.tsx  |   5 +-
 .../src/components/user_dashboard.tsx         |  34 ++-
 .../components/view_logs/RequestLogsTable.tsx |   1 +
 .../src/components/view_logs/index.tsx        |   9 +-
 14 files changed, 315 insertions(+), 58 deletions(-)
 create mode 100644 tests/e2e/ui/tests/tables/tableScrolling.spec.ts

diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts
new file mode 100644
index 00000000000..f51c772185c
--- /dev/null
+++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts
@@ -0,0 +1,239 @@
+import { test, expect, type APIRequestContext, type Locator, type Page as PlaywrightPage } from "@playwright/test";
+import { ADMIN_STORAGE_PATH } from "../../constants";
+import { Page } from "../../fixtures/pages";
+import { navigateToPage } from "../../helpers/navigation";
+import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
+
+/**
+ * LIT-4738 table scrolling. On the paginated pages the app shell <main> is the only page scroller
+ * and must never overflow: rows scroll inside the table body under a header that stays put, and the
+ * pagination footer sits at the bottom of the page instead of below the fold or inside a clipped
+ * box. Pages that keep plain page scrolling must never paint rows past a fixed-height ancestor.
+ * The viewport is pinned so "more rows than fit" means the same thing on every machine.
+ */
+
+const VIEWPORT = { width: 1280, height: 720 };
+const SEED_ROWS = 40;
+const LOG_ROWS = 20;
+const BODY_SCROLL_PX = 500;
+/** p-8 on Keys and Teams, p-6 on Logs: the footer may sit at most one page padding above the edge. */
+const MAX_FOOTER_GAP_PX = 40;
+
+const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+
+const adminHeaders = (): Record<string, string> => ({
+  Authorization: `Bearer ${masterKey()}`,
+});
+
+/** Keys and Teams render a <main> of their own inside the app shell's, which comes first in document order. */
+const pageScroller = (page: PlaywrightPage): Locator => page.locator("main").first();
+
+/** Tabs keep every panel mounted, so a bare test id can match a hidden table; scope to the visible one. */
+const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
+
+/** The page's data table; Model Hub also renders a plain links table above it, which this skips. */
+const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
+
+const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
+
+interface BoxMetrics {
+  top: number;
+  bottom: number;
+  scrollHeight: number;
+  clientHeight: number;
+  scrollWidth: number;
+  clientWidth: number;
+}
+
+const metrics = (locator: Locator): Promise<BoxMetrics> =>
+  locator.evaluate((el) => {
+    const rect = el.getBoundingClientRect();
+    return {
+      top: rect.top,
+      bottom: rect.bottom,
+      scrollHeight: el.scrollHeight,
+      clientHeight: el.clientHeight,
+      scrollWidth: el.scrollWidth,
+      clientWidth: el.clientWidth,
+    };
+  });
+
+async function postOk(
+  request: APIRequestContext,
+  path: string,
+  data: Record<string, unknown>,
+): Promise<Record<string, any>> {
+  const res = await request.post(path, { headers: adminHeaders(), data });
+  expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true);
+  return (await res.json()) as Record<string, any>;
+}
+
+/** One request at a time: a burst of forty management calls starves the proxy's transaction pool. */
+async function oneAtATime<T>(count: number, call: (index: number) => Promise<T>): Promise<T[]> {
+  const results: T[] = [];
+  for (let i = 0; i < count; i++) {
+    results.push(await call(i));
+  }
+  return results;
+}
+
+async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise<void> {
+  await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count);
+}
+
+async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"): Promise<void> {
+  await visibleTestId(page, "pagination-page-size").click();
+  await page.getByRole("option", { name: size, exact: true }).click();
+}
+
+/**
+ * The page scroller stays put, the table body is what scrolls, the header does not move while the
+ * body scrolls, and the pagination footer sits at the bottom of the page.
+ */
+async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
+  const scroller = await metrics(pageScroller(page));
+  const body = visibleTestId(page, "data-table-scroller");
+  const bodyBefore = await metrics(body);
+  const headBefore = await metrics(visibleTestId(page, "data-table-head"));
+  const footer = await metrics(visibleDataTable(page));
+
+  expect(scroller.scrollHeight, "page scroller must not overflow vertically").toBe(scroller.clientHeight);
+  expect(scroller.scrollWidth, "page scroller must not overflow horizontally").toBe(scroller.clientWidth);
+  expect(bodyBefore.scrollHeight, "table body must be the element that scrolls").toBeGreaterThan(
+    bodyBefore.clientHeight,
+  );
+  expect(footer.bottom, "pagination footer must be inside the page").toBeLessThanOrEqual(scroller.bottom);
+  expect(scroller.bottom - footer.bottom, "pagination footer must sit at the bottom of the page").toBeLessThanOrEqual(
+    MAX_FOOTER_GAP_PX,
+  );
+
+  await body.evaluate((el, px) => {
+    el.scrollTop = px;
+  }, BODY_SCROLL_PX);
+  await expect.poll(() => body.evaluate((el) => el.scrollTop)).toBeGreaterThan(0);
+  const headAfter = await metrics(visibleTestId(page, "data-table-head"));
+  expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top));
+}
+
+/**
+ * Every row must sit inside each ancestor up to the nearest one that really scrolls vertically; a
+ * fixed-height box that neither grows nor scrolls lets rows paint past its bottom edge.
+ */
+const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise<string[]> =>
+  visibleDataTable(page)
+    .locator("table")
+    .evaluate((table) => {
+      const scrollsVertically = (el: Element): boolean =>
+        /auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1;
+      const describe = (el: Element): string =>
+        `<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`;
+      return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => {
+        const rowBottom = row.getBoundingClientRect().bottom;
+        const spills: string[] = [];
+        for (let el = row.parentElement; el && el !== document.body && !scrollsVertically(el); el = el.parentElement) {
+          const bottom = el.getBoundingClientRect().bottom;
+          if (rowBottom > bottom + 1) {
+            spills.push(
+              `row ${index} bottom ${Math.round(rowBottom)} past ${describe(el)} bottom ${Math.round(bottom)}`,
+            );
+          }
+        }
+        return spills;
+      });
+    });
+
+type Cleanup = (request: APIRequestContext) => Promise<unknown>;
+const cleanups: Cleanup[] = [];
+
+test.describe("Admin tables scroll inside the page", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
+
+  test.afterEach(async ({ request }) => {
+    for (const cleanup of cleanups.splice(0)) {
+      // Teardown must never turn a passing test red or mask a real failure.
+      await cleanup(request).catch(() => {});
+    }
+  });
+
+  test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
+    page,
+    request,
+  }) => {
+    const suffix = uniqueSuffix();
+    const created = await oneAtATime(SEED_ROWS, (i) =>
+      postOk(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
+    );
+    cleanups.push((r) => r.post("/key/delete", { headers: adminHeaders(), data: { keys: created.map((k) => k.key) } }));
+
+    await navigateToPage(page, Page.ApiKeys);
+    await expectRowsAtLeast(page, SEED_ROWS);
+    await expectBodyIsTheOnlyScroller(page);
+  });
+
+  test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const created = await oneAtATime(SEED_ROWS, (i) =>
+      postOk(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
+    );
+    cleanups.push((r) =>
+      r.post("/team/delete", { headers: adminHeaders(), data: { team_ids: created.map((t) => t.team_id) } }),
+    );
+
+    await navigateToPage(page, Page.Teams);
+    await expectRowsAtLeast(page, SEED_ROWS);
+    await expectBodyIsTheOnlyScroller(page);
+  });
+
+  test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({
+    page,
+    request,
+  }) => {
+    const suffix = uniqueSuffix();
+    const ids = await oneAtATime(LOG_ROWS, (i) =>
+      sendChatCompletion(request, { model: CHAT_MODEL_A, prompt: `scroll ${suffix} ${i}` }),
+    );
+    await waitForSpendLog(request, ids[ids.length - 1]);
+
+    await navigateToPage(page, Page.Logs);
+    await expect(visibleTestId(page, "datatable-search")).toBeVisible({ timeout: 20_000 });
+    await setRowsPerPage(page, "25");
+    await expectRowsAtLeast(page, LOG_ROWS);
+    await expectBodyIsTheOnlyScroller(page);
+  });
+
+  test("Tags: no row paints past the box it lives in", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`);
+    await oneAtATime(SEED_ROWS, (i) => postOk(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }));
+    cleanups.push((r) =>
+      oneAtATime(SEED_ROWS, (i) => r.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } })),
+    );
+
+    await navigateToPage(page, Page.TagManagement);
+    await expectRowsAtLeast(page, SEED_ROWS);
+    expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+  });
+
+  test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => {
+    const suffix = uniqueSuffix();
+    const created = await oneAtATime(SEED_ROWS, (i) =>
+      postOk(request, "/model/new", {
+        model_name: `e2e-scroll-model-${suffix}-${i}`,
+        litellm_params: {
+          model: "openai/fake-gpt-4",
+          api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
+          api_key: "fake-key",
+        },
+      }),
+    );
+    cleanups.push((r) =>
+      oneAtATime(SEED_ROWS, (i) =>
+        r.post("/model/delete", { headers: adminHeaders(), data: { id: created[i].model_info.id } }),
+      ),
+    );
+
+    await navigateToPage(page, Page.ModelHubTable);
+    await expectRowsAtLeast(page, SEED_ROWS);
+    expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx
index 488190a0fdf..feea01f19ca 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/TagTable.tsx
@@ -41,6 +41,7 @@ const TagTable: React.FC<TagTableProps> = ({ data, onEdit, onDelete, onSelectTag
       data={data}
       columns={columns}
       getRowId={(tag, index) => tag.name || String(index)}
+      fillHeight
       sortingMode="client"
       sorting={sorting}
       onSortingChange={setSorting}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx
index a04afcd6d45..583492bf837 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/index.tsx
@@ -126,7 +126,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
   }, [accessToken]);
 
   return (
-    <div className="mx-4 h-[75vh]">
+    <div className="mx-4 h-full">
       {selectedTagId ? (
         <TagInfoView
           tagId={selectedTagId}
@@ -139,7 +139,7 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
           editTag={editTag}
         />
       ) : (
-        <div className="mt-2 h-[75vh] w-full gap-2 p-8">
+        <div className="flex h-full w-full flex-col p-8 pt-10">
           <div className="mt-2 mb-4 flex w-full items-center justify-between">
             <h1>Tag Management</h1>
             <div className="flex items-center space-x-2">
@@ -162,23 +162,21 @@ const TagManagement: React.FC<TagProps> = ({ accessToken, userID, userRole }) =>
             </p>
           </div>
 
-          <Button className="mb-4" onClick={() => setIsCreateModalVisible(true)}>
+          <Button className="mb-4 self-start" onClick={() => setIsCreateModalVisible(true)}>
             + Create New Tag
           </Button>
 
-          <div className="mt-2 grid h-[75vh] w-full grid-cols-1 gap-2 pt-2 pb-2">
-            <div>
-              <TagTable
-                data={tags}
-                isLoading={isLoadingTags}
-                onEdit={(tag) => {
-                  setSelectedTagId(tag.name);
-                  setEditTag(true);
-                }}
-                onDelete={handleDelete}
-                onSelectTag={setSelectedTagId}
-              />
-            </div>
+          <div className="mt-2 flex min-h-0 flex-1 flex-col">
+            <TagTable
+              data={tags}
+              isLoading={isLoadingTags}
+              onEdit={(tag) => {
+                setSelectedTagId(tag.name);
+                setEditTag(true);
+              }}
+              onDelete={handleDelete}
+              onSelectTag={setSelectedTagId}
+            />
           </div>
 
           {/* Create Tag Modal */}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
index f74f444ce26..1745c710c51 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/vector-stores/_components/index.tsx
@@ -137,8 +137,8 @@ const VectorStoreManagement: React.FC<VectorStoreProps> = ({ accessToken, userID
       />
     </div>
   ) : (
-    <div className="mx-4 h-[75vh]">
-      <div className="gap-2 p-8 h-[75vh] w-full mt-2">
+    <div className="mx-4">
+      <div className="gap-2 p-8 w-full mt-2">
         <div className="flex justify-between mt-2 w-full items-center mb-4">
           <h1 className="text-xl font-semibold tracking-tight text-foreground">Vector Store Management</h1>
           <div className="flex items-center space-x-2">
diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx
index 475e3dcd70b..74f44db5cb9 100644
--- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx
+++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx
@@ -400,7 +400,7 @@ const ModelHubTable: React.FC<ModelHubTableProps> = ({ accessToken, publicPage,
   }
 
   return (
-    <div className="mx-4 h-[75vh]">
+    <div className="mx-4">
       {publicPage == false ? (
         <div className="w-full m-2 mt-2 p-8">
           {/* Header with Title, Description and URL */}
diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx
index ef58237a6aa..4be91f22339 100644
--- a/ui/litellm-dashboard/src/components/Teams.tsx
+++ b/ui/litellm-dashboard/src/components/Teams.tsx
@@ -559,6 +559,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
     {
       key: "your-teams",
       label: "Your Teams",
+      className: "flex min-h-0 flex-1 flex-col",
       children: (
         <>
           <TeamsTable
@@ -608,6 +609,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
     {
       key: "available-teams",
       label: "Available Teams",
+      className: "min-h-0 flex-1 overflow-y-auto",
       children: <AvailableTeamsPanel accessToken={accessToken} userID={userID} />,
     },
     ...(isProxyAdminRole(userRole || "")
@@ -615,6 +617,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
           {
             key: "default-settings",
             label: "Default Team Settings",
+            className: "min-h-0 flex-1 overflow-y-auto",
             children: <TeamSSOSettings accessToken={accessToken} userID={userID || ""} userRole={userRole || ""} />,
           },
         ]
@@ -622,7 +625,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
   ];
 
   return (
-    <main className={selectedTeamId ? "px-12 py-6" : "p-8"}>
+    <main className={selectedTeamId ? "px-12 py-6" : "flex h-full flex-col p-8"}>
       {selectedTeamId ? (
         <TeamInfoView
           teamId={selectedTeamId}
@@ -642,7 +645,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
           premiumUser={premiumUser}
         />
       ) : (
-        <Tabs defaultValue={tabItems[0].key} className="gap-6">
+        <Tabs defaultValue={tabItems[0].key} className="min-h-0 flex-1 gap-6">
           <PageHeader
             icon={<Users />}
             title="Teams"
@@ -674,7 +677,7 @@ const Teams: React.FC<TeamProps> = ({ accessToken, userID, userRole, premiumUser
             )}
           />
           {tabItems.map((item) => (
-            <TabsContent key={item.key} value={item.key}>
+            <TabsContent key={item.key} value={item.key} className={item.className}>
               {item.children}
             </TabsContent>
           ))}
diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx
index 0b38d51e36a..3fb19e522f3 100644
--- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx
+++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx
@@ -164,7 +164,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet
       isLoading={isLoading}
       loadingMessage="Loading teams..."
       noDataMessage="No teams found"
-      maxBodyHeight="calc(75vh - 210px)"
+      fillHeight
       size="compact"
       toolbar={(table) => (
         <>
diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
index ebedf57af45..f28ae27b6bc 100644
--- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx
@@ -256,7 +256,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
   }
 
   return (
-    <div className="flex h-full flex-col gap-6 overflow-hidden">
+    <div className="flex min-h-0 flex-1 flex-col gap-6">
       <PageHeader
         icon={<KeyRound />}
         title="Virtual Keys"
@@ -283,7 +283,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) {
         isLoading={isLoading}
         loadingMessage="Loading keys..."
         noDataMessage="No keys found"
-        maxBodyHeight="calc(75vh - 210px)"
+        fillHeight
         size="compact"
         toolbar={(table) => (
           <>
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
index f502ea77127..daf20d927e7 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
@@ -613,8 +613,12 @@ describe("DataTable layout", () => {
 
   it("makes the header sticky and constrains body height when maxBodyHeight is set", () => {
     render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameEmailColumns} maxBodyHeight={240} />);
-    expect(screen.getByTestId("data-table-head")).toHaveClass("sticky");
-    expect(screen.getByTestId("data-table-scroller")).toHaveStyle({ maxHeight: "240px" });
+    const scroller = screen.getByTestId("data-table-scroller");
+    expect(scroller).toHaveStyle({ maxHeight: "240px" });
+    expect(scroller).toHaveClass("overflow-auto");
+    // As in fill mode: the Table primitive's own overflow container would otherwise capture the sticky header.
+    expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible");
+    expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background");
   });
 
   it("caps fillHeight at the parent's height instead of stretching to it, so a short table stays short", () => {
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
index 60267606951..f04430c1698 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
@@ -59,18 +59,28 @@ const noop = () => {};
 /**
  * Height-filling mode. The table still sizes to its rows; the parent's height is only a ceiling, so
  * a short table keeps its footer under the last row and a long one scrolls its rows instead of the
- * page. `table-container` is the Table primitive's own overflow-x wrapper; left as a scroll box it
- * captures the sticky header and the header scrolls away with the rows. And rows pass under that
- * header, which the semi-transparent header row tint alone would not hide.
+ * page.
  */
 const FILL_CLASSES = {
   outer: "flex max-h-full min-h-0 flex-col",
   frame: "flex min-h-0 flex-col",
-  body: "min-h-0 [&_[data-slot=table-container]]:overflow-visible",
+  body: "min-h-0",
+} as const;
+
+const NO_FILL_CLASSES = { outer: "", frame: "", body: "" } as const;
+
+/**
+ * Sticky header, in both fill and maxBodyHeight mode. `table-container` is the Table primitive's own
+ * overflow-x wrapper; left as a scroll box it captures the sticky header and the header scrolls away
+ * with the rows. And rows pass under that header, which the semi-transparent header row tint alone
+ * would not hide.
+ */
+const STICKY_CLASSES = {
+  body: "[&_[data-slot=table-container]]:overflow-visible",
   header: "bg-background",
 } as const;
 
-const NO_FILL_CLASSES = { outer: "", frame: "", body: "", header: "" } as const;
+const NO_STICKY_CLASSES = { body: "", header: "" } as const;
 
 function columnDefId<TData, TValue>(column: ColumnDef<TData, TValue>): string | undefined {
   if ("id" in column && typeof column.id === "string") {
@@ -533,6 +543,7 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
   const visibleColumnCount = table.getVisibleLeafColumns().length;
   const stickyHeader = maxBodyHeight !== undefined || fillHeight;
   const fill = fillHeight ? FILL_CLASSES : NO_FILL_CLASSES;
+  const sticky = stickyHeader ? STICKY_CLASSES : NO_STICKY_CLASSES;
   const tableStyle = enableColumnResizing ? { width: table.getTotalSize(), minWidth: "100%" } : undefined;
 
   const renderPagination = (): React.ReactNode => {
@@ -593,13 +604,13 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
         {toolbar !== undefined && <div className="shrink-0 border-b border-border px-4 py-3">{toolbar(table)}</div>}
         <div
           data-testid="data-table-scroller"
-          className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", fill.body)}
+          className={cn(stickyHeader ? "overflow-auto" : "overflow-x-auto", sticky.body, fill.body)}
           style={maxBodyHeight !== undefined ? { maxHeight: maxBodyHeight } : undefined}
         >
           <TableRoot className={enableColumnResizing ? "table-fixed" : ""} style={tableStyle}>
             <TableHeader
               data-testid="data-table-head"
-              className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", fill.header)}
+              className={cn(stickyHeader ? "sticky top-0 z-sticky" : "", sticky.header)}
             >
               {table.getHeaderGroups().map((headerGroup) => (
                 <TableRow key={headerGroup.id} className="bg-muted/50">
diff --git a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
index aa9df4a0319..b0380255b95 100644
--- a/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
+++ b/ui/litellm-dashboard/src/components/team/TeamVirtualKeysTable.tsx
@@ -443,7 +443,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
   }, []);
 
   return (
-    <div className="w-full h-full overflow-hidden">
+    <div className="w-full">
       {selectedKey ? (
         <KeyInfoView
           keyId={selectedKey.token}
@@ -453,7 +453,7 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
           onDelete={refetch}
         />
       ) : (
-        <div className="py-4 flex-1 overflow-hidden">
+        <div className="py-4">
           <DataTable
             data={displayKeys}
             columns={columns}
@@ -471,7 +471,6 @@ export function TeamVirtualKeysTable({ teamId, teamAlias, organization }: TeamVi
             columnResizeMode="onChange"
             isLoading={isLoading || isFetching}
             loadingMessage="Loading keys..."
-            maxBodyHeight="75vh"
             size="compact"
             toolbar={(table) => (
               <>
diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx
index 614b3b7af62..dcadc141103 100644
--- a/ui/litellm-dashboard/src/components/user_dashboard.tsx
+++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx
@@ -216,24 +216,22 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer";
 
   return (
-    <main className="h-[75vh] p-8">
-      <div className="flex h-full flex-col">
-        <VirtualKeysTable
-          headerActions={
-            canCreateKey ? (
-              <CreateKey
-                key={selectedTeam ? selectedTeam.team_id : null}
-                team={selectedTeam as Team | null}
-                teams={teams as Team[]}
-                data={keys}
-                addKey={addKey}
-                autoOpenCreate={autoOpenCreate}
-                prefillData={prefillData}
-              />
-            ) : undefined
-          }
-        />
-      </div>
+    <main className="flex h-full flex-col p-8">
+      <VirtualKeysTable
+        headerActions={
+          canCreateKey ? (
+            <CreateKey
+              key={selectedTeam ? selectedTeam.team_id : null}
+              team={selectedTeam as Team | null}
+              teams={teams as Team[]}
+              data={keys}
+              addKey={addKey}
+              autoOpenCreate={autoOpenCreate}
+              prefillData={prefillData}
+            />
+          ) : undefined
+        }
+      />
     </main>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
index 17caa4466fa..db108ee681c 100644
--- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsTable.tsx
@@ -86,6 +86,7 @@ export function RequestLogsTable({
       data={data}
       columns={columns}
       getRowId={(row) => row.request_id}
+      fillHeight
       sortingMode="server"
       sorting={sorting}
       onSortingChange={onSortingChange}
diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx
index f20285dffa0..aadf90fad6c 100644
--- a/ui/litellm-dashboard/src/components/view_logs/index.tsx
+++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx
@@ -27,6 +27,9 @@ const AUDIT_LOGS_TAB: LogsTab = { id: "audit logs", label: "Audit Logs" };
 const DELETED_KEYS_TAB: LogsTab = { id: "deleted keys", label: "Deleted Keys" };
 const DELETED_TEAMS_TAB: LogsTab = { id: "deleted teams", label: "Deleted Teams" };
 
+const tabContentClassName = (tabId: LogsTabId): string =>
+  tabId === REQUEST_LOGS_TAB.id ? "flex min-h-0 flex-1 flex-col" : "min-h-0 flex-1 overflow-y-auto";
+
 export default function SpendLogsTable({ accessToken, token, userRole, userID, premiumUser }: SpendLogsTableProps) {
   const [activeTab, setActiveTab] = useState<LogsTabId>(REQUEST_LOGS_TAB.id);
   const canViewAuditLogs = useCan("viewAuditLogs");
@@ -78,8 +81,8 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
   };
 
   return (
-    <div className="box-border w-full overflow-x-hidden p-6">
-      <Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)}>
+    <div className="flex h-full w-full flex-col p-6">
+      <Tabs value={activeTab} onValueChange={(value) => setActiveTab(value as LogsTabId)} className="min-h-0 flex-1">
         <TabsList variant="line">
           {tabs.map((tab) => (
             <TabsTrigger key={tab.id} value={tab.id} className="flex-none">
@@ -88,7 +91,7 @@ export default function SpendLogsTable({ accessToken, token, userRole, userID, p
           ))}
         </TabsList>
         {tabs.map((tab) => (
-          <TabsContent key={tab.id} value={tab.id} keepMounted>
+          <TabsContent key={tab.id} value={tab.id} keepMounted className={tabContentClassName(tab.id)}>
             {renderPanel(tab.id)}
           </TabsContent>
         ))}

From e26d607f5df343cfd62077da46382398c543679d Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Thu, 3 Sep 2026 17:22:47 -0700
Subject: [PATCH 27/38] feat(ui): configure auto-router affinity idle TTL
 (#39679)

---
 .../complexity_router/README.md               |  7 ++
 .../components/add_model/AffinityControls.tsx | 66 ++++++++++++++-----
 .../add_model/ComplexityRouterConfig.test.tsx | 40 +++++++++++
 .../add_model/ComplexityRouterConfig.tsx      |  2 +
 .../add_model/add_auto_router_tab.test.tsx    |  7 +-
 .../add_model/add_auto_router_tab.tsx         |  1 +
 .../build_complexity_router_config.test.ts    |  9 +++
 .../build_complexity_router_config.ts         |  6 ++
 ...d_updated_complexity_router_config.test.ts | 38 +++++++++++
 .../edit_auto_router_modal.test.tsx           | 43 ++++++++++++
 .../edit_auto_router_modal.tsx                |  8 +++
 .../src/lib/autorouter_presets.test.ts        | 10 +++
 .../src/lib/autorouter_presets.ts             |  1 +
 13 files changed, 220 insertions(+), 18 deletions(-)

diff --git a/litellm/router_strategy/complexity_router/README.md b/litellm/router_strategy/complexity_router/README.md
index e3da70f50fe..afa27719064 100644
--- a/litellm/router_strategy/complexity_router/README.md
+++ b/litellm/router_strategy/complexity_router/README.md
@@ -190,6 +190,9 @@ model_list:
 
         # Let that replacement also override a kept session pin, for image turns only (default: false)
         modality_pin_override: true
+
+        # Refreshes on every pin reuse, so this is idle time rather than total session length (default: 3600)
+        session_affinity_ttl_seconds: 300
 ```
 
 ## Usage
@@ -240,6 +243,10 @@ affinity write happens upstream of the gate and stores the session's own model,
 turn replays the original pin and the override is never pinned in its place. It does nothing
 unless `modality_routing` is also on.
 
+### Session pin retention
+
+`session_affinity_ttl_seconds` is the idle window for both the model pin selected by session affinity and the deployment pin. Every request that reuses a pin refreshes its TTL, so a session actively sending requests stays pinned. After the window passes with no pin reuse, the next request classifies again and creates a fresh pin. Omit the setting to track the default of 3600 seconds.
+
 ### Heuristic-first chaining
 
 `classifier_type: heuristic_first` runs the local scorer on every request and only calls the LLM
diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx
index 4d9e2122739..9022d424369 100644
--- a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx
@@ -1,26 +1,58 @@
 import React from "react";
 
+import { Input } from "@/components/ui/input";
 import { Switch } from "@/components/ui/switch";
 
 import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig";
-import { DEFAULT_DEPLOYMENT_AFFINITY } from "./ComplexityRouterConfig";
+import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY_TTL_SECONDS } from "./ComplexityRouterConfig";
 
 export const AffinityControls: React.FC<{
   value: ComplexityRouterConfigValue;
   onChange: (value: ComplexityRouterConfigValue) => void;
-}> = ({ value, onChange }) => (
-  <>
-    <div className="flex items-center gap-2 mb-2">
-      <Switch
-        checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
-        onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
-        aria-label="Pin a session to one deployment per model group"
-      />
-      <strong className="font-semibold">Pin a session to one deployment per model group</strong>
-    </div>
-    <span className="block text-xs text-muted-foreground">
-      Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
-      load-balance every turn.
-    </span>
-  </>
-);
+}> = ({ value, onChange }) => {
+  const [ttlDraft, setTtlDraft] = React.useState<string | null>(null);
+  const commitTtl = (raw: string) => {
+    setTtlDraft(null);
+    if (raw.trim() === "") {
+      onChange({ ...value, session_affinity_ttl_seconds: undefined });
+      return;
+    }
+    const parsed = Number(raw);
+    if (!Number.isFinite(parsed)) return;
+    onChange({ ...value, session_affinity_ttl_seconds: Math.max(1, Math.round(parsed)) });
+  };
+
+  return (
+    <>
+      <div className="flex items-center gap-2 mb-2">
+        <Switch
+          checked={value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY}
+          onCheckedChange={(deploymentAffinity) => onChange({ ...value, deployment_affinity: deploymentAffinity })}
+          aria-label="Pin a session to one deployment per model group"
+        />
+        <strong className="font-semibold">Pin a session to one deployment per model group</strong>
+      </div>
+      <span className="block text-xs mb-3 text-muted-foreground">
+        Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to
+        load-balance every turn.
+      </span>
+      <div style={{ maxWidth: 320 }}>
+        <label className="block text-sm font-medium mb-1" htmlFor="session-affinity-ttl">
+          How long a pin survives idle (seconds)
+        </label>
+        <Input
+          id="session-affinity-ttl"
+          inputMode="numeric"
+          value={ttlDraft ?? value.session_affinity_ttl_seconds ?? ""}
+          placeholder={String(DEFAULT_SESSION_AFFINITY_TTL_SECONDS)}
+          onChange={(event) => setTtlDraft(event.target.value)}
+          onBlur={(event) => commitTtl(event.target.value)}
+        />
+        <span className="block text-xs mt-1 text-muted-foreground">
+          Refreshes after every request that reuses a pin. Empty tracks the backend default of{" "}
+          {DEFAULT_SESSION_AFFINITY_TTL_SECONDS} seconds.
+        </span>
+      </div>
+    </>
+  );
+};
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
index 02cb0543e17..bdca5205b2e 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.test.tsx
@@ -947,6 +947,46 @@ describe("ComplexityRouterConfig affinity panel", () => {
 
     expect(screen.getByRole("switch", { name: "Pin a session to one deployment per model group" })).not.toBeChecked();
   });
+
+  it("writes an idle TTL on blur and keeps the partial input as a draft while typing", () => {
+    const onChange = vi.fn();
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
+    expect(ttl).toHaveAttribute("placeholder", "3600");
+    fireEvent.change(ttl, { target: { value: "300" } });
+    expect(onChange).not.toHaveBeenCalled();
+    fireEvent.blur(ttl);
+
+    expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 300 });
+  });
+
+  it("clearing the idle TTL returns the router to its backend default", () => {
+    const onChange = vi.fn();
+    const value = { ...defaultValue, session_affinity_ttl_seconds: 300 };
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} value={value} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
+    expect(ttl).toHaveValue("300");
+    fireEvent.change(ttl, { target: { value: "" } });
+    fireEvent.blur(ttl);
+
+    expect(onChange).toHaveBeenCalledWith({ ...value, session_affinity_ttl_seconds: undefined });
+  });
+
+  it("clamps a non-positive idle TTL to the backend's minimum", () => {
+    const onChange = vi.fn();
+    renderWithProviders(<ComplexityRouterConfig {...baseProps} onChange={onChange} />);
+    fireEvent.click(screen.getByText("Advanced: Affinity"));
+
+    const ttl = screen.getByLabelText("How long a pin survives idle (seconds)");
+    fireEvent.change(ttl, { target: { value: "0" } });
+    fireEvent.blur(ttl);
+
+    expect(onChange).toHaveBeenCalledWith({ ...defaultValue, session_affinity_ttl_seconds: 1 });
+  });
 });
 
 describe("ComplexityRouterConfig default model", () => {
diff --git a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
index 791fd1fcfbb..06363830d64 100644
--- a/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/ComplexityRouterConfig.tsx
@@ -58,6 +58,7 @@ export const DEFAULT_CLASSIFIER_CONTEXT_WINDOW_SIZE = 3;
 export const DEFAULT_CLASSIFIER_CONTEXT_BUDGET_CHARS = 8000;
 export const MIN_QUOTED_CONTEXT_TURN_CHARS = 120;
 export const DEFAULT_SESSION_AFFINITY = false;
+export const DEFAULT_SESSION_AFFINITY_TTL_SECONDS = 3600;
 export const DEFAULT_DEPLOYMENT_AFFINITY = true;
 
 export type ClassificationMode = "every_request" | "user_turn";
@@ -411,6 +412,7 @@ export interface ComplexityRouterConfigValue {
   hybrid_boundary_margin?: number;
   classification_mode?: ClassificationMode;
   session_affinity?: boolean;
+  session_affinity_ttl_seconds?: number;
   modality_routing?: boolean;
   modality_pin_override?: boolean;
   deployment_affinity?: boolean;
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
index 836b4c6ff3d..dfef8171c51 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx
@@ -493,7 +493,7 @@ describe("AddAutoRouterTab", () => {
     });
   });
 
-  it("carries session affinity turned on through to the create payload", async () => {
+  it("carries session affinity turned on and its idle window through to the create payload", async () => {
     const user = userEvent.setup();
     vi.mocked(getMissingTiersError).mockReturnValue(null);
 
@@ -503,12 +503,17 @@ describe("AddAutoRouterTab", () => {
     expandDetailedConfiguration();
     await user.click(screen.getByText("Advanced: Classification Method"));
     await user.click(await screen.findByRole("radio", { name: /Once per session/ }));
+    await user.click(screen.getByText("Advanced: Affinity"));
+    const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
+    fireEvent.change(ttl, { target: { value: "300" } });
+    fireEvent.blur(ttl);
 
     await user.click(screen.getByRole("button", { name: /add auto router/i }));
 
     await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled());
     expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({
       session_affinity: true,
+      session_affinity_ttl_seconds: 300,
     });
   });
 
diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
index c335beade91..1a1725b8dd0 100644
--- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx
@@ -388,6 +388,7 @@ const AddAutoRouterTab: React.FC<AddAutoRouterTabProps> = ({
     reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score,
     enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation,
     contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer,
+    sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds,
   };
 
   const submitRecommendedRouter = async (name: string) => {
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
index 53ab859baa0..9ee555f5dd2 100644
--- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
+++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts
@@ -73,6 +73,15 @@ describe("buildComplexityRouterConfig", () => {
     expect(config.context_window_escalation_buffer).toBe(0.9);
   });
 
+  it("omits session_affinity_ttl_seconds when untouched, so the router tracks the backend default", () => {
+    expect(buildComplexityRouterConfig(baseParams)).not.toHaveProperty("session_affinity_ttl_seconds");
+  });
+
+  it("emits an explicit session affinity idle window", () => {
+    const config = buildComplexityRouterConfig({ ...baseParams, sessionAffinityTtlSeconds: 300 });
+    expect(config.session_affinity_ttl_seconds).toBe(300);
+  });
+
   it("trims escalation keywords and drops blank entries", () => {
     const config = buildComplexityRouterConfig({
       ...baseParams,
diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
index b8af55e8c8a..956e593a234 100644
--- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
+++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts
@@ -138,6 +138,7 @@ export interface BuildComplexityRouterConfigParams {
   tierModelParams?: TierModelParamsByTier;
   enableContextWindowEscalation?: boolean;
   contextWindowEscalationBuffer?: number;
+  sessionAffinityTtlSeconds?: number;
 }
 
 /**
@@ -175,6 +176,7 @@ export interface ComplexityRouterConfigPayload {
   hybrid_boundary_margin?: number;
   classification_mode: ClassificationMode;
   session_affinity: boolean;
+  session_affinity_ttl_seconds?: number;
   deployment_affinity: boolean;
   modality_routing: boolean;
   modality_pin_override: boolean;
@@ -456,6 +458,7 @@ export const buildComplexityRouterConfig = ({
   tierModelParams,
   enableContextWindowEscalation,
   contextWindowEscalationBuffer,
+  sessionAffinityTtlSeconds,
 }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => {
   const serializedTierModelConfigs = customTierSet
     ? serializeTierModelConfigs(
@@ -522,6 +525,9 @@ export const buildComplexityRouterConfig = ({
     ...(contextWindowEscalationBuffer !== undefined && {
       context_window_escalation_buffer: contextWindowEscalationBuffer,
     }),
+    ...(sessionAffinityTtlSeconds !== undefined && {
+      session_affinity_ttl_seconds: sessionAffinityTtlSeconds,
+    }),
     ...scorerKnobs,
   };
   if (!customTierSet) return payload;
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
index f24f3033901..877b35199a4 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts
@@ -257,6 +257,43 @@ describe("buildUpdatedComplexityRouterConfig session affinity", () => {
   });
 });
 
+describe("buildUpdatedComplexityRouterConfig session affinity ttl", () => {
+  it("writes an edited idle window", () => {
+    const result = buildUpdatedComplexityRouterConfig(STORED, { ...FORM_VALUE, session_affinity_ttl_seconds: 300 });
+    expect(result.session_affinity_ttl_seconds).toBe(300);
+  });
+
+  it("carries a stored idle window through an untouched open-and-save", () => {
+    const stored = { ...STORED, session_affinity_ttl_seconds: 900 };
+    const result = buildUpdatedComplexityRouterConfig(stored, hydrateComplexityRouterConfig(stored, undefined));
+    expect(result.session_affinity_ttl_seconds).toBe(900);
+  });
+
+  it("drops the key when the field is cleared, so the router goes back to tracking the backend default", () => {
+    const result = buildUpdatedComplexityRouterConfig(
+      { ...STORED, session_affinity_ttl_seconds: 900 },
+      { ...FORM_VALUE, session_affinity_ttl_seconds: undefined },
+    );
+    expect(result).not.toHaveProperty("session_affinity_ttl_seconds");
+  });
+
+  it("keeps the idle window on a custom tier set, whose deployment pin still uses it", () => {
+    const result = buildUpdatedComplexityRouterConfig(STORED, {
+      ...FORM_VALUE,
+      session_affinity_ttl_seconds: 300,
+      custom_tier_set: {
+        tiers: [
+          { id: "a", name: "CASUAL", definition: "small talk", models: ["gpt-4o-mini"] },
+          { id: "b", name: "AUDIT", definition: "security review", models: ["o1"] },
+        ],
+        fallback_tier_id: "a",
+      },
+    });
+    expect(result.session_affinity).toBe(false);
+    expect(result.session_affinity_ttl_seconds).toBe(300);
+  });
+});
+
 describe("buildUpdatedComplexityRouterConfig modality pin override", () => {
   it("writes modality_pin_override explicitly both ways", () => {
     expect(
@@ -529,6 +566,7 @@ describe("managed keys survive an untouched open-and-save", () => {
     classifier_fallback: "default_model",
     classification_mode: "user_turn",
     session_affinity: true,
+    session_affinity_ttl_seconds: 300,
     modality_routing: true,
     modality_pin_override: true,
     deployment_affinity: false,
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
index 6067e72e547..970bcaa545f 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.test.tsx
@@ -550,6 +550,49 @@ describe("EditAutoRouterModal deployment affinity", () => {
     expect(savedConfig().deployment_affinity).toBe(false);
   });
 
+  it("preserves an idle TTL through an untouched save", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    expect(await screen.findByLabelText("How long a pin survives idle (seconds)")).toHaveValue("300");
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
+  });
+
+  it("persists an edited idle TTL", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig(STORED_CONFIG);
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
+    fireEvent.change(ttl, { target: { value: "300" } });
+    fireEvent.blur(ttl);
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig().session_affinity_ttl_seconds).toBe(300);
+  });
+
+  it("removes the idle TTL when cleared", async () => {
+    const user = userEvent.setup();
+    renderWithStoredConfig({ ...STORED_CONFIG, session_affinity_ttl_seconds: 300 });
+
+    await user.click(await screen.findByText("Advanced: Affinity"));
+    const ttl = await screen.findByLabelText("How long a pin survives idle (seconds)");
+    fireEvent.change(ttl, { target: { value: "" } });
+    fireEvent.blur(ttl);
+
+    await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(savedConfig()).not.toHaveProperty("session_affinity_ttl_seconds");
+  });
+
   // modality_pin_override is a managed key, so the modal rewrites it from form state on save. A
   // hydration gap would silently turn a stored override off on the next untouched save.
   it("shows a stored modality_pin_override=true as on and preserves it through an untouched save", async () => {
diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
index 54766df93ec..ea1e5cba6a3 100644
--- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
+++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx
@@ -104,6 +104,7 @@ export interface StoredComplexityRouterConfig {
   dimension_weights?: unknown;
   reasoning_override_min_score?: unknown;
   session_affinity?: unknown;
+  session_affinity_ttl_seconds?: unknown;
   modality_routing?: unknown;
   modality_pin_override?: unknown;
   deployment_affinity?: unknown;
@@ -182,6 +183,11 @@ export const hydrateComplexityRouterConfig = (
     reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score),
     session_affinity:
       typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY,
+    session_affinity_ttl_seconds:
+      typeof parsedConfig.session_affinity_ttl_seconds === "number" &&
+      Number.isFinite(parsedConfig.session_affinity_ttl_seconds)
+        ? parsedConfig.session_affinity_ttl_seconds
+        : undefined,
     modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false,
     modality_pin_override:
       typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false,
@@ -224,6 +230,7 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([
   "hybrid_boundary_margin",
   "classification_mode",
   "session_affinity",
+  "session_affinity_ttl_seconds",
   "modality_routing",
   "modality_pin_override",
   "deployment_affinity",
@@ -322,6 +329,7 @@ export const buildUpdatedComplexityRouterConfig = (
     classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns,
     classifierFallback: value.classifier_fallback,
     sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY,
+    sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds,
     modalityRouting: value.modality_routing ?? false,
     modalityPinOverride: value.modality_pin_override ?? false,
     deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
index 3dd9911c794..ffede7b2a6b 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts
@@ -83,9 +83,19 @@ describe("autorouter_presets", () => {
       expect(config.tier_boundaries).toBeUndefined();
       expect(config.token_thresholds).toBeUndefined();
       expect(config.dimension_weights).toBeUndefined();
+      expect(config.session_affinity_ttl_seconds).toBeUndefined();
     }
   });
 
+  it("carries a preset's session affinity idle window into the prefilled form state", () => {
+    const config = getPresetByKey("anthropic_family")!.complexity_router_config;
+    const prefill = buildPresetPrefill({ ...config, session_affinity_ttl_seconds: 300 }, groupsOnly([]));
+    expect(prefill.complexityRouterConfig.session_affinity_ttl_seconds).toBe(300);
+    expect(
+      buildPresetPrefill(config, groupsOnly([])).complexityRouterConfig.session_affinity_ttl_seconds,
+    ).toBeUndefined();
+  });
+
   it("keeps the model-family presets on the heuristic classifier", () => {
     for (const key of ["anthropic_family", "gemini_family", "openai_family"]) {
       expect(getPresetByKey(key)!.complexity_router_config.classifier_type).toBe("heuristic");
diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
index ff482bd23b8..b01108f1631 100644
--- a/ui/litellm-dashboard/src/lib/autorouter_presets.ts
+++ b/ui/litellm-dashboard/src/lib/autorouter_presets.ts
@@ -284,6 +284,7 @@ export const buildPresetPrefill = (
       classifier_context_include_assistant_turns: config.classifier_context_include_assistant_turns,
       classification_mode: config.classification_mode ?? DEFAULT_CLASSIFICATION_MODE,
       session_affinity: config.session_affinity ?? DEFAULT_SESSION_AFFINITY,
+      session_affinity_ttl_seconds: config.session_affinity_ttl_seconds,
       deployment_affinity: config.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY,
       modality_routing: config.modality_routing ?? false,
       modality_pin_override: config.modality_pin_override ?? false,

From ff97e71652b19e84df692b03af76fb9f22c77709 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Thu, 3 Sep 2026 17:25:26 -0700
Subject: [PATCH 28/38] refactor(ui): type and de-mutate the table scrolling
 spec, drop CSS narration

DataTable loses the comment that narrated its sticky header classes. The
table scrolling e2e spec now types every management API response it
reads, seeds rows through an immutable reduce instead of pushing into
arrays, and deletes what it seeded in each test's finally block instead
of draining a shared mutable list in afterEach.

Refs LIT-4738

Claude-Session: https://claude.ai/code/session_018yW93iDaEMhoQUXcYjus7D
---
 .../ui/tests/tables/tableScrolling.spec.ts    | 179 ++++++++----------
 .../shared/DataTable/DataTable.test.tsx       |   1 -
 .../components/shared/DataTable/DataTable.tsx |   6 -
 3 files changed, 79 insertions(+), 107 deletions(-)

diff --git a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts
index f51c772185c..5c7438cfd11 100644
--- a/tests/e2e/ui/tests/tables/tableScrolling.spec.ts
+++ b/tests/e2e/ui/tests/tables/tableScrolling.spec.ts
@@ -4,37 +4,23 @@ import { Page } from "../../fixtures/pages";
 import { navigateToPage } from "../../helpers/navigation";
 import { CHAT_MODEL_A, masterKey, sendChatCompletion, waitForSpendLog } from "../../helpers/traffic";
 
-/**
- * LIT-4738 table scrolling. On the paginated pages the app shell <main> is the only page scroller
- * and must never overflow: rows scroll inside the table body under a header that stays put, and the
- * pagination footer sits at the bottom of the page instead of below the fold or inside a clipped
- * box. Pages that keep plain page scrolling must never paint rows past a fixed-height ancestor.
- * The viewport is pinned so "more rows than fit" means the same thing on every machine.
- */
-
 const VIEWPORT = { width: 1280, height: 720 };
 const SEED_ROWS = 40;
 const LOG_ROWS = 20;
 const BODY_SCROLL_PX = 500;
-/** p-8 on Keys and Teams, p-6 on Logs: the footer may sit at most one page padding above the edge. */
 const MAX_FOOTER_GAP_PX = 40;
 
-const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+interface GeneratedKey {
+  key: string;
+}
 
-const adminHeaders = (): Record<string, string> => ({
-  Authorization: `Bearer ${masterKey()}`,
-});
+interface CreatedTeam {
+  team_id: string;
+}
 
-/** Keys and Teams render a <main> of their own inside the app shell's, which comes first in document order. */
-const pageScroller = (page: PlaywrightPage): Locator => page.locator("main").first();
-
-/** Tabs keep every panel mounted, so a bare test id can match a hidden table; scope to the visible one. */
-const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
-
-/** The page's data table; Model Hub also renders a plain links table above it, which this skips. */
-const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
-
-const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
+interface CreatedModel {
+  model_info: { id: string };
+}
 
 interface BoxMetrics {
   top: number;
@@ -45,6 +31,18 @@ interface BoxMetrics {
   clientWidth: number;
 }
 
+const uniqueSuffix = (): string => `${Date.now()}-${Math.random().toString(36).slice(2, 8)}`;
+
+const adminHeaders = (): Record<string, string> => ({ Authorization: `Bearer ${masterKey()}` });
+
+const appShellMain = (page: PlaywrightPage): Locator => page.locator("main").first();
+
+const visibleTestId = (page: PlaywrightPage, id: string): Locator => page.getByTestId(id).filter({ visible: true });
+
+const visibleDataTable = (page: PlaywrightPage): Locator => visibleTestId(page, "data-table-root").first();
+
+const visibleRows = (page: PlaywrightPage): Locator => visibleDataTable(page).locator("tbody tr");
+
 const metrics = (locator: Locator): Promise<BoxMetrics> =>
   locator.evaluate((el) => {
     const rect = el.getBoundingClientRect();
@@ -58,24 +56,17 @@ const metrics = (locator: Locator): Promise<BoxMetrics> =>
     };
   });
 
-async function postOk(
-  request: APIRequestContext,
-  path: string,
-  data: Record<string, unknown>,
-): Promise<Record<string, any>> {
+async function postOk<T>(request: APIRequestContext, path: string, data: Record<string, unknown>): Promise<T> {
   const res = await request.post(path, { headers: adminHeaders(), data });
   expect(res.ok(), `POST ${path} failed (${res.status()}): ${await res.text()}`).toBe(true);
-  return (await res.json()) as Record<string, any>;
+  return (await res.json()) as T;
 }
 
-/** One request at a time: a burst of forty management calls starves the proxy's transaction pool. */
-async function oneAtATime<T>(count: number, call: (index: number) => Promise<T>): Promise<T[]> {
-  const results: T[] = [];
-  for (let i = 0; i < count; i++) {
-    results.push(await call(i));
-  }
-  return results;
-}
+const oneAtATime = <T>(count: number, call: (index: number) => Promise<T>): Promise<readonly T[]> =>
+  Array.from({ length: count }, (_, i) => i).reduce<Promise<readonly T[]>>(
+    async (previous, i) => [...(await previous), await call(i)],
+    Promise.resolve([]),
+  );
 
 async function expectRowsAtLeast(page: PlaywrightPage, count: number): Promise<void> {
   await expect.poll(() => visibleRows(page).count(), { timeout: 30_000 }).toBeGreaterThanOrEqual(count);
@@ -86,12 +77,8 @@ async function setRowsPerPage(page: PlaywrightPage, size: "25" | "50" | "100"):
   await page.getByRole("option", { name: size, exact: true }).click();
 }
 
-/**
- * The page scroller stays put, the table body is what scrolls, the header does not move while the
- * body scrolls, and the pagination footer sits at the bottom of the page.
- */
 async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
-  const scroller = await metrics(pageScroller(page));
+  const scroller = await metrics(appShellMain(page));
   const body = visibleTestId(page, "data-table-scroller");
   const bodyBefore = await metrics(body);
   const headBefore = await metrics(visibleTestId(page, "data-table-head"));
@@ -115,73 +102,61 @@ async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void>
   expect(Math.round(headAfter.top), "header must stay put while the body scrolls").toBe(Math.round(headBefore.top));
 }
 
-/**
- * Every row must sit inside each ancestor up to the nearest one that really scrolls vertically; a
- * fixed-height box that neither grows nor scrolls lets rows paint past its bottom edge.
- */
 const rowsPaintingPastAnAncestor = (page: PlaywrightPage): Promise<string[]> =>
   visibleDataTable(page)
     .locator("table")
     .evaluate((table) => {
       const scrollsVertically = (el: Element): boolean =>
         /auto|scroll/.test(getComputedStyle(el).overflowY) && el.scrollHeight > el.clientHeight + 1;
+      const boxesUpToTheScroller = (el: Element | null): Element[] =>
+        el === null || el === document.body || scrollsVertically(el)
+          ? []
+          : [el, ...boxesUpToTheScroller(el.parentElement)];
       const describe = (el: Element): string =>
         `<${el.tagName.toLowerCase()} class="${el.getAttribute("class") ?? ""}">`;
       return Array.from(table.querySelectorAll("tbody tr")).flatMap((row, index) => {
         const rowBottom = row.getBoundingClientRect().bottom;
-        const spills: string[] = [];
-        for (let el = row.parentElement; el && el !== document.body && !scrollsVertically(el); el = el.parentElement) {
-          const bottom = el.getBoundingClientRect().bottom;
-          if (rowBottom > bottom + 1) {
-            spills.push(
-              `row ${index} bottom ${Math.round(rowBottom)} past ${describe(el)} bottom ${Math.round(bottom)}`,
-            );
-          }
-        }
-        return spills;
+        return boxesUpToTheScroller(row.parentElement)
+          .filter((box) => rowBottom > box.getBoundingClientRect().bottom + 1)
+          .map(
+            (box) =>
+              `row ${index} bottom ${Math.round(rowBottom)} past ${describe(box)} bottom ${Math.round(box.getBoundingClientRect().bottom)}`,
+          );
       });
     });
 
-type Cleanup = (request: APIRequestContext) => Promise<unknown>;
-const cleanups: Cleanup[] = [];
-
 test.describe("Admin tables scroll inside the page", () => {
   test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
 
-  test.afterEach(async ({ request }) => {
-    for (const cleanup of cleanups.splice(0)) {
-      // Teardown must never turn a passing test red or mask a real failure.
-      await cleanup(request).catch(() => {});
-    }
-  });
-
   test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
     page,
     request,
   }) => {
     const suffix = uniqueSuffix();
-    const created = await oneAtATime(SEED_ROWS, (i) =>
-      postOk(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
+    const keys = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<GeneratedKey>(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
     );
-    cleanups.push((r) => r.post("/key/delete", { headers: adminHeaders(), data: { keys: created.map((k) => k.key) } }));
-
-    await navigateToPage(page, Page.ApiKeys);
-    await expectRowsAtLeast(page, SEED_ROWS);
-    await expectBodyIsTheOnlyScroller(page);
+    try {
+      await navigateToPage(page, Page.ApiKeys);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      await expectBodyIsTheOnlyScroller(page);
+    } finally {
+      await request.post("/key/delete", { headers: adminHeaders(), data: { keys: keys.map((k) => k.key) } });
+    }
   });
 
   test("Teams: rows scroll under a sticky header and the page itself never scrolls", async ({ page, request }) => {
     const suffix = uniqueSuffix();
-    const created = await oneAtATime(SEED_ROWS, (i) =>
-      postOk(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
+    const teams = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<CreatedTeam>(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
     );
-    cleanups.push((r) =>
-      r.post("/team/delete", { headers: adminHeaders(), data: { team_ids: created.map((t) => t.team_id) } }),
-    );
-
-    await navigateToPage(page, Page.Teams);
-    await expectRowsAtLeast(page, SEED_ROWS);
-    await expectBodyIsTheOnlyScroller(page);
+    try {
+      await navigateToPage(page, Page.Teams);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      await expectBodyIsTheOnlyScroller(page);
+    } finally {
+      await request.post("/team/delete", { headers: adminHeaders(), data: { team_ids: teams.map((t) => t.team_id) } });
+    }
   });
 
   test("Request Logs: rows scroll under a sticky header and the page itself never scrolls", async ({
@@ -204,20 +179,24 @@ test.describe("Admin tables scroll inside the page", () => {
   test("Tags: no row paints past the box it lives in", async ({ page, request }) => {
     const suffix = uniqueSuffix();
     const names = Array.from({ length: SEED_ROWS }, (_, i) => `e2e-scroll-tag-${suffix}-${i}`);
-    await oneAtATime(SEED_ROWS, (i) => postOk(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }));
-    cleanups.push((r) =>
-      oneAtATime(SEED_ROWS, (i) => r.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } })),
+    await oneAtATime(SEED_ROWS, (i) =>
+      postOk<unknown>(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }),
     );
-
-    await navigateToPage(page, Page.TagManagement);
-    await expectRowsAtLeast(page, SEED_ROWS);
-    expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    try {
+      await navigateToPage(page, Page.TagManagement);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    } finally {
+      await oneAtATime(SEED_ROWS, (i) =>
+        request.post("/tag/delete", { headers: adminHeaders(), data: { name: names[i] } }),
+      );
+    }
   });
 
   test("Model Hub: no row paints past the box it lives in", async ({ page, request }) => {
     const suffix = uniqueSuffix();
-    const created = await oneAtATime(SEED_ROWS, (i) =>
-      postOk(request, "/model/new", {
+    const models = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<CreatedModel>(request, "/model/new", {
         model_name: `e2e-scroll-model-${suffix}-${i}`,
         litellm_params: {
           model: "openai/fake-gpt-4",
@@ -226,14 +205,14 @@ test.describe("Admin tables scroll inside the page", () => {
         },
       }),
     );
-    cleanups.push((r) =>
-      oneAtATime(SEED_ROWS, (i) =>
-        r.post("/model/delete", { headers: adminHeaders(), data: { id: created[i].model_info.id } }),
-      ),
-    );
-
-    await navigateToPage(page, Page.ModelHubTable);
-    await expectRowsAtLeast(page, SEED_ROWS);
-    expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    try {
+      await navigateToPage(page, Page.ModelHubTable);
+      await expectRowsAtLeast(page, SEED_ROWS);
+      expect(await rowsPaintingPastAnAncestor(page)).toEqual([]);
+    } finally {
+      await oneAtATime(SEED_ROWS, (i) =>
+        request.post("/model/delete", { headers: adminHeaders(), data: { id: models[i].model_info.id } }),
+      );
+    }
   });
 });
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
index daf20d927e7..149554a3ac3 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
@@ -616,7 +616,6 @@ describe("DataTable layout", () => {
     const scroller = screen.getByTestId("data-table-scroller");
     expect(scroller).toHaveStyle({ maxHeight: "240px" });
     expect(scroller).toHaveClass("overflow-auto");
-    // As in fill mode: the Table primitive's own overflow container would otherwise capture the sticky header.
     expect(scroller).toHaveClass("[&_[data-slot=table-container]]:overflow-visible");
     expect(screen.getByTestId("data-table-head")).toHaveClass("sticky", "bg-background");
   });
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
index f04430c1698..17a5fe42d1c 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
@@ -69,12 +69,6 @@ const FILL_CLASSES = {
 
 const NO_FILL_CLASSES = { outer: "", frame: "", body: "" } as const;
 
-/**
- * Sticky header, in both fill and maxBodyHeight mode. `table-container` is the Table primitive's own
- * overflow-x wrapper; left as a scroll box it captures the sticky header and the header scrolls away
- * with the rows. And rows pass under that header, which the semi-transparent header row tint alone
- * would not hide.
- */
 const STICKY_CLASSES = {
   body: "[&_[data-slot=table-container]]:overflow-visible",
   header: "bg-background",

From 1add1b4655007159afc5b75699a05fb782f59c4a 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 17:25:37 -0700
Subject: [PATCH 29/38] perf(mcp): cache SSO identity assertion reads on the
 ID-JAG path (#39348)

* perf(mcp): cache SSO identity assertion reads on the ID-JAG path

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

* fix(mcp): guard sso assertion cache against stale relogin reads

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

* fix(mcp): keep sso assertion cache entries and generation markers in separate namespaces

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

* test(router): rename duplicate get_configured_mode test so ruff F811 passes

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

* fix(mcp): use a process-wide epoch for sso assertion cache invalidation

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

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/constants.py                          |   1 +
 .../sso_assertion_store.py                    |  74 ++++++++-
 .../test_sso_assertion_store.py               | 144 +++++++++++++++++-
 3 files changed, 207 insertions(+), 12 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index cd72adc3db5..da731cb5eb2 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -151,6 +151,7 @@ DEFAULT_SEMANTIC_GUARD_SIMILARITY_THRESHOLD = float(os.getenv("DEFAULT_SEMANTIC_
 MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS: Final = int(os.getenv("MCP_OAUTH2_TOKEN_EXPIRY_BUFFER_SECONDS", "60"))
 MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE", "200"))
 MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL: Final = int(os.getenv("MCP_OAUTH2_TOKEN_CACHE_DEFAULT_TTL", "3600"))
+MCP_SSO_ASSERTION_CACHE_TTL_SECONDS: Final = int(os.getenv("MCP_SSO_ASSERTION_CACHE_TTL_SECONDS", "60"))
 
 # Default npm cache directory for STDIO MCP servers.
 # npm/npx needs a writable cache dir; in containers the default (~/.npm)
diff --git a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py
index 8a0d41584f9..6552008ca54 100644
--- a/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py
+++ b/litellm/proxy/_experimental/mcp_server/outbound_credentials/sso_assertion_store.py
@@ -11,7 +11,8 @@ being registered, so a gateway with no EMA upstream never stores bearer material
 The row is one encrypted payload per user, latest login wins. ``expires_at`` mirrors the
 id_token ``exp`` claim and is judged by the reader, never enforced by deletion here: an
 expired assertion with a refresh token is still renewable, and the DB row is the source of
-truth, the same contract as the per-user OAuth credential store.
+truth, the same contract as the per-user OAuth credential store. Reads use a per-process cache with
+TTL ``MCP_SSO_ASSERTION_CACHE_TTL_SECONDS``; invalidation also guards against stale in-flight reads.
 """
 
 from __future__ import annotations
@@ -24,6 +25,8 @@ import jwt
 from pydantic import BaseModel, ConfigDict, SecretStr, TypeAdapter, ValidationError
 
 from litellm._logging import verbose_proxy_logger
+from litellm.caching.in_memory_cache import InMemoryCache
+from litellm.constants import MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE, MCP_SSO_ASSERTION_CACHE_TTL_SECONDS
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
@@ -45,6 +48,46 @@ class SSOIdentityAssertion(BaseModel):
     expires_at: datetime | None = None
 
 
+class SSOAssertionCache:
+    """Process-local read cache. ``invalidate`` bumps a process-wide epoch so a fetch that started
+    before a login cannot repopulate the old assertion after it."""
+
+    def __init__(self, ttl_seconds: int = MCP_SSO_ASSERTION_CACHE_TTL_SECONDS) -> None:
+        self._entries = InMemoryCache(
+            max_size_in_memory=MCP_OAUTH2_TOKEN_CACHE_MAX_SIZE,
+            default_ttl=ttl_seconds,
+        )
+        self._epoch: int = 0
+
+    def epoch(self) -> int:
+        return self._epoch
+
+    def get(self, user_id: str) -> SSOIdentityAssertion | None:
+        cached: Final = self._entries.get_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id
+        )
+        return cached if isinstance(cached, SSOIdentityAssertion) else None
+
+    def set_if_unchanged(self, user_id: str, assertion: SSOIdentityAssertion, seen_epoch: int) -> None:
+        if self._epoch != seen_epoch:
+            return
+        self._entries.set_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id, assertion
+        )
+
+    def invalidate(self, user_id: str) -> None:
+        self._epoch += 1
+        self._entries.delete_cache(  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+            user_id
+        )
+
+    def flush(self) -> None:
+        self._entries.flush_cache()  # pyright: ignore[reportUnknownMemberType]  # InMemoryCache is untyped
+
+
+_ASSERTION_CACHE: Final = SSOAssertionCache()
+
+
 class _IdTokenClaims(BaseModel):
     exp: float | None = None
     iss: str | None = None
@@ -107,7 +150,9 @@ async def ema_assertion_retention_enabled() -> bool:
     return row is not None
 
 
-async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAssertion) -> None:
+async def persist_sso_identity_assertion(
+    user_id: str, assertion: SSOIdentityAssertion, cache: SSOAssertionCache = _ASSERTION_CACHE
+) -> None:
     from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper  # noqa: PLC0415  # runtime global
     from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # runtime global
 
@@ -127,11 +172,10 @@ async def persist_sso_identity_assertion(user_id: str, assertion: SSOIdentityAss
             "update": {"assertion_b64": encoded},
         },
     )
+    cache.invalidate(user_id)
 
 
-async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | None:
-    """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
-    rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
+async def _read_assertion_from_db(user_id: str) -> SSOIdentityAssertion | None:
     from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper  # noqa: PLC0415  # runtime global
     from litellm.proxy.proxy_server import prisma_client  # noqa: PLC0415  # runtime global
 
@@ -160,6 +204,21 @@ async def fetch_sso_identity_assertion(user_id: str) -> SSOIdentityAssertion | N
     )
 
 
+async def fetch_sso_identity_assertion(
+    user_id: str, cache: SSOAssertionCache = _ASSERTION_CACHE
+) -> SSOIdentityAssertion | None:
+    """The stored assertion for ``user_id``, or ``None`` when absent, undecryptable (salt-key
+    rotation), or unparseable. Expiry is not judged here; the reader owns that policy."""
+    cached: Final = cache.get(user_id)
+    if cached is not None:
+        return cached
+    seen_epoch: Final = cache.epoch()
+    assertion: Final = await _read_assertion_from_db(user_id)
+    if assertion is not None:
+        cache.set_if_unchanged(user_id, assertion, seen_epoch)
+    return assertion
+
+
 class AssertionStoreUnavailable(Exception):
     """Raised by ``fetch`` when the backing store is unreachable (e.g. the DB is down).
 
@@ -189,9 +248,12 @@ class DbSSOAssertionStore:
     from credential resolution and from the upstream-401 retry.
     """
 
+    def __init__(self, cache: SSOAssertionCache = _ASSERTION_CACHE) -> None:
+        self._cache = cache
+
     async def fetch(self, user_id: str) -> SSOIdentityAssertion | None:
         try:
-            return await fetch_sso_identity_assertion(user_id)
+            return await fetch_sso_identity_assertion(user_id, cache=self._cache)
         except Exception as exc:  # noqa: BLE001  # any driver/storage failure is an outage, not an absence
             raise AssertionStoreUnavailable(str(exc)) from exc
 
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py
index 7b82e004f37..5d6d47b8c38 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/outbound_credentials/test_sso_assertion_store.py
@@ -7,6 +7,7 @@ round-trips exactly, a store failure never escapes into the login path, and a sa
 rotation re-encrypts stored rows like the sibling per-user credential tables.
 """
 
+import asyncio
 import json
 import os
 import time
@@ -16,8 +17,10 @@ import jwt as pyjwt
 import pytest
 
 from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_store import (
+    _ASSERTION_CACHE,
     AssertionStoreUnavailable,
     DbSSOAssertionStore,
+    SSOAssertionCache,
     assertion_from_sso_login,
     ema_assertion_retention_enabled,
     fetch_sso_identity_assertion,
@@ -25,7 +28,7 @@ from litellm.proxy._experimental.mcp_server.outbound_credentials.sso_assertion_s
     retain_sso_identity_assertion_for_ema,
     rotate_sso_identity_assertions_master_key,
 )
-from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
+from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper
 from litellm.types.mcp import MCPAuth
 
 SALT_KEY = "test-salt-key-for-sso-assertion-tests-1234"
@@ -38,6 +41,11 @@ def _set_salt_key(monkeypatch):
     monkeypatch.setenv("LITELLM_SALT_KEY", SALT_KEY)
 
 
+@pytest.fixture(autouse=True)
+def _flush_assertion_cache():
+    _ASSERTION_CACHE.flush()
+
+
 def _make_id_token(exp_offset: int = 3600, iss: str = ISSUER) -> str:
     return pyjwt.encode(
         {"iss": iss, "sub": "u1", "exp": int(time.time()) + exp_offset},
@@ -52,9 +60,7 @@ def _make_prisma(stored: dict, db_has_id_jag_server: bool = False):
     ``db_has_id_jag_server`` drives the retention gate's authoritative DB fallback;
     it is wired explicitly so the gate never reads a truthy bare MagicMock."""
     prisma = MagicMock()
-    prisma.db.litellm_mcpservertable.find_first = AsyncMock(
-        return_value=MagicMock() if db_has_id_jag_server else None
-    )
+    prisma.db.litellm_mcpservertable.find_first = AsyncMock(return_value=MagicMock() if db_has_id_jag_server else None)
 
     async def _upsert(where, data):
         stored[where["user_id"]] = data["update"]["assertion_b64"]
@@ -235,6 +241,78 @@ async def test_persist_overwrites_previous_login():
     assert fetched.refresh_token is not None
 
 
+@pytest.mark.asyncio
+async def test_fetch_serves_second_read_from_cache_without_db_read():
+    stored = {}
+    prisma = _make_prisma(stored)
+    cache = SSOAssertionCache()
+    token = _make_id_token()
+    assertion = assertion_from_sso_login(token, "rt_1")
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        await persist_sso_identity_assertion("user-a", assertion, cache=cache)
+        first = await fetch_sso_identity_assertion("user-a", cache=cache)
+        second = await fetch_sso_identity_assertion("user-a", cache=cache)
+    assert first is not None
+    assert second is not None
+    assert first.id_token.get_secret_value() == token
+    assert second.id_token.get_secret_value() == token
+    prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_persist_busts_cache_so_relogin_is_visible_immediately():
+    stored = {}
+    prisma = _make_prisma(stored)
+    cache = SSOAssertionCache()
+    first_token = _make_id_token(exp_offset=100)
+    second_token = _make_id_token(exp_offset=7200)
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
+        first = await fetch_sso_identity_assertion("user-a", cache=cache)
+        await persist_sso_identity_assertion("user-a", assertion_from_sso_login(second_token, "rt_new"), cache=cache)
+        second = await fetch_sso_identity_assertion("user-a", cache=cache)
+    assert first is not None
+    assert second is not None
+    assert first.id_token.get_secret_value() == first_token
+    assert second.id_token.get_secret_value() == second_token
+
+
+@pytest.mark.asyncio
+async def test_fetch_racing_a_relogin_does_not_cache_the_previous_assertion():
+    stored = {}
+    prisma = _make_prisma(stored)
+    cache = SSOAssertionCache()
+    first_token = _make_id_token(exp_offset=100)
+    second_token = _make_id_token(exp_offset=7200)
+    db_read_started = asyncio.Event()
+    relogin_done = asyncio.Event()
+    unpaused_find_unique = prisma.db.litellm_ssoidentityassertion.find_unique
+
+    async def _paused_find_unique(where):
+        row = await unpaused_find_unique(where=where)
+        db_read_started.set()
+        await relogin_done.wait()
+        return row
+
+    prisma.db.litellm_ssoidentityassertion.find_unique = _paused_find_unique
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
+        racing_fetch = asyncio.create_task(fetch_sso_identity_assertion("user-a", cache=cache))
+        await db_read_started.wait()
+        await persist_sso_identity_assertion(
+            "user-a",
+            assertion_from_sso_login(second_token, "rt_new"),
+            cache=cache,
+        )
+        relogin_done.set()
+        raced = await racing_fetch
+        after = await fetch_sso_identity_assertion("user-a", cache=cache)
+    assert raced is not None
+    assert after is not None
+    assert raced.id_token.get_secret_value() == first_token
+    assert after.id_token.get_secret_value() == second_token
+
+
 @pytest.mark.asyncio
 async def test_fetch_missing_row_returns_none():
     prisma = _make_prisma({})
@@ -242,6 +320,18 @@ async def test_fetch_missing_row_returns_none():
         assert await fetch_sso_identity_assertion("nobody") is None
 
 
+@pytest.mark.asyncio
+async def test_fetch_does_not_cache_a_missing_row():
+    prisma = _make_prisma({})
+    cache = SSOAssertionCache()
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        first = await fetch_sso_identity_assertion("nobody", cache=cache)
+        second = await fetch_sso_identity_assertion("nobody", cache=cache)
+    assert first is None
+    assert second is None
+    assert prisma.db.litellm_ssoidentityassertion.find_unique.await_count == 2
+
+
 @pytest.mark.asyncio
 async def test_fetch_undecryptable_row_returns_none():
     prisma = _make_prisma({"user-a": "not-an-encrypted-blob"})
@@ -251,13 +341,37 @@ async def test_fetch_undecryptable_row_returns_none():
 
 @pytest.mark.asyncio
 async def test_fetch_unparseable_payload_returns_none():
-    from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
-
     prisma = _make_prisma({"user-a": encrypt_value_helper("]]not json")})
     with patch("litellm.proxy.proxy_server.prisma_client", prisma):
         assert await fetch_sso_identity_assertion("user-a") is None
 
 
+@pytest.mark.asyncio
+async def test_cached_assertion_expires_after_ttl():
+    stored = {}
+    prisma = _make_prisma(stored)
+    cache = SSOAssertionCache(ttl_seconds=1)
+    first_token = _make_id_token(exp_offset=100)
+    second_token = _make_id_token(exp_offset=7200)
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        await persist_sso_identity_assertion("user-a", assertion_from_sso_login(first_token, None), cache=cache)
+        first = await fetch_sso_identity_assertion("user-a", cache=cache)
+        await persist_sso_identity_assertion(
+            "user-a",
+            assertion_from_sso_login(second_token, "rt_new"),
+            cache=SSOAssertionCache(),
+        )
+        cached = await fetch_sso_identity_assertion("user-a", cache=cache)
+        time.sleep(1.1)
+        expired = await fetch_sso_identity_assertion("user-a", cache=cache)
+    assert first is not None
+    assert cached is not None
+    assert expired is not None
+    assert first.id_token.get_secret_value() == first_token
+    assert cached.id_token.get_secret_value() == first_token
+    assert expired.id_token.get_secret_value() == second_token
+
+
 @pytest.mark.asyncio
 async def test_retain_noop_when_no_id_jag_server():
     stored = {}
@@ -357,6 +471,24 @@ async def test_db_store_converts_a_driver_failure_into_assertion_store_unavailab
             await DbSSOAssertionStore().fetch("alice")
 
 
+@pytest.mark.asyncio
+async def test_db_store_uses_injected_cache():
+    stored = {}
+    prisma = _make_prisma(stored)
+    cache = SSOAssertionCache()
+    token = _make_id_token()
+    with patch("litellm.proxy.proxy_server.prisma_client", prisma):  # test-quality-ok: fake DB seam has no injection boundary
+        await persist_sso_identity_assertion("alice", assertion_from_sso_login(token, None), cache=cache)
+        store = DbSSOAssertionStore(cache=cache)
+        first = await store.fetch("alice")
+        second = await store.fetch("alice")
+    assert first is not None
+    assert second is not None
+    assert first.id_token.get_secret_value() == token
+    assert second.id_token.get_secret_value() == token
+    prisma.db.litellm_ssoidentityassertion.find_unique.assert_awaited_once()
+
+
 @pytest.mark.asyncio
 async def test_db_store_returns_none_for_a_user_with_no_stored_assertion():
     """An absent row stays an absence, not an outage, so a user who never signed in still gets the

From 5dd3fdbc3dce6cac86b6a5bde930b0fd59d8d301 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 17:31:05 -0700
Subject: [PATCH 30/38] fix(router): evict stale global pattern_router entries
 on upsert/delete (#39664)

* fix(router): evict stale global pattern_router entries on upsert/delete

upsert_deployment and delete_deployment cleaned team_pattern_routers but left
the outgoing deployment in the global pattern_router, so wildcard requests kept
round-robining onto the stale entry after a PATCH /model/{id}/update.

Fixes #29064

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

* test(router): dedupe test_get_configured_mode_reads_deployment_model_info name

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

* fix(router): restore global pattern_router eviction dropped by previous commit

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

---------

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 litellm/router.py                 |  1 +
 tests/test_litellm/test_router.py | 34 +++++++++++++++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/litellm/router.py b/litellm/router.py
index dea9aa62729..6da201725b6 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -9546,6 +9546,7 @@ class Router:
             public_model_name for _, public_model_name in self.team_model_to_deployment_indices
         )
 
+        self.pattern_router.remove_deployment(model_id)
         for team_id in list(self.team_pattern_routers.keys()):
             team_pattern_router = self.team_pattern_routers[team_id]
             team_pattern_router.remove_deployment(model_id)
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index f7f0d79b4fd..5fc96bcfbb1 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -5101,6 +5101,40 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment():
     )
 
 
+def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete():
+    """
+    Regression for #29064: upsert_deployment removed the old deployment from
+    model_list but left it in the global pattern_router, so wildcard requests
+    round-robined between the stale and the corrected deployment.
+    """
+    from litellm.types.router import Deployment, LiteLLM_Params
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "openai/*",
+                "litellm_params": {"model": "openai/openai/*", "api_key": "sk-old"},
+                "model_info": {"id": "global-wildcard"},
+            }
+        ]
+    )
+
+    router.upsert_deployment(
+        Deployment(
+            model_name="openai/*",
+            litellm_params=LiteLLM_Params(model="openai/*", api_key="sk-new"),
+            model_info={"id": "global-wildcard"},
+        )
+    )
+
+    matches = router.pattern_router.route("openai/gpt-5.2")
+    assert matches is not None
+    assert [m["litellm_params"]["api_key"] for m in matches] == ["sk-new"]
+
+    router.delete_deployment(id="global-wildcard")
+    assert router.pattern_router.patterns == {}
+
+
 def test_pattern_match_router_remove_deployment():
     """
     remove_deployment must drop only the deployment with the given model id and

From fe770700f4cf3e02999d33912c36f6779fe2ca43 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 17:31:58 -0700
Subject: [PATCH 31/38] fix(caching): keep a node timeout from forcing a
 cluster-wide topology reinit on redis-py 8.x (#39349)

* fix(caching): keep node timeout from forcing cluster reinit

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

* docs(caching): describe the 8.x timeout-tolerant wrapper in the module docstring

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

* style(caching): format redis cluster isolation wrapper

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

* fix(caching): keep concurrent reinit requests when tolerating a node timeout

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

* test(caching): cover redis cluster redirect branches

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

* fix(caching): let overlapping tolerated timeouts release their own reinit requests

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>
---
 .../caching/redis_cluster_node_isolation.py   |  93 ++++--
 .../test_redis_cluster_node_isolation.py      | 280 ++++++++++++++++--
 2 files changed, 336 insertions(+), 37 deletions(-)

diff --git a/litellm/caching/redis_cluster_node_isolation.py b/litellm/caching/redis_cluster_node_isolation.py
index ae8c78709d9..0035801018b 100644
--- a/litellm/caching/redis_cluster_node_isolation.py
+++ b/litellm/caching/redis_cluster_node_isolation.py
@@ -19,13 +19,13 @@ connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-cov
 retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
 topology changed.
 
-redis-py 8.x fixed this upstream with gentler machinery than this override's
-``node.disconnect()`` (which also kills connections other coroutines are mid-operation
-on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
-killed connection): it marks in-use connections for reconnect only after their current
-operation completes, disconnects only the idle pooled ones, and defers reinitialization
-to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
-recovery API, the factory returns the base ``RedisCluster`` unmodified.
+redis-py 8.x recovers connections per-connection, so the copied override is not used. Upstream
+still flips the shared ``_initialize`` flag on any node's timeout, funneling every concurrent
+caller through the reinit lock and, if ``CLUSTER SLOTS`` lands on the slow node, into a full
+teardown. For those versions the factory returns a thin wrapper around upstream's
+``_execute_command`` that clears the flag again after an isolated timeout (a ConnectionError,
+a third consecutive timeout on the same node, or a concurrent request from any other command
+or ``aclose()`` still reinits).
 """
 
 import asyncio
@@ -44,6 +44,8 @@ class _ClusterNodeAttrs(Protocol):
     mode; typing ``target_node`` as this Protocol at the one boundary keeps the override's
     own logic fully typed without a banned ``typing.cast``."""
 
+    name: str
+
     async def execute_command(
         self,
         *args: object,
@@ -78,18 +80,20 @@ class _ClusterAttrs(Protocol):
 #: this override can't see (Python won't error -- it'll just run our now-stale copy), so
 #: construction logs a loud warning rather than silently trusting an unverified copy.
 _VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
+_CONSECUTIVE_TIMEOUTS_BEFORE_REINIT: Final = 3
 
 
-def get_litellm_async_redis_cluster_class(
+def get_litellm_async_redis_cluster_class(  # noqa: C901  # supports redis-py version-specific cluster implementations
     cluster_node_class: type | None = None,
+    base_cluster_class: type | None = None,
 ) -> type["_AsyncRedisClusterType"]:
-    """Returns the base ``RedisCluster`` when the installed redis-py already recovers a
-    node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
-    subclass with the per-node isolation fix for older versions whose upstream branch
-    tears down the whole cluster client.
+    """Returns a timeout-tolerant ``RedisCluster`` subclass when installed redis-py already
+    recovers node-level connections per-connection (8.x+), else builds the ``RedisCluster``
+    subclass with the per-node isolation fix for older versions whose upstream branch tears
+    down the whole cluster client.
 
-    ``cluster_node_class`` exists for dependency injection in tests; production callers
-    leave it unset and the installed ``ClusterNode`` is used.
+    ``cluster_node_class`` and ``base_cluster_class`` exist for dependency injection in tests;
+    production callers leave them unset and the installed redis-py classes are used.
 
     Imported lazily because this module is reachable from a base ``import litellm`` while
     redis is not a base dependency. Cheap to call repeatedly: the underlying redis
@@ -118,13 +122,68 @@ def get_litellm_async_redis_cluster_class(
     from redis.exceptions import TimeoutError as _RedisTimeoutError
 
     node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
+    base_class: Final = base_cluster_class if base_cluster_class is not None else _BaseAsyncRedisCluster
     if hasattr(node_class, "update_active_connections_for_reconnect"):
         verbose_logger.debug(
-            "redis-py %s recovers a node-level connection error per-connection upstream; "
-            "using the base RedisCluster without litellm's node-isolation override.",
+            "redis-py %s recovers node connections per-connection upstream; using "
+            "LiteLLM's timeout-tolerant RedisCluster wrapper.",
             redis.__version__,
         )
-        return _BaseAsyncRedisCluster
+
+        class LiteLLMAsyncRedisClusterTimeoutTolerant(
+            base_class  # pyright: ignore[reportGeneralTypeIssues, reportUntypedBaseClass]  # the injected base class is selected at runtime
+        ):
+            def __init__(
+                self,
+                *args: object,
+                **kwargs: object,  # kwargs-ok: passes redis-py's constructor kwargs through untouched
+            ) -> None:
+                self._litellm_initialize = False
+                self._litellm_reinit_requests = 0
+                self._litellm_tolerated_timeouts = 0
+                super().__init__(*args, **kwargs)
+                self._litellm_consecutive_timeouts: dict[  # mutable-ok: per-node counter updated on the command hot path
+                    str, int
+                ] = {}
+
+            @property
+            def _initialize(self) -> bool:
+                return self._litellm_initialize
+
+            @_initialize.setter
+            def _initialize(self, value: bool) -> None:
+                if value:
+                    self._litellm_reinit_requests += 1
+                self._litellm_initialize = value
+
+            async def _execute_command(
+                self,
+                target_node: _ClusterNodeAttrs,
+                *args: object,
+                **kwargs: object,  # kwargs-ok: matches redis-py's own command dispatch signature
+            ) -> object:
+                outstanding_before: Final = self._litellm_reinit_requests - self._litellm_tolerated_timeouts
+                pending_before: Final = self._litellm_initialize
+                try:
+                    result: Final = await super()._execute_command(target_node, *args, **kwargs)
+                except _RedisTimeoutError:
+                    timeouts: Final = self._litellm_consecutive_timeouts.get(target_node.name, 0) + 1
+                    if timeouts >= _CONSECUTIVE_TIMEOUTS_BEFORE_REINIT:
+                        self._litellm_consecutive_timeouts.pop(target_node.name, None)
+                        raise
+                    self._litellm_consecutive_timeouts[target_node.name] = timeouts
+                    self._litellm_tolerated_timeouts += 1
+                    if (
+                        not pending_before
+                        and self._litellm_reinit_requests - self._litellm_tolerated_timeouts == outstanding_before
+                    ):
+                        self._initialize = False
+                    raise
+                if self._litellm_consecutive_timeouts:
+                    self._litellm_consecutive_timeouts.pop(target_node.name, None)
+                return result
+
+        return LiteLLMAsyncRedisClusterTimeoutTolerant
 
     if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
         verbose_logger.warning(
diff --git a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py
index c16ceec8c31..9ebc3673ab6 100644
--- a/tests/test_litellm/caching/test_redis_cluster_node_isolation.py
+++ b/tests/test_litellm/caching/test_redis_cluster_node_isolation.py
@@ -5,20 +5,26 @@ CLIENT PAUSE) showed 100% of concurrent commands to the other two, untouched nod
 stalling for the full pause duration before this fix, and zero after -- these tests pin
 the same behavior at the unit level so it can run without a live Redis Cluster."""
 
+import asyncio
 from typing import TYPE_CHECKING
-from unittest.mock import AsyncMock
+from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
 from redis.exceptions import (
+    AskError,
     BusyLoadingError,
     ClusterDownError,
+    ClusterError,
     MaxConnectionsError,
     MovedError,
+    TryAgainError,
 )
 from redis.exceptions import (
     ConnectionError as RedisConnectionError,
 )
-from redis.exceptions import TimeoutError as RedisTimeoutError
+from redis.exceptions import (
+    TimeoutError as RedisTimeoutError,
+)
 
 from litellm.caching.redis_cluster_node_isolation import (
     get_litellm_async_redis_cluster_class,
@@ -39,10 +45,35 @@ class _NodeClassWithoutPerConnectionRecovery:
 class _FakeClusterNode:
     def __init__(self, name: str, raises: Exception | None = None, response: object = None) -> None:
         self.name = name
-        self.execute_command = AsyncMock(side_effect=raises, return_value=response)
+
+        async def execute_command(*args: object, **kwargs: object) -> object:
+            await asyncio.sleep(0)
+            if raises is not None:
+                raise raises
+            return response
+
+        self.execute_command = AsyncMock(side_effect=execute_command)
         self.disconnect = AsyncMock()
 
 
+class _Fake8xRedisCluster:
+    def __init__(self) -> None:
+        self._initialize = False
+
+    async def _execute_command(
+        self, target_node: _FakeClusterNode, *args: object, **kwargs: object
+    ) -> object:
+        try:
+            return await target_node.execute_command(*args, **kwargs)
+        except (RedisConnectionError, RedisTimeoutError):
+            self._initialize = True
+            await asyncio.sleep(0)
+            raise
+
+    async def aclose(self) -> None:
+        self._initialize = True
+
+
 class _FakeNodesManager:
     def __init__(self, node_to_return: _FakeClusterNode) -> None:
         self._moved_exception: object = None
@@ -68,31 +99,198 @@ def _build_cluster_instance() -> "_AsyncRedisClusterType":
     return instance
 
 
-def test_per_connection_recovery_redis_py_gets_the_unmodified_upstream_class() -> None:
-    """Regression (redis-py 8.x): when upstream ClusterNode already recovers a node-level
-    connection error per-connection, the factory must NOT install the copied override,
-    whose node.disconnect() also kills connections other coroutines are mid-operation on."""
-    from redis.asyncio.cluster import RedisCluster
-
+def _build_8x_cluster_instance() -> _Fake8xRedisCluster:
     cluster_cls = get_litellm_async_redis_cluster_class(
-        cluster_node_class=_NodeClassWithPerConnectionRecovery
+        cluster_node_class=_NodeClassWithPerConnectionRecovery,
+        base_cluster_class=_Fake8xRedisCluster,
+    )
+    return cluster_cls()
+
+
+def test_unverified_redis_version_logs_warning(caplog: pytest.LogCaptureFixture) -> None:
+    import redis
+
+    with patch.object(redis, "__version__", "8.0.1"):
+        get_litellm_async_redis_cluster_class(cluster_node_class=_NodeClassWithoutPerConnectionRecovery)
+
+    assert "not in the set this cluster-teardown-storm fix was verified against" in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_single_timeout_does_not_request_topology_reinit() -> None:
+    error = RedisTimeoutError("timeout")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command.side_effect = error
+    instance = _build_8x_cluster_instance()
+
+    with pytest.raises(RedisTimeoutError) as exc_info:
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert exc_info.value is error
+    assert instance._initialize is False
+
+
+@pytest.mark.asyncio
+async def test_connection_error_preserves_upstream_topology_reinit() -> None:
+    error = RedisConnectionError("connection error")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command.side_effect = error
+    instance = _build_8x_cluster_instance()
+
+    with pytest.raises(RedisConnectionError) as exc_info:
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert exc_info.value is error
+    assert instance._initialize is True
+
+
+@pytest.mark.asyncio
+async def test_three_consecutive_timeouts_request_topology_reinit_and_reset_counter() -> None:
+    errors = [
+        RedisTimeoutError("timeout-1"),
+        RedisTimeoutError("timeout-2"),
+        RedisTimeoutError("timeout-3"),
+    ]
+    fourth_error = RedisTimeoutError("timeout-4")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command.side_effect = [*errors, fourth_error]
+    instance = _build_8x_cluster_instance()
+
+    for error in errors:
+        with pytest.raises(RedisTimeoutError) as exc_info:
+            await instance._execute_command(target_node, "GET", "k")
+        assert exc_info.value is error
+
+    assert instance._initialize is True
+    instance._initialize = False
+
+    with pytest.raises(RedisTimeoutError) as exc_info:
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert exc_info.value is fourth_error
+    assert instance._initialize is False
+
+
+@pytest.mark.asyncio
+async def test_success_resets_consecutive_timeout_counter() -> None:
+    errors = [RedisTimeoutError("timeout-1"), RedisTimeoutError("timeout-2")]
+    final_error = RedisTimeoutError("timeout-3")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command.side_effect = [*errors, b"value", final_error]
+    instance = _build_8x_cluster_instance()
+
+    for error in errors:
+        with pytest.raises(RedisTimeoutError) as exc_info:
+            await instance._execute_command(target_node, "GET", "k")
+        assert exc_info.value is error
+        assert instance._initialize is False
+
+    result = await instance._execute_command(target_node, "GET", "k")
+    assert result == b"value"
+    assert instance._initialize is False
+
+    with pytest.raises(RedisTimeoutError) as exc_info:
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert exc_info.value is final_error
+    assert instance._initialize is False
+
+
+@pytest.mark.asyncio
+async def test_timeout_counters_are_per_node() -> None:
+    node_a_errors = [RedisTimeoutError("node-a-1"), RedisTimeoutError("node-a-2")]
+    node_b_error = RedisTimeoutError("node-b-1")
+    node_a = _FakeClusterNode("node-a")
+    node_b = _FakeClusterNode("node-b")
+    node_a.execute_command.side_effect = node_a_errors
+    node_b.execute_command.side_effect = node_b_error
+    instance = _build_8x_cluster_instance()
+
+    for target_node, error in (
+        (node_a, node_a_errors[0]),
+        (node_b, node_b_error),
+        (node_a, node_a_errors[1]),
+    ):
+        with pytest.raises(RedisTimeoutError) as exc_info:
+            await instance._execute_command(target_node, "GET", "k")
+        assert exc_info.value is error
+
+    assert instance._initialize is False
+
+
+@pytest.mark.asyncio
+async def test_timeout_does_not_clear_concurrent_topology_reinit_request() -> None:
+    error = RedisTimeoutError("timeout")
+    instance = _build_8x_cluster_instance()
+
+    async def request_reinit(*args: object, **kwargs: object) -> object:
+        await instance.aclose()
+        raise error
+
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command.side_effect = request_reinit
+
+    with pytest.raises(RedisTimeoutError) as exc_info:
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert exc_info.value is error
+    assert instance._initialize is True
+
+
+@pytest.mark.asyncio
+async def test_tolerated_timeout_does_not_erase_concurrent_connection_error_reinit() -> None:
+    instance = _build_8x_cluster_instance()
+    failing_node = _FakeClusterNode("node-a", raises=RedisConnectionError("gone"))
+    slow_node = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow"))
+
+    results = await asyncio.gather(
+        instance._execute_command(failing_node, "GET", "a"),
+        instance._execute_command(slow_node, "GET", "b"),
+        return_exceptions=True,
     )
 
-    assert cluster_cls is RedisCluster
+    assert isinstance(results[0], RedisConnectionError)
+    assert isinstance(results[1], RedisTimeoutError)
+    assert instance._initialize is True
 
 
-def test_pre_recovery_redis_py_still_gets_the_node_isolation_override() -> None:
-    """Old redis-py (5.x) responds to a node-level error with a full-cluster aclose(),
-    so those versions must keep litellm's per-node isolation override."""
-    from redis.asyncio.cluster import RedisCluster
+@pytest.mark.asyncio
+async def test_overlapping_tolerated_timeouts_do_not_request_topology_reinit() -> None:
+    instance = _build_8x_cluster_instance()
+    node_a = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow-a"))
+    node_b = _FakeClusterNode("node-b", raises=RedisTimeoutError("slow-b"))
 
-    cluster_cls = get_litellm_async_redis_cluster_class(
-        cluster_node_class=_NodeClassWithoutPerConnectionRecovery
+    results = await asyncio.gather(
+        instance._execute_command(node_a, "GET", "a"),
+        instance._execute_command(node_b, "GET", "b"),
+        return_exceptions=True,
     )
 
-    assert cluster_cls is not RedisCluster
-    assert issubclass(cluster_cls, RedisCluster)
-    assert "_execute_command" in cluster_cls.__dict__
+    assert all(isinstance(result, RedisTimeoutError) for result in results)
+    assert instance._initialize is False
+
+
+@pytest.mark.asyncio
+async def test_tolerated_timeout_does_not_clear_pending_reinit() -> None:
+    instance = _build_8x_cluster_instance()
+    instance._initialize = True
+    target_node = _FakeClusterNode("node-a", raises=RedisTimeoutError("slow"))
+
+    with pytest.raises(RedisTimeoutError):
+        await instance._execute_command(target_node, "GET", "k")
+
+    assert instance._initialize is True
+
+
+@pytest.mark.asyncio
+async def test_success_returns_value_without_topology_reinit() -> None:
+    target_node = _FakeClusterNode("node-a", response=b"value")
+    instance = _build_8x_cluster_instance()
+
+    result = await instance._execute_command(target_node, "GET", "k")
+
+    assert result == b"value"
+    assert instance._initialize is False
 
 
 @pytest.mark.asyncio
@@ -110,6 +308,48 @@ async def test_node_level_error_resets_only_that_node_not_the_whole_client(error
     instance.aclose.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_moved_error_retries_without_full_reinit_before_threshold() -> None:
+    moved_error = MovedError("1 127.0.0.1:7001")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command = AsyncMock(side_effect=[moved_error, b"value"])
+    instance = _build_cluster_instance()
+    instance.RedisClusterRequestTTL = 2
+    instance.nodes_manager = _FakeNodesManager(node_to_return=target_node)
+    instance._determine_slot = AsyncMock(return_value=0)
+
+    result = await instance._execute_command(target_node, "GET", "k")
+
+    assert result == b"value"
+    assert instance.nodes_manager._moved_exception is moved_error
+    instance.aclose.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_ask_error_sends_asking_and_retries_on_redirected_node() -> None:
+    ask_error = AskError("0 127.0.0.1:7001")
+    target_node = _FakeClusterNode("node-a")
+    target_node.execute_command = AsyncMock(side_effect=[ask_error, None, b"value"])
+    instance = _build_cluster_instance()
+    instance.RedisClusterRequestTTL = 2
+    instance.get_node = Mock(return_value=target_node)
+
+    result = await instance._execute_command(target_node, "GET", "k")
+
+    assert result == b"value"
+    instance.get_node.assert_called_once_with(node_name="127.0.0.1:7001")
+
+
+@pytest.mark.asyncio
+async def test_try_again_error_exhausts_ttl() -> None:
+    target_node = _FakeClusterNode("node-a", raises=TryAgainError("try again"))
+    instance = _build_cluster_instance()
+    instance.RedisClusterRequestTTL = 2
+
+    with pytest.raises(ClusterError):
+        await instance._execute_command(target_node, "GET", "k")
+
+
 @pytest.mark.asyncio
 async def test_successful_command_touches_neither_disconnect_nor_aclose() -> None:
     target_node = _FakeClusterNode("node-a", response=b"v")

From a06d63f99e9286d43716677c6a6766b2ab92d42a Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Thu, 3 Sep 2026 17:33:31 -0700
Subject: [PATCH 32/38] fix(logging): blocked requests no longer report
 guardrail_status=success in multi-guardrail configs (#39596)

* fix(logging): aggregate guardrail_status by severity across guardrail entries

A pre_call guardrail that passed (e.g. hide-secrets recording a mask)
appends its entry before a later guardrail's block, and the first-wins
reader reported the blocked request as guardrail_status=success in
StandardLoggingPayload.status_fields. Take the most severe status
across all entries instead: guardrail_intervened >
guardrail_failed_to_respond > success > not_run.

* refactor(logging): express guardrail status severity as an immutable order

Replace the precedence dict and rebinding loop with a severity-ordered
tuple and a max() aggregation, per the repo's no-mutation and
mutable-collection lint gates; parametrize the severity test cases.
No behavior change.

* style(logging): apply ruff format to entries binding
---
 litellm/litellm_core_utils/litellm_logging.py | 32 +++++--
 .../test_tracing_guardrails.py                | 88 +++++++++++++++++++
 2 files changed, 111 insertions(+), 9 deletions(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 51d6858b7dc..17a19f05fa3 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -10,7 +10,7 @@ import subprocess
 import sys
 import time
 import traceback
-from collections.abc import Callable, Mapping, Sequence
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from datetime import datetime as dt_object
 from functools import lru_cache
 from types import MappingProxyType, TracebackType
@@ -5878,14 +5878,28 @@ def _get_status_fields(
     #########################################################
     # Map - guardrail_information.guardrail_status to guardrail_status
     #########################################################
-    guardrail_status: GuardrailStatus = "not_run"
-    if guardrail_information and isinstance(guardrail_information, list):
-        for information in guardrail_information:
-            if isinstance(information, dict):
-                raw_status = information.get("guardrail_status", "not_run")
-                if raw_status != "not_run":
-                    guardrail_status = GUARDRAIL_STATUS_MAP.get(raw_status, "not_run")
-                    break
+    # Severity order, least severe first. The status aggregates across ALL
+    # guardrail entries rather than taking the first non-"not_run" one: a
+    # pre_call guardrail that passed (e.g. a mask) records its entry before a
+    # later guardrail's block, and first-wins would report a blocked request
+    # as "success".
+    GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = (
+        "not_run",
+        "success",
+        "guardrail_failed_to_respond",
+        "guardrail_intervened",
+    )
+    entries: Final[Sequence[object]] = guardrail_information if isinstance(guardrail_information, list) else ()
+    raw_statuses: Final[Iterator[object]] = (
+        entry.get("guardrail_status", "not_run") for entry in entries if isinstance(entry, dict)
+    )
+    # A guardrail is free to write any value here, and an unhashable one would
+    # raise TypeError on the mapping lookup and drop the whole payload.
+    guardrail_status: Final[GuardrailStatus] = max(
+        (GUARDRAIL_STATUS_MAP.get(raw_status, "not_run") for raw_status in raw_statuses if isinstance(raw_status, str)),
+        key=GUARDRAIL_STATUS_SEVERITY.index,
+        default="not_run",
+    )
 
     return StandardLoggingPayloadStatusFields(llm_api_status=llm_api_status, guardrail_status=guardrail_status)
 
diff --git a/tests/guardrails_tests/test_tracing_guardrails.py b/tests/guardrails_tests/test_tracing_guardrails.py
index bd8b7bad33f..ac85803ba39 100644
--- a/tests/guardrails_tests/test_tracing_guardrails.py
+++ b/tests/guardrails_tests/test_tracing_guardrails.py
@@ -806,3 +806,91 @@ def test_guardrail_status_fields_computation():
     )
     assert status_fields_no_guardrail.get("llm_api_status") == "success"
     assert status_fields_no_guardrail.get("guardrail_status") == "not_run"
+
+
+@pytest.mark.parametrize(
+    "status, guardrail_information, expected_guardrail_status",
+    [
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="pre_call_success_before_blocker",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "guardrail_intervened"},
+                {"guardrail_status": "success"},
+            ],
+            "guardrail_intervened",
+            id="blocker_before_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "guardrail_failed_to_respond"},
+            ],
+            "guardrail_failed_to_respond",
+            id="failure_outranks_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "guardrail_failed_to_respond"},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="intervention_outranks_failure",
+        ),
+        pytest.param(
+            "success",
+            [
+                {"guardrail_status": "success"},
+                {"guardrail_status": "success"},
+            ],
+            "success",
+            id="all_success_stays_success",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": "some_new_status"},
+                {"guardrail_status": "blocked"},
+            ],
+            "guardrail_intervened",
+            id="unknown_status_does_not_mask_blocker",
+        ),
+        pytest.param(
+            "failure",
+            [
+                {"guardrail_status": {"unhashable": True}},
+                {"guardrail_status": "guardrail_intervened"},
+            ],
+            "guardrail_intervened",
+            id="unhashable_status_is_skipped",
+        ),
+    ],
+)
+def test_guardrail_status_fields_severity_across_entries(
+    status, guardrail_information, expected_guardrail_status
+):
+    """
+    A blocked request must never be reported as a guardrail success.
+
+    With multiple guardrails on one request (e.g. a pre_call mask that passes,
+    then a post_call guardrail that blocks), entries are recorded in execution
+    order, so the earlier "success" entry must not shadow the later
+    "guardrail_intervened" entry: the aggregate takes the most severe status,
+    regardless of entry order.
+    """
+    from litellm.litellm_core_utils.litellm_logging import _get_status_fields
+
+    fields = _get_status_fields(
+        status=status, guardrail_information=guardrail_information, error_str=None
+    )
+    assert fields.get("guardrail_status") == expected_guardrail_status

From bd10977a9ab84fefd020cc8f3ec235213049ff24 Mon Sep 17 00:00:00 2001
From: tin-berri <tin@berri.ai>
Date: Thu, 3 Sep 2026 17:44:07 -0700
Subject: [PATCH 33/38] fix(snowflake): normalize Cortex Claude request shapes
 (#39453)

* fix(snowflake): normalize Cortex Claude request shapes

Co-authored-by: Kamron Javaherpour <kamron@kargo.com>

Co-authored-by: Oleksandr Kononov <oleks.konov@kargo.com>

* style(snowflake): format Cortex request transformations

* fix(snowflake): annotate Cortex wire payloads

* fix(snowflake): route Cortex content through the shared Anthropic converters

* fix(snowflake): surface Cortex prompt-cache usage and thinking blocks

Parse Cortex's Anthropic-dialect responses and SSE with Anthropic's own parser so cache_creation/cache_read counts, thinking blocks and signatures reach the caller. Restore thinking for every Claude model: Cortex documents extended thinking broadly and only adaptive thinking is 4.6-gated.

* fix(snowflake): echo signed thinking blocks on every assistant turn

The reference converter extends signed thinking blocks on each assistant turn, not just tool-call turns, so a replayed thinking-plus-text response keeps its signed block. Content-less thinking turns send no empty text block.

* fix(snowflake): preserve thinking list content

---------

Co-authored-by: Oleksandr Kononov <oleks.konov@kargo.com>
---
 litellm/llms/snowflake/chat/transformation.py | 383 ++++++++--------
 .../test_snowflake_chat_transformation.py     | 416 ++++++++++++++++--
 .../test_snowflake_native_endpoints.py        |  60 ++-
 3 files changed, 652 insertions(+), 207 deletions(-)

diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py
index 0968185b084..c64fc583edc 100644
--- a/litellm/llms/snowflake/chat/transformation.py
+++ b/litellm/llms/snowflake/chat/transformation.py
@@ -8,23 +8,31 @@ Routes to native Cortex REST API endpoints based on model:
 Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api
 """
 
+import copy
 import json
+import re
 from collections.abc import Mapping, Sequence
 from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict
 
 import httpx
 from typing_extensions import ReadOnly
 
-from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk
+from litellm.litellm_core_utils.prompt_templates.factory import (
+    anthropic_process_openai_file_message,
+    convert_to_anthropic_tool_result,
+    create_anthropic_image_param,
+    select_anthropic_content_block_type_for_file,
+)
+from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser
+from litellm.llms.anthropic.chat.transformation import AnthropicConfig
+from litellm.llms.anthropic.common_utils import normalize_cache_control_in_anthropic_payload
+from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolMessage
 from litellm.types.utils import (
-    ChatCompletionMessageToolCall,
-    ChatCompletionUsageBlock,
     Choices,
-    Function,
     GenericStreamingChunk,
     Message,
     ModelResponse,
-    Usage,
+    ModelResponseStream,
 )
 
 from ...base_llm.base_model_iterator import BaseModelResponseIterator
@@ -93,6 +101,103 @@ def _is_claude_model(model: str) -> bool:
     return any(name.startswith(p) for p in _CLAUDE_MODEL_PREFIXES)
 
 
+def _convert_image_url_to_anthropic(block: Mapping[str, object]) -> object:
+    """One OpenAI ``image_url`` block in the native shape Cortex accepts.
+
+    Cortex documents base64 sources only, so remote URLs are inlined the way every
+    other base64-only Anthropic dialect (Bedrock invoke, Vertex) inlines them, and
+    pdf/text data URIs become document blocks rather than malformed image blocks.
+    """
+    image_url: Final = block.get("image_url")
+    url: Final = image_url if isinstance(image_url, str) else _image_url_field(image_url, "url")
+    if not url:
+        return block
+
+    converted: Final = (
+        anthropic_process_openai_file_message({"type": "file", "file": {"file_data": url}})
+        if select_anthropic_content_block_type_for_file(_data_uri_media_type(url)) == "document"
+        else create_anthropic_image_param(
+            image_url if isinstance(image_url, dict) else url,  # mutable-ok: caller's JSON block
+            format=_image_url_field(image_url, "format"),
+            is_bedrock_invoke=True,
+        )
+    )
+    cache_control: Final = block.get("cache_control")
+    if cache_control is None:
+        return converted
+    return {**converted, "cache_control": cache_control}  # mutable-ok: JSON wire block
+
+
+def _image_url_field(image_url: object, key: str) -> str | None:
+    value: Final = image_url.get(key) if isinstance(image_url, dict) else None
+    return value if isinstance(value, str) else None
+
+
+def _data_uri_media_type(url: str) -> str:
+    match: Final = re.match(r"data:([^;,]+)", url)
+    return match.group(1) if match else ""
+
+
+def _convert_image_url_blocks_to_anthropic(content: object) -> object:
+    if not isinstance(content, list):
+        return content
+    return [  # mutable-ok: JSON wire blocks
+        _convert_image_url_to_anthropic(block)
+        if isinstance(block, Mapping) and block.get("type") == "image_url"
+        else block
+        for block in content
+    ]
+
+
+def _convert_tool_result_to_anthropic(
+    content: object, tool_call_id: str, cache_control: object
+) -> Mapping[str, object]:
+    """The Anthropic ``tool_result`` block for one OpenAI tool message.
+
+    Delegating to the shared converter keeps image, document and per-block cache
+    breakpoints identical to every other Anthropic dialect; only the plain-string
+    and non-list shapes it does not model are handled here.
+    """
+    if not isinstance(content, list):
+        plain: Final[dict[str, object]] = {  # mutable-ok: JSON wire block
+            "type": "tool_result",
+            "tool_use_id": tool_call_id,
+            "content": content if isinstance(content, str) else json.dumps(content),
+        }
+        return {**plain, "cache_control": cache_control} if cache_control is not None else plain
+    converted: Final = convert_to_anthropic_tool_result(
+        ChatCompletionToolMessage(role="tool", tool_call_id=tool_call_id, content=content),
+        force_base64=True,
+    )
+    if cache_control is None:
+        return converted
+    return {**converted, "cache_control": cache_control}  # mutable-ok: JSON wire block
+
+
+def _signed_thinking_blocks(msg: object) -> list[dict[str, object]]:  # mutable-ok: JSON wire blocks
+    """The assistant turn's thinking blocks that can legally be echoed back.
+
+    Only signed blocks round-trip: Cortex rejects a thinking block whose signature is
+    missing, which is what an unsigned block from a non-thinking turn would produce.
+    """
+    blocks: Final = msg.get("thinking_blocks") if isinstance(msg, dict) else getattr(msg, "thinking_blocks", None)
+    if not isinstance(blocks, list):
+        return []  # mutable-ok: JSON wire blocks
+    return [  # mutable-ok: JSON wire blocks
+        dict(block)
+        for block in blocks
+        if isinstance(block, Mapping) and (block.get("signature") or block.get("type") == "redacted_thinking")
+    ]
+
+
+def _clean_input_schema(schema: object) -> object:  # mutable-ok: JSON schema copy
+    return (
+        {key: value for key, value in schema.items() if key != "$schema"}
+        if isinstance(schema, Mapping)
+        else schema  # mutable-ok: JSON schema copy
+    )  # mutable-ok: JSON schema copy
+
+
 class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
     """
     Snowflake Cortex REST API — unified provider.
@@ -178,7 +283,7 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                 if "description" in func:
                     anthropic_tool["description"] = func["description"]
                 if "parameters" in func:
-                    anthropic_tool["input_schema"] = func["parameters"]
+                    anthropic_tool["input_schema"] = _clean_input_schema(func["parameters"])
                 else:
                     anthropic_tool["input_schema"] = {
                         "type": "object",
@@ -186,10 +291,16 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                     }
                 anthropic_tools.append(anthropic_tool)
             else:
-                anthropic_tools.append(tool)
+                anthropic_tools.append(
+                    {**tool, "input_schema": _clean_input_schema(tool["input_schema"])}  # mutable-ok: JSON wire tool
+                    if "input_schema" in tool
+                    else tool
+                )
         return anthropic_tools
 
-    def _extract_system_and_messages(self, messages: list[AllMessageValues]) -> tuple[str | None, list[dict]]:
+    def _extract_system_and_messages(  # mutable-ok: JSON wire messages
+        self, messages: list[AllMessageValues]
+    ) -> tuple[list[dict] | None, list[dict]]:
         """
         Split messages into system prompt and conversation turns for Anthropic format.
 
@@ -197,26 +308,39 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         - assistant messages with tool_calls → tool_use content blocks
         - tool role messages → user role with tool_result content blocks
         """
-        system_parts: Final[list[str]] = []
-        conversation: Final[list[dict]] = []
+        system_parts: Final[list[dict]] = []  # mutable-ok: JSON wire messages
+        conversation: Final[list[dict]] = []  # mutable-ok: JSON wire messages
 
         for msg in messages:
             if isinstance(msg, dict):
                 role = msg.get("role", "")
                 content: Any = msg.get("content", "")
+                msg_cache_control: object = msg.get("cache_control")
             else:
                 role = getattr(msg, "role", "")
                 content = getattr(msg, "content", "")
+                msg_cache_control = getattr(msg, "cache_control", None)
 
             if role == "system":
                 if isinstance(content, str) and content:
-                    system_parts.append(content)
+                    system_parts.append({"type": "text", "text": content})  # mutable-ok: JSON wire system block
                 elif isinstance(content, list):
-                    system_parts.append("\n".join(b.get("text", "") for b in content if b.get("type") == "text"))
+                    system_parts.extend(
+                        {  # mutable-ok: JSON wire system block
+                            "type": "text",
+                            "text": block.get("text", ""),
+                            **(
+                                {"cache_control": block["cache_control"]} if "cache_control" in block else {}
+                            ),  # mutable-ok: JSON wire block
+                        }
+                        for block in content
+                        if isinstance(block, Mapping) and block.get("type") == "text"
+                    )
             elif role == "assistant":
                 tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else getattr(msg, "tool_calls", None)
+                thinking_blocks = _signed_thinking_blocks(msg)
                 if tool_calls:
-                    content_blocks: list[dict[str, object]] = []
+                    content_blocks: list[dict[str, object]] = list(thinking_blocks)  # mutable-ok: JSON wire blocks
                     if content:
                         content_blocks.append({"type": "text", "text": content})
                     for tc in tool_calls:
@@ -239,18 +363,26 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                             }
                         )
                     conversation.append({"role": "assistant", "content": content_blocks})
+                elif thinking_blocks:
+                    thinking_content = (
+                        [
+                            *thinking_blocks,
+                            *copy.deepcopy(content),
+                        ]
+                        if isinstance(content, list)
+                        else [*thinking_blocks, *([{"type": "text", "text": content}] if content else [])]
+                    )  # rebind-ok: loop-local normalized content
+                    conversation.append({"role": "assistant", "content": thinking_content})
                 else:
                     conversation.append({"role": "assistant", "content": content})
             elif role == "tool":
-                tool_call_id = (
+                tool_call_id_value = (
                     msg.get("tool_call_id", "") if isinstance(msg, dict) else getattr(msg, "tool_call_id", "")
                 )
-                tool_content = content if isinstance(content, str) else json.dumps(content)
-                tool_result_block = {
-                    "type": "tool_result",
-                    "tool_use_id": tool_call_id,
-                    "content": tool_content,
-                }
+                tool_call_id = (
+                    tool_call_id_value if isinstance(tool_call_id_value, str) else ""
+                )  # rebind-ok: normalized loop value
+                tool_result_block = _convert_tool_result_to_anthropic(content, tool_call_id, msg_cache_control)
                 if (
                     conversation
                     and conversation[-1]["role"] == "user"
@@ -260,11 +392,18 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
                 ):
                     conversation[-1]["content"].append(tool_result_block)
                 else:
-                    conversation.append({"role": "user", "content": [tool_result_block]})
+                    conversation.append(
+                        {"role": "user", "content": [tool_result_block]}  # mutable-ok: JSON wire message
+                    )  # mutable-ok: JSON wire message
             else:
-                conversation.append({"role": role, "content": content})
+                conversation.append(  # mutable-ok: JSON wire message
+                    {  # mutable-ok: JSON wire message
+                        "role": role,
+                        "content": _convert_image_url_blocks_to_anthropic(content),
+                    }  # mutable-ok: JSON wire message
+                )
 
-        system: Final[str | None] = "\n\n".join(system_parts) if system_parts else None
+        system: Final[list[dict] | None] = system_parts if system_parts else None  # mutable-ok: JSON wire messages
         return system, conversation
 
     def transform_request(
@@ -339,7 +478,9 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         extra_body: dict,
     ) -> dict:
         """Anthropic Messages format for /messages endpoint."""
-        system, conversation = self._extract_system_and_messages(messages)
+        passthrough_system: Final = optional_params.pop("system", None)
+        extracted_system, conversation = self._extract_system_and_messages(messages)
+        system: Final = passthrough_system if passthrough_system is not None else extracted_system
 
         if "tools" in optional_params:
             optional_params["tools"] = self._transform_tools_to_anthropic(optional_params["tools"])
@@ -353,16 +494,19 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
 
         model_name: Final = model.removeprefix("snowflake/")
 
-        body: Final[dict[str, object]] = {
-            "model": model_name,
-            "messages": conversation,
-            "stream": stream,
-            **optional_params,
-            **extra_body,
-        }
-
+        body: Final[dict[str, object]] = normalize_cache_control_in_anthropic_payload(  # mutable-ok: JSON wire body
+            {  # mutable-ok: JSON wire body
+                "model": model_name,
+                "messages": conversation,
+                "stream": stream,
+                **optional_params,
+                **extra_body,  # mutable-ok: JSON wire body
+            }
+        )
         if system is not None:
-            body["system"] = system
+            body["system"] = normalize_cache_control_in_anthropic_payload(  # mutable-ok: JSON wire payload
+                {"system": system}  # mutable-ok: JSON wire payload
+            )["system"]
 
         if "max_tokens" not in body:
             body["max_tokens"] = 4096  # reasonable default; Anthropic API max varies by model
@@ -435,23 +579,10 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             additional_args={"complete_input_dict": request_data},
         )
 
-        text_content = ""
-        tool_calls: Final = []
-
-        for block in response_json.get("content", []):
-            if block.get("type") == "text":
-                text_content += block.get("text", "")
-            elif block.get("type") == "tool_use":
-                tool_calls.append(
-                    ChatCompletionMessageToolCall(
-                        id=block.get("id", ""),
-                        type="function",
-                        function=Function(
-                            name=block.get("name", ""),
-                            arguments=json.dumps(block.get("input", {})),
-                        ),
-                    )
-                )
+        anthropic_config: Final = AnthropicConfig()
+        text_content, _, thinking_blocks, reasoning_content, tool_calls, _, _, _ = (
+            anthropic_config.extract_response_content(completion_response=dict(response_json))
+        )
 
         _stop_reason_map: Final = {
             "end_turn": "stop",
@@ -461,9 +592,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
         }
         finish_reason: Final = _stop_reason_map.get(response_json.get("stop_reason", "end_turn"), "stop")
 
-        message: Final = Message(content=text_content or None, role="assistant")
-        if tool_calls:
-            message.tool_calls = tool_calls
+        message: Final = Message(
+            content=text_content or None,
+            role="assistant",
+            tool_calls=tool_calls or None,
+            thinking_blocks=thinking_blocks,
+            reasoning_content=reasoning_content,
+        )
 
         choice: Final = Choices(
             finish_reason=finish_reason,
@@ -471,11 +606,13 @@ class SnowflakeConfig(SnowflakeBaseConfig, OpenAIGPTConfig):
             message=message,
         )
 
-        usage_data: Final = response_json.get("usage", {})
-        usage: Final = Usage(
-            prompt_tokens=usage_data.get("input_tokens", 0),
-            completion_tokens=usage_data.get("output_tokens", 0),
-            total_tokens=usage_data.get("input_tokens", 0) + usage_data.get("output_tokens", 0),
+        # Cortex reports prompt-cache creation/read counts alongside input_tokens; the
+        # shared calculator folds them into prompt_tokens_details so cached input is
+        # visible and billed at its own rate.
+        usage: Final = anthropic_config.calculate_usage(
+            usage_object=response_json.get("usage", {}),
+            reasoning_content=reasoning_content,
+            completion_response=dict(response_json),
         )
 
         model_response.choices = [choice]
@@ -516,15 +653,19 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
         json_mode: bool | None = False,
     ):
         super().__init__(streaming_response=streaming_response, sync_stream=sync_stream)
-        self._tool_index = 0
-        self._tool_id = ""
-        self._tool_name = ""
-        self._input_tokens = 0
+        # Cortex streams the Anthropic SSE dialect on /messages, so its events are parsed
+        # by Anthropic's own parser: thinking deltas, signatures and prompt-cache usage
+        # all arrive the way they do on every other Anthropic-dialect provider.
+        self._anthropic_parser: Final = AnthropicStreamParser(
+            streaming_response=streaming_response,
+            sync_stream=sync_stream,
+            json_mode=json_mode,
+        )
 
-    def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
+    def chunk_parser(self, chunk: dict) -> GenericStreamingChunk | ModelResponseStream:
         if "choices" in chunk:
             return self._parse_openai_chunk(chunk)
-        return self._parse_anthropic_chunk(chunk)
+        return self._anthropic_parser.chunk_parser(chunk)
 
     def _parse_openai_chunk(self, chunk: dict) -> GenericStreamingChunk:
         choices: Final = chunk.get("choices", [])
@@ -566,117 +707,3 @@ class SnowflakeStreamingHandler(BaseModelResponseIterator):
             index=choice.get("index", 0),
             tool_use=tool_use,
         )
-
-    def _parse_anthropic_chunk(self, chunk: dict) -> GenericStreamingChunk:
-        event_type: Final = chunk.get("type", "")
-
-        if event_type == "message_start":
-            message: Final = chunk.get("message", {})
-            usage_data = message.get("usage", {})
-            self._input_tokens = usage_data.get("input_tokens", 0)
-            return GenericStreamingChunk(
-                text="",
-                is_finished=False,
-                finish_reason="",
-                usage=None,
-                index=0,
-                tool_use=None,
-            )
-
-        elif event_type == "content_block_delta":
-            delta = chunk.get("delta", {})
-            delta_type: Final = delta.get("type", "")
-
-            if delta_type == "text_delta":
-                return GenericStreamingChunk(
-                    text=delta.get("text", ""),
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=None,
-                )
-            elif delta_type == "input_json_delta":
-                return GenericStreamingChunk(
-                    text="",
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=ChatCompletionToolCallChunk(
-                        id=self._tool_id,
-                        type="function",
-                        function={
-                            "name": self._tool_name,
-                            "arguments": delta.get("partial_json", ""),
-                        },
-                        index=self._tool_index,
-                    ),
-                )
-
-        elif event_type == "content_block_start":
-            content_block: Final = chunk.get("content_block", {})
-            if content_block.get("type") == "tool_use":
-                self._tool_id = content_block.get("id", "")
-                self._tool_name = content_block.get("name", "")
-                self._tool_index = chunk.get("index", 0)
-                return GenericStreamingChunk(
-                    text="",
-                    is_finished=False,
-                    finish_reason="",
-                    usage=None,
-                    index=chunk.get("index", 0),
-                    tool_use=ChatCompletionToolCallChunk(
-                        id=self._tool_id,
-                        type="function",
-                        function={"name": self._tool_name, "arguments": ""},
-                        index=self._tool_index,
-                    ),
-                )
-
-        elif event_type == "message_delta":
-            delta = chunk.get("delta", {})
-            stop_reason: Final = delta.get("stop_reason", "")
-            usage_data = chunk.get("usage", {})
-            _stop_map: Final = {
-                "end_turn": "stop",
-                "max_tokens": "length",
-                "tool_use": "tool_calls",
-                "stop_sequence": "stop",
-            }
-            usage = None
-            if usage_data or self._input_tokens:
-                output_t: Final = usage_data.get("output_tokens", 0)
-                input_t: Final = self._input_tokens or usage_data.get("input_tokens", 0)
-                usage = ChatCompletionUsageBlock(
-                    prompt_tokens=input_t,
-                    completion_tokens=output_t,
-                    total_tokens=input_t + output_t,
-                )
-            return GenericStreamingChunk(
-                text="",
-                is_finished=True,
-                finish_reason=_stop_map.get(stop_reason, "stop"),
-                usage=usage,
-                index=0,
-                tool_use=None,
-            )
-
-        elif event_type == "message_stop":
-            return GenericStreamingChunk(
-                text="",
-                is_finished=True,
-                finish_reason="stop",
-                usage=None,
-                index=0,
-                tool_use=None,
-            )
-
-        return GenericStreamingChunk(
-            text="",
-            is_finished=False,
-            finish_reason="",
-            usage=None,
-            index=0,
-            tool_use=None,
-        )
diff --git a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
index a182656e4a8..25a961c3413 100644
--- a/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
+++ b/tests/test_litellm/llms/snowflake/chat/test_snowflake_chat_transformation.py
@@ -17,7 +17,7 @@ import pytest
 import litellm
 from litellm import completion, acompletion
 from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
-from litellm.llms.snowflake.chat.transformation import SnowflakeConfig
+from litellm.llms.snowflake.chat.transformation import SnowflakeConfig, SnowflakeStreamingHandler
 from litellm.types.utils import ModelResponse
 
 
@@ -114,8 +114,7 @@ class TestSnowflakeToolTransformation:
             )
 
             assert transformed_request["tool_choice"] == value, (
-                f"tool_choice='{value}' should pass through unchanged, "
-                f"got {transformed_request['tool_choice']}"
+                f"tool_choice='{value}' should pass through unchanged, got {transformed_request['tool_choice']}"
             )
 
     def test_transform_response_with_tool_calls(self):
@@ -159,9 +158,7 @@ class TestSnowflakeToolTransformation:
             headers={"Content-Type": "application/json"},
         )
 
-        model_response = ModelResponse(
-            choices=[litellm.Choices(index=0, message=litellm.Message())]
-        )
+        model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
 
         logging_obj = MagicMock()
 
@@ -232,9 +229,7 @@ class TestSnowflakeToolTransformation:
             headers={"Content-Type": "application/json"},
         )
 
-        model_response = ModelResponse(
-            choices=[litellm.Choices(index=0, message=litellm.Message())]
-        )
+        model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
 
         logging_obj = MagicMock()
 
@@ -280,9 +275,7 @@ class TestSnowflakeToolTransformation:
             headers={"Content-Type": "application/json"},
         )
 
-        model_response = ModelResponse(
-            choices=[litellm.Choices(index=0, message=litellm.Message())]
-        )
+        model_response = ModelResponse(choices=[litellm.Choices(index=0, message=litellm.Message())])
 
         logging_obj = MagicMock()
 
@@ -300,10 +293,7 @@ class TestSnowflakeToolTransformation:
 
         # Verify standard response works
         assert isinstance(result, ModelResponse)
-        assert (
-            result.choices[0].message.content
-            == "Hello! I'm doing well, thank you for asking."
-        )
+        assert result.choices[0].message.content == "Hello! I'm doing well, thank you for asking."
 
     def test_get_supported_openai_params_includes_tools(self):
         """
@@ -318,6 +308,385 @@ class TestSnowflakeToolTransformation:
         assert "max_tokens" in supported_params
 
 
+class TestSnowflakeCortexClaudeFixes:
+    def setup_method(self):
+        self.config = SnowflakeConfig()
+
+    @staticmethod
+    def _transform(messages, optional_params=None):
+        return SnowflakeConfig().transform_request(
+            model="snowflake/claude-sonnet-4-6",
+            messages=messages,
+            optional_params=optional_params or {},
+            litellm_params={},
+            headers={},
+        )
+
+    def test_thinking_is_offered_on_every_claude_model(self):
+        """Cortex documents extended thinking (budget_tokens) for Claude generally, so a
+        4.6-only gate would silently drop it on the models that do support it."""
+        for model in (
+            "snowflake/claude-sonnet-4-6",
+            "snowflake/claude-sonnet-4-5",
+            "snowflake/claude-3-7-sonnet",
+            "snowflake/claude-4-opus",
+        ):
+            assert "thinking" in self.config.get_supported_openai_params(model), model
+        assert "thinking" not in self.config.get_supported_openai_params("snowflake/llama3.1-70b")
+
+    def test_system_blocks_preserve_cache_control_and_strip_ttl(self):
+        body = self._transform(
+            [
+                {
+                    "role": "system",
+                    "content": [
+                        {
+                            "type": "text",
+                            "text": "You are helpful",
+                            "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                        }
+                    ],
+                },
+                {"role": "user", "content": "hi"},
+            ]
+        )
+        assert body["system"] == [{"type": "text", "text": "You are helpful", "cache_control": {"type": "ephemeral"}}]
+
+    def test_direct_system_param_is_normalized(self):
+        body = self._transform(
+            [{"role": "user", "content": "hi"}],
+            {"system": [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral", "ttl": "1h"}}]},
+        )
+        assert body["system"] == [{"type": "text", "text": "direct", "cache_control": {"type": "ephemeral"}}]
+
+    def test_message_and_tool_cache_control_are_normalized(self):
+        body = self._transform(
+            [
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": "hi", "cache_control": {"type": "ephemeral", "ttl": "1h"}}],
+                }
+            ],
+            {
+                "tools": [
+                    {
+                        "name": "f",
+                        "input_schema": {"type": "object", "properties": {}},
+                        "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                    }
+                ]
+            },
+        )
+        assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
+        assert body["tools"][0]["cache_control"] == {"type": "ephemeral"}
+
+    def test_extra_body_message_override_is_normalized(self):
+        body = self._transform(
+            [{"role": "user", "content": "original"}],
+            {
+                "extra_body": {
+                    "messages": [
+                        {
+                            "role": "user",
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": "override",
+                                    "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                                }
+                            ],
+                        }
+                    ]
+                }
+            },
+        )
+        assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral"}
+
+    def test_image_blocks_are_converted_to_anthropic_source(self):
+        body = self._transform(
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {
+                            "type": "image_url",
+                            "image_url": {"url": "data:image/png;base64,ZmFrZQ==", "format": "image/jpeg"},
+                        }
+                    ],
+                }
+            ]
+        )
+        assert body["messages"][0]["content"] == [
+            {"type": "image", "source": {"type": "base64", "media_type": "image/jpeg", "data": "ZmFrZQ=="}}
+        ]
+
+    def test_tool_result_image_list_is_converted(self):
+        body = self._transform(
+            [
+                {"role": "user", "content": "look"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
+                    ],
+                },
+                {
+                    "role": "tool",
+                    "tool_call_id": "call_1",
+                    "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}}],
+                },
+            ]
+        )
+        assert body["messages"][2]["content"][0]["content"] == [
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}}
+        ]
+
+    def test_tool_result_preserves_cache_control(self):
+        """A cache breakpoint the bridge puts on a tool message must survive onto the tool_result."""
+        for tool_content in ("done", [{"type": "text", "text": "done"}]):
+            body = self._transform(
+                [
+                    {"role": "user", "content": "look"},
+                    {
+                        "role": "tool",
+                        "tool_call_id": "call_1",
+                        "content": tool_content,
+                        "cache_control": {"type": "ephemeral", "ttl": "1h"},
+                    },
+                ]
+            )
+            tool_result = body["messages"][1]["content"][0]
+            assert tool_result["cache_control"] == {"type": "ephemeral"}, tool_content
+
+    def test_pdf_data_uri_becomes_a_document_block(self):
+        """A bridged pdf data URI is a document block; forwarding it as an image is malformed."""
+        body = self._transform(
+            [
+                {
+                    "role": "user",
+                    "content": [
+                        {"type": "image_url", "image_url": {"url": "data:application/pdf;base64,ZmFrZQ=="}},
+                    ],
+                }
+            ]
+        )
+        assert body["messages"][0]["content"] == [
+            {
+                "type": "document",
+                "source": {"type": "base64", "media_type": "application/pdf", "data": "ZmFrZQ=="},
+            }
+        ]
+
+    def test_multipart_tool_result_preserves_text_and_converts_image(self):
+        body = self._transform(
+            [
+                {"role": "user", "content": "look"},
+                {
+                    "role": "tool",
+                    "tool_call_id": "call_1",
+                    "content": [
+                        {"type": "text", "text": "first"},
+                        {"type": "image_url", "image_url": {"url": "data:image/png;base64,ZmFrZQ=="}},
+                        {"type": "text", "text": "last"},
+                    ],
+                },
+            ]
+        )
+        assert body["messages"][1]["content"][0]["content"] == [
+            {"type": "text", "text": "first"},
+            {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "ZmFrZQ=="}},
+            {"type": "text", "text": "last"},
+        ]
+
+    def test_plain_text_tool_result_remains_string(self):
+        body = self._transform(
+            [{"role": "user", "content": "look"}, {"role": "tool", "tool_call_id": "call_1", "content": "done"}]
+        )
+        assert body["messages"][1]["content"][0]["content"] == "done"
+
+    def test_anthropic_tool_schema_strips_only_top_level_schema_key(self):
+        tools = [
+            {
+                "name": "f",
+                "input_schema": {"$schema": "schema", "type": "object", "properties": {"$schema": {"type": "string"}}},
+            }
+        ]
+        body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
+        schema = body["tools"][0]["input_schema"]
+        assert "$schema" not in schema
+        assert "$schema" in schema["properties"]
+
+    def test_tool_schema_strips_only_top_level_schema_key(self):
+        tools = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "f",
+                    "parameters": {
+                        "$schema": "https://json-schema.org/draft/2020-12/schema",
+                        "type": "object",
+                        "properties": {"$schema": {"type": "string"}},
+                    },
+                },
+            }
+        ]
+        body = self._transform([{"role": "user", "content": "hi"}], {"tools": tools})
+        schema = body["tools"][0]["input_schema"]
+        assert "$schema" not in schema
+        assert "$schema" in schema["properties"]
+
+    def test_streaming_tool_identity_is_emitted_only_on_start(self):
+        handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
+        start = handler.chunk_parser(
+            {
+                "type": "content_block_start",
+                "index": 0,
+                "content_block": {"type": "tool_use", "id": "tool_1", "name": "read"},
+            }
+        )
+        first_delta = handler.chunk_parser(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "input_json_delta", "partial_json": '{"path":'},
+            }
+        )
+        second_delta = handler.chunk_parser(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "input_json_delta", "partial_json": '"/tmp"}'},
+            }
+        )
+
+        def _tool_call(chunk):
+            return chunk.choices[0].delta.tool_calls[0]
+
+        assert _tool_call(start).id == "tool_1"
+        assert _tool_call(start).function.name == "read"
+        assert _tool_call(first_delta).id is None
+        assert _tool_call(first_delta).function.name is None
+        assert _tool_call(second_delta).id is None
+        assert _tool_call(second_delta).function.name is None
+        assert _tool_call(first_delta).function.arguments == '{"path":'
+        assert _tool_call(second_delta).function.arguments == '"/tmp"}'
+
+    def test_signed_thinking_blocks_lead_the_assistant_turn(self):
+        """Multi-turn tool use with thinking only works if the signed block is echoed back first."""
+        body = self._transform(
+            [
+                {"role": "user", "content": "hi"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "thinking_blocks": [
+                        {"type": "thinking", "thinking": "391", "signature": "Eto"},
+                        {"type": "thinking", "thinking": "unsigned"},
+                    ],
+                    "tool_calls": [
+                        {"id": "call_1", "type": "function", "function": {"name": "read", "arguments": "{}"}}
+                    ],
+                },
+            ]
+        )
+        blocks = body["messages"][1]["content"]
+        assert blocks[0] == {"type": "thinking", "thinking": "391", "signature": "Eto"}
+        assert [b["type"] for b in blocks] == ["thinking", "tool_use"]
+
+    def test_signed_thinking_blocks_lead_a_plain_text_assistant_turn(self):
+        """A thinking response without a tool call must also round-trip on the next request."""
+        body = self._transform(
+            [
+                {"role": "user", "content": "hi"},
+                {
+                    "role": "assistant",
+                    "content": "391",
+                    "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
+                },
+                {"role": "user", "content": "continue"},
+            ]
+        )
+        assert body["messages"][1] == {
+            "role": "assistant",
+            "content": [
+                {"type": "thinking", "thinking": "391", "signature": "Eto"},
+                {"type": "text", "text": "391"},
+            ],
+        }
+
+    def test_signed_thinking_blocks_preserve_list_content(self):
+        """Cached assistant text reaches this transform as a content list, not a string."""
+        body = self._transform(
+            [
+                {"role": "user", "content": "hi"},
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}}],
+                    "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
+                },
+                {"role": "user", "content": "continue"},
+            ]
+        )
+        assert body["messages"][1] == {
+            "role": "assistant",
+            "content": [
+                {"type": "thinking", "thinking": "391", "signature": "Eto"},
+                {"type": "text", "text": "391", "cache_control": {"type": "ephemeral"}},
+            ],
+        }
+
+    def test_thinking_only_assistant_turn_sends_no_empty_text_block(self):
+        """Anthropic-shaped APIs reject empty text blocks, so a content-less thinking turn is thinking only."""
+        body = self._transform(
+            [
+                {"role": "user", "content": "hi"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "thinking_blocks": [{"type": "thinking", "thinking": "391", "signature": "Eto"}],
+                },
+                {"role": "user", "content": "continue"},
+            ]
+        )
+        assert body["messages"][1]["content"] == [{"type": "thinking", "thinking": "391", "signature": "Eto"}]
+
+    def test_streaming_surfaces_thinking_and_prompt_cache_usage(self):
+        """Cortex streams thinking deltas, signatures and cache counts; all must reach the caller."""
+        handler = SnowflakeStreamingHandler(streaming_response=[], sync_stream=True)
+        handler.chunk_parser(
+            {
+                "type": "message_start",
+                "message": {"usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323}},
+            }
+        )
+        thinking = handler.chunk_parser(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "thinking_delta", "thinking": "391"},
+            }
+        )
+        signature = handler.chunk_parser(
+            {
+                "type": "content_block_delta",
+                "index": 0,
+                "delta": {"type": "signature_delta", "signature": "Eto"},
+            }
+        )
+        final = handler.chunk_parser(
+            {
+                "type": "message_delta",
+                "delta": {"stop_reason": "end_turn"},
+                "usage": {"output_tokens": 8, "cache_read_input_tokens": 1323},
+            }
+        )
+
+        assert thinking.choices[0].delta.reasoning_content == "391"
+        assert signature.choices[0].delta.thinking_blocks[0]["signature"] == "Eto"
+        assert final.usage.prompt_tokens_details.cached_tokens == 1323
+
+
 class TestSnowFlakeCompletion:
     model_name = "mistral"
 
@@ -380,10 +749,7 @@ class TestSnowFlakeCompletion:
         # PAT key was used
         post_kwargs = mock_post.call_args_list[-1][1]
         assert "xxxxx" in post_kwargs["headers"]["Authorization"]
-        assert (
-            post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"]
-            == "PROGRAMMATIC_ACCESS_TOKEN"
-        )
+        assert post_kwargs["headers"]["X-Snowflake-Authorization-Token-Type"] == "PROGRAMMATIC_ACCESS_TOKEN"
 
         # account id was used
         assert "AAAA-BBBB" in post_kwargs["url"]
@@ -495,9 +861,7 @@ class TestSnowflakeChatCompletion:
                 )
                 mock_post.assert_called_once()
         else:
-            with patch.object(
-                AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp
-            ) as mock_post:
+            with patch.object(AsyncHTTPHandler, "post", new_callable=AsyncMock, return_value=mock_resp) as mock_post:
                 response = asyncio.run(
                     acompletion(
                         model="snowflake/mistral-7b",
@@ -580,8 +944,4 @@ class TestSnowflakeChatCompletion:
             chunks_received = asyncio.run(_run())
 
         assert len(chunks_received) > 0
-        content = "".join(
-            c.choices[0].delta.content
-            for c in chunks_received
-            if c.choices[0].delta.content
-        )
+        content = "".join(c.choices[0].delta.content for c in chunks_received if c.choices[0].delta.content)
diff --git a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py
index fb21e2e6f6b..7970f7771fc 100644
--- a/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py
+++ b/tests/test_litellm/llms/snowflake/test_snowflake_native_endpoints.py
@@ -338,7 +338,7 @@ class TestAnthropicConfigRequest:
             litellm_params={},
             headers={},
         )
-        assert body["system"] == "You are helpful."
+        assert body["system"] == [{"type": "text", "text": "You are helpful."}]
         assert all(m["role"] != "system" for m in body["messages"])
         assert body["messages"][0] == {"role": "user", "content": "Hello"}
 
@@ -422,6 +422,64 @@ class TestAnthropicConfigResponse:
         assert result.usage.completion_tokens == 5
         assert result.usage.total_tokens == 15
 
+    def test_prompt_cache_usage_is_surfaced(self):
+        """Cortex reports cache creation/read counts; dropping them hides caching and bills cached input at full price."""
+        raw = httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "model": "claude-sonnet-4-6",
+                "content": [{"type": "text", "text": "hi"}],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 18, "cache_creation_input_tokens": 1323, "cache_read_input_tokens": 0},
+            },
+        )
+        result = self.cfg.transform_response(
+            model="snowflake/claude-sonnet-4-6",
+            raw_response=raw,
+            model_response=ModelResponse(),
+            logging_obj=_mock_logging(),
+            request_data={},
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+        assert result.usage.prompt_tokens == 1341
+        assert result.usage.prompt_tokens_details.cache_creation_tokens == 1323
+        assert result.usage.prompt_tokens_details.cached_tokens == 0
+
+    def test_thinking_block_and_signature_are_preserved(self):
+        """The signature must survive so a client can echo the thinking block on the next turn."""
+        raw = httpx.Response(
+            200,
+            json={
+                "id": "msg_1",
+                "model": "claude-sonnet-4-6",
+                "content": [
+                    {"type": "thinking", "thinking": "391", "signature": "Eto"},
+                    {"type": "text", "text": "391"},
+                ],
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 10, "output_tokens": 5},
+            },
+        )
+        result = self.cfg.transform_response(
+            model="snowflake/claude-sonnet-4-6",
+            raw_response=raw,
+            model_response=ModelResponse(),
+            logging_obj=_mock_logging(),
+            request_data={},
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            encoding=None,
+        )
+        message = result.choices[0].message
+        assert message.content == "391"
+        assert message.reasoning_content == "391"
+        assert message.thinking_blocks[0]["signature"] == "Eto"
+
     def test_stop_reason_end_turn_maps_to_stop(self):
         raw = _make_anthropic_response()
         result = self.cfg.transform_response(

From aec083cdac6fa55b950a705d71198604c7763baa 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 18:19:04 -0700
Subject: [PATCH 34/38] feat(proxy): per-worker admission control that rejects
 excess requests with 503 (#39352)

* feat(proxy): reject excess per-worker requests with 503

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

* refactor(proxy): drop redundant suppressions in admission middleware

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

* fix(proxy): exempt the /metrics/ redirect target from admission control

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

* ci: allowlist live Granian saturation benchmark

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

* fix(proxy): regenerate dashboard API types

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

* fix(proxy): normalize root_path for admission exemptions, validate settings, inject state

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

* test(proxy): cover prometheus metric factory, lifespan scope, and prefix lookalike paths

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

* fix(proxy): queue behind pending waiters, cache admission settings parsing, log invalid limits

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

* refactor(proxy): simplify invalid admission settings handling

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>
---
 .github/ci-coverage-allowlist.yml             |   5 +
 litellm/proxy/_types.py                       |   9 +
 .../health_endpoints/_health_endpoints.py     |  19 +-
 .../admission_control_middleware.py           | 315 ++++++++++++++
 litellm/proxy/proxy_server.py                 |  13 +
 .../test_granian_admission_saturation.py      | 150 +++++++
 .../health_endpoints/test_health_endpoints.py |  12 +
 .../test_admission_control_middleware.py      | 402 ++++++++++++++++++
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  16 +
 9 files changed, 940 insertions(+), 1 deletion(-)
 create mode 100644 litellm/proxy/middleware/admission_control_middleware.py
 create mode 100644 tests/load_tests/test_granian_admission_saturation.py
 create mode 100644 tests/test_litellm/proxy/middleware/test_admission_control_middleware.py

diff --git a/.github/ci-coverage-allowlist.yml b/.github/ci-coverage-allowlist.yml
index 32232de381c..0da07038152 100644
--- a/.github/ci-coverage-allowlist.yml
+++ b/.github/ci-coverage-allowlist.yml
@@ -79,6 +79,11 @@ test_paths:
       - tests/load_tests/test_otel_load_test.py
       - tests/load_tests/test_vertex_embeddings_load_test.py
       - tests/load_tests/test_vertex_load_tests.py
+  - reason: >-
+      Env-gated saturation benchmark requires a live proxy and provider credentials, so it is run
+      locally rather than in pull-request jobs
+    paths:
+      - tests/load_tests/test_granian_admission_saturation.py
   - reason: >-
       A local-only agent rig: test_a2a_completion_bridge.py needs a LangGraph server on
       localhost:2024 and test_a2a.py drives a live A2A endpoint, so neither can run in a
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 5d5a25e7cd6..1a807fd39bb 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2404,6 +2404,15 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
     """
 
     completion_model: str | None = Field(None, description="proxy level default model for all chat completion calls")
+    max_in_flight_requests_per_worker: int | None = Field(
+        None, gt=0, description="maximum concurrent requests handled by each worker"
+    )
+    max_queued_requests_per_worker: int | None = Field(
+        None, ge=0, description="maximum requests waiting for a worker slot"
+    )
+    admission_queue_timeout_seconds: float = Field(
+        1.0, gt=0, description="maximum time a request waits for a worker slot"
+    )
     plugins: list[PluginConfig] | None = Field(
         None, description="external services registered as embeddable UI plugins"
     )
diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py
index 8b57bdca2fe..65d0ec8c0dc 100644
--- a/litellm/proxy/health_endpoints/_health_endpoints.py
+++ b/litellm/proxy/health_endpoints/_health_endpoints.py
@@ -50,6 +50,9 @@ from litellm.proxy.health_check import (
     perform_health_check,
     run_with_timeout,
 )
+from litellm.proxy.middleware.admission_control_middleware import (
+    get_admission_control_stats,
+)
 from litellm.proxy.middleware.in_flight_requests_middleware import (
     get_in_flight_requests,
 )
@@ -63,6 +66,13 @@ from litellm.secret_managers.main import get_secret_bool
 #### Health ENDPOINTS ####
 
 
+class _HealthBacklogResponse(TypedDict):
+    in_flight_requests: ReadOnly[int]
+    admitted_requests: ReadOnly[int]
+    queued_requests: ReadOnly[int]
+    rejected_requests: ReadOnly[int]
+
+
 def _reject_os_environ_references(params: dict) -> None:
     """
     Validate that the provided params do not contain any ``os.environ/``
@@ -1759,7 +1769,14 @@ async def health_backlog():
     for the event loop to get to them, adding latency before LiteLLM even starts
     its own timer.
     """
-    return {"in_flight_requests": get_in_flight_requests()}
+    stats: Final = get_admission_control_stats()
+    response: Final[_HealthBacklogResponse] = {
+        "in_flight_requests": get_in_flight_requests(),
+        "admitted_requests": stats.admitted,
+        "queued_requests": stats.queued,
+        "rejected_requests": stats.rejected_total,
+    }
+    return response
 
 
 @router.get(
diff --git a/litellm/proxy/middleware/admission_control_middleware.py b/litellm/proxy/middleware/admission_control_middleware.py
new file mode 100644
index 00000000000..aa62ef9e3bf
--- /dev/null
+++ b/litellm/proxy/middleware/admission_control_middleware.py
@@ -0,0 +1,315 @@
+import asyncio
+import os
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import Annotated, Final, Protocol, TypeAlias, runtime_checkable
+
+from pydantic import Field, TypeAdapter, ValidationError
+from starlette.responses import JSONResponse
+from starlette.types import ASGIApp, Receive, Scope, Send
+
+from litellm._logging import verbose_proxy_logger
+
+_EXEMPT_PATHS: Final[frozenset[str]] = frozenset(
+    {
+        "/health/liveliness",
+        "/health/liveness",
+        "/health/readiness",
+        "/health/readiness/details",
+        "/health/backlog",
+        "/health/drain",
+        "/metrics",
+        "/metrics/",
+    }
+)
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlSettings:
+    max_in_flight_requests: int
+    max_queued_requests: int
+    queue_timeout_seconds: float
+
+
+AdmissionControlSettingsGetter: TypeAlias = Callable[[], AdmissionControlSettings | None]  # mutable-ok: Callable params
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlStats:
+    admitted: int
+    queued: int
+    rejected_total: int
+
+
+@runtime_checkable
+class _Gauge(Protocol):
+    def inc(self, amount: float = 1) -> None: ...
+
+    def dec(self, amount: float = 1) -> None: ...
+
+
+@runtime_checkable
+class _CounterChild(Protocol):
+    def inc(self, amount: float = 1) -> None: ...
+
+
+@runtime_checkable
+class _Counter(Protocol):
+    def labels(self, reason: str) -> _CounterChild: ...
+
+
+@dataclass(frozen=True, slots=True)
+class AdmissionControlMetrics:
+    admitted_gauge: _Gauge
+    queued_gauge: _Gauge
+    rejected_counter: _Counter
+
+
+AdmissionControlMetricsFactory: TypeAlias = Callable[[], AdmissionControlMetrics | None]  # mutable-ok: Callable params
+
+
+class AdmissionControlState:
+    """Per-process admission counters and the in-flight semaphore shared by one worker's requests."""
+
+    def __init__(self, metrics_factory: AdmissionControlMetricsFactory) -> None:
+        self._metrics_factory = metrics_factory
+        self._metrics: AdmissionControlMetrics | None = None
+        self._metrics_init_attempted = False
+        self._admitted = 0
+        self._queued = 0
+        self._rejected_total = 0
+        self._semaphore: asyncio.Semaphore | None = None
+        self._semaphore_loop: asyncio.AbstractEventLoop | None = None
+
+    def get_stats(self) -> AdmissionControlStats:
+        return AdmissionControlStats(
+            admitted=self._admitted,
+            queued=self._queued,
+            rejected_total=self._rejected_total,
+        )
+
+    def get_semaphore(self, max_in_flight_requests: int) -> asyncio.Semaphore:
+        loop: Final = asyncio.get_running_loop()
+        if self._semaphore_loop is not loop:
+            self._semaphore = asyncio.Semaphore(max_in_flight_requests)
+            self._semaphore_loop = loop
+        semaphore: Final = self._semaphore
+        if semaphore is None:
+            raise RuntimeError("Admission control semaphore was not initialized")
+        return semaphore
+
+    def record_admission(self) -> None:
+        self._admitted += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.admitted_gauge.inc()
+
+    def record_release(self) -> None:
+        self._admitted -= 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.admitted_gauge.dec()
+
+    def record_queue(self) -> None:
+        self._queued += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.queued_gauge.inc()
+
+    def record_dequeue(self) -> None:
+        self._queued -= 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.queued_gauge.dec()
+
+    def record_rejection(self, reason: str) -> None:
+        self._rejected_total += 1
+        metrics: Final = self._get_metrics()
+        if metrics is not None:
+            metrics.rejected_counter.labels(reason=reason).inc()
+
+    def _get_metrics(self) -> AdmissionControlMetrics | None:
+        if not self._metrics_init_attempted:
+            self._metrics_init_attempted = True
+            self._metrics = self._metrics_factory()
+        return self._metrics
+
+
+class AdmissionControlMiddleware:
+    def __init__(
+        self,
+        app: ASGIApp,
+        get_settings: AdmissionControlSettingsGetter,
+        state: AdmissionControlState,
+    ) -> None:
+        self.app = app
+        self.get_settings = get_settings
+        self.state = state
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        settings: Final = self.get_settings()
+        if settings is None or _get_route_path(scope) in _EXEMPT_PATHS:
+            await self.app(scope, receive, send)
+            return
+
+        state: Final = self.state
+        semaphore: Final = state.get_semaphore(settings.max_in_flight_requests)
+        if not semaphore.locked():
+            await semaphore.acquire()
+            state.record_admission()
+        elif state.get_stats().queued >= settings.max_queued_requests:
+            state.record_rejection("queue_full")
+            await _overloaded_response(state)(scope, receive, send)
+            return
+        else:
+            state.record_queue()
+            try:
+                await asyncio.wait_for(
+                    semaphore.acquire(),
+                    timeout=settings.queue_timeout_seconds,
+                )
+            except asyncio.TimeoutError:
+                state.record_dequeue()
+                state.record_rejection("queue_timeout")
+                await _overloaded_response(state)(scope, receive, send)
+                return
+            except asyncio.CancelledError:
+                state.record_dequeue()
+                raise
+            state.record_dequeue()
+            state.record_admission()
+
+        try:
+            await self.app(scope, receive, send)
+        finally:
+            semaphore.release()
+            state.record_release()
+
+
+def _get_route_path(scope: Scope) -> str:
+    """Strip the ASGI root_path (SERVER_ROOT_PATH) the same way Starlette does before route matching."""
+    path: Final[str] = scope["path"]
+    root_path: Final[str] = scope.get("root_path", "")
+    if not root_path or not path.startswith(root_path):
+        return path
+    if path == root_path:
+        return ""
+    if path[len(root_path)] == "/":
+        return path[len(root_path) :]
+    return path
+
+
+def _create_gauge(gauge_type: Callable[..., object], name: str, description: str) -> _Gauge:
+    metric: Final = (
+        gauge_type(name, description, multiprocess_mode="livesum")
+        if "PROMETHEUS_MULTIPROC_DIR" in os.environ
+        else gauge_type(name, description)
+    )
+    if not isinstance(metric, _Gauge):
+        raise TypeError("Admission gauge has an unexpected type")
+    return metric
+
+
+def create_prometheus_admission_metrics() -> AdmissionControlMetrics | None:
+    try:
+        from prometheus_client import Counter, Gauge
+
+        return AdmissionControlMetrics(
+            admitted_gauge=_create_gauge(
+                Gauge,
+                "litellm_admission_admitted_requests",
+                "Number of requests admitted by this worker",
+            ),
+            queued_gauge=_create_gauge(
+                Gauge,
+                "litellm_admission_queued_requests",
+                "Number of requests queued by this worker",
+            ),
+            rejected_counter=Counter(  # mutable-ok: Prometheus requires runtime Counter construction
+                "litellm_admission_rejected_requests_total",
+                "Number of requests rejected by this worker",
+                labelnames=("reason",),
+            ),
+        )
+    except (ImportError, ValueError):
+        return None
+
+
+admission_control_state: Final = AdmissionControlState(create_prometheus_admission_metrics)
+
+
+def get_admission_control_stats() -> AdmissionControlStats:
+    return admission_control_state.get_stats()
+
+
+_PositiveInt: TypeAlias = Annotated[int, Field(gt=0)]
+_NonNegativeInt: TypeAlias = Annotated[int, Field(ge=0)]
+_PositiveFloat: TypeAlias = Annotated[float, Field(gt=0)]
+_AdmissionControlRaw: TypeAlias = int | float | str | None
+
+
+def _hashable(value: object) -> _AdmissionControlRaw:
+    return value if value is None or isinstance(value, (int, float, str)) else repr(value)
+
+
+_POSITIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_PositiveInt)
+_NON_NEGATIVE_INT_ADAPTER: Final[TypeAdapter[int]] = TypeAdapter(_NonNegativeInt)
+_POSITIVE_FLOAT_ADAPTER: Final[TypeAdapter[float]] = TypeAdapter(_PositiveFloat)
+
+
+@lru_cache(maxsize=16)
+def _parse_admission_control_settings(
+    max_in_flight_raw: _AdmissionControlRaw,
+    max_queued_raw: _AdmissionControlRaw,
+    queue_timeout_raw: _AdmissionControlRaw,
+) -> AdmissionControlSettings | None:
+    try:
+        max_in_flight: Final = _POSITIVE_INT_ADAPTER.validate_python(max_in_flight_raw)
+        max_queued: Final = (
+            max_in_flight if max_queued_raw is None else _NON_NEGATIVE_INT_ADAPTER.validate_python(max_queued_raw)
+        )
+        queue_timeout: Final = _POSITIVE_FLOAT_ADAPTER.validate_python(queue_timeout_raw)
+    except ValidationError as exc:
+        verbose_proxy_logger.error(
+            "Ignoring invalid admission control settings, per-worker admission control is disabled: %s",
+            exc,
+        )
+        return None
+    return AdmissionControlSettings(
+        max_in_flight_requests=max_in_flight,
+        max_queued_requests=max_queued,
+        queue_timeout_seconds=queue_timeout,
+    )
+
+
+def get_admission_control_settings(settings: Mapping[str, object]) -> AdmissionControlSettings | None:
+    max_in_flight_raw: Final = settings.get("max_in_flight_requests_per_worker")
+    if max_in_flight_raw is None:
+        return None
+    return _parse_admission_control_settings(
+        _hashable(max_in_flight_raw),
+        _hashable(settings.get("max_queued_requests_per_worker")),
+        _hashable(settings.get("admission_queue_timeout_seconds", 1.0)),
+    )
+
+
+def _overloaded_response(state: AdmissionControlState) -> JSONResponse:
+    stats: Final = state.get_stats()
+    return JSONResponse(
+        status_code=503,
+        headers={"retry-after": "1"},  # mutable-ok: Starlette expects a plain headers mapping
+        content={  # mutable-ok: Starlette serializes a plain response mapping
+            "error": {  # mutable-ok: nested response mapping
+                "message": (
+                    f"Worker at capacity: {stats.admitted} in-flight, {stats.queued} queued requests. Retry later."
+                ),
+                "type": "overloaded_error",
+                "code": "503",
+            }
+        },
+    )
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index c43cc510990..83f63c15529 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -583,6 +583,11 @@ try:
 except ImportError:
     build_billing_metrics_recorder = None
     shutdown_billing_metrics_recorder = None
+from litellm.proxy.middleware.admission_control_middleware import (
+    AdmissionControlMiddleware,
+    admission_control_state,
+    get_admission_control_settings,
+)
 from litellm.proxy.middleware.in_flight_requests_middleware import (
     InFlightRequestsMiddleware,
 )
@@ -16502,6 +16507,9 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
     {
         "max_parallel_requests": "Integer",
         "global_max_parallel_requests": "Integer",
+        "max_in_flight_requests_per_worker": "Integer",
+        "max_queued_requests_per_worker": "Integer",
+        "admission_queue_timeout_seconds": "Float",
         "max_request_size_mb": "Integer",
         "max_batch_file_size_mb": "Integer",
         "max_file_size_mb": "Integer",
@@ -18177,6 +18185,11 @@ app.add_middleware(
     get_max_request_size_mb=lambda: general_settings.get("max_request_size_mb"),
     is_request_size_limit_enabled=lambda: premium_user is True,
 )
+app.add_middleware(
+    AdmissionControlMiddleware,
+    get_settings=lambda: get_admission_control_settings(general_settings),
+    state=admission_control_state,
+)
 
 
 async def _stream_mcp_asgi_response(handle_fn, scope: dict, receive) -> "StreamingResponse":
diff --git a/tests/load_tests/test_granian_admission_saturation.py b/tests/load_tests/test_granian_admission_saturation.py
new file mode 100644
index 00000000000..b42c06037c2
--- /dev/null
+++ b/tests/load_tests/test_granian_admission_saturation.py
@@ -0,0 +1,150 @@
+import asyncio
+import os
+import socket
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Final
+
+import httpx
+import pytest
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("LITELLM_RUN_SATURATION_BENCHMARK") != "1",
+    reason="set LITELLM_RUN_SATURATION_BENCHMARK=1 to run the saturation benchmark",
+)
+
+
+def _free_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
+        listener.bind(("127.0.0.1", 0))
+        return int(listener.getsockname()[1])
+
+
+def _percentile(values: list[float], percentile: float) -> float:
+    return sorted(values)[min(int(len(values) * percentile), len(values) - 1)]
+
+
+@pytest.mark.asyncio
+async def test_granian_admission_control_saturation(tmp_path: Path) -> None:
+    fake_port: Final = _free_port()
+    proxy_port: Final = _free_port()
+    fake_script: Final = Path(__file__).parents[1] / "_fake_openai_endpoint_server.py"
+    config_path: Final = tmp_path / "saturation_config.yaml"
+    config_path.write_text(
+        f"""model_list:
+  - model_name: slow-endpoint
+    litellm_params:
+      model: openai/slow-endpoint
+      api_base: http://127.0.0.1:{fake_port}/v1
+general_settings:
+  master_key: sk-saturation
+  max_in_flight_requests_per_worker: 8
+  max_queued_requests_per_worker: 8
+  admission_queue_timeout_seconds: 0.5
+"""
+    )
+    fake_process: Final = subprocess.Popen(
+        [sys.executable, str(fake_script), "--host", "127.0.0.1", "--port", str(fake_port)],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    try:
+        proxy_process: Final = subprocess.Popen(
+            [
+                sys.executable,
+                "-m",
+                "litellm.proxy.proxy_cli",
+                "--config",
+                str(config_path),
+                "--run_granian",
+                "--num_workers",
+                "1",
+                "--port",
+                str(proxy_port),
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        try:
+            async with httpx.AsyncClient(base_url=f"http://127.0.0.1:{proxy_port}") as client:
+                deadline: Final = time.monotonic() + 60
+                while time.monotonic() < deadline:
+                    try:
+                        response: Final = await client.get("/health/liveliness", timeout=2)
+                        if response.status_code == 200:
+                            break
+                    except httpx.HTTPError:
+                        pass
+                    await asyncio.sleep(0.25)
+                else:
+                    raise AssertionError("Granian proxy did not become healthy")
+
+                liveness_latencies: Final[list[float]] = []
+                stop_sampling: Final = asyncio.Event()
+
+                async def sample_liveness() -> None:
+                    while not stop_sampling.is_set():
+                        start: Final = time.perf_counter()
+                        try:
+                            response = await client.get("/health/liveliness", timeout=2)
+                            response.raise_for_status()
+                            liveness_latencies.append(time.perf_counter() - start)
+                        except httpx.HTTPError:
+                            pass
+                        await asyncio.sleep(0.05)
+
+                async def send_completion() -> tuple[int, float, bool]:
+                    start: Final = time.perf_counter()
+                    response = await client.post(
+                        "/chat/completions",
+                        headers={"Authorization": "Bearer sk-saturation"},
+                        json={
+                            "model": "slow-endpoint",
+                            "messages": [{"role": "user", "content": "hello"}],
+                        },
+                        timeout=10,
+                    )
+                    return response.status_code, time.perf_counter() - start, "retry-after" in response.headers
+
+                sampler: Final = asyncio.create_task(sample_liveness())
+                results: Final = await asyncio.gather(*(send_completion() for _ in range(200)))
+                stop_sampling.set()
+                await sampler
+
+            statuses: Final = [result[0] for result in results]
+            latencies: Final = [result[1] for result in results]
+            rejected: Final = [result for result in results if result[0] == 503]
+            assert set(statuses) <= {200, 503}
+            assert rejected
+            assert all(result[2] for result in rejected)
+            assert _percentile(latencies, 0.99) < 5
+            assert liveness_latencies
+            assert _percentile(liveness_latencies, 0.95) < 0.5
+
+            duration: Final = max(latencies)
+            print(
+                "\nmetric          value\n"
+                f"rps             {len(results) / duration:.2f}\n"
+                f"200 count       {statuses.count(200)}\n"
+                f"503 count       {statuses.count(503)}\n"
+                f"p50             {_percentile(latencies, 0.50):.3f}s\n"
+                f"p95             {_percentile(latencies, 0.95):.3f}s\n"
+                f"p99             {_percentile(latencies, 0.99):.3f}s\n"
+                f"liveness p95    {_percentile(liveness_latencies, 0.95):.3f}s"
+            )
+        finally:
+            proxy_process.terminate()
+            try:
+                proxy_process.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                proxy_process.kill()
+                proxy_process.wait()
+    finally:
+        fake_process.terminate()
+        try:
+            fake_process.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            fake_process.kill()
+            fake_process.wait()
diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
index e3f71692c78..0e90c107865 100644
--- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
+++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
@@ -1238,6 +1238,18 @@ def test_health_liveness_endpoint(proxy_client):
     print(f"\n/health/liveness response time: {duration_ms:.2f}ms")
 
 
+def test_health_backlog_includes_admission_control_stats(proxy_client):
+    response = proxy_client.get("/health/backlog")
+
+    assert response.status_code == 200, response.text
+    assert set(response.json()) == {
+        "in_flight_requests",
+        "admitted_requests",
+        "queued_requests",
+        "rejected_requests",
+    }
+
+
 def test_health_readiness(proxy_client):
     """
     Test /health/readiness endpoint.
diff --git a/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py
new file mode 100644
index 00000000000..f1ca13daa03
--- /dev/null
+++ b/tests/test_litellm/proxy/middleware/test_admission_control_middleware.py
@@ -0,0 +1,402 @@
+import asyncio
+import json
+from typing import Final
+
+import pytest
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+
+from litellm.proxy.middleware.admission_control_middleware import (
+    AdmissionControlMetrics,
+    AdmissionControlMiddleware,
+    AdmissionControlSettings,
+    AdmissionControlState,
+    AdmissionControlStats,
+    _parse_admission_control_settings,
+    create_prometheus_admission_metrics,
+    get_admission_control_settings,
+)
+
+
+@pytest.fixture
+def state() -> AdmissionControlState:
+    return AdmissionControlState(lambda: None)
+
+
+async def _call(
+    middleware: AdmissionControlMiddleware,
+    path: str = "/",
+    root_path: str = "",
+) -> tuple[Message, ...]:
+    messages: Final[list[Message]] = []
+
+    async def receive() -> Message:
+        return {"type": "http.request", "body": b"", "more_body": False}
+
+    async def send(message: Message) -> None:
+        messages.append(message)
+
+    scope: Final[Scope] = {
+        "type": "http",
+        "path": path,
+        "root_path": root_path,
+        "method": "GET",
+        "headers": [],
+    }
+    await middleware(scope, receive, send)
+    return tuple(messages)
+
+
+def _handler_with_release(
+    started: asyncio.Event,
+    release: asyncio.Event,
+) -> ASGIApp:
+    async def handler(scope: Scope, receive: Receive, send: Send) -> None:
+        started.set()
+        await release.wait()
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok", "more_body": False})
+
+    return handler
+
+
+def test_is_not_base_http_middleware() -> None:
+    assert not issubclass(AdmissionControlMiddleware, BaseHTTPMiddleware)
+
+
+@pytest.mark.asyncio
+async def test_capacity_rejects_excess_and_releases_queued_request(state: AdmissionControlState) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    middleware: Final = AdmissionControlMiddleware(
+        _handler_with_release(started, release),
+        lambda: AdmissionControlSettings(1, 1, 1.0),
+        state,
+    )
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    second: Final = asyncio.create_task(_call(middleware))
+    await asyncio.sleep(0)
+    assert state.get_stats().queued == 1
+
+    third: Final = await _call(middleware)
+    assert third[0]["status"] == 503
+    headers: Final = dict(third[0]["headers"])
+    assert headers[b"retry-after"] == b"1"
+    assert headers[b"content-type"] == b"application/json"
+    assert json.loads(third[1]["body"])["error"] == {
+        "message": "Worker at capacity: 1 in-flight, 1 queued requests. Retry later.",
+        "type": "overloaded_error",
+        "code": "503",
+    }
+    assert state.get_stats().rejected_total == 1
+
+    release.set()
+    assert (await first)[0]["status"] == 200
+    assert (await second)[0]["status"] == 200
+    assert state.get_stats() == AdmissionControlStats(0, 0, 1)
+
+
+@pytest.mark.asyncio
+async def test_pending_waiter_is_not_skipped_after_admission_is_released(state: AdmissionControlState) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    third_trigger: Final = asyncio.Event()
+    middleware: Final = AdmissionControlMiddleware(
+        _handler_with_release(started, release),
+        lambda: AdmissionControlSettings(1, 2, 1.0),
+        state,
+    )
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    second: Final = asyncio.create_task(_call(middleware))
+    await asyncio.sleep(0)
+
+    async def call_third() -> tuple[Message, ...]:
+        await third_trigger.wait()
+        return await _call(middleware)
+
+    third: Final = asyncio.create_task(call_third())
+    await asyncio.sleep(0)
+    release.set()
+    third_trigger.set()
+    await asyncio.sleep(0)
+
+    assert state.get_stats().queued == 2
+    await asyncio.gather(first, second, third)
+
+
+@pytest.mark.asyncio
+async def test_queue_timeout_rejects_and_decrements_queue(state: AdmissionControlState) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    middleware: Final = AdmissionControlMiddleware(
+        _handler_with_release(started, release),
+        lambda: AdmissionControlSettings(1, 1, 0.05),
+        state,
+    )
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    start_time: Final = asyncio.get_running_loop().time()
+    second: Final = await _call(middleware)
+    elapsed: Final = asyncio.get_running_loop().time() - start_time
+
+    assert second[0]["status"] == 503
+    assert elapsed < 0.5
+    assert state.get_stats().queued == 0
+    assert state.get_stats().rejected_total == 1
+    release.set()
+    await first
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("root_path", "probe_path"),
+    (
+        ("", "/health/liveliness"),
+        ("/proxy", "/proxy/health/liveliness"),
+        ("/proxy", "/proxy/metrics"),
+    ),
+)
+async def test_exempt_path_passes_through_when_saturated(
+    state: AdmissionControlState,
+    root_path: str,
+    probe_path: str,
+) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+
+    async def handler(scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["path"] == "/":
+            started.set()
+            await release.wait()
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok", "more_body": False})
+
+    middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    health: Final = await _call(middleware, probe_path, root_path)
+    assert health[0]["status"] == 200
+    blocked: Final = await _call(middleware, "/proxy/v1/chat/completions", root_path)
+    assert blocked[0]["status"] == 503
+    lookalike: Final = await _call(middleware, "/proxyhealth/liveliness", "/proxy")
+    assert lookalike[0]["status"] == 503
+    release.set()
+    await first
+
+
+@pytest.mark.asyncio
+async def test_non_http_scope_passes_through_when_saturated(state: AdmissionControlState) -> None:
+    seen: Final[list[str]] = []
+
+    async def handler(scope: Scope, receive: Receive, send: Send) -> None:
+        seen.append(scope["type"])
+
+    middleware: Final = AdmissionControlMiddleware(handler, lambda: AdmissionControlSettings(1, 0, 1.0), state)
+    state.record_admission()
+
+    async def receive() -> Message:
+        return {"type": "lifespan.startup"}
+
+    async def send(message: Message) -> None:
+        return None
+
+    await middleware({"type": "lifespan"}, receive, send)
+    assert seen == ["lifespan"]
+
+
+@pytest.mark.asyncio
+async def test_none_settings_does_not_limit_concurrency() -> None:
+    active: Final = [0]
+    peak: Final = [0]
+    all_started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+
+    async def handler(scope: Scope, receive: Receive, send: Send) -> None:
+        active[0] += 1
+        peak[0] = max(peak[0], active[0])
+        if active[0] == 3:
+            all_started.set()
+        await release.wait()
+        active[0] -= 1
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"ok", "more_body": False})
+
+    middleware: Final = AdmissionControlMiddleware(handler, lambda: None, AdmissionControlState(lambda: None))
+    requests: Final = tuple(asyncio.create_task(_call(middleware)) for _ in range(3))
+    await all_started.wait()
+    assert peak[0] == 3
+    release.set()
+    results: Final = await asyncio.gather(*requests)
+    assert tuple(result[0]["status"] for result in results) == (200, 200, 200)
+
+
+@pytest.mark.asyncio
+async def test_cancelling_queued_request_does_not_leak_counter(state: AdmissionControlState) -> None:
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    middleware: Final = AdmissionControlMiddleware(
+        _handler_with_release(started, release),
+        lambda: AdmissionControlSettings(1, 1, 1.0),
+        state,
+    )
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    queued: Final = asyncio.create_task(_call(middleware))
+    await asyncio.sleep(0)
+    queued.cancel()
+    with pytest.raises(asyncio.CancelledError):
+        await queued
+    assert state.get_stats().queued == 0
+    release.set()
+    await first
+
+
+@pytest.mark.asyncio
+async def test_streaming_response_holds_admission_until_final_body(state: AdmissionControlState) -> None:
+    first_chunk_sent: Final = asyncio.Event()
+    finish_stream: Final = asyncio.Event()
+
+    async def handler(scope: Scope, receive: Receive, send: Send) -> None:
+        await send({"type": "http.response.start", "status": 200, "headers": []})
+        await send({"type": "http.response.body", "body": b"first", "more_body": True})
+        first_chunk_sent.set()
+        await finish_stream.wait()
+        await send({"type": "http.response.body", "body": b"last", "more_body": False})
+
+    middleware: Final = AdmissionControlMiddleware(
+        handler,
+        lambda: AdmissionControlSettings(1, 1, 1.0),
+        state,
+    )
+    first: Final = asyncio.create_task(_call(middleware))
+    await first_chunk_sent.wait()
+    second: Final = asyncio.create_task(_call(middleware))
+    await asyncio.sleep(0)
+    assert not second.done()
+    assert state.get_stats().queued == 1
+    finish_stream.set()
+    assert (await first)[0]["status"] == 200
+    assert (await second)[0]["status"] == 200
+    assert state.get_stats().admitted == 0
+    assert state.get_stats().queued == 0
+
+
+class _FakeGauge:
+    def __init__(self) -> None:
+        self.value = 0.0
+
+    def inc(self, amount: float = 1) -> None:
+        self.value += amount
+
+    def dec(self, amount: float = 1) -> None:
+        self.value -= amount
+
+
+class _FakeCounter:
+    def __init__(self) -> None:
+        self.by_reason: Final[dict[str, _FakeGauge]] = {}
+
+    def labels(self, reason: str) -> _FakeGauge:
+        return self.by_reason.setdefault(reason, _FakeGauge())
+
+
+@pytest.mark.asyncio
+async def test_metrics_track_admitted_queued_and_rejected() -> None:
+    admitted: Final = _FakeGauge()
+    queued: Final = _FakeGauge()
+    rejected: Final = _FakeCounter()
+    state: Final = AdmissionControlState(
+        lambda: AdmissionControlMetrics(admitted_gauge=admitted, queued_gauge=queued, rejected_counter=rejected)
+    )
+    started: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    middleware: Final = AdmissionControlMiddleware(
+        _handler_with_release(started, release),
+        lambda: AdmissionControlSettings(1, 1, 0.05),
+        state,
+    )
+
+    first: Final = asyncio.create_task(_call(middleware))
+    await started.wait()
+    second: Final = asyncio.create_task(_call(middleware))
+    await asyncio.sleep(0)
+    assert (admitted.value, queued.value) == (1.0, 1.0)
+    await _call(middleware)
+    assert rejected.by_reason["queue_full"].value == 1.0
+    await second
+    assert rejected.by_reason["queue_timeout"].value == 1.0
+    release.set()
+    await first
+    assert (admitted.value, queued.value) == (0.0, 0.0)
+
+
+def test_create_prometheus_admission_metrics_registers_named_metrics() -> None:
+    from prometheus_client import REGISTRY
+
+    metrics: Final = create_prometheus_admission_metrics()
+    if metrics is not None:
+        metrics.admitted_gauge.inc()
+        metrics.queued_gauge.inc()
+        metrics.rejected_counter.labels(reason="queue_full").inc()
+        assert REGISTRY.get_sample_value("litellm_admission_admitted_requests") == 1.0
+        assert REGISTRY.get_sample_value("litellm_admission_queued_requests") == 1.0
+    assert REGISTRY.get_sample_value("litellm_admission_rejected_requests_total", {"reason": "queue_full"}) is not None
+    assert create_prometheus_admission_metrics() is None
+
+
+@pytest.mark.parametrize(
+    ("settings", "expected"),
+    (
+        ({}, None),
+        ({"max_in_flight_requests_per_worker": None}, None),
+        ({"max_in_flight_requests_per_worker": 0}, None),
+        ({"max_in_flight_requests_per_worker": "many"}, None),
+        ({"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": -1}, None),
+        ({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": 0}, None),
+        ({"max_in_flight_requests_per_worker": 3, "admission_queue_timeout_seconds": -0.5}, None),
+        (
+            {"max_in_flight_requests_per_worker": 3, "max_queued_requests_per_worker": 0},
+            AdmissionControlSettings(3, 0, 1.0),
+        ),
+        (
+            {"max_in_flight_requests_per_worker": 3},
+            AdmissionControlSettings(3, 3, 1.0),
+        ),
+        (
+            {
+                "max_in_flight_requests_per_worker": 3,
+                "max_queued_requests_per_worker": 5,
+                "admission_queue_timeout_seconds": 0.25,
+            },
+            AdmissionControlSettings(3, 5, 0.25),
+        ),
+    ),
+)
+def test_get_admission_control_settings(
+    settings: dict[str, object],
+    expected: AdmissionControlSettings | None,
+) -> None:
+    assert get_admission_control_settings(settings) == expected
+
+
+def test_invalid_admission_control_settings_logs_once(caplog: pytest.LogCaptureFixture) -> None:
+    _parse_admission_control_settings.cache_clear()
+    caplog.set_level("ERROR")
+    settings: Final = {"max_in_flight_requests_per_worker": [1]}
+
+    assert get_admission_control_settings(settings) is None
+    assert get_admission_control_settings(settings) is None
+
+    messages: Final = tuple(
+        record.message
+        for record in caplog.records
+        if record.message.startswith("Ignoring invalid admission control settings")
+    )
+    assert len(messages) == 1
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index 6d2ff431137..3dcfeb64866 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -25501,6 +25501,12 @@ export interface components {
          * @description Documents all the fields supported by `general_settings` in config.yaml
          */
         ConfigGeneralSettings: {
+            /**
+             * Admission Queue Timeout Seconds
+             * @description maximum time a request waits for a worker slot
+             * @default 1
+             */
+            admission_queue_timeout_seconds: number;
             /**
              * Alert To Webhook Url
              * @description Mapping of alert type to webhook url. e.g. `alert_to_webhook_url: {'budget_alerts': 'https://nothooks.slack.com/services/T00000000/B00000000/XXXXXXXXXXXXXXXXXXXXXXXX'}`
@@ -25709,11 +25715,21 @@ export interface components {
              * @description max file size in MB for /v1/files uploads, for any purpose, if a file is larger than this size it will be rejected before being forwarded to the provider
              */
             max_file_size_mb?: number | null;
+            /**
+             * Max In Flight Requests Per Worker
+             * @description maximum concurrent requests handled by each worker
+             */
+            max_in_flight_requests_per_worker?: number | null;
             /**
              * Max Parallel Requests
              * @description maximum parallel requests for each api key
              */
             max_parallel_requests?: number | null;
+            /**
+             * Max Queued Requests Per Worker
+             * @description maximum requests waiting for a worker slot
+             */
+            max_queued_requests_per_worker?: number | null;
             /**
              * Max Request Size Mb
              * @description max request size in MB, if a request is larger than this size it will be rejected

From 4a537e2c19da060321df8cf93a7d4268492a9f0f 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 18:23:37 -0700
Subject: [PATCH 35/38] fix(proxy): emit SSE keepalives on queue, rag, azure
 passthrough, usage chat and policy enrich streams (#39273)

Co-authored-by: yassin <yassin@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../policy_endpoints/endpoints.py             |  15 +-
 .../usage_endpoints/ai_usage_chat.py          |   4 +-
 .../usage_endpoints/endpoints.py              |  19 ++-
 .../llm_passthrough_endpoints.py              | 137 +++++++++++-------
 litellm/proxy/proxy_server.py                 |  33 +++--
 litellm/proxy/rag_endpoints/endpoints.py      |  88 ++++++-----
 .../policy_endpoints/test_endpoints.py        |  68 +++++++++
 .../usage_endpoints/test_ai_usage_chat.py     |  58 ++++++++
 .../test_llm_pass_through_endpoints.py        |  90 ++++++++++++
 .../proxy_server/test_streaming_helpers.py    |  89 ++++++++++++
 .../proxy/rag_endpoints/test_rag_endpoints.py |  95 +++++++++++-
 11 files changed, 588 insertions(+), 108 deletions(-)

diff --git a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py
index f58f3722741..69356922ea1 100644
--- a/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py
+++ b/litellm/proxy/management_endpoints/policy_endpoints/endpoints.py
@@ -12,7 +12,7 @@ All /policy management endpoints
 import copy
 import json
 import os
-from collections.abc import AsyncIterator
+from collections.abc import AsyncGenerator, AsyncIterator
 from typing import TYPE_CHECKING, Final, Literal, cast
 
 from fastapi import APIRouter, Depends, HTTPException, Request
@@ -20,6 +20,7 @@ from fastapi.responses import Response, StreamingResponse
 from pydantic import BaseModel, Field
 from typing_extensions import TypedDict
 
+import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
     COMPETITOR_LLM_TEMPERATURE,
@@ -32,6 +33,10 @@ from litellm.llms.openai.chat.guardrail_translation.handler import (
 )
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.sse_keepalive import (
+    SSE_COMMENT_PING,
+    wrap_sse_stream_with_keepalive_pings,
+)
 from litellm.proxy.guardrails.guardrail_hooks.custom_code import (
     RESPONSE_REJECTION_GUARDRAIL_CODE,
     CustomCodeGuardrail,
@@ -811,7 +816,7 @@ async def _stream_competitor_events(
     llm_enrichment: dict,
     brand_name: str,
     model: str,
-) -> AsyncIterator[str]:
+) -> AsyncGenerator[str, None]:
     """Stream competitor names as SSE events, then emit a final 'done' event."""
     competitors: Final[list[str]] = list(data.competitors or [])
 
@@ -883,7 +888,11 @@ async def enrich_policy_template_stream(
     model: Final = data.model or DEFAULT_COMPETITOR_DISCOVERY_MODEL
 
     return StreamingResponse(
-        _stream_competitor_events(data, template, llm_enrichment, brand_name, model),
+        wrap_sse_stream_with_keepalive_pings(
+            _stream_competitor_events(data, template, llm_enrichment, brand_name, model),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            ping_chunk=SSE_COMMENT_PING,
+        ),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
     )
diff --git a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py
index 9d5ddda017a..7ef496f94e3 100644
--- a/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py
+++ b/litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py
@@ -4,7 +4,7 @@ usage/spend data by querying the aggregated daily activity endpoints.
 """
 
 import json
-from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, Sequence
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence
 from datetime import date
 from typing import Any, Final, Literal, Protocol, cast, overload
 
@@ -543,7 +543,7 @@ async def stream_usage_ai_chat(
     model: str | None = None,
     user_id: str | None = None,
     is_admin: bool = False,
-) -> AsyncIterator[str]:
+) -> AsyncGenerator[str, None]:
     """Stream SSE events: status → tool_call → chunk → done."""
     resolved_model: Final = (model or "").strip() or DEFAULT_COMPETITOR_DISCOVERY_MODEL
     truncated: Final = messages[-MAX_CHAT_MESSAGES:] if len(messages) > MAX_CHAT_MESSAGES else messages
diff --git a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py
index b7b0ae2d8e5..d1e92d0c7e4 100644
--- a/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py
+++ b/litellm/proxy/management_endpoints/usage_endpoints/endpoints.py
@@ -10,8 +10,13 @@ from fastapi import APIRouter, Depends, Request
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
 
+import litellm
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
+from litellm.proxy.common_utils.sse_keepalive import (
+    SSE_COMMENT_PING,
+    wrap_sse_stream_with_keepalive_pings,
+)
 
 router: Final = APIRouter()
 
@@ -56,11 +61,15 @@ async def usage_ai_chat(
     messages: Final = [{"role": m.role, "content": m.content} for m in data.messages]
 
     return StreamingResponse(
-        stream_usage_ai_chat(
-            messages=messages,
-            model=data.model,
-            user_id=user_id,
-            is_admin=is_admin,
+        wrap_sse_stream_with_keepalive_pings(
+            stream_usage_ai_chat(
+                messages=messages,
+                model=data.model,
+                user_id=user_id,
+                is_admin=is_admin,
+            ),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            ping_chunk=SSE_COMMENT_PING,
         ),
         media_type="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
index 688123c9d41..6b1d6405a6a 100644
--- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
+++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
@@ -9,10 +9,11 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
 from __future__ import annotations
 
 import hmac
+import inspect
 import json
 import os
 import re
-from collections.abc import Callable, Mapping
+from collections.abc import AsyncGenerator, Callable, Mapping
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Annotated, Final, cast
 
@@ -32,6 +33,7 @@ from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
 from litellm.llms.anthropic.common_utils import AnthropicModelInfo
 from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
 from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
+from litellm.passthrough.main import AsyncPassthroughStreamingResponse
 from litellm.proxy._types import *
 from litellm.proxy.auth.handle_jwt import JWTHandler
 from litellm.proxy.auth.route_checks import RouteChecks
@@ -40,6 +42,7 @@ from litellm.proxy.auth.user_api_key_auth import (
     user_api_key_auth,
     user_api_key_auth_websocket,
 )
+from litellm.proxy.common_request_processing import open_sse_before_first_byte
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
@@ -47,6 +50,9 @@ from litellm.proxy.common_utils.http_parsing_utils import (
     get_form_data,
     get_request_body,
 )
+from litellm.proxy.common_utils.sse_keepalive import (
+    wrap_passthrough_sse_bytes_with_keepalive_pings,
+)
 from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
 from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
     HttpPassThroughEndpointHelpers,
@@ -1478,6 +1484,74 @@ def is_azure_ai_search_service_level_index_create(method: str, endpoint: str) ->
     return path == "indexes" or path.endswith("/indexes")
 
 
+async def _relay_upstream_bytes(upstream: AsyncGenerator[bytes, bytes]) -> AsyncGenerator[bytes, None]:
+    try:
+        async for chunk in upstream:
+            yield chunk
+    finally:
+        await upstream.aclose()
+
+
+async def _relay_azure_router_model(
+    llm_router: litellm.Router,
+    model: str,
+    endpoint: str,
+    request: Request,
+    request_body: Mapping[str, object],
+    is_streaming_request: bool,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> Response:
+    result: Final = await llm_router.allm_passthrough_route(
+        model=model,
+        method=request.method,
+        endpoint=endpoint,
+        request_query_params=request.query_params,
+        request_headers=_safe_get_request_headers(request),
+        stream=is_streaming_request,
+        content=None,
+        data=None,
+        files=None,
+        json=(request_body if request.headers.get("content-type") == "application/json" else None),
+        params=None,
+        headers=None,
+        cookies=None,
+        litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict),
+    )
+
+    if not is_streaming_request:
+        upstream: Final = cast(httpx.Response, result)
+        return Response(
+            content=await upstream.aread(),
+            status_code=upstream.status_code,
+            headers=HttpPassThroughEndpointHelpers.get_response_headers(headers=upstream.headers, custom_headers=None),
+        )
+
+    if inspect.isasyncgen(result):
+        sse_headers: Final = {"content-type": "text/event-stream"}
+        return StreamingResponse(
+            content=wrap_passthrough_sse_bytes_with_keepalive_pings(
+                stream=_relay_upstream_bytes(result),
+                ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+                upstream_headers=sse_headers,
+            ),
+            status_code=200,
+            headers=sse_headers,
+        )
+
+    upstream_stream: Final = cast(AsyncPassthroughStreamingResponse, result)
+    return StreamingResponse(
+        content=wrap_passthrough_sse_bytes_with_keepalive_pings(
+            stream=_relay_upstream_bytes(upstream_stream),
+            ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
+            upstream_headers=upstream_stream.headers,
+        ),
+        status_code=upstream_stream.status_code,
+        headers=HttpPassThroughEndpointHelpers.get_response_headers(
+            headers=upstream_stream.headers, custom_headers=None
+        ),
+    )
+
+
 @router.api_route(
     "/azure_ai/{endpoint:path}",
     methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
@@ -1528,55 +1602,18 @@ async def azure_proxy_route(
             if is_router_model:
                 request_body = await get_request_body(request)
                 is_streaming_request = is_passthrough_request_streaming(request_body)
-                result = await llm_router.allm_passthrough_route(
-                    model=part,
-                    method=request.method,
-                    endpoint=endpoint,
-                    request_query_params=request.query_params,
-                    request_headers=_safe_get_request_headers(request),
-                    stream=is_streaming_request,
-                    content=None,
-                    data=None,
-                    files=None,
-                    json=(request_body if request.headers.get("content-type") == "application/json" else None),
-                    params=None,
-                    headers=None,
-                    cookies=None,
-                    litellm_metadata=get_passthrough_router_request_metadata(user_api_key_dict),
-                )
-
-                if is_streaming_request:
-                    # Check if result is an async generator (from _async_streaming)
-                    import inspect
-
-                    if inspect.isasyncgen(result):
-                        # Result is already an async generator, use it directly
-                        return StreamingResponse(
-                            content=result,
-                            status_code=200,
-                            headers={"content-type": "text/event-stream"},
-                        )
-                    else:
-                        # Result is an httpx.Response, use aiter_bytes()
-                        result = cast(httpx.Response, result)
-                        return StreamingResponse(
-                            content=result.aiter_bytes(),
-                            status_code=result.status_code,
-                            headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                                headers=result.headers,
-                                custom_headers=None,
-                            ),
-                        )
-
-                # Non-streaming response
-                result = cast(httpx.Response, result)
-                content = await result.aread()
-                return Response(
-                    content=content,
-                    status_code=result.status_code,
-                    headers=HttpPassThroughEndpointHelpers.get_response_headers(
-                        headers=result.headers,
-                        custom_headers=None,
+                return await open_sse_before_first_byte(
+                    _relay_azure_router_model(
+                        llm_router=llm_router,
+                        model=part,
+                        endpoint=endpoint,
+                        request=request,
+                        request_body=request_body,
+                        is_streaming_request=is_streaming_request,
+                        user_api_key_dict=user_api_key_dict,
+                    ),
+                    ping_interval_seconds=(
+                        litellm.sse_keepalive_ping_interval_seconds if is_streaming_request else None
                     ),
                 )
             elif is_vector_store_index:
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 83f63c15529..1f39a78e12a 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -15238,20 +15238,33 @@ async def async_queue_request(
 
         if llm_router is None:
             raise HTTPException(status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value})
-
-        response: Final = await llm_router.schedule_acompletion(**data)
+        router: Final = llm_router
 
         if "stream" in data and data["stream"] is True:  # use generate_responses to stream responses
-            return StreamingResponse(
-                async_data_generator(
-                    user_api_key_dict=user_api_key_dict,
-                    response=response,
-                    request_data=data,
-                    request=request,
-                ),
-                media_type="text/event-stream",
+
+            async def produce_queue_stream() -> StreamingResponse:
+                return StreamingResponse(
+                    async_data_generator(
+                        user_api_key_dict=user_api_key_dict,
+                        response=await router.schedule_acompletion(**data),
+                        request_data=data,
+                        request=request,
+                    ),
+                    media_type="text/event-stream",
+                )
+
+            async def audit_late_failure(exc: Exception) -> HTTPException | None:
+                return await proxy_logging_obj.post_call_failure_hook(
+                    user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=data
+                )
+
+            return await open_sse_before_first_byte(
+                produce_queue_stream(),
+                ping_interval_seconds=ttft_keepalive_interval(data, router),
+                on_late_failure=audit_late_failure,
             )
 
+        response: Final = await router.schedule_acompletion(**data)
         fastapi_response.headers.update({"x-litellm-priority": str(data["priority"])})
         return response
     except Exception as e:
diff --git a/litellm/proxy/rag_endpoints/endpoints.py b/litellm/proxy/rag_endpoints/endpoints.py
index c8c6c505375..d6a402e1860 100644
--- a/litellm/proxy/rag_endpoints/endpoints.py
+++ b/litellm/proxy/rag_endpoints/endpoints.py
@@ -23,11 +23,14 @@ from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH
 from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
     LiteLLM_ManagedVectorStore,
 )
-from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
 from litellm.proxy._types import *
 from litellm.proxy.auth.auth_utils import is_request_body_safe
 from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
-from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
+from litellm.proxy.common_request_processing import (
+    ProxyBaseLLMRequestProcessing,
+    open_sse_before_first_byte,
+    ttft_keepalive_interval,
+)
 from litellm.proxy.common_utils.http_parsing_utils import (
     _read_request_body,
     _safe_get_request_headers,
@@ -48,6 +51,7 @@ from litellm.proxy.vector_store_endpoints.utils import (
     assert_user_can_access_vector_store_id,
 )
 from litellm.repositories.table_repositories import ManagedVectorStoresRepository
+from litellm.types.utils import ModelResponse
 
 if TYPE_CHECKING:
     from litellm.proxy.utils import PrismaClient
@@ -756,43 +760,53 @@ async def rag_query(
             merged_retrieval_config.get("custom_llm_provider"),
         )
 
-        # Call query
-        response: Final = await litellm.aquery(
-            model=model,
-            messages=messages,
-            retrieval_config=merged_retrieval_config,
-            vector_store_params=store_data,
-            rerank=rerank,
-            stream=stream,
-            router=llm_router,
-            **request_data,
-        )
-
-        hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
-        custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers(
-            user_api_key_dict=user_api_key_dict,
-            call_id=hidden_params.get("litellm_call_id", None) or "",
-            model_id=hidden_params.get("model_id", None) or "",
-            cache_key=hidden_params.get("cache_key", None) or "",
-            api_base=hidden_params.get("api_base", None) or "",
-            version=version,
-            response_cost=hidden_params.get("response_cost", None),
-            request_data=request_data,
-        )
-
-        if isinstance(response, CustomStreamWrapper):
-            return StreamingResponse(
-                select_data_generator(
-                    response=response,
-                    user_api_key_dict=user_api_key_dict,
-                    request_data=request_data,
-                    request=request,
-                ),
-                media_type="text/event-stream",
-                headers=custom_headers,
+        async def query() -> ModelResponse:
+            return await litellm.aquery(
+                model=model,
+                messages=messages,
+                retrieval_config=merged_retrieval_config,
+                vector_store_params=store_data,
+                rerank=rerank,
+                stream=stream,
+                router=llm_router,
+                **request_data,
             )
 
-        fastapi_response.headers.update(custom_headers)
+        def custom_headers_for(response: ModelResponse) -> Mapping[str, str]:
+            hidden_params: Final = getattr(response, "_hidden_params", {}) or {}
+            return ProxyBaseLLMRequestProcessing.get_custom_headers(
+                user_api_key_dict=user_api_key_dict,
+                call_id=hidden_params.get("litellm_call_id", None) or "",
+                model_id=hidden_params.get("model_id", None) or "",
+                cache_key=hidden_params.get("cache_key", None) or "",
+                api_base=hidden_params.get("api_base", None) or "",
+                version=version,
+                response_cost=hidden_params.get("response_cost", None),
+                request_data=request_data,
+            )
+
+        if stream:
+
+            async def produce_stream() -> StreamingResponse:
+                response: Final = await query()
+                return StreamingResponse(
+                    select_data_generator(
+                        response=response,
+                        user_api_key_dict=user_api_key_dict,
+                        request_data=request_data,
+                        request=request,
+                    ),
+                    media_type="text/event-stream",
+                    headers=custom_headers_for(response),
+                )
+
+            return await open_sse_before_first_byte(
+                produce_stream(),
+                ping_interval_seconds=ttft_keepalive_interval(data, llm_router),
+            )
+
+        response: Final = await query()
+        fastapi_response.headers.update(custom_headers_for(response))
         return response
 
     except HTTPException:
diff --git a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py
index 4e063dd0c5b..86f9aeafc08 100644
--- a/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/policy_endpoints/test_endpoints.py
@@ -225,3 +225,71 @@ def test_compute_overall_action_all_passed():
 
 def test_compute_overall_action_empty():
     assert _compute_overall_action([]) == "passed"
+
+
+class TestEnrichPolicyTemplateStreamKeepalive:
+    async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]:
+        import asyncio
+        from unittest.mock import MagicMock
+
+        import litellm
+        import litellm.proxy.management_endpoints.policy_endpoints.endpoints as policy_endpoints
+        import litellm.proxy.proxy_server as proxy_server
+        from fastapi.responses import StreamingResponse
+        from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+        from litellm.proxy.management_endpoints.policy_endpoints.endpoints import (
+            EnrichTemplateRequest,
+            enrich_policy_template_stream,
+        )
+
+        monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval)
+
+        async def _name_chunks():
+            await asyncio.sleep(delay)
+            chunk = MagicMock()
+            chunk.choices = [MagicMock()]
+            chunk.choices[0].delta.content = "Rival Air\n"
+            yield chunk
+
+        class SlowRouter:
+            async def acompletion(self, **kwargs):
+                return _name_chunks()
+
+        async def _no_variations(competitors, model):
+            return {}
+
+        monkeypatch.setattr(proxy_server, "llm_router", SlowRouter())
+        monkeypatch.setattr(policy_endpoints, "_generate_competitor_variations", _no_variations)
+
+        response = await enrich_policy_template_stream(
+            data=EnrichTemplateRequest(
+                template_id="competitor-mention-detection",
+                parameters={"brand_name": "Acme"},
+                model="gpt-5.4-mini",
+            ),
+            request=MagicMock(),
+            user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+        )
+        assert isinstance(response, StreamingResponse)
+        chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator]
+        return chunks, dict(response.headers)
+
+    @pytest.mark.asyncio
+    async def test_endpoint_pings_while_competitor_discovery_is_still_running(self, monkeypatch):
+        chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05)
+
+        assert headers["content-type"].startswith("text/event-stream")
+        assert headers["cache-control"] == "no-cache"
+        assert headers["x-accel-buffering"] == "no"
+        assert chunks[0] == b": ping\n\n"
+        assert chunks.count(b": ping\n\n") >= 3
+        assert b'data: {"type": "competitor", "name": "Rival Air"}\n\n' in chunks
+        assert chunks[-1].startswith(b'data: {"type": "done"')
+
+    @pytest.mark.asyncio
+    async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch):
+        chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15)
+
+        assert b": ping\n\n" not in chunks
+        assert chunks[0] == b'data: {"type": "competitor", "name": "Rival Air"}\n\n'
+        assert chunks[-1].startswith(b'data: {"type": "done"')
diff --git a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py
index 3a32b3cc128..e5616a1975d 100644
--- a/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py
+++ b/tests/test_litellm/proxy/management_endpoints/usage_endpoints/test_ai_usage_chat.py
@@ -466,3 +466,61 @@ class TestUsageAiChatServiceAccountGuard:
                 is_admin=False,
             )
         assert "Endpoint-level guard missing" in str(exc_info.value)
+
+
+class TestUsageAiChatKeepalive:
+    async def _collect_endpoint_body(self, monkeypatch, interval, delay=0.3) -> tuple[list[bytes], dict]:
+        import asyncio
+
+        import litellm
+        from fastapi.responses import StreamingResponse
+        from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
+        from litellm.proxy.management_endpoints.usage_endpoints.endpoints import (
+            ChatMessage,
+            UsageAIChatRequest,
+            usage_ai_chat,
+        )
+
+        monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval)
+
+        async def slow_acompletion(**kwargs):
+            await asyncio.sleep(delay)
+            response = MagicMock()
+            response.choices = [MagicMock()]
+            response.choices[0].message.tool_calls = None
+            response.choices[0].message.content = "Total spend is $50.25"
+            return response
+
+        with patch(  # test-quality-ok: the stream calls the module-level litellm.acompletion directly; no injection seam
+            "litellm.proxy.management_endpoints.usage_endpoints.ai_usage_chat.litellm.acompletion",
+            new=AsyncMock(side_effect=slow_acompletion),
+        ):
+            response = await usage_ai_chat(
+                data=UsageAIChatRequest(messages=[ChatMessage(role="user", content="hi")], model="gpt-4o-mini"),
+                request=MagicMock(),
+                user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
+            )
+            assert isinstance(response, StreamingResponse)
+            chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator]
+        return chunks, dict(response.headers)
+
+    @pytest.mark.asyncio
+    async def test_endpoint_pings_while_the_planning_completion_is_still_running(self, monkeypatch):
+        chunks, headers = await self._collect_endpoint_body(monkeypatch, interval=0.05)
+
+        assert headers["content-type"].startswith("text/event-stream")
+        assert headers["cache-control"] == "no-cache"
+        assert headers["x-accel-buffering"] == "no"
+        assert chunks[0].startswith(b'data: {"type": "status"')
+        assert chunks[1] == b": ping\n\n"
+        assert chunks.count(b": ping\n\n") >= 3
+        assert b'"content": "Total spend is $50.25"' in b"".join(chunks)
+        assert chunks[-1] == b'data: {"type": "done"}\n\n'
+
+    @pytest.mark.asyncio
+    async def test_endpoint_stream_is_untouched_while_keepalives_are_unconfigured(self, monkeypatch):
+        chunks, _ = await self._collect_endpoint_body(monkeypatch, interval=None, delay=0.15)
+
+        assert b": ping\n\n" not in chunks
+        assert chunks[0].startswith(b'data: {"type": "status"')
+        assert chunks[-1] == b'data: {"type": "done"}\n\n'
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
index 5154f738e9a..acb45038df0 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
@@ -5116,3 +5116,93 @@ class TestAzureRouterModelStreamingDispatch:
         assert result.status_code == 200
         body = b"".join([chunk async for chunk in result.body_iterator])
         assert body == upstream_body
+
+
+class TestAzureRouterModelStreamingKeepalive:
+    async def _dispatch(self, monkeypatch, interval, headers_delay=0.0, body_delay=0.0) -> StreamingResponse:
+        import asyncio
+
+        import litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints as ep
+        import litellm.proxy.proxy_server as proxy_server
+        from litellm.passthrough.main import AsyncPassthroughStreamingResponse
+
+        monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval)
+
+        class _StallingBody(httpx.AsyncByteStream):
+            async def __aiter__(self):
+                await asyncio.sleep(body_delay)
+                yield b"data: hello\n\n"
+
+        async def _upstream_response() -> httpx.Response:
+            await asyncio.sleep(headers_delay)
+            return httpx.Response(
+                200,
+                headers={"content-type": "text/event-stream", "x-upstream": "kept"},
+                stream=_StallingBody(),
+                request=httpx.Request("POST", "https://my-azure.openai.azure.com/openai/deployments/gpt-5/x"),
+            )
+
+        logging_obj = MagicMock()
+        logging_obj.async_flush_passthrough_collected_chunks = AsyncMock()
+
+        class StreamingRouter:
+            async def allm_passthrough_route(self, **kwargs):
+                return await AsyncPassthroughStreamingResponse(
+                    response=_upstream_response(),
+                    litellm_logging_obj=logging_obj,
+                    provider_config=MagicMock(),
+                )
+
+        async def fake_get_request_body(_request):
+            return {"model": "gpt-5", "stream": True}
+
+        monkeypatch.setattr(proxy_server, "llm_router", StreamingRouter())
+        monkeypatch.setattr(ep, "get_request_body", fake_get_request_body)
+        monkeypatch.setattr(ep, "is_passthrough_request_using_router_model", lambda *a, **k: True)
+
+        request = MagicMock(spec=Request)
+        request.method = "POST"
+        request.headers = {"content-type": "application/json"}
+        request.query_params = {}
+
+        result = await azure_proxy_route(
+            endpoint="openai/deployments/gpt-5/chat/completions",
+            request=request,
+            fastapi_response=MagicMock(spec=Response),
+            user_api_key_dict=UserAPIKeyAuth(api_key="hashed-token"),
+        )
+        assert isinstance(result, StreamingResponse)
+        return result
+
+    @pytest.mark.asyncio
+    async def test_pings_while_upstream_headers_are_still_pending(self, monkeypatch):
+        result = await self._dispatch(monkeypatch, interval=0.05, headers_delay=0.3)
+
+        chunks = [chunk async for chunk in result.body_iterator]
+
+        assert result.status_code == 200
+        assert result.headers["x-accel-buffering"] == "no"
+        assert chunks[0] == b": ping\n\n"
+        assert chunks.count(b": ping\n\n") >= 3
+        assert b"".join(chunks).endswith(b"data: hello\n\n")
+
+    @pytest.mark.asyncio
+    async def test_pings_while_upstream_body_is_still_pending(self, monkeypatch):
+        result = await self._dispatch(monkeypatch, interval=0.05, body_delay=0.3)
+
+        chunks = [chunk async for chunk in result.body_iterator]
+
+        assert result.status_code == 200
+        assert result.headers["x-upstream"] == "kept"
+        assert chunks[0] == b": ping\n\n"
+        assert chunks.count(b": ping\n\n") >= 3
+        assert chunks[-1] == b"data: hello\n\n"
+
+    @pytest.mark.asyncio
+    async def test_relays_upstream_bytes_untouched_while_keepalives_are_unconfigured(self, monkeypatch):
+        result = await self._dispatch(monkeypatch, interval=None, headers_delay=0.15, body_delay=0.15)
+
+        chunks = [chunk async for chunk in result.body_iterator]
+
+        assert result.headers["x-upstream"] == "kept"
+        assert chunks == [b"data: hello\n\n"]
diff --git a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
index fdaad567f95..87e10ce7e8d 100644
--- a/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
+++ b/tests/test_litellm/proxy/proxy_server/test_streaming_helpers.py
@@ -1819,3 +1819,92 @@ async def test_run_thread_stream_is_untouched_while_keepalives_are_unconfigured(
 
     assert not any(chunk.startswith(": ping") for chunk in chunks)
     assert chunks[-1] == "data: [DONE]\n\n"
+
+
+# ---------------------------------------------------------------------------
+# async_queue_request: SSE keepalives during the time-to-first-token
+# ---------------------------------------------------------------------------
+
+
+async def _queue_streaming(monkeypatch, interval, delay=0.3, fails_with=None):
+    _patch_logging_flags(monkeypatch)
+    monkeypatch.setattr(litellm, "sse_keepalive_ping_interval_seconds", interval)
+
+    router = MagicMock()
+    router.get_model_list.return_value = []
+
+    async def _schedule_after_the_scheduler_queue_drains(**kwargs):
+        await asyncio.sleep(delay)
+        if fails_with is not None:
+            raise fails_with
+        return _async_iter([_simple_chunk(content="queued reply")])
+
+    router.schedule_acompletion = _schedule_after_the_scheduler_queue_drains
+    monkeypatch.setattr(ps, "llm_router", router)
+
+    request = MagicMock()
+    request.url = "http://testserver/queue/chat/completions"
+    request.method = "POST"
+    request.headers = {}
+    request.json = AsyncMock(
+        return_value={
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "hi"}],
+            "priority": 0,
+            "stream": True,
+        }
+    )
+    request.is_disconnected = AsyncMock(return_value=False)
+
+    return await ps.async_queue_request(
+        request=request,
+        fastapi_response=Response(),
+        user_api_key_dict=_user_auth(),
+    )
+
+
+@pytest.mark.asyncio
+async def test_queue_request_pings_while_the_scheduler_is_still_waiting(monkeypatch):
+    response = await _queue_streaming(monkeypatch, interval=0.05)
+
+    assert isinstance(response, StreamingResponse)
+    assert response.headers["x-accel-buffering"] == "no"
+    chunks = [chunk async for chunk in response.body_iterator]
+
+    assert chunks[0] == b": ping\n\n"
+    assert chunks.count(b": ping\n\n") >= 3
+    assert b'"content":"queued reply"' in chunks[-2]
+    assert chunks[-1] == b"data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_queue_request_audits_a_failure_that_arrives_after_the_first_ping(monkeypatch):
+    audited = []
+
+    async def _record_failure(*, user_api_key_dict, original_exception, request_data, **kwargs):
+        audited.append(original_exception)
+        return None
+
+    monkeypatch.setattr(ps.proxy_logging_obj, "post_call_failure_hook", _record_failure)
+
+    boom = RuntimeError("scheduler died after the wire was already open")
+    response = await _queue_streaming(monkeypatch, interval=0.05, fails_with=boom)
+
+    assert isinstance(response, StreamingResponse)
+    chunks = [chunk async for chunk in response.body_iterator]
+
+    assert chunks[0] == b": ping\n\n"
+    assert audited == [boom]
+    assert json.loads(chunks[-2].removeprefix(b"data: "))["error"]["code"] == "500"
+    assert chunks[-1] == b"data: [DONE]\n\n"
+
+
+@pytest.mark.asyncio
+async def test_queue_request_stream_is_untouched_while_keepalives_are_unconfigured(monkeypatch):
+    response = await _queue_streaming(monkeypatch, interval=None, delay=0.15)
+
+    assert isinstance(response, StreamingResponse)
+    chunks = [chunk if isinstance(chunk, bytes) else chunk.encode() async for chunk in response.body_iterator]
+
+    assert not any(chunk.startswith(b": ping") for chunk in chunks)
+    assert chunks[-1] == b"data: [DONE]\n\n"
diff --git a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py
index a176e91eaa4..832435711c6 100644
--- a/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py
+++ b/tests/test_litellm/proxy/rag_endpoints/test_rag_endpoints.py
@@ -11,7 +11,6 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import pytest
 from fastapi.testclient import TestClient
 
-
 from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.proxy_server import app
@@ -324,6 +323,100 @@ def test_rag_query_stream_returns_event_stream(client_internal_user):
     assert "data: [DONE]" in response.text
 
 
+def test_rag_query_stream_pings_while_retrieval_is_still_running(client_internal_user, monkeypatch):
+    import asyncio
+
+    import litellm as litellm_module
+
+    monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 0.05)
+
+    async def slow_aquery(**kwargs):
+        await asyncio.sleep(0.3)
+        return await litellm_module.acompletion(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "What is the codename?"}],
+            mock_response="The codename is AZURE-FALCON-42.",
+            stream=True,
+            api_key="test-key",
+        )
+
+    with (
+        patch(  # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam
+            "litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
+            new=AsyncMock(side_effect=slow_aquery),
+        ),
+        patch("litellm.vector_store_registry", None),  # test-quality-ok: proxy module global, no injection seam
+        patch("litellm.proxy.proxy_server.prisma_client", None),  # test-quality-ok: proxy module global, no injection seam
+    ):
+        response = client_internal_user.post(
+            "/v1/rag/query",
+            json={
+                "model": "gpt-4o-mini",
+                "messages": [{"role": "user", "content": "What is the codename?"}],
+                "retrieval_config": {
+                    "vector_store_id": "vs_test_123",
+                    "custom_llm_provider": "openai",
+                },
+                "stream": True,
+            },
+        )
+
+    assert response.status_code == 200, response.text
+    assert response.headers.get("content-type", "").startswith("text/event-stream")
+    assert response.headers["x-accel-buffering"] == "no"
+    assert response.text.startswith(": ping\n\n")
+    assert response.text.count(": ping\n\n") >= 3
+    assert '"object":"chat.completion.chunk"' in response.text
+    assert response.text.endswith("data: [DONE]\n\n")
+
+
+def test_rag_query_stream_keeps_response_headers_when_retrieval_beats_the_keepalive(
+    client_internal_user, monkeypatch
+):
+    import litellm as litellm_module
+
+    monkeypatch.setattr(litellm_module, "sse_keepalive_ping_interval_seconds", 5)
+
+    async def fast_aquery(**kwargs):
+        response = await litellm_module.acompletion(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "What is the codename?"}],
+            mock_response="The codename is AZURE-FALCON-42.",
+            stream=True,
+            api_key="test-key",
+        )
+        response._hidden_params["response_cost"] = 3.45e-06
+        return response
+
+    with (
+        patch(  # test-quality-ok: the handler calls the module-level litellm.aquery directly; no injection seam
+            "litellm.proxy.rag_endpoints.endpoints.litellm.aquery",
+            new=AsyncMock(side_effect=fast_aquery),
+        ),
+        patch("litellm.vector_store_registry", None),  # test-quality-ok: proxy module global, no injection seam
+        patch("litellm.proxy.proxy_server.prisma_client", None),  # test-quality-ok: proxy module global, no injection seam
+    ):
+        response = client_internal_user.post(
+            "/v1/rag/query",
+            json={
+                "model": "gpt-4o-mini",
+                "messages": [{"role": "user", "content": "What is the codename?"}],
+                "retrieval_config": {
+                    "vector_store_id": "vs_test_123",
+                    "custom_llm_provider": "openai",
+                },
+                "stream": True,
+            },
+        )
+
+    assert response.status_code == 200, response.text
+    assert response.headers.get("content-type", "").startswith("text/event-stream")
+    assert response.headers.get("x-litellm-response-cost") == "3.45e-06"
+    assert not response.text.startswith(": ping")
+    assert '"object":"chat.completion.chunk"' in response.text
+    assert response.text.endswith("data: [DONE]\n\n")
+
+
 def test_rag_query_merges_managed_store_params(client_internal_user):
     """
     Regression: /v1/rag/query must consult the managed vector store registry

From 4e18c0f63a33bd5289dbb42c89bb22c276a809c1 Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Thu, 3 Sep 2026 18:29:32 -0700
Subject: [PATCH 36/38] fix(azure): restrict the storage credential chain to
 deployment identities (#39637)

* fix(azure): restrict the storage credential chain to deployment identities

The keyless Azure Storage path walks the full DefaultAzureCredential chain, so a
proxy with no storage service principal authenticates as whichever identity the
host happens to carry: an operator's az login on a workstation, or the
AZURE_CLIENT_ID/AZURE_CLIENT_SECRET service principal set for Azure OpenAI.
Neither is the identity granted Storage Blob Data Contributor.

Narrow the chain to workload identity and managed identity, the two credentials
a deployment legitimately holds. Azure OpenAI, Postgres IAM auth and the other
callers of get_azure_ad_token_provider keep the full chain.

* test(azure): read the credential chain off the mock instead of an accumulator

* chore: drop a stray launch traceback committed at the repo root

* fix(azure): let the storage chain reach a system assigned managed identity

DefaultAzureCredential keeps one managed identity link and pins it to
AZURE_CLIENT_ID, so a host that sets that variable for Azure OpenAI and runs as
a system assigned identity never got asked for a storage token. Build the chain
from the three credentials a deployment can carry instead of subtracting the
ones it cannot.
---
 .../azure_storage/azure_storage.py            |   2 +-
 .../get_azure_ad_token_provider.py            |  24 +++
 .../get_azure_ad_token_provider.py            |   1 +
 .../azure_storage/test_azure_storage.py       |  21 ++-
 .../test_get_azure_ad_token_provider.py       | 138 ++++++++++++++++++
 5 files changed, 184 insertions(+), 2 deletions(-)

diff --git a/litellm/integrations/azure_storage/azure_storage.py b/litellm/integrations/azure_storage/azure_storage.py
index 16ef6920114..13058bf4f22 100644
--- a/litellm/integrations/azure_storage/azure_storage.py
+++ b/litellm/integrations/azure_storage/azure_storage.py
@@ -36,7 +36,7 @@ AZURE_STORAGE_TOKEN_SCOPE: Final = "https://storage.azure.com/.default"
 def _cached_credential_chain_token_provider() -> Callable[[], str]:
     return get_azure_ad_token_provider(
         azure_scope=AZURE_STORAGE_TOKEN_SCOPE,
-        azure_credential=AzureCredentialType.DefaultAzureCredential,
+        azure_credential=AzureCredentialType.DeploymentIdentityCredential,
     )
 
 
diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py
index c2dc09bc65d..5d056ea3fe0 100644
--- a/litellm/secret_managers/get_azure_ad_token_provider.py
+++ b/litellm/secret_managers/get_azure_ad_token_provider.py
@@ -57,9 +57,11 @@ def get_azure_ad_token_provider(
     from azure import identity
     from azure.identity import (
         CertificateCredential,
+        ChainedTokenCredential,
         ClientSecretCredential,
         DefaultAzureCredential,
         ManagedIdentityCredential,
+        WorkloadIdentityCredential,
         get_bearer_token_provider,
     )
 
@@ -101,6 +103,28 @@ def get_azure_ad_token_provider(
         # DefaultAzureCredential doesn't require explicit environment variables
         # It automatically discovers credentials from the environment (managed identity, CLI, etc.)
         credential = DefaultAzureCredential()
+    elif cred == AzureCredentialType.DeploymentIdentityCredential:
+        # DefaultAzureCredential cannot express this: excluding its developer credentials still
+        # leaves one managed identity link, which AZURE_CLIENT_ID pins to a user assigned identity,
+        # so a host running as a system assigned identity never gets asked
+        workload_client_id: Final = os.environ.get("AZURE_CLIENT_ID")
+        workload_tenant_id: Final = os.environ.get("AZURE_TENANT_ID")
+        workload_token_file: Final = os.environ.get("AZURE_FEDERATED_TOKEN_FILE")
+        credential = ChainedTokenCredential(
+            *(
+                (
+                    WorkloadIdentityCredential(
+                        client_id=workload_client_id,
+                        tenant_id=workload_tenant_id,
+                        token_file_path=workload_token_file,
+                    ),
+                )
+                if workload_client_id and workload_tenant_id and workload_token_file
+                else ()
+            ),
+            *((ManagedIdentityCredential(client_id=workload_client_id),) if workload_client_id else ()),
+            ManagedIdentityCredential(),
+        )
     else:
         cred_cls: Final = getattr(identity, cred)
         credential = cred_cls()
diff --git a/litellm/types/secret_managers/get_azure_ad_token_provider.py b/litellm/types/secret_managers/get_azure_ad_token_provider.py
index 5d2f7409f95..6b4700d081d 100644
--- a/litellm/types/secret_managers/get_azure_ad_token_provider.py
+++ b/litellm/types/secret_managers/get_azure_ad_token_provider.py
@@ -6,3 +6,4 @@ class AzureCredentialType(str, Enum):
     ManagedIdentityCredential = "ManagedIdentityCredential"
     CertificateCredential = "CertificateCredential"
     DefaultAzureCredential = "DefaultAzureCredential"
+    DeploymentIdentityCredential = "DeploymentIdentityCredential"
diff --git a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py
index a96eae0f9c3..6e1dab4a71a 100644
--- a/tests/test_litellm/integrations/azure_storage/test_azure_storage.py
+++ b/tests/test_litellm/integrations/azure_storage/test_azure_storage.py
@@ -42,6 +42,7 @@ def workload_identity_env_vars(monkeypatch):
         "AZURE_STORAGE_ENDPOINT_SUFFIX",
         "AZURE_CLIENT_SECRET",
         "AZURE_CREDENTIAL",
+        "AZURE_TOKEN_CREDENTIALS",
         "AZURE_SCOPE",
     ):
         monkeypatch.delenv(unset, raising=False)
@@ -206,10 +207,28 @@ def test_default_chain_provider_is_storage_scoped_and_built_once_per_process():
     assert first() == "chain-token"
     mock_builder.assert_called_once_with(
         azure_scope="https://storage.azure.com/.default",
-        azure_credential=AzureCredentialType.DefaultAzureCredential,
+        azure_credential=AzureCredentialType.DeploymentIdentityCredential,
     )
 
 
+def test_storage_chain_reaches_only_the_identities_a_deployment_carries(workload_identity_env_vars):
+    """
+    The chain runs on a server, where a developer sign-in is a person and not the deployment, so
+    the storage token must come from workload identity or managed identity or from nothing
+    """
+    _cached_credential_chain_token_provider.cache_clear()
+    with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "chain-token") as bearer:
+        _cached_credential_chain_token_provider()
+    _cached_credential_chain_token_provider.cache_clear()
+
+    bearer.assert_called_once()
+    with bearer.call_args.args[0] as chain:
+        assert {type(link).__name__ for link in chain.credentials} == {
+            "WorkloadIdentityCredential",
+            "ManagedIdentityCredential",
+        }
+
+
 @pytest.mark.asyncio
 async def test_chain_tokens_are_read_from_the_provider_on_every_refresh(
     workload_identity_env_vars,
diff --git a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py
index c9ec22ab0df..4bc7c21d8d2 100644
--- a/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py
+++ b/tests/test_litellm/secret_managers/test_get_azure_ad_token_provider.py
@@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch
 # Adds the grandparent directory to sys.path to allow importing project modules
 
 import pytest
+from azure.core.exceptions import ClientAuthenticationError
 
 from litellm.secret_managers.get_azure_ad_token_provider import (
     get_azure_ad_token_provider,
@@ -16,6 +17,143 @@ from litellm.types.secret_managers.get_azure_ad_token_provider import (
 )
 
 
+class TestDeploymentIdentityCredential:
+    @staticmethod
+    def _chain_for(credential_type):
+        with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer:
+            get_azure_ad_token_provider(
+                azure_scope="https://storage.azure.com/.default",
+                azure_credential=credential_type,
+            )
+        bearer.assert_called_once()
+        with bearer.call_args.args[0] as chain:
+            return {type(link).__name__ for link in chain.credentials}
+
+    @staticmethod
+    def _managed_identity_client_ids(credential_type):
+        with patch("azure.identity.get_bearer_token_provider", return_value=lambda: "token") as bearer:
+            get_azure_ad_token_provider(
+                azure_scope="https://storage.azure.com/.default",
+                azure_credential=credential_type,
+            )
+        with bearer.call_args.args[0] as chain:
+            return [
+                (link._credential._settings or {}).get("client_id")
+                for link in chain.credentials
+                if type(link).__name__ == "ManagedIdentityCredential"
+            ]
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "workload-identity-client-id",
+            "AZURE_TENANT_ID": "workload-identity-tenant-id",
+            "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token",
+        },
+        clear=True,
+    )
+    def test_deployment_identity_reaches_workload_and_managed_identity_only(self):
+        assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == {
+            "WorkloadIdentityCredential",
+            "ManagedIdentityCredential",
+        }
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "workload-identity-client-id",
+            "AZURE_TENANT_ID": "workload-identity-tenant-id",
+            "AZURE_FEDERATED_TOKEN_FILE": "/var/run/secrets/azure/tokens/azure-identity-token",
+            "AZURE_TOKEN_CREDENTIALS": "dev",
+        },
+        clear=True,
+    )
+    def test_deployment_identity_survives_a_developer_only_token_credentials_setting(self):
+        """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, which is every
+        credential this chain drops, so the deployment's own identity has to win over it"""
+        assert self._chain_for(AzureCredentialType.DeploymentIdentityCredential) == {
+            "WorkloadIdentityCredential",
+            "ManagedIdentityCredential",
+        }
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "azure-openai-client-id",
+            "AZURE_CLIENT_SECRET": "azure-openai-client-secret",
+            "AZURE_TENANT_ID": "azure-openai-tenant-id",
+        },
+        clear=True,
+    )
+    def test_default_azure_credential_keeps_its_full_chain(self):
+        """Azure OpenAI callers pass DefaultAzureCredential and must be unaffected by the
+        narrowing that the storage callback asks for"""
+        full_chain = self._chain_for(AzureCredentialType.DefaultAzureCredential)
+
+        assert "EnvironmentCredential" in full_chain
+        assert "AzureCliCredential" in full_chain
+        assert "EnvironmentCredential" not in self._chain_for(
+            AzureCredentialType.DeploymentIdentityCredential
+        )
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "azure-openai-client-id",
+            "AZURE_CLIENT_SECRET": "azure-openai-client-secret",
+            "AZURE_TENANT_ID": "azure-openai-tenant-id",
+        },
+        clear=True,
+    )
+    def test_deployment_identity_refuses_to_mint_a_token_for_a_configured_service_principal(self):
+        """A host carrying only an Azure OpenAI client secret must get no token at all, and the
+        refusal must name the identities that were actually tried"""
+        provider = get_azure_ad_token_provider(
+            azure_scope="https://storage.azure.com/.default",
+            azure_credential=AzureCredentialType.DeploymentIdentityCredential,
+        )
+
+        with pytest.raises(ClientAuthenticationError) as refusal:
+            provider()
+
+        assert "ManagedIdentityCredential" in str(refusal.value)
+        assert "EnvironmentCredential" not in str(refusal.value)
+        assert "AzureCliCredential" not in str(refusal.value)
+        assert "azure-openai-client-secret" not in str(refusal.value)
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "azure-openai-client-id",
+            "AZURE_CLIENT_SECRET": "azure-openai-client-secret",
+            "AZURE_TENANT_ID": "azure-openai-tenant-id",
+        },
+        clear=True,
+    )
+    def test_deployment_identity_still_reaches_a_system_assigned_managed_identity(self):
+        """AZURE_CLIENT_ID names one identity for the whole proxy, and pointing it at Azure OpenAI
+        must not hide the system assigned identity the host runs as"""
+        client_ids = self._managed_identity_client_ids(AzureCredentialType.DeploymentIdentityCredential)
+
+        assert "azure-openai-client-id" in client_ids
+        assert None in client_ids
+
+    @patch.dict(
+        os.environ,
+        {
+            "AZURE_CLIENT_ID": "user-assigned-identity-client-id",
+            "AZURE_TOKEN_CREDENTIALS": "dev",
+        },
+        clear=True,
+    )
+    def test_deployment_identity_keeps_the_user_assigned_identity_under_a_dev_only_setting(self):
+        """AZURE_TOKEN_CREDENTIALS=dev asks the SDK for developer credentials only, and the
+        identity a host actually runs as has to survive that"""
+        assert "user-assigned-identity-client-id" in self._managed_identity_client_ids(
+            AzureCredentialType.DeploymentIdentityCredential
+        )
+
+
 class TestGetAzureAdTokenProvider:
     @patch.dict(
         os.environ,

From 7ae352e5cf0cfd8ae7e5f2c153938c19c1e9c773 Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Fri, 4 Sep 2026 01:36:46 +0000
Subject: [PATCH 37/38] fix(model_checks): drop wildcard routes like bedrock/*
 from /v1/models (#31731)

* fix: remove wildcard routes from /v1/models response

Wildcard routes like bedrock/* were leaking into the /v1/models response
because _get_wildcard_models only removed them from unique_models in the
fallback branches (no router or no deployment), but not when the router
had a matching deployment. Now wildcards are always removed from the base
list; they are only re-added to the result when return_wildcard_routes=True
is explicitly passed.

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

* refactor(model_checks): collapse wildcard expansion branches and tighten regression tests

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

---------

Co-authored-by: Krrish Dholakia <krrishdholakia@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: yassin <yassin@berri.ai>
---
 litellm/proxy/auth/model_checks.py            | 31 +++-----
 .../proxy/auth/test_model_checks.py           | 78 +++++++++++++++++++
 2 files changed, 88 insertions(+), 21 deletions(-)

diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py
index 1625198892f..de2ca4762f1 100644
--- a/litellm/proxy/auth/model_checks.py
+++ b/litellm/proxy/auth/model_checks.py
@@ -387,33 +387,22 @@ def _get_wildcard_models(
     all_wildcard_models: Final = []
     for model in unique_models:
         if _check_wildcard_routing(model=model):
-            if return_wildcard_routes:  # will add the wildcard route to the list eg: anthropic/*.
+            if return_wildcard_routes:
                 all_wildcard_models.append(model)
 
-            ## get litellm params from model
-            if llm_router is not None:
-                model_list = llm_router.get_model_list(model_name=model, team_id=team_id)
-                if model_list:
-                    for router_model in model_list:
-                        wildcard_models = get_known_models_from_wildcard(
+            models_to_remove.add(model)
+
+            model_list = llm_router.get_model_list(model_name=model, team_id=team_id) if llm_router else None
+            if model_list:
+                for router_model in model_list:
+                    all_wildcard_models.extend(
+                        get_known_models_from_wildcard(
                             wildcard_model=model,
                             litellm_params=LiteLLM_Params(**router_model["litellm_params"]),
                         )
-                        all_wildcard_models.extend(wildcard_models)
-                else:
-                    # Router has no deployment for this wildcard (e.g., BYOK team models)
-                    # Fall back to expanding from known provider models
-                    wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)
-                    if wildcard_models:
-                        models_to_remove.add(model)
-                        all_wildcard_models.extend(wildcard_models)
+                    )
             else:
-                # get all known provider models
-                wildcard_models = get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)
-
-                if wildcard_models:
-                    models_to_remove.add(model)
-                    all_wildcard_models.extend(wildcard_models)
+                all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None))
 
     for model in models_to_remove:
         unique_models.remove(model)
diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py
index 62073f4bf51..d58683fd1e5 100644
--- a/tests/test_litellm/proxy/auth/test_model_checks.py
+++ b/tests/test_litellm/proxy/auth/test_model_checks.py
@@ -754,6 +754,84 @@ def test_expand_wildcard_invalid_litellm_params_passthrough():
     assert result == [deployment]
 
 
+def test_get_complete_model_list_excludes_wildcard_routes_by_default():
+    """Regression (LIT-4108): a wildcard with a matching router deployment leaked into /v1/models."""
+    from litellm import Router
+    from litellm.proxy.auth.model_checks import get_complete_model_list
+
+    router = Router(
+        model_list=[
+            {
+                "model_name": "bedrock/*",
+                "litellm_params": {"model": "bedrock/*"},
+            },
+            {
+                "model_name": "gpt-4",
+                "litellm_params": {"model": "openai/gpt-4"},
+            },
+        ]
+    )
+
+    result = get_complete_model_list(
+        key_models=[],
+        team_models=[],
+        proxy_model_list=["bedrock/*", "gpt-4"],
+        user_model=None,
+        infer_model_from_keys=False,
+        return_wildcard_routes=False,
+        llm_router=router,
+    )
+
+    assert "bedrock/*" not in result
+    assert "gpt-4" in result
+    assert any(m.startswith("bedrock/") for m in result)
+
+
+def test_get_complete_model_list_excludes_wildcard_routes_without_router():
+    from litellm.proxy.auth.model_checks import get_complete_model_list
+
+    result = get_complete_model_list(
+        key_models=[],
+        team_models=[],
+        proxy_model_list=["bedrock/*", "gpt-4"],
+        user_model=None,
+        infer_model_from_keys=False,
+        return_wildcard_routes=False,
+        llm_router=None,
+    )
+
+    assert "bedrock/*" not in result
+    assert "gpt-4" in result
+    assert any(m.startswith("bedrock/") for m in result)
+
+
+def test_get_complete_model_list_includes_wildcard_routes_when_requested():
+    from litellm import Router
+    from litellm.proxy.auth.model_checks import get_complete_model_list
+
+    router = Router(
+        model_list=[
+            {
+                "model_name": "bedrock/*",
+                "litellm_params": {"model": "bedrock/*"},
+            },
+        ]
+    )
+
+    result = get_complete_model_list(
+        key_models=[],
+        team_models=[],
+        proxy_model_list=["bedrock/*"],
+        user_model=None,
+        infer_model_from_keys=False,
+        return_wildcard_routes=True,
+        llm_router=router,
+    )
+
+    assert result.count("bedrock/*") == 1
+    assert any(m.startswith("bedrock/") and m != "bedrock/*" for m in result)
+
+
 def test_add_known_models_refreshes_models_by_provider_for_wildcard_expansion():
     """models_by_provider was a frozen import-time snapshot of set unions, so cost map
     reloads (which call add_known_models) never reached wildcard expansion until a

From b7f53ce9a9af08e920bdd4d991b0013e4f52abb6 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 3 Sep 2026 18:39:37 -0700
Subject: [PATCH 38/38] fix(mcp): pre-flight the ID-JAG credential at the
 transport edge (#35392)

---
 .../mcp_server/mcp_server_manager.py          |  40 +++--
 .../proxy/_experimental/mcp_server/server.py  |  14 +-
 .../mcp_server/test_mcp_server.py             | 123 +++++++++++++
 .../mcp_server/test_mcp_server_manager.py     | 167 +++++++++++++++++-
 4 files changed, 324 insertions(+), 20 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index 0c3932d2ba1..ce4928ff83d 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -3638,30 +3638,48 @@ class MCPServerManager:
         user_api_key_auth: UserAPIKeyAuth | None,
         raw_headers: Mapping[str, str] | None = None,
     ) -> None:
-        """Run the OBO exchange for a caller-supplied subject at the transport edge.
+        """Mint an exchange-backed server's upstream credential at the transport edge.
 
         Single-server routes call this before the MCP session opens, where an HTTP status and
         ``WWW-Authenticate`` still reach the client. A rejected subject raises the RFC 9728
         challenge and any other ``CredError`` maps onto its public HTTP status, so an exchange
         failure surfaces as a failure instead of the session continuing into an empty tool list.
         A successful exchange is cached by the exchanger, so the session's list/call reuses it.
+
+        Each mode pre-flights only where it would resolve the subject the session goes on to use,
+        which is what keeps the pre-flight from reaching a verdict the session would contradict.
+        ``oauth2_token_exchange`` mints from the caller's inbound bearer, so without one there is
+        nothing to exchange and the missing-subject case stays the preemptive challenge's job.
+        ``oauth2_id_jag`` is the mirror image: tool listing resolves it from the identity assertion
+        captured for this user at SSO login and never from the inbound bearer, so the pre-flight is
+        faithful exactly when no identity bearer was sent (a LiteLLM key in ``Authorization`` is not one),
+        and a caller that did send one is passed through
+        untouched rather than judged against a subject the listing will not use. That store-sourced
+        case is the one whose missing-assertion 412 and store-outage 503 the session cannot report.
+        Only OBO has a discovery challenge to raise; ID-JAG's failures are plain statuses whose body
+        already names what the user has to do, so they map through ``raise_public`` as at egress.
         """
-        if server.auth_type != MCPAuth.oauth2_token_exchange:
-            return
-        if not self._extract_bearer_token(oauth2_headers, None):
-            return
-        resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
-        spec: Final = to_server_spec(resolved_server)
-        if spec is None or not isinstance(spec.config, TokenExchangeConfig):
-            return
         subject_token: Final = self._extract_subject_token(oauth2_headers, raw_headers, user_api_key_auth)
-        if subject_token is None:
+        match server.auth_type:
+            case MCPAuth.oauth2_token_exchange:
+                if not self._extract_bearer_token(oauth2_headers, None):
+                    return
+            case MCPAuth.oauth2_id_jag:
+                if subject_token is not None:
+                    return
+            case _:
+                return
+        resolved_server: Final = await self.ensure_oauth_metadata_discovered(server)
+        spec: Final = _to_server_spec_fail_closed(resolved_server)
+        if spec is None or not isinstance(spec.config, (TokenExchangeConfig, IdJagConfig)):
+            return
+        if subject_token is None and isinstance(spec.config, TokenExchangeConfig):
             raise_token_exchange_challenge(resolved_server, root_path=get_server_root_path())
         match await self._cred_provider.resolve_credentials(to_subject(user_api_key_auth, subject_token), spec):
             case Ok(_):
                 return
             case Error(err):
-                if err.tag == "unauthorized":
+                if err.tag == "unauthorized" and isinstance(spec.config, TokenExchangeConfig):
                     raise_token_exchange_challenge(
                         resolved_server,
                         root_path=get_server_root_path(),
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index bb075220530..3d7c947a913 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -3851,15 +3851,15 @@ if MCP_AVAILABLE:
 
                 raise_token_exchange_challenge(server, root_path=get_server_root_path())
 
-            # token_exchange (OBO) with a subject present: run the exchange here at the transport
-            # edge, so a rejected subject raises the RFC 9728 challenge (and a gateway fault its
-            # public status) instead of the session opening and list_tools masking the failure as
-            # an empty tool list. Gated to single-server routes; the multi-server aggregate keeps
-            # absorbing per-server auth failures so one bad server cannot 401 the whole connect.
+            # Exchange-backed modes (token_exchange's OBO mint, id_jag's stored-assertion mint): run
+            # the exchange here at the transport edge, so a rejected subject raises the RFC 9728
+            # challenge and any other failure its public status, instead of the session opening and
+            # list_tools masking it as an empty tool list. The manager owns which modes pre-flight
+            # and what each mints from. Gated to single-server routes the key may reach; the
+            # multi-server aggregate keeps absorbing per-server auth failures so one bad server
+            # cannot 401 the whole connect.
             if (
                 server
-                and server.auth_type == MCPAuth.oauth2_token_exchange
-                and oauth2_headers
                 and len(mcp_servers or []) == 1
                 and server.server_id
                 in frozenset(
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 0fd35e674b7..9a6815a61e5 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -8201,6 +8201,129 @@ class TestPreemptive401ModeAware:
         await self._run(delegate, self.LITELLM_KEY_HEADERS, has_stored_token=False)
 
 
+class TestSingleServerPreflightReachesIdJag:
+    """The connect-time preflight is what turns a credential failure into an HTTP status the client
+    can read. An oauth2_id_jag server has to reach it: its subject comes from the assertion stored at
+    SSO login, so the failure is decided before any IdP call and there is nothing later in the session
+    that can report it (tools/list degrades to an empty list, tools/call to 'tool not found')."""
+
+    def _id_jag_server(self) -> MCPServer:
+        return MCPServer(
+            server_id="id-idjag",
+            name="idjag",
+            alias="idjag",
+            server_name="idjag",
+            url="https://idjag.test/mcp",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_id_jag,
+            client_id="gateway-client",
+            client_secret="gateway-secret",
+            token_exchange_endpoint="https://org-idp.test/oauth2/token",
+            id_jag_resource_token_endpoint="https://resource-as.test/oauth2/token",
+            mcp_info={"server_name": "idjag"},
+        )
+
+    async def _run(self, server: MCPServer, mcp_servers: list[str], preflight: AsyncMock) -> None:
+        from litellm.proxy._experimental.mcp_server import server as server_module
+
+        with (
+            patch.object(  # test-quality-ok: route wiring must use the manager's configured server
+                server_module.global_mcp_server_manager,
+                "get_mcp_server_by_name",
+                return_value=server,
+            ),
+            patch.object(  # test-quality-ok: route wiring must invoke the manager preflight
+                server_module.global_mcp_server_manager,
+                "preflight_token_exchange",
+                preflight,
+            ),
+            patch.object(  # test-quality-ok: allowed-set resolution needs the DB; the test controls its answer
+                server_module, "_get_allowed_mcp_servers", AsyncMock(return_value=[server])
+            ),
+        ):
+            await server_module._raise_preemptive_401_for_unauthenticated_servers(
+                scope={"type": "http", "method": "POST", "path": "/mcp/idjag", "headers": []},
+                mcp_servers=mcp_servers,
+                oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"},
+                mcp_server_auth_headers=None,
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"),
+                client_ip=None,
+            )
+
+    @pytest.mark.asyncio
+    async def test_id_jag_single_server_route_surfaces_the_preflight_status(self):
+        """The 412 the preflight raises must propagate out of connect, not be swallowed."""
+        server = self._id_jag_server()
+        preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion"))
+
+        with pytest.raises(HTTPException) as exc:
+            await self._run(server, ["idjag"], preflight)
+
+        assert exc.value.status_code == 412
+        assert preflight.await_args.kwargs["server"] is server
+
+    @pytest.mark.asyncio
+    async def test_token_exchange_without_a_bearer_still_challenges_and_never_pre_flights(self):
+        """The already-shipped OBO path must be untouched by the call site dropping its mode test.
+        A token_exchange server with no inbound bearer has nothing to exchange, so it still gets the
+        RFC 9728 discovery challenge from the block above and the preflight is never reached; pushing
+        a subject-less exchange through the resolver would turn that challenge into some other status
+        and strand a client that only had to SSO and retry."""
+        from litellm.proxy._experimental.mcp_server import server as server_module
+
+        token_exchange = MCPServer(
+            server_id="id-obo",
+            name="obo",
+            alias="obo",
+            server_name="obo",
+            url="https://obo.test/mcp",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_token_exchange,
+            token_exchange_endpoint="https://idp.test/oauth2/token",
+            client_id="cid",
+            client_secret="csec",
+            mcp_info={"server_name": "obo"},
+        )
+        preflight = AsyncMock()
+
+        with (
+            patch.object(  # test-quality-ok: route wiring must use the manager's configured server
+                server_module.global_mcp_server_manager,
+                "get_mcp_server_by_name",
+                return_value=token_exchange,
+            ),
+            patch.object(  # test-quality-ok: route wiring must invoke the manager preflight
+                server_module.global_mcp_server_manager,
+                "preflight_token_exchange",
+                preflight,
+            ),
+            pytest.raises(HTTPException) as exc,
+        ):
+            await server_module._raise_preemptive_401_for_unauthenticated_servers(
+                scope={"type": "http", "method": "POST", "path": "/mcp/obo", "headers": []},
+                mcp_servers=["obo"],
+                oauth2_headers=None,
+                mcp_server_auth_headers=None,
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"),
+                client_ip=None,
+            )
+
+        assert exc.value.status_code == 401
+        headers = exc.value.headers or {}
+        assert "resource_metadata" in (headers.get("WWW-Authenticate") or headers.get("www-authenticate") or "")
+        preflight.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_id_jag_multi_server_route_still_absorbs_the_failure(self):
+        """The aggregate contract is unchanged: with more than one target the preflight does not run,
+        so one server with no stored assertion cannot fail the whole connect."""
+        preflight = AsyncMock(side_effect=HTTPException(status_code=412, detail="no stored assertion"))
+
+        await self._run(self._id_jag_server(), ["idjag", "other"], preflight)
+
+        preflight.assert_not_awaited()
+
+
 def _make_obo_server(alias: str) -> MCPServer:
     return MCPServer(
         server_id=f"id-{alias}",
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
index 482f779bcb8..02b1a19081a 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
@@ -2616,6 +2616,165 @@ class TestMCPServerManager:
         await manager.preflight_token_exchange(server=server, oauth2_headers=None, user_api_key_auth=None)
         assert resolved == ["good-subject"]
 
+    def _id_jag_server(self, server_id: str) -> "MCPServer":
+        return MCPServer(
+            server_id=server_id,
+            name=f"{server_id}-server",
+            url="https://up.example.com/mcp",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.oauth2_id_jag,
+            client_id="gateway-client",
+            client_secret="gateway-secret",
+            token_exchange_endpoint="https://org-idp.example/oauth2/token",
+            id_jag_resource_token_endpoint="https://resource-as.example/oauth2/token",
+        )
+
+    @pytest.mark.asyncio
+    async def test_preflight_id_jag_surfaces_missing_assertion_as_a_plain_412(self):
+        """ID-JAG's missing/expired-assertion precondition must reach the client as a 412 whose body
+        names the fix, at the transport edge. Without the preflight the session opens and the caller
+        gets a 200 with an empty tool list and then 'tool not found', which is not what happened.
+        412 is a precondition, not an RFC 9728 discovery challenge, so it carries no
+        WWW-Authenticate: there is nothing for the client to discover and retry against."""
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError
+
+        summary = (
+            "ID-JAG requires an IdP identity assertion for this user and none is stored. "
+            "Sign in through LiteLLM SSO so the gateway captures one."
+        )
+
+        class _FakeProvider:
+            async def resolve_credentials(self, subject, server):
+                return Error(CredError.of_precondition_required(summary))
+
+        manager = MCPServerManager(cred_provider=_FakeProvider())
+
+        with pytest.raises(HTTPException) as exc_info:
+            await manager.preflight_token_exchange(
+                server=self._id_jag_server("id-jag-preflight-412"),
+                oauth2_headers=None,
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"),
+            )
+        assert exc_info.value.status_code == 412
+        assert summary in exc_info.value.detail
+        assert not (exc_info.value.headers or {})
+
+    @pytest.mark.asyncio
+    async def test_preflight_id_jag_surfaces_assertion_store_outage_as_503(self):
+        """A store outage is the other failure the session would swallow, and it is a different
+        answer than 412: the user has nothing to fix by signing in again."""
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError
+
+        class _FakeProvider:
+            async def resolve_credentials(self, subject, server):
+                return Error(CredError.of_upstream_unavailable("assertion store unreachable"))
+
+        manager = MCPServerManager(cred_provider=_FakeProvider())
+
+        with pytest.raises(HTTPException) as exc_info:
+            await manager.preflight_token_exchange(
+                server=self._id_jag_server("id-jag-preflight-503"),
+                oauth2_headers=None,
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"),
+            )
+        assert exc_info.value.status_code == 503
+
+    @pytest.mark.asyncio
+    async def test_preflight_id_jag_preflights_litellm_key_and_skips_identity_bearer(self):
+        """ID-JAG preflights when Authorization carries a LiteLLM key, but skips a caller identity
+        bearer that the session passes through unchanged."""
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.httpx_auth import (
+            StaticHeaderAuth,
+        )
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Ok
+
+        subjects = []
+
+        class _FakeProvider:
+            async def resolve_credentials(self, subject, server):
+                subjects.append(
+                    (
+                        subject.subject_id,
+                        subject.inbound_token.get_secret_value() if subject.inbound_token else None,
+                    )
+                )
+                return Ok(StaticHeaderAuth("Bearer minted-id-jag", header_name="Authorization"))
+
+        manager = MCPServerManager(cred_provider=_FakeProvider())
+        caller = UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1")
+
+        await manager.preflight_token_exchange(
+            server=self._id_jag_server("id-jag-preflight-key"),
+            oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"},
+            raw_headers={"authorization": "Bearer sk-litellm-virtual-key"},
+            user_api_key_auth=caller,
+        )
+        assert subjects == [("u-1", None)]
+
+        await manager.preflight_token_exchange(
+            server=self._id_jag_server("id-jag-preflight-identity"),
+            oauth2_headers={"Authorization": "Bearer caller-idp-id-token"},
+            raw_headers={
+                "x-litellm-api-key": "Bearer sk-admission-key",
+                "authorization": "Bearer caller-idp-id-token",
+            },
+            user_api_key_auth=UserAPIKeyAuth(api_key="hashed-key", user_id="u-1"),
+        )
+        assert subjects == [("u-1", None)]
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "server_fields",
+        [
+            {"auth_type": MCPAuth.none},
+            {"auth_type": MCPAuth.api_key, "authentication_token": "static-upstream-key"},
+            {"auth_type": MCPAuth.bearer_token, "authentication_token": "static-upstream-key"},
+            {
+                "auth_type": MCPAuth.oauth2,
+                "oauth2_flow": "client_credentials",
+                "client_id": "cid",
+                "client_secret": "csec",
+                "token_url": "https://idp.example.com/token",
+            },
+            {"auth_type": MCPAuth.true_passthrough},
+        ],
+    )
+    async def test_preflight_resolves_nothing_for_a_mode_that_does_not_pre_flight(self, server_fields):
+        """The manager is the only thing deciding which modes pre-flight, so it has to reject every
+        other mode itself. The single-server call site no longer tests the mode before calling, so a
+        mode that falls through here would start resolving its credential a second time, at connect,
+        for flows that never had a connect-time resolution at all."""
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.result import Error
+        from litellm.proxy._experimental.mcp_server.outbound_credentials.types import CredError
+
+        calls = []
+
+        class _FakeProvider:
+            async def resolve_credentials(self, subject, server):
+                calls.append(server.server_id)
+                return Error(CredError.of_misconfigured("the preflight must never get here"))
+
+        manager = MCPServerManager(cred_provider=_FakeProvider())
+        server = MCPServer(
+            server_id="not-pre-flighted",
+            name="not-pre-flighted-server",
+            url="https://up.example.com/mcp",
+            transport=MCPTransport.http,
+            **server_fields,
+        )
+
+        assert (
+            await manager.preflight_token_exchange(
+                server=server,
+                oauth2_headers={"Authorization": "Bearer sk-litellm-virtual-key"},
+                user_api_key_auth=UserAPIKeyAuth(api_key="sk-litellm-virtual-key", user_id="u-1"),
+            )
+            is None
+        )
+        assert calls == []
+
     @pytest.mark.asyncio
     @pytest.mark.parametrize(
         "authorization",
@@ -2638,7 +2797,9 @@ class TestMCPServerManager:
         resolved: Final[list[str | None]] = []
 
         class _FakeProvider:
-            async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]:
+            async def resolve_credentials(
+                self, subject: Subject, server: ServerSpec
+            ) -> Ok[StaticHeaderAuth, CredError]:
                 resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None)
                 return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))
 
@@ -2664,7 +2825,9 @@ class TestMCPServerManager:
         resolved: Final[list[str | None]] = []
 
         class _FakeProvider:
-            async def resolve_credentials(self, subject: Subject, server: ServerSpec) -> Ok[StaticHeaderAuth, CredError]:
+            async def resolve_credentials(
+                self, subject: Subject, server: ServerSpec
+            ) -> Ok[StaticHeaderAuth, CredError]:
                 resolved.append(subject.inbound_token.get_secret_value() if subject.inbound_token else None)
                 return Ok(StaticHeaderAuth("Bearer MINTED", header_name="Authorization"))