Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_idjag_sso_provider_gap

.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

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

enterprise/litellm_enterprise/proxy/audit_logging_endpoints.py

@@ -48,6 +48,18 @@ def _build_json_field_or_condition(json_key: str, value: str) -> dict[str, objec
     }
 
 
+def _build_search_condition(search: str) -> dict[str, object]:
+    """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": search},
+            {"changed_by_api_key": search},
+        )
+    }
+
+
 @router.get(
     "/audit",
     tags=["Audit Logging"],
@@ -83,6 +95,10 @@ 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",
+    ),
     # Sorting parameters
     sort_by: str | None = Query(
         None,
@@ -118,6 +134,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,14 +147,14 @@ 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]] = (
         {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(
@@ -195,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})

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):

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

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(

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)
@@ -432,6 +433,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

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,
     )
 
 

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)
 

litellm/litellm_core_utils/llm_cost_calc/utils.py

@@ -415,40 +415,64 @@ 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:
+    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 +482,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 +1067,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 +1098,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 +1113,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 +1180,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 +1191,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 +1250,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 +1308,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 +1327,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 +1338,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 +1348,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
 

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]:
     """

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)

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)

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,
-        )

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -3656,30 +3656,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(),

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
 

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(

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"
     )

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,
         )
 
 

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)

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(

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
@@ -20,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,
@@ -74,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):
@@ -119,8 +120,56 @@ 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 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 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:
+        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 +379,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 +440,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(TeamRepository(prisma_client).table, records)
+    return [_record_to_response(r, assigned_team_ids=attached[r.access_group_id]) for r in records]
 
 
 @router.get(
@@ -411,7 +462,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(TeamRepository(prisma_client).table, (record,))
+    return _record_to_response(record, assigned_team_ids=attached[record.access_group_id])
 
 
 @router.put(
@@ -461,8 +513,10 @@ 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 [])
+            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

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 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"),
@@ -5943,6 +5948,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 +6168,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": 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 +6193,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 +6283,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 +6294,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 +6334,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 +6373,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

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"},
     )

litellm/proxy/management_endpoints/scim/scim_v2.py

@@ -1123,13 +1123,15 @@ 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,
         )
 
-        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
 
@@ -1699,7 +1701,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,
@@ -1717,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)
@@ -1771,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()
             )

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

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"},

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]}

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",
+            }
+        },
+    )

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:
@@ -1659,6 +1696,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 +1709,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 +1943,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 +1954,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 +2012,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 +2106,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 +2118,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)

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,
 )
@@ -15233,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:
@@ -16502,6 +16520,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 +18198,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":

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:

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],
@@ -2229,6 +2285,31 @@ 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."""
+    raw: Final = f"${next_param_index}"
+    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 = {raw} OR team_id = {raw} OR "user" = {raw} OR end_user = {raw} '
+        f"OR session_id = {raw} OR model_id = {raw})))"
+    )
+    return _SpendLogSearchCondition(sql=sql, params=(search, start_date, end_date))
+
+
 @router.get(
     "/spend/logs/v2",
     tags=["Budget & Spend Tracking"],
@@ -2329,6 +2410,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 (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"
+        ),
+    ),
 ):
     """
     View spend logs with pagination support.
@@ -2392,8 +2481,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 +2667,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 +2675,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 +2764,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 +2804,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
@@ -3214,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):
@@ -3244,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]] = {}

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"
 

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,
@@ -9545,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)
@@ -11762,7 +11764,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 +11774,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 +11808,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 +11857,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 +11910,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 +11940,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(

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

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()

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 or case-insensitive alias substring."""
+
+    OR: ReadOnly[tuple[KeyTokenWhere, KeyAliasContainsWhere]]
 
 
 class BulkUpdateKeyRequestItem(BaseModel):

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"

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):

ruff-strict-budget.json

@@ -96,7 +96,7 @@
     "limit": 10
   },
   "DTZ007": {
-    "limit": 17
+    "limit": 6
   },
   "DTZ011": {
     "limit": 3

tests/e2e/ui/tests/tables/tableScrolling.spec.ts

@@ -0,0 +1,218 @@
+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";
+
+const VIEWPORT = { width: 1280, height: 720 };
+const SEED_ROWS = 40;
+const LOG_ROWS = 20;
+const BODY_SCROLL_PX = 500;
+const MAX_FOOTER_GAP_PX = 40;
+
+interface GeneratedKey {
+  key: string;
+}
+
+interface CreatedTeam {
+  team_id: string;
+}
+
+interface CreatedModel {
+  model_info: { id: string };
+}
+
+interface BoxMetrics {
+  top: number;
+  bottom: number;
+  scrollHeight: number;
+  clientHeight: number;
+  scrollWidth: number;
+  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();
+    return {
+      top: rect.top,
+      bottom: rect.bottom,
+      scrollHeight: el.scrollHeight,
+      clientHeight: el.clientHeight,
+      scrollWidth: el.scrollWidth,
+      clientWidth: el.clientWidth,
+    };
+  });
+
+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 T;
+}
+
+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);
+}
+
+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();
+}
+
+async function expectBodyIsTheOnlyScroller(page: PlaywrightPage): Promise<void> {
+  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"));
+  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));
+}
+
+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;
+        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)}`,
+          );
+      });
+    });
+
+test.describe("Admin tables scroll inside the page", () => {
+  test.use({ storageState: ADMIN_STORAGE_PATH, viewport: VIEWPORT });
+
+  test("Virtual Keys: rows scroll under a sticky header and the page itself never scrolls", async ({
+    page,
+    request,
+  }) => {
+    const suffix = uniqueSuffix();
+    const keys = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<GeneratedKey>(request, "/key/generate", { key_alias: `e2e-scroll-key-${suffix}-${i}` }),
+    );
+    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 teams = await oneAtATime(SEED_ROWS, (i) =>
+      postOk<CreatedTeam>(request, "/team/new", { team_alias: `e2e-scroll-team-${suffix}-${i}` }),
+    );
+    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 ({
+    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<unknown>(request, "/tag/new", { name: names[i], description: "LIT-4738 scroll" }),
+    );
+    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 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",
+          api_base: `http://127.0.0.1:${process.env.MOCK_LLM_PORT ?? "8090"}/v1`,
+          api_key: "fake-key",
+        },
+      }),
+    );
+    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 } }),
+      );
+    }
+  });
+});

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@");

tests/enterprise/litellm_enterprise/proxy/test_audit_logging_endpoints.py

@@ -1,5 +1,6 @@
 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 +9,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 +133,45 @@ 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_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"}

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

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

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()

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

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")

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,

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,258 @@ 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():
+    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():
+    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():
+    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():
+    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():
+    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():
+    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():
+    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"

tests/test_litellm/llms/dashscope/test_dashscope_cost_calculator.py

@@ -649,6 +649,90 @@ 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):
+        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):
+        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):
+        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."""

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()

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)

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(

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

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}",

tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py

@@ -2706,6 +2706,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",
@@ -2728,7 +2887,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"))
 
@@ -2754,7 +2915,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"))
 

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"]

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

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.

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"')

tests/test_litellm/proxy/management_endpoints/scim/test_scim_v2_endpoints.py

@@ -19,11 +19,13 @@ from litellm.proxy._types import (
     NewUserResponse,
     ProxyErrorTypes,
     ProxyException,
+    UserAPIKeyAuth,
 )
 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 +39,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 +306,85 @@ 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
+    assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+
+
+@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
+    assert new_user_mock.call_args.kwargs["user_api_key_dict"] == UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
+
+
 @pytest.mark.asyncio
 async def test_scim_create_user_respects_default_role_set_via_ui(mocker, monkeypatch):
     """
@@ -1176,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"""

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,45 @@ 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, 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)
@@ -373,6 +409,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 +1058,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 +1092,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 +1137,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 +1161,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 +1190,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 +1207,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 +1238,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"]), team_to_remove])
     mock_team_table.find_unique = AsyncMock(return_value=team_to_remove)
 
     resp = client.put(
@@ -1145,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 = (

tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py

@@ -6347,6 +6347,93 @@ 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_ors_token_and_alias_contains():
+    """
+    LIT-4741: `search` matches a key by its alias (case-insensitive contains) OR by
+    its ID (the token column), with the pasted value used verbatim.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _build_key_filter_conditions,
+    )
+
+    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.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="member-key-id",
+            )
+        )
+    )
+
+    assert where.get("AND"), f"expected top-level AND, got: {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."""
+    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="key-id-123",
+    )
+
+    where = json.loads(json.dumps(mock_find_many.call_args.kwargs["where"]))
+    assert _search_clause("key-id-123", "key-id-123") in where["AND"], f"search not in Prisma where: {where}"
+
+
 @pytest.mark.asyncio
 async def test_generate_key_negative_max_budget():
     """
@@ -14870,6 +14957,16 @@ async def test_list_keys_non_admin_cannot_opt_into_substring():
     assert kwargs["user_id"] == "alice"
 
 
+@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="key-id-123")
+    assert kwargs["search"] == "key-id-123"
+    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)

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'

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(

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

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,126 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
         assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
 
 
+class TestVertexPassthroughDefaultLocationOnShortRoutes:
+    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]],
+        forwarder: Mock,
+    ) -> None:
+        from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
+            PassthroughEndpointRouter,
+        )
+
+        async def receive():
+            return {"type": "http.request", "body": b"{}", "more_body": False}
+
+        request: Final = Request(
+            {
+                "type": "http",
+                "method": "POST",
+                "path": f"/vertex_ai/{endpoint}",
+                "headers": headers,
+                "query_string": b"",
+            },
+            receive=receive,
+        )
+        router: Final = PassthroughEndpointRouter()
+        if default_config is not None:
+            router.set_default_vertex_config(dict(default_config))
+        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: Final = Mock()
+        mock_credentials.token = "test-token"
+        caller: Final = UserAPIKeyAuth(api_key="test-key")
+        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", new=forwarder),
+            mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
+        ):
+            await vertex_proxy_route(
+                endpoint=endpoint,
+                request=request,
+                fastapi_response=Response(),
+                user_api_key_dict=caller,
+            )
+
+    @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
+    ):
+        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,
+        )
+        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(
+        ("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):
+        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:
     """The operable index is only the segment right after ``indexes``.
 
@@ -5009,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"]

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",

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"

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

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
 
@@ -58,6 +56,24 @@ 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 = \$\1 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
+    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):
     """
     Rebuild the Prisma-style ``where`` dict the filter_fns below expect from the
@@ -77,6 +93,16 @@ 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 = (int(g) for g in search_clause.groups())
+        where["search"] = {
+            "value": params[raw_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 +115,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 +2378,208 @@ 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,
+    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="key-hash-7", start_date=start, end_date=end, next_param_index=3
+    )
+
+    assert condition.sql == (
+        "(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 == ("key-hash-7", 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": "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},
+        {**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"}),
+        ("hashed-7", {"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
@@ -3095,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}")
@@ -3265,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,
                 },
             ]
 
@@ -3620,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")
 
@@ -3668,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={
@@ -3679,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]
@@ -3694,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"""
@@ -4602,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 []
 
 
@@ -6351,3 +6734,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)

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,

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,

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
@@ -4943,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
@@ -7362,14 +7554,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():
@@ -12701,33 +12893,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

type-discipline-budget.json

@@ -3,7 +3,7 @@
     "limit": 22328
   },
   "LIT002": {
-    "limit": 26760
+    "limit": 26750
   },
   "LIT003": {
     "limit": 261
@@ -27,10 +27,10 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16470
+    "limit": 16468
   },
   "LIT011": {
-    "limit": 5516
+    "limit": 5514
   },
   "LIT012": {
     "limit": 4489

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();
   });

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)}
             />

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: "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("pasted-key-id");
+    expect(callUrl.searchParams.has("key_alias")).toBe(false);
+    expect(callUrl.searchParams.has("key_hash")).toBe(false);
+  });
 });
 
 describe("useDeletedKeys", () => {

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,

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");
 

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}

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);
 

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,
       });

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}

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 */}

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">

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 */}

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>
           ))}

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) => (
         <>

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();

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,
@@ -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) => (
           <>
@@ -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)}

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>
+    </>
+  );
+};

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", () => {

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;

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,
     });
   });
 

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) => {

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,

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;