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

# Conflicts:
#	tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py

litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql

@@ -0,0 +1,2 @@
+-- AlterTable
+ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;

litellm-proxy-extras/litellm_proxy_extras/schema.prisma

@@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable {
   delegate_auth_to_upstream Boolean @default(false)
   oauth_passthrough Boolean @default(false)
   dcr_bridge Boolean?
+  per_server_oauth_discovery Boolean @default(false)
   is_byok               Boolean  @default(false)
   byok_description      String[] @default([])
   byok_api_key_help_url String?

litellm/integrations/azure_sentinel/azure_sentinel.py

@@ -16,23 +16,32 @@ import asyncio
 import os
 import time
 import traceback
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from types import MappingProxyType
-from typing import Final
+from typing import Final, TypeVar
 from urllib.parse import urlparse
 
 from litellm._logging import verbose_logger
+from litellm.integrations.batch_utils import (
+    BatchSendCancelled,
+    send_batch_with_413_split,
+    undelivered_after_http_error,
+)
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.llms.custom_httpx.http_handler import (
+    MaskedHTTPStatusError,
     get_async_httpx_client,
     httpxSpecialProvider,
 )
+from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
 from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
 DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com"
 DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default"
 
+_QueuedPayload = TypeVar("_QueuedPayload", StandardLoggingPayload, StandardAuditLogPayload)
+
 MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType(
     {
         "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE,
@@ -153,6 +162,8 @@ class AzureSentinelLogger(CustomBatchLogger):
         asyncio.create_task(self.periodic_flush())
         self.log_queue: list[StandardLoggingPayload] = []
         self.audit_log_queue: list[StandardAuditLogPayload] = []
+        self.logs_awaiting_retry = False
+        self.audit_logs_awaiting_retry = False
 
     @staticmethod
     def _normalize_authority_host(authority_host: str) -> str:
@@ -245,8 +256,8 @@ class AzureSentinelLogger(CustomBatchLogger):
 
             self.log_queue.append(standard_logging_payload)
 
-            if len(self.log_queue) >= self.batch_size:
-                await self.async_send_batch()
+            if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
+                await self._threshold_send_logs()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@@ -275,8 +286,8 @@ class AzureSentinelLogger(CustomBatchLogger):
 
             self.log_queue.append(standard_logging_payload)
 
-            if len(self.log_queue) >= self.batch_size:
-                await self.async_send_batch()
+            if len(self.log_queue) >= self.batch_size and not self.logs_awaiting_retry:
+                await self._threshold_send_logs()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Layer Error - %s\n%s", e, traceback.format_exc())
@@ -298,12 +309,24 @@ class AzureSentinelLogger(CustomBatchLogger):
 
             self.audit_log_queue.append(audit_log)
 
-            if len(self.audit_log_queue) >= self.batch_size:
-                await self.async_send_audit_batch()
+            if len(self.audit_log_queue) >= self.batch_size and not self.audit_logs_awaiting_retry:
+                await self._threshold_send_audit_logs()
 
         except Exception as e:
             verbose_logger.exception("Azure Sentinel Audit Log Layer Error - %s\n%s", e, traceback.format_exc())
 
+    async def _threshold_send_logs(self) -> None:
+        async with self.flush_lock:
+            if self.logs_awaiting_retry or len(self.log_queue) < self.batch_size:
+                return
+            await self.async_send_batch()
+
+    async def _threshold_send_audit_logs(self) -> None:
+        async with self.flush_lock:
+            if self.audit_logs_awaiting_retry or len(self.audit_log_queue) < self.batch_size:
+                return
+            await self.async_send_audit_batch()
+
     async def async_send_batch(self):
         """
         Sends the batch of logs to Azure Monitor Logs Ingestion API
@@ -311,67 +334,110 @@ class AzureSentinelLogger(CustomBatchLogger):
         Raises:
             Raises a NON Blocking verbose_logger.exception if an error occurs
         """
-        await self._async_send_batch_to_api(
-            log_queue=self.log_queue,
-            api_endpoint=self.api_endpoint,
-            log_type="logs",
-        )
+        batch_to_send: Final = tuple(self.log_queue)
+        self.log_queue = []  # mutable-ok: queue ownership is detached before the async send
+        try:
+            undelivered: Final = await self._async_send_batch_to_api(
+                log_queue=batch_to_send,
+                api_endpoint=self.api_endpoint,
+                log_type="logs",
+            )
+        except BatchSendCancelled as cancelled:
+            self.log_queue = self._requeue(cancelled.undelivered, self.log_queue, "logs")
+            self.logs_awaiting_retry = bool(self.log_queue)
+            raise asyncio.CancelledError() from cancelled
+        except asyncio.CancelledError:
+            self.log_queue = self._requeue(batch_to_send, self.log_queue, "logs")
+            self.logs_awaiting_retry = bool(self.log_queue)
+            raise
+        self.log_queue = self._requeue(undelivered, self.log_queue, "logs")
+        self.logs_awaiting_retry = bool(undelivered) and bool(self.log_queue)
 
     async def async_send_audit_batch(self):
         """
         Sends the batch of audit logs to Azure Monitor Logs Ingestion API
         """
-        await self._async_send_batch_to_api(
-            log_queue=self.audit_log_queue,
-            api_endpoint=self.audit_api_endpoint,
-            log_type="audit logs",
+        batch_to_send: Final = tuple(self.audit_log_queue)
+        self.audit_log_queue = []  # mutable-ok: queue ownership is detached before the async send
+        try:
+            undelivered: Final = await self._async_send_batch_to_api(
+                log_queue=batch_to_send,
+                api_endpoint=self.audit_api_endpoint,
+                log_type="audit logs",
+            )
+        except BatchSendCancelled as cancelled:
+            self.audit_log_queue = self._requeue(cancelled.undelivered, self.audit_log_queue, "audit logs")
+            self.audit_logs_awaiting_retry = bool(self.audit_log_queue)
+            raise asyncio.CancelledError() from cancelled
+        except asyncio.CancelledError:
+            self.audit_log_queue = self._requeue(batch_to_send, self.audit_log_queue, "audit logs")
+            self.audit_logs_awaiting_retry = bool(self.audit_log_queue)
+            raise
+        self.audit_log_queue = self._requeue(undelivered, self.audit_log_queue, "audit logs")
+        self.audit_logs_awaiting_retry = bool(undelivered) and bool(self.audit_log_queue)
+
+    def _requeue(
+        self,
+        undelivered: tuple[_QueuedPayload, ...],
+        queue: list[_QueuedPayload],
+        log_type: str,
+    ) -> list[_QueuedPayload]:
+        merged: Final = [*undelivered, *queue]  # mutable-ok: queue trimming returns a mutable logger queue
+        overflow: Final = len(merged) - self.max_queue_size
+        if overflow <= 0:
+            return merged
+
+        verbose_logger.warning(
+            "Azure Sentinel: %s queue exceeded max_queue_size=%s, dropped %s oldest records",
+            log_type,
+            self.max_queue_size,
+            overflow,
         )
+        return merged[overflow:]
 
     async def _async_send_batch_to_api(
         self,
-        log_queue: list[StandardLoggingPayload | StandardAuditLogPayload],
+        log_queue: tuple[_QueuedPayload, ...],
         api_endpoint: str,
         log_type: str,
-    ) -> None:
+    ) -> tuple[_QueuedPayload, ...]:
+        if not log_queue:
+            return ()
+
+        verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
         try:
-            if not log_queue:
-                return
-
-            verbose_logger.debug("Azure Sentinel - about to flush %s %s", len(log_queue), log_type)
-
-            # Get OAuth2 token
             bearer_token: Final = await self._get_oauth_token()
+        except MaskedHTTPStatusError as e:
+            return undelivered_after_http_error(log_queue, e.status_code, "Azure Sentinel OAuth token", str(e))
+        except Exception as e:
+            verbose_logger.exception("Azure Sentinel Error getting OAuth token - %s", e)
+            return tuple(log_queue)
 
-            # Convert log queue to JSON array format expected by Logs Ingestion API
-            # Each log entry should be a JSON object in the array
-            body: Final = safe_dumps(log_queue)
+        headers: Final = {
+            "Authorization": f"Bearer {bearer_token}",
+            "Content-Type": "application/json",
+        }
 
-            # Set headers for Logs Ingestion API
-            headers: Final = {
-                "Authorization": f"Bearer {bearer_token}",
-                "Content-Type": "application/json",
-            }
-
-            # Send the request
-            response = await self.async_httpx_client.post(url=api_endpoint, data=body.encode("utf-8"), headers=headers)
-
-            if response.status_code not in [200, 204]:
-                verbose_logger.error(
-                    "Azure Sentinel API error: status_code=%s, response=%s",
-                    response.status_code,
-                    response.text,
-                )
-                raise Exception(f"Failed to send logs to Azure Sentinel: {response.status_code} - {response.text}")
-
-            verbose_logger.debug(
-                "Azure Sentinel: Response from API status_code: %s",
-                response.status_code,
+        async def _send_batch(batch: Sequence[_QueuedPayload]):
+            body: Final = safe_dumps(batch)
+            return await self.async_httpx_client.post(
+                url=api_endpoint,
+                data=body.encode("utf-8"),
+                headers=headers,
             )
 
-        except Exception as e:
-            verbose_logger.exception("Azure Sentinel Error sending batch API - %s\n%s", e, traceback.format_exc())
-        finally:
-            log_queue.clear()
+        return await send_batch_with_413_split(
+            batch=log_queue,
+            send_batch=_send_batch,
+            exceeds_limits=lambda batch: (
+                len(batch) > self.batch_size
+                or len(safe_dumps(batch).encode("utf-8")) > AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
+            ),
+            success_status_codes=frozenset({200, 204}),
+            integration_name="Azure Sentinel",
+            drop_error_message="Azure Sentinel API Error - Payload too large for a single record",
+            non_success_handler=undelivered_after_http_error,
+        )
 
     async def flush_queue(self):
         if self.flush_lock is None:

litellm/integrations/batch_utils.py

@@ -0,0 +1,160 @@
+import asyncio
+from collections.abc import Awaitable, Callable, Sequence
+from typing import Final, Generic, TypeVar
+
+import httpx
+
+from litellm._logging import verbose_logger
+from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
+
+_BatchItem = TypeVar("_BatchItem")
+
+_RETRYABLE_CLIENT_STATUS_CODES: Final = frozenset({408, 429})
+
+
+def is_retryable_status(status_code: int) -> bool:
+    return not 400 <= status_code < 500 or status_code in _RETRYABLE_CLIENT_STATUS_CODES
+
+
+def undelivered_after_http_error(
+    batch: Sequence[_BatchItem],
+    status_code: int,
+    integration_name: str,
+    detail: str,
+) -> tuple[_BatchItem, ...]:
+    """The records to requeue after a non-2xx: all of them on a status a retry can clear, none on
+    a 4xx that would only repeat, since retaining those retries a misconfiguration forever."""
+    if is_retryable_status(status_code):
+        verbose_logger.error(
+            "%s API error: status_code=%s, will retry %s records - %s",
+            integration_name,
+            status_code,
+            len(batch),
+            detail,
+        )
+        return tuple(batch)
+    verbose_logger.error(
+        "%s API error: status_code=%s is not retryable, dropped %s records - %s",
+        integration_name,
+        status_code,
+        len(batch),
+        detail,
+    )
+    return ()
+
+
+def requeue_after_http_error(
+    batch: Sequence[_BatchItem],
+    status_code: int,
+    integration_name: str,
+    detail: str,
+) -> tuple[_BatchItem, ...]:
+    verbose_logger.error(
+        "%s API error: status_code=%s, will retry %s records - %s",
+        integration_name,
+        status_code,
+        len(batch),
+        detail,
+    )
+    return tuple(batch)
+
+
+class BatchSendCancelled(asyncio.CancelledError, Generic[_BatchItem]):
+    """Cancellation of a batch send, carrying only the records the destination never accepted.
+
+    A batch split under the size cap is delivered in pieces, so requeueing all of it after a
+    cancellation partway through would send the accepted pieces a second time.
+    """
+
+    def __init__(self, undelivered: tuple[_BatchItem, ...]) -> None:
+        super().__init__()
+        self.undelivered: Final = undelivered
+
+
+async def _keep_the_remainder_on_cancel(
+    send: Awaitable[tuple[_BatchItem, ...]],
+    remainder: Sequence[_BatchItem],
+) -> tuple[_BatchItem, ...]:
+    try:
+        return await send
+    except BatchSendCancelled as cancelled:
+        raise BatchSendCancelled((*cancelled.undelivered, *remainder)) from cancelled
+
+
+async def send_batch_with_413_split(
+    batch: Sequence[_BatchItem],
+    send_batch: Callable[[Sequence[_BatchItem]], Awaitable[httpx.Response]],
+    exceeds_limits: Callable[[Sequence[_BatchItem]], bool],
+    success_status_codes: frozenset[int],
+    integration_name: str,
+    drop_error_message: str,
+    non_success_handler: Callable[
+        [Sequence[_BatchItem], int, str, str], tuple[_BatchItem, ...]
+    ] = requeue_after_http_error,
+) -> tuple[_BatchItem, ...]:
+    async def _halve() -> tuple[_BatchItem, ...]:
+        midpoint: Final = len(batch) // 2
+        left_batch: Final = batch[:midpoint]
+        right_batch: Final = batch[midpoint:]
+        left_undelivered: Final = await _keep_the_remainder_on_cancel(
+            send_batch_with_413_split(
+                batch=left_batch,
+                send_batch=send_batch,
+                exceeds_limits=exceeds_limits,
+                success_status_codes=success_status_codes,
+                integration_name=integration_name,
+                drop_error_message=drop_error_message,
+                non_success_handler=non_success_handler,
+            ),
+            right_batch,
+        )
+        if left_undelivered:
+            return (*left_undelivered, *right_batch)
+        return await send_batch_with_413_split(
+            batch=right_batch,
+            send_batch=send_batch,
+            exceeds_limits=exceeds_limits,
+            success_status_codes=success_status_codes,
+            integration_name=integration_name,
+            drop_error_message=drop_error_message,
+            non_success_handler=non_success_handler,
+        )
+
+    async def _handle_413() -> tuple[_BatchItem, ...]:
+        if len(batch) == 1:
+            verbose_logger.error(drop_error_message)
+            return ()
+        return await _halve()
+
+    if not batch:
+        return ()
+
+    try:
+        oversized: Final = exceeds_limits(batch)
+    except Exception as e:  # noqa: BLE001  # any record that cannot be serialized is isolated and dropped alone
+        if len(batch) > 1:
+            return await _halve()
+        verbose_logger.exception("%s dropped a record that cannot be serialized - %s", integration_name, e)
+        return ()
+    if oversized and len(batch) > 1:
+        return await _halve()
+
+    try:
+        response: Final = await send_batch(batch)
+    except MaskedHTTPStatusError as e:
+        if e.status_code == 413:
+            return await _handle_413()
+        return non_success_handler(batch, e.status_code, integration_name, str(e))
+    except asyncio.CancelledError as cancelled:
+        raise BatchSendCancelled(tuple(batch)) from cancelled
+    except Exception as e:
+        verbose_logger.exception("%s Error sending batch API - %s", integration_name, e)
+        return tuple(batch)
+
+    if response.status_code == 413:
+        return await _handle_413()
+    if response.status_code not in success_status_codes:
+        return non_success_handler(batch, response.status_code, integration_name, response.text)
+
+    verbose_logger.debug("%s delivered %s records, status_code=%s", integration_name, len(batch), response.status_code)
+    return ()

litellm/integrations/datadog/datadog.py

@@ -29,6 +29,7 @@ from typing_extensions import ReadOnly, TypedDict
 import litellm
 from litellm._logging import verbose_logger
 from litellm._uuid import uuid
+from litellm.integrations.batch_utils import BatchSendCancelled, requeue_after_http_error, send_batch_with_413_split
 from litellm.integrations.custom_batch_logger import CustomBatchLogger
 from litellm.integrations.datadog.datadog_handler import (
     get_datadog_base_url_from_env,
@@ -43,7 +44,6 @@ from litellm.integrations.datadog.datadog_mock_client import (
 )
 from litellm.litellm_core_utils.dd_tracing import tracer
 from litellm.llms.custom_httpx.http_handler import (
-    MaskedHTTPStatusError,
     _get_httpx_client,
     get_async_httpx_client,
     httpxSpecialProvider,
@@ -396,6 +396,9 @@ class DataDogLogger(
             if self.is_mock_mode:
                 verbose_logger.debug("[DATADOG MOCK] Batch of %s events successfully mocked", len(batch_to_send))
 
+        except BatchSendCancelled as cancelled:
+            self.log_queue = list(cancelled.undelivered) + self.log_queue  # mutable-ok: logger queue remains appendable
+            raise asyncio.CancelledError() from cancelled
         except Exception as e:
             self.log_queue = batch_to_send + self.log_queue
             verbose_logger.exception("Datadog Error sending batch API - %s\n%s", e, traceback.format_exc())
@@ -413,53 +416,16 @@ class DataDogLogger(
         that could not be delivered because of a non-413 (transient) error, so the caller
         re-queues only those and never the events already accepted by Datadog.
         """
-        pending: Final[list[list]] = [batch]
-        while pending:
-            chunk = pending.pop()
-            if not chunk:
-                continue
-            if len(chunk) > 1 and self._exceeds_intake_limits(chunk):
-                mid = len(chunk) // 2
-                pending.append(chunk[mid:])
-                pending.append(chunk[:mid])
-                continue
-            try:
-                response = await self.async_send_compressed_data(chunk)
-            except Exception as e:
-                if isinstance(e, MaskedHTTPStatusError) and e.status_code == 413:
-                    response = e.response
-                else:
-                    verbose_logger.exception("Datadog Error sending batch API - %s", e)
-                    return self._undelivered(chunk, pending)
-
-            if response.status_code == 413:
-                if len(chunk) == 1:
-                    verbose_logger.error(DD_ERRORS.DATADOG_413_ERROR.value)
-                    continue
-                mid = len(chunk) // 2
-                pending.append(chunk[mid:])
-                pending.append(chunk[:mid])
-                continue
-
-            if response.status_code != 202:
-                verbose_logger.error(
-                    "Datadog: unexpected response status_code=%s, text=%s",
-                    response.status_code,
-                    response.text,
-                )
-                return self._undelivered(chunk, pending)
-
-            verbose_logger.debug(
-                "Datadog: delivered %s events, status_code=%s, text=%s",
-                len(chunk),
-                response.status_code,
-                response.text,
-            )
-        return []
-
-    @staticmethod
-    def _undelivered(chunk: list, pending: list[list]) -> list:
-        return chunk + [event for remaining in reversed(pending) for event in remaining]
+        undelivered: Final = await send_batch_with_413_split(
+            batch=batch,
+            send_batch=self.async_send_compressed_data,
+            exceeds_limits=self._exceeds_intake_limits,
+            success_status_codes=frozenset({202}),
+            integration_name="Datadog",
+            drop_error_message=DD_ERRORS.DATADOG_413_ERROR.value,
+            non_success_handler=requeue_after_http_error,
+        )
+        return list(undelivered)  # mutable-ok: caller prepends records to the logger queue
 
     @staticmethod
     def _exceeds_intake_limits(chunk: Sequence[DatadogPayload]) -> bool:
@@ -606,7 +572,7 @@ class DataDogLogger(
         )
         return dd_payload
 
-    async def async_send_compressed_data(self, data: list) -> Response:
+    async def async_send_compressed_data(self, data: Sequence[DatadogPayload]) -> Response:
         """
         Async helper to send compressed data to datadog self.intake_url
 

litellm/llms/anthropic/common_utils.py

@@ -1641,6 +1641,25 @@ def flatten_unencrypted_web_search_results_in_anthropic_messages(  # mutable-ok:
     return [_flatten_web_search_results_in_message(m) for m in messages]  # mutable-ok: JSON wire format
 
 
+def _without_provider_specific_fields(block: object) -> object:
+    if not isinstance(block, dict) or "provider_specific_fields" not in block:
+        return block
+    return {k: v for k, v in block.items() if k != "provider_specific_fields"}  # mutable-ok: JSON wire format
+
+
+def _strip_provider_specific_fields_in_message(message: object) -> object:
+    if not isinstance(message, dict) or not isinstance(message.get("content"), list):
+        return message
+    content: Final = [_without_provider_specific_fields(b) for b in message["content"]]  # mutable-ok: JSON wire format
+    return {**message, "content": content}  # mutable-ok: JSON wire format
+
+
+def strip_provider_specific_fields_from_anthropic_messages(
+    messages: Sequence[object],
+) -> Sequence[object]:
+    return [_strip_provider_specific_fields_in_message(m) for m in messages]  # mutable-ok: JSON wire format
+
+
 def _normalized_cache_control(cache_control: object) -> dict[str, str] | None:  # mutable-ok: JSON wire format
     if not isinstance(cache_control, Mapping):
         return None

litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py

@@ -1346,7 +1346,7 @@ class LiteLLMAnthropicMessagesAdapter:
                     # Add provider_specific_fields if signature is present
                     if provider_specific_fields:
                         tool_use_block.provider_specific_fields = provider_specific_fields
-                    new_content.append(tool_use_block.model_dump())
+                    new_content.append(tool_use_block.model_dump(exclude_none=True))
 
         return new_content
 

litellm/llms/anthropic/experimental_pass_through/messages/handler.py

@@ -18,6 +18,7 @@ from litellm.llms.anthropic.common_utils import (
     flatten_unencrypted_web_search_results_in_anthropic_messages,
     sanitize_tool_use_ids_in_anthropic_messages,
     strip_empty_content_blocks_from_anthropic_messages,
+    strip_provider_specific_fields_from_anthropic_messages,
 )
 from litellm.llms.base_llm.anthropic_messages.transformation import (
     BaseAnthropicMessagesConfig,
@@ -650,7 +651,7 @@ def anthropic_messages_handler(
 
     return base_llm_http_handler.anthropic_messages_handler(
         model=model,
-        messages=messages,
+        messages=strip_provider_specific_fields_from_anthropic_messages(messages),
         anthropic_messages_provider_config=anthropic_messages_provider_config,
         anthropic_messages_optional_request_params=dict(anthropic_messages_optional_request_params),
         _is_async=is_async,

litellm/llms/anthropic/experimental_pass_through/responses_adapters/transformation.py

@@ -647,7 +647,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
                         id=item.call_id or item.id or "",
                         name=item.name,
                         input=input_data,
-                    ).model_dump()
+                    ).model_dump(exclude_none=True)
                 )
                 stop_reason = "tool_use"
 
@@ -676,7 +676,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
                             id=item.get("call_id") or item.get("id", ""),
                             name=item.get("name", ""),
                             input=input_data,
-                        ).model_dump()
+                        ).model_dump(exclude_none=True)
                     )
                     stop_reason = "tool_use"
 

litellm/model_prices_and_context_window_backup.json

@@ -12238,6 +12238,48 @@
         "output_cost_per_token": 2.65e-06,
         "supports_pdf_input": true
     },
+    "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0": {
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 8172,
+        "max_tokens": 8172,
+        "mode": "embedding",
+        "input_cost_per_token": 1.62e-07,
+        "input_cost_per_image": 7.2e-05,
+        "input_cost_per_video_per_second": 0.00084,
+        "input_cost_per_audio_per_second": 0.000168,
+        "output_cost_per_token": 0.0,
+        "output_vector_size": 3072,
+        "supports_embedding_image_input": true,
+        "supports_image_input": true,
+        "supports_video_input": true,
+        "supports_audio_input": true
+    },
+    "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 300000,
+        "max_output_tokens": 10000,
+        "max_tokens": 10000,
+        "mode": "chat",
+        "output_cost_per_token": 2.88e-07,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_response_schema": true,
+        "supports_vision": true
+    },
+    "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": {
+        "input_cost_per_token": 4.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 10000,
+        "max_tokens": 10000,
+        "mode": "chat",
+        "output_cost_per_token": 1.68e-07,
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_response_schema": true
+    },
     "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": {
         "input_cost_per_token": 9.6e-07,
         "litellm_provider": "bedrock",
@@ -43648,6 +43690,23 @@
         "input_cost_per_token_batches": 1.65e-06,
         "output_cost_per_token_batches": 8.25e-06
     },
+    "us-gov.anthropic.claude-3-haiku-20240307-v1:0": {
+        "deprecation_date": "2026-09-10",
+        "input_cost_per_token": 3e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 4096,
+        "max_tokens": 4096,
+        "mode": "chat",
+        "output_cost_per_token": 1.5e-06,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "cache_read_input_token_cost": 3e-08,
+        "cache_creation_input_token_cost": 3.75e-07
+    },
     "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": {
         "cache_creation_input_token_cost": 4.5e-06,
         "cache_creation_input_token_cost_above_1hr": 7.2e-06,
@@ -43742,6 +43801,160 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "us-gov.anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "us-gov.anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
+    "us-gov.nvidia.nemotron-nano-3-30b": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.88e-07,
+        "source": "https://aws.amazon.com/bedrock/pricing/",
+        "supports_function_calling": true,
+        "supports_native_structured_output": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.nvidia.nemotron-nano-12b-v2": {
+        "input_cost_per_token": 2.4e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 7.2e-07,
+        "supports_system_messages": true,
+        "supports_vision": true
+    },
+    "us-gov.nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
+    "us-gov.nvidia.nemotron-super-3-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 7.8e-07,
+        "source": "https://aws.amazon.com/bedrock/pricing/",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.openai.gpt-oss-20b-1:0": {
+        "input_cost_per_token": 8.4e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3.6e-07,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.openai.gpt-oss-120b-1:0": {
+        "input_cost_per_token": 1.8e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 7.2e-07,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.xai.grok-4.6": {
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_prompt_caching": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "au.anthropic.claude-haiku-4-5-20251001-v1:0": {
         "cache_creation_input_token_cost": 1.375e-06,
         "cache_creation_input_token_cost_above_1hr": 2.2e-06,
@@ -59547,6 +59760,16 @@
         "supports_system_messages": true,
         "supports_vision": true
     },
+    "bedrock/us-gov-west-1/nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
     "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": {
         "input_cost_per_token": 1.8e-07,
         "litellm_provider": "bedrock",
@@ -59651,6 +59874,70 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "bedrock/us-gov-west-1/anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
     "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": {
         "input_cost_per_token": 7.2e-08,
         "litellm_provider": "bedrock",
@@ -59676,6 +59963,16 @@
         "supports_system_messages": true,
         "supports_vision": true
     },
+    "bedrock/us-gov-east-1/nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
     "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": {
         "input_cost_per_token": 1.8e-07,
         "litellm_provider": "bedrock",
@@ -59780,6 +60077,70 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "bedrock/us-gov-east-1/anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "bedrock/us-gov-east-1/anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
     "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": {
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
@@ -59894,6 +60255,120 @@
         "output_cost_per_token": 3e-06,
         "cache_read_input_token_cost": 2.4e-07
     },
+    "bedrock_mantle/us-gov-west-1/xai.grok-4.6": {
+        "use_openai_responses_path": true,
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": {
+        "input_cost_per_token": 4.8e-08,
+        "output_cost_per_token": 9.6e-08,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": {
+        "input_cost_per_token": 1.56e-07,
+        "output_cost_per_token": 4.8e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": {
+        "input_cost_per_token": 1.68e-07,
+        "output_cost_per_token": 4.8e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": {
+        "input_cost_per_token": 8.4e-08,
+        "output_cost_per_token": 3.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "output_cost_per_token": 7.2e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
     "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": {
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
@@ -59921,6 +60396,63 @@
         "cache_read_input_token_cost": 3.3e-07,
         "output_cost_per_token": 1.98e-05
     },
+    "bedrock_mantle/us-gov-east-1/xai.grok-4.6": {
+        "use_openai_responses_path": true,
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": {
+        "input_cost_per_token": 8.4e-08,
+        "output_cost_per_token": 3.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "output_cost_per_token": 7.2e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
     "azure/us-gov/gpt-5.1": {
         "cache_read_input_token_cost": 1.71875e-07,
         "default_reasoning_effort": "none",

litellm/models/mcp_server.py

@@ -98,6 +98,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase):
     delegate_auth_to_upstream: bool = False
     oauth_passthrough: bool = False
     dcr_bridge: bool | None = None
+    per_server_oauth_discovery: bool = False
     is_byok: bool = False
     byok_description: list[str] = Field(default_factory=list)
     byok_api_key_help_url: str | None = None

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -197,7 +197,7 @@ def _gateway_dcr_challenge_target(
     if targets is None:
         return None
     server: Final = global_mcp_server_manager.get_mcp_server_by_name(targets[0], client_ip=client_ip)
-    if server is None or not server.is_gateway_managed_oauth2:
+    if server is None or not server.advertises_gateway_authorization_server:
         return None
     return targets[0]
 

litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py

@@ -1481,11 +1481,19 @@ async def _persist_dcr_client_registration(
         )
         updated_row: Final = await update_mcp_server(
             prisma_client=prisma_client,
-            data=UpdateMCPServerRequest(
-                server_id=mcp_server.server_id,
-                credentials=credentials,
-                oauth2_flow="authorization_code",
-                **({"token_url": mcp_server.token_url} if mcp_server.token_url else {}),
+            data=(
+                UpdateMCPServerRequest(
+                    server_id=mcp_server.server_id,
+                    credentials=credentials,
+                    oauth2_flow="authorization_code",
+                    token_url=mcp_server.token_url,
+                )
+                if mcp_server.token_url
+                else UpdateMCPServerRequest(
+                    server_id=mcp_server.server_id,
+                    credentials=credentials,
+                    oauth2_flow="authorization_code",
+                )
             ),
             touched_by="mcp_oauth_dcr",
         )
@@ -2367,7 +2375,7 @@ async def _build_oauth_protected_resource_response(
     if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange:
         _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource")
 
-    if explicitly_named and mcp_server is not None and mcp_server.is_gateway_managed_oauth2:
+    if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server:
         return {
             "authorization_servers": [f"{request_base_url}/mcp"],
             "resource": resource_url,

litellm/proxy/_experimental/mcp_server/mcp_server_manager.py

@@ -155,6 +155,7 @@ from litellm.proxy._types import (
     MCPTransportType,
     SpecialMCPServerNames,
     UserAPIKeyAuth,
+    is_per_server_oauth_discovery_eligible,
 )
 from litellm.proxy.auth.ip_address_utils import IPAddressUtils
 from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
@@ -344,6 +345,7 @@ class MCPServerConfig(TypedDict, total=False):
     token_endpoint_auth_method: MCPTokenEndpointAuthMethod
     scopes: str | Sequence[str]
     dcr_bridge: object
+    per_server_oauth_discovery: ReadOnly[object]
     extra_headers: _StringList
     allowed_tools: _StringList
     disallowed_tools: _StringList
@@ -414,6 +416,31 @@ def _blank_to_none(value: str | None) -> str | None:
     return value.strip() or None
 
 
+def _config_per_server_oauth_discovery(
+    server_config: MCPServerConfig,
+    server_ref: str,
+    auth_type: MCPAuthType | None,
+    oauth2_flow: object,
+) -> bool:
+    match server_config.get("per_server_oauth_discovery", False):
+        case bool() as enabled:
+            pass
+        case other:
+            raise ValueError(
+                f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery must be a boolean "
+                f"(got {other!r})."
+            )
+    relay_eligible: Final = is_per_server_oauth_discovery_eligible(
+        auth_type, oauth2_flow, server_config.get("delegate_auth_to_upstream", False)
+    )
+    if enabled and not relay_eligible:
+        raise ValueError(
+            f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery is only supported for "
+            "auth_type oauth2 with oauth2_flow authorization_code and without delegate_auth_to_upstream."
+        )
+    return enabled
+
+
 def _pinned_config_server_id(raw_server_id: object, server_name: str) -> str | None:
     """Return the ``server_id`` an admin pinned for this config.yaml server, or ``None`` when absent.
 
@@ -2307,6 +2334,9 @@ class MCPServerManager:
                 )
 
             config_dcr_bridge = server_config.get("dcr_bridge", None)
+            config_per_server_oauth_discovery = _config_per_server_oauth_discovery(
+                server_config, server_name or server_id, auth_type, config_oauth2_flow
+            )
             if config_dcr_bridge is not None and not isinstance(config_dcr_bridge, bool):
                 raise ValueError(
                     f"Invalid config for MCP server '{server_name or server_id}': dcr_bridge "
@@ -2378,6 +2408,7 @@ class MCPServerManager:
                 delegate_auth_to_upstream=bool(server_config.get("delegate_auth_to_upstream", False)),
                 oauth_passthrough=bool(server_config.get("oauth_passthrough", False)),
                 dcr_bridge=config_dcr_bridge,
+                per_server_oauth_discovery=config_per_server_oauth_discovery,
                 # AWS SigV4 fields
                 aws_access_key_id=server_config.get("aws_access_key_id", None),
                 aws_secret_access_key=server_config.get("aws_secret_access_key", None),
@@ -2903,6 +2934,7 @@ class MCPServerManager:
             delegate_auth_to_upstream=bool(getattr(mcp_server, "delegate_auth_to_upstream", False)),
             oauth_passthrough=bool(getattr(mcp_server, "oauth_passthrough", False)),
             dcr_bridge=getattr(mcp_server, "dcr_bridge", None),
+            per_server_oauth_discovery=bool(getattr(mcp_server, "per_server_oauth_discovery", False)),
             created_at=getattr(mcp_server, "created_at", None),
             updated_at=getattr(mcp_server, "updated_at", None),
             tool_name_to_display_name=_deserialize_json_dict(getattr(mcp_server, "tool_name_to_display_name", None)),
@@ -6692,6 +6724,7 @@ class MCPServerManager:
             registration_url=server.configured_registration_url or server.registration_url,
             oauth2_flow=server.oauth2_flow,
             dcr_bridge=server.dcr_bridge,
+            per_server_oauth_discovery=server.per_server_oauth_discovery,
             token_exchange_endpoint=server.token_exchange_endpoint,
             audience=server.audience,
             subject_token_type=server.subject_token_type,
@@ -6810,6 +6843,7 @@ class MCPServerManager:
             delegate_auth_to_upstream=server.delegate_auth_to_upstream,
             oauth_passthrough=getattr(server, "oauth_passthrough", False),
             dcr_bridge=server.dcr_bridge,
+            per_server_oauth_discovery=server.per_server_oauth_discovery,
             is_byok=server.is_byok,
             byok_description=server.byok_description,
             byok_api_key_help_url=server.byok_api_key_help_url,

litellm/proxy/_experimental/mcp_server/utils.py

@@ -752,7 +752,7 @@ def interpolate_headers(headers: Mapping[str, str], variables: Mapping[str, str]
 def build_env_var_setup_url(server_id: str) -> str:
     """The frontend URL where a user can fill in their per-user env vars."""
     base: Final = os.environ.get("PROXY_BASE_URL", "").rstrip("/")
-    path: Final = f"/ui/?page=mcp-servers&fill_env_vars={quote(server_id, safe='')}"
+    path: Final = f"/ui/mcp-servers?fill_env_vars={quote(server_id, safe='')}"
     return f"{base}{path}" if base else path
 
 

litellm/proxy/_lazy_openapi_snapshot.json

@@ -9504,6 +9504,18 @@
               "description": "Name of the guardrail in guardrails.ai",
               "title": "Guard Name"
             },
+            "inspect_embeddings": {
+              "anyOf": [
+                {
+                  "type": "boolean"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "When True, the Aim and Cato Networks guardrails send /embeddings `input` to the vendor as user messages. Off by default because embedding input is documents being indexed, not a conversation.",
+              "title": "Inspect Embeddings"
+            },
             "keyword_redaction_tag": {
               "anyOf": [
                 {
@@ -11261,6 +11273,12 @@
               ],
               "description": "Threshold configuration for Lakera guardrail categories"
             },
+            "ccr_retrieval": {
+              "default": true,
+              "description": "Inject the Headroom retrieval tool for hashes declared by the compression service.",
+              "title": "Ccr Retrieval",
+              "type": "boolean"
+            },
             "checks": {
               "anyOf": [
                 {
@@ -11656,6 +11674,18 @@
               "description": "Include scanner category summaries in responses (sets `plr_scanners` header).",
               "title": "Include Scanners"
             },
+            "inspect_embeddings": {
+              "anyOf": [
+                {
+                  "type": "boolean"
+                },
+                {
+                  "type": "null"
+                }
+              ],
+              "description": "When True, the Aim and Cato Networks guardrails send /embeddings `input` to the vendor as user messages. Off by default because embedding input is documents being indexed, not a conversation.",
+              "title": "Inspect Embeddings"
+            },
             "is_detector_server": {
               "anyOf": [
                 {
@@ -16325,6 +16355,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -17948,6 +17983,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -18829,6 +18869,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -20838,6 +20883,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -22312,6 +22362,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -22832,6 +22887,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {
@@ -25341,6 +25401,11 @@
               "title": "Oauth Passthrough",
               "type": "boolean"
             },
+            "per_server_oauth_discovery": {
+              "default": false,
+              "title": "Per Server Oauth Discovery",
+              "type": "boolean"
+            },
             "registration_url": {
               "anyOf": [
                 {

litellm/proxy/_types.py

@@ -1379,6 +1379,35 @@ def _dcr_bridge_auth_type_error(auth_type: object) -> ValueError:
     )
 
 
+def _per_server_oauth_discovery_error() -> ValueError:
+    return ValueError(
+        "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow "
+        "authorization_code and without delegate_auth_to_upstream."
+    )
+
+
+def is_per_server_oauth_discovery_eligible(
+    auth_type: object, oauth2_flow: object, delegate_auth_to_upstream: object
+) -> bool:
+    return auth_type == MCPAuth.oauth2 and oauth2_flow == "authorization_code" and not delegate_auth_to_upstream
+
+
+def _reject_unsupported_per_server_oauth_discovery(values: object, require_auth_type: bool) -> None:
+    """Partial updates may omit eligibility fields; those are checked against the stored row by the
+    update endpoint. Every field the payload does carry must be eligible on its own."""
+    if not isinstance(values, dict) or not values.get("per_server_oauth_discovery"):
+        return
+    auth_type_ok: Final = values.get("auth_type") == MCPAuth.oauth2 or (
+        not require_auth_type and "auth_type" not in values
+    )
+    oauth2_flow_ok: Final = values.get("oauth2_flow") == "authorization_code" or (
+        not require_auth_type and "oauth2_flow" not in values
+    )
+    if auth_type_ok and oauth2_flow_ok and not values.get("delegate_auth_to_upstream"):
+        return
+    raise _per_server_oauth_discovery_error()
+
+
 class NewMCPServerRequest(LiteLLMPydanticObjectBase):
     server_id: str | None = None
     server_name: str | None = None
@@ -1420,6 +1449,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
     delegate_auth_to_upstream: bool = False
     oauth_passthrough: bool = False
     dcr_bridge: bool | None = None
+    per_server_oauth_discovery: bool = False
     is_byok: bool = False
     byok_description: list[str] = Field(default_factory=list)
     byok_api_key_help_url: str | None = None
@@ -1484,6 +1514,12 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase):
             return values
         raise _dcr_bridge_auth_type_error(auth_type)
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object:
+        _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=True)
+        return values
+
 
 class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
     server_id: str
@@ -1526,6 +1562,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
     delegate_auth_to_upstream: bool = False
     oauth_passthrough: bool = False
     dcr_bridge: bool | None = None
+    per_server_oauth_discovery: bool = False
     is_byok: bool = False
     byok_description: list[str] = Field(default_factory=list)
     byok_api_key_help_url: str | None = None
@@ -1570,6 +1607,12 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase):
             return values
         raise _dcr_bridge_auth_type_error(auth_type)
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object:
+        _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=False)
+        return values
+
 
 from litellm.models.mcp_server import (  # noqa: E402
     LiteLLM_MCPServerTable as LiteLLM_MCPServerTable,

litellm/proxy/client/cli/README.md

@@ -478,6 +478,7 @@ Launch a coding agent with all of its LLM traffic routed through your LiteLLM pr
 lite claude
 lite codex
 lite opencode
+lite pi
 ```
 
 Anything you type after the agent name is forwarded to it untouched, so the usual flags keep working:
@@ -491,17 +492,19 @@ Each command resolves your LiteLLM key (logging in via SSO when none is stored a
 
 The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. It also gets `CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY=1` (again unless you already set it) so Claude Code v2.1.129+ fills its `/model` picker from the proxy's `/v1/models`; Claude Code only lists entries whose id contains `claude` or `anthropic`, and older versions ignore the variable. Export it as `0` to turn discovery off. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). OpenCode additionally gets `OPENCODE_CONFIG_CONTENT` holding a generated `litellm` provider (`@ai-sdk/openai-compatible`, the proxy `/v1` URL, `{env:OPENAI_API_KEY}`) with one model entry per chat model your key can see on `/v1/models`, so its model picker mirrors the proxy without a hand-maintained `opencode.json`; OpenCode merges that over your own config files, and if you already export `OPENCODE_CONFIG_CONTENT` yours is left alone. When the list cannot be fetched, `lite opencode` says so on stderr and launches anyway.
 
+pi ignores base-URL environment variables entirely, so `lite pi` (kept out of the `lite --help` command listing for now, but fully functional) wires it up differently: before handoff it fetches the models your key can use from the proxy's `/v1/models` (plus each model's context window and output cap from `/model_group/info`, when available) and syncs them into a `litellm` provider entry in pi's `~/.pi/agent/models.json` (honoring `PI_CODING_AGENT_DIR`), then starts pi on that provider's first model via an injected `--model litellm/<id>`. Only that one provider entry is rewritten; the rest of the file, including any other custom providers, is left alone. The entry references the key as `$LITELLM_PROXY_API_KEY`, which the wrapper exports for the session, so the token itself never lands on disk and plain `pi` outside the wrapper simply shows the litellm models as unavailable. Your own flags come after the injected pin, so `lite pi --model litellm/<other-id>` wins, and inside the TUI the `/model` picker lists every synced litellm model.
+
 Options (these belong to the wrapper, so put them before the agent's own flags):
 
 - `--skip-verify`: Skip the pre-launch key check (useful offline or with non-standard auth).
 
-To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model` or `lite codex -m my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints.
+To pin the model, pass the agent's own model flag (for example `lite claude --model my-proxy-model`, `lite codex -m my-proxy-model`, or `lite pi --model my-proxy-model`), or export the variable the agent reads (`ANTHROPIC_MODEL` / `ANTHROPIC_SMALL_FAST_MODEL` for Claude Code); the wrapper preserves anything you already have set. Whatever model the agent ends up requesting must exist on the proxy, since requests land on the proxy's `/v1/messages` (Anthropic) or `/v1/chat/completions` and `/v1/responses` (OpenAI) endpoints.
 
 #### About the `lite login` credential
 
 The token minted by `lite login` is a short-lived, per-session agent credential, not a managed virtual key. It is scoped to the user and team you authenticated as, inherits that user's and team's models and budgets, and is enforced on the proxy exactly like a virtual key on the same team (guardrails, routing, logging, spend). Spend is tracked against the shared team and user budgets, so running several agents (or logging in more than once) does not hand each session its own separate budget; they all draw down the same team/user allowance. There is no separate per-session cap, so sustained agent use is not capped at a small chat-session limit.
 
-The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, and `lite opencode` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
+The credential is short-lived by design (default 24h, configurable via `LITELLM_CLI_JWT_EXPIRATION_HOURS`); run `lite login` again to refresh it, which also re-reads your latest team and user settings. It does not appear in the Keys UI and cannot be rotated or revoked mid-session. `lite auth print-token` (usable as Claude Code's `apiKeyHelper`) prints it while it's still fresh and fails once it expires -- there is no silent renewal, so a long-running session needs a fresh `lite login` once a day. `lite claude`, `lite codex`, `lite opencode`, and `lite pi` work with it on a default deployment; `EXPERIMENTAL_UI_LOGIN` is not required. `lite login --pkce` is the exception to the daily re-login: it signs in through your system browser with OAuth authorization code and PKCE and stores a refresh token next to the key, so every `lite` command and `lite auth print-token` renew the key on their own shortly before it expires, `lite whoami` shows when the current key expires, and `lite logout` revokes the refresh token on the proxy (it needs a proxy that serves `/.well-known/litellm-cli-auth`; see [Browser sign-in with PKCE](https://docs.litellm.ai/docs/proxy/cli_sso#browser-sign-in-with-pkce)). When a renewal is refused, for example after a `lite logout` run from another copy of the credential, the command prints why on stderr and, once the key has run out, tells you to run `lite login --pkce` again. Only the holder can end a `--pkce` session early, with `lite logout`; an admin has no button for it, but every renewal re-reads the user on the proxy, so deactivating the user or removing them from the team makes the next renewal fail and the key runs out within `LITELLM_CLI_JWT_EXPIRATION_HOURS`. On a proxy with more than one worker or replica, configure Redis (`litellm_settings.cache` with Redis `cache_params`, or `general_settings.coordination_redis`) so a refresh token stays single-use and `lite logout` holds on every worker; without Redis each worker keeps its own record. If you need a long-lived, rotatable key that shows up in the Keys UI, create a dedicated virtual key in the dashboard and pass it via `--api-key` or `LITELLM_PROXY_API_KEY` instead.
 
 ### Route Every Claude Code Session Through the Proxy
 

litellm/proxy/client/cli/commands/agents.py

@@ -5,7 +5,7 @@ import sys
 from collections.abc import Callable, Mapping, Sequence
 from dataclasses import dataclass
 from types import MappingProxyType
-from typing import Final
+from typing import Final, TypeAlias
 
 import click
 import requests
@@ -13,6 +13,15 @@ from pydantic import BaseModel, TypeAdapter, ValidationError
 
 from .auth import CliContextObj, context_secret_vault, get_stored_api_key, login
 from .cmd_quoting import quote_for_cmd
+from .pi import (
+    LITELLM_PROXY_API_KEY_ENV,
+    PI_PROVIDER_NAME,
+    PiSyncError,
+    fetch_model_ids,
+    fetch_model_limits,
+    models_json_path,
+    sync_models_json,
+)
 
 ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL"
 ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN"
@@ -32,19 +41,24 @@ _SKIP_VERIFY_FLAG: Final = "--skip-verify"
 
 PROFILE_ANTHROPIC: Final = "anthropic"
 PROFILE_OPENAI: Final = "openai"
+PROFILE_LITELLM: Final = "litellm"
 
 _KNOWN_AGENTS: Final[dict[str, tuple[str, frozenset[str]]]] = {
     "claude": ("Claude Code", frozenset({PROFILE_ANTHROPIC})),
     "codex": ("Codex", frozenset({PROFILE_OPENAI})),
     "opencode": ("OpenCode", frozenset({PROFILE_OPENAI})),
+    "pi": ("pi", frozenset({PROFILE_LITELLM})),
 }
 
 _INSTALL_DOCS: Final[dict[str, str]] = {
     "claude": "https://docs.claude.com/en/docs/claude-code/setup",
     "codex": "https://developers.openai.com/codex/cli",
     "opencode": "https://opencode.ai/docs",
+    "pi": "https://pi.dev",
 }
 
+_HIDDEN_AGENTS: Final = frozenset({"pi"})
+
 CODEX_PROXY_PROVIDER: Final = "litellm"
 
 
@@ -81,6 +95,8 @@ def build_agent_env(
     the environment is left alone. CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY
     defaults to 1 so Claude Code (v2.1.129+) fills its /model picker from the
     proxy's /v1/models; likewise left alone when already set.
+    pi ignores both base URL variables and instead resolves $LITELLM_PROXY_API_KEY
+    from its synced models.json provider entry.
     """
     env: Final = dict(base_env)
     root: Final = base_url.rstrip("/")
@@ -95,6 +111,8 @@ def build_agent_env(
     if PROFILE_OPENAI in profiles:
         env[OPENAI_BASE_URL_ENV] = root + "/v1"
         env[OPENAI_API_KEY_ENV] = api_key
+    if PROFILE_LITELLM in profiles:
+        env[LITELLM_PROXY_API_KEY_ENV] = api_key
     return env
 
 
@@ -130,6 +148,40 @@ _PROXY_ARGS: Final[dict[str, Callable[[str], list[str]]]] = {
 }
 
 
+def prepare_pi(
+    base_url: str,
+    api_key: str,
+    base_env: Mapping[str, str],
+    *,
+    get: Callable[..., requests.Response] = requests.get,
+) -> tuple[str, ...]:
+    """Sync the proxy's model list into pi's models.json before handoff.
+
+    pi has no base-URL env vars, so this file is the only way to point it at the
+    proxy. Only the litellm provider entry is touched; the synced entry references
+    the key as $LITELLM_PROXY_API_KEY, which build_agent_env exports. The returned
+    --model pin is needed because pi ignores a bare --provider when picking the
+    interactive startup model; a user-supplied --model comes later in argv and wins.
+    """
+    ids: Final = fetch_model_ids(base_url, api_key, get=get)
+    if isinstance(ids, PiSyncError):
+        raise AgentRunError(ids.message)
+    limits: Final = fetch_model_limits(base_url, api_key, get=get)
+    path: Final = models_json_path(base_env)
+    error: Final = sync_models_json(path, base_url, ids, limits)
+    if error is not None:
+        raise AgentRunError(error.message)
+    click.echo(f"litellm: synced {len(ids)} proxy models into {path}")
+    return ("--model", f"{PI_PROVIDER_NAME}/{ids[0]}")
+
+
+_Preparer: TypeAlias = Callable[[str, str, Mapping[str, str]], Sequence[str]]
+
+_PREPARERS: Final[Mapping[str, _Preparer]] = MappingProxyType(
+    {"pi": prepare_pi}  # mutable-ok: MappingProxyType freezes the provider registry
+)
+
+
 def agent_launch_args(command: str, base_url: str) -> list[str]:
     """Extra CLI args an agent needs to actually honor the proxy.
 
@@ -407,15 +459,14 @@ def run_agent(
     warn: Callable[[str], None] = _warn,
     launcher: Callable[[str, Sequence[str], Mapping[str, str]], None] = _hand_off,
     reattach_terminal: Callable[[], None] | None = None,
+    preparers: Mapping[str, _Preparer] = MappingProxyType(_PREPARERS),
 ) -> None:
     """Validate, wire the environment, and hand off to the agent.
 
-    On success this never returns: POSIX replaces the current process, Windows
-    waits on the agent and exits with its status. Raises AgentRunError for
-    missing binaries, an unreachable proxy, or a rejected key. The model list is
-    synced only once the key check passed, so an unreachable proxy costs one
-    timeout rather than two, and --skip-verify keeps the launch fully offline.
-    reattach_terminal, when given, runs just before handoff to restore stdin.
+    On success this replaces the current process and never returns. Raises
+    AgentRunError for missing binaries, an unreachable proxy, a rejected key, or
+    a failed pre-launch config sync (pi). reattach_terminal, when given, runs
+    just before handoff to restore stdin.
     """
     if not command:
         raise AgentRunError("Nothing to run.")
@@ -435,13 +486,16 @@ def run_agent(
     if isinstance(synced, ModelSyncSkipped):
         warn(f"litellm: not syncing {display_name} models from the proxy: {synced.reason}")
 
+    prepare: Final = preparers.get(os.path.basename(command[0]))
+    prepared_args: Final = tuple(prepare(base_url, api_key, env_before_sync)) if prepare is not None else ()
+
     env: Final = MappingProxyType(
         {
             **build_agent_env(env_before_sync, base_url, api_key, profiles),
             **(_NO_EXTRA_ENV if isinstance(synced, ModelSyncSkipped) else synced),
         }
     )
-    extra_args: Final = agent_launch_args(command[0], base_url)
+    extra_args: Final = (*agent_launch_args(command[0], base_url), *prepared_args)
     if reattach_terminal is not None:
         reattach_terminal()
     launcher(binary, [command[0], *extra_args, *command[1:]], env)
@@ -501,6 +555,7 @@ def _make_agent_command(binary: str, display_name: str) -> click.Command:
         name=binary,
         context_settings={"ignore_unknown_options": True},
         short_help=f"Run {display_name} through your LiteLLM proxy",
+        hidden=binary in _HIDDEN_AGENTS,
     )
     @click.option("--skip-verify", is_flag=True, default=False, help=_SKIP_VERIFY_HELP)
     @click.argument("args", nargs=-1, type=click.UNPROCESSED)
@@ -533,6 +588,7 @@ __all__ = [
     "build_agent_env",
     "opencode_model_sync_env",
     "opencode_provider_config",
+    "prepare_pi",
     "resolve_api_key",
     "run_agent",
     "verify_proxy_key",

litellm/proxy/client/cli/commands/pi.py

@@ -0,0 +1,210 @@
+"""Sync a LiteLLM provider into pi's models.json.
+
+pi ignores ANTHROPIC_BASE_URL/OPENAI_BASE_URL, so `lite pi` routes it through the
+proxy by writing a provider entry instead. The key is stored as a $-reference so
+the short-lived login token never lands on disk.
+"""
+
+import json
+import os
+import tempfile
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from types import MappingProxyType
+from typing import Final
+
+import requests
+from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError
+
+PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR"
+PI_PROVIDER_NAME: Final = "litellm"
+LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY"
+
+
+@dataclass(frozen=True, slots=True)
+class PiSyncError:
+    message: str
+
+
+@dataclass(frozen=True, slots=True)
+class ModelLimits:
+    context_window: int | None
+    max_tokens: int | None
+
+
+class _Model(BaseModel):
+    id: str
+
+
+class _ModelList(BaseModel):
+    data: tuple[_Model, ...]
+
+
+class _ModelGroup(BaseModel):
+    model_group: str
+    max_input_tokens: float | None = None
+    max_output_tokens: float | None = None
+
+
+class _ModelGroupList(BaseModel):
+    data: tuple[_ModelGroup, ...]
+
+
+def fetch_model_ids(
+    base_url: str,
+    api_key: str,
+    *,
+    get: Callable[..., requests.Response] = requests.get,
+) -> tuple[str, ...] | PiSyncError:
+    url: Final = base_url.rstrip("/") + "/v1/models"
+    try:
+        resp: Final = get(
+            url,
+            headers={"Authorization": f"Bearer {api_key}"},  # mutable-ok: requests headers require a dict
+            timeout=10,
+        )
+    except requests.RequestException as e:
+        return PiSyncError(f"Could not list models from the proxy: {e}")
+    if resp.status_code != 200:
+        return PiSyncError(f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot build pi's model list.")
+    try:
+        listing: Final = _ModelList.model_validate(resp.json())
+    except (ValueError, ValidationError) as e:
+        return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}")
+    ids: Final = tuple(dict.fromkeys(model.id for model in listing.data))
+    if not ids:
+        return PiSyncError("The proxy returned no models for your key, so pi would have nothing to run.")
+    return ids
+
+
+_NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({})
+
+
+def fetch_model_limits(
+    base_url: str,
+    api_key: str,
+    *,
+    get: Callable[..., requests.Response] = requests.get,
+) -> Mapping[str, ModelLimits]:
+    """Best effort: pi falls back to its own defaults for models without limits,
+    so an unavailable /model_group/info must not block the launch."""
+    url: Final = base_url.rstrip("/") + "/model_group/info"
+    try:
+        resp: Final = get(
+            url,
+            headers={"Authorization": f"Bearer {api_key}"},  # mutable-ok: requests headers require a dict
+            timeout=10,
+        )
+        if resp.status_code != 200:
+            return _NO_LIMITS
+        listing: Final = _ModelGroupList.model_validate(resp.json())
+    except (requests.RequestException, ValueError, ValidationError):
+        return _NO_LIMITS
+    return MappingProxyType(
+        {
+            group.model_group: ModelLimits(
+                context_window=int(group.max_input_tokens) if group.max_input_tokens else None,
+                max_tokens=int(group.max_output_tokens) if group.max_output_tokens else None,
+            )
+            for group in listing.data
+        }
+    )
+
+
+def models_json_path(env: Mapping[str, str]) -> Path:
+    override: Final = env.get(PI_CONFIG_DIR_ENV)
+    root: Final = Path(override) if override else Path.home() / ".pi" / "agent"
+    return root / "models.json"
+
+
+def _model_entry(
+    model_id: str, limits: Mapping[str, ModelLimits]
+) -> dict[str, JsonValue]:  # mutable-ok: JSON object is serialized
+    limit: Final = limits.get(model_id)
+    context: Final[dict[str, JsonValue]] = (  # mutable-ok: JSON field
+        {"contextWindow": limit.context_window} if limit and limit.context_window else {}  # mutable-ok: JSON field
+    )
+    output: Final[dict[str, JsonValue]] = (  # mutable-ok: JSON field
+        {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {}
+    )  # mutable-ok: JSON field
+    return {"id": model_id, **context, **output}  # mutable-ok: JSON serialization requires a mutable object
+
+
+def provider_block(
+    base_url: str,
+    model_ids: tuple[str, ...],
+    limits: Mapping[str, ModelLimits] = _NO_LIMITS,
+) -> dict[str, JsonValue]:  # mutable-ok: JSON object is serialized
+    """openai-completions is the one API shape every LiteLLM model serves.
+
+    Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which
+    breaks compaction thresholds and over-asks models with smaller output caps.
+    """
+    return {  # mutable-ok: JSON serialization requires a mutable object
+        "baseUrl": base_url.rstrip("/") + "/v1",
+        "api": "openai-completions",
+        "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}",
+        "models": [_model_entry(model_id, limits) for model_id in model_ids],  # mutable-ok: JSON array
+    }
+
+
+_MODELS_FILE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue])
+
+
+def sync_models_json(
+    path: Path,
+    base_url: str,
+    model_ids: tuple[str, ...],
+    limits: Mapping[str, ModelLimits] = _NO_LIMITS,
+) -> PiSyncError | None:
+    """Replace only the litellm provider entry, leaving the rest of the file intact."""
+    try:
+        current: Final = (  # mutable-ok: JSON object default
+            _MODELS_FILE_ADAPTER.validate_json(path.read_text()) if path.exists() else {}
+        )
+    except (OSError, ValidationError) as e:
+        return PiSyncError(f"Could not read {path} as a JSON object: {e}. Fix or move the file, then retry.")
+    existing_providers: Final = current.get("providers", {})  # mutable-ok: JSON object default
+    if not isinstance(existing_providers, dict):
+        return PiSyncError(f'"providers" in {path} is not an object; fix or move the file, then retry.')
+    updated: Final = {  # mutable-ok: JSON serialization requires a mutable object
+        **current,
+        "providers": {  # mutable-ok: JSON serialization requires a mutable object
+            **existing_providers,
+            PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits),
+        },
+    }
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+    except OSError as e:
+        return PiSyncError(f"Could not write {path}: {e}")
+    try:
+        fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp")
+    except OSError as e:
+        return PiSyncError(f"Could not write {path}: {e}")
+    try:
+        with os.fdopen(fd, "w") as file:
+            file.write(json.dumps(updated, indent=2) + "\n")
+        os.replace(tmp_name, path)
+    except OSError as e:
+        try:
+            os.unlink(tmp_name)
+        except FileNotFoundError:
+            pass
+        return PiSyncError(f"Could not write {path}: {e}")
+    return None
+
+
+__all__ = (
+    "LITELLM_PROXY_API_KEY_ENV",
+    "PI_CONFIG_DIR_ENV",
+    "PI_PROVIDER_NAME",
+    "ModelLimits",
+    "PiSyncError",
+    "fetch_model_ids",
+    "fetch_model_limits",
+    "models_json_path",
+    "provider_block",
+    "sync_models_json",
+)

litellm/proxy/guardrails/_content_utils.py

@@ -8,7 +8,7 @@ skip the other shapes — these helpers normalise that so every hook sees
 every text fragment.
 """
 
-from collections.abc import Callable, Iterator, Mapping
+from collections.abc import Callable, Iterator, Mapping, Sequence
 from typing import Any, Final
 
 # Call types whose body carries free-form chat / prompt text that
@@ -33,6 +33,22 @@ def is_text_content_call_type(call_type: str) -> bool:
     return call_type in TEXT_CONTENT_CALL_TYPES
 
 
+# Call types whose request body carries no conversation at all. Embeddings carry
+# ``input`` — documents being indexed, not a prompt — which
+# :func:`build_inspection_messages` would lift into synthetic chat messages.
+#
+# Deny-list on purpose: ``TEXT_CONTENT_CALL_TYPES`` above omits conversational
+# call types (``anthropic_messages``, ``responses``, ``call_mcp_tool``), so a
+# blocking guardrail gated on that allow-list would stop inspecting real chat
+# traffic. Testing this instead leaves an unrecognised call type inspected.
+NON_CONVERSATIONAL_CALL_TYPES: Final[frozenset[str]] = frozenset({"embedding", "aembedding"})
+
+
+def is_non_conversational_call_type(call_type: str) -> bool:
+    """Return True if ``call_type``'s body carries no conversation to inspect."""
+    return call_type in NON_CONVERSATIONAL_CALL_TYPES
+
+
 TEXT_PART_TYPES: Final[frozenset[str]] = frozenset(
     {"text", "input_text", "output_text", "summary_text", "reasoning_text"}
 )
@@ -196,7 +212,17 @@ def walk_user_text(data: dict[str, Any], visit: Callable[[str], str]) -> int:
     return visited
 
 
-def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: list[dict[str, Any]]) -> None:
+def is_string_batch_input(data: Mapping[str, object]) -> bool:
+    """Return True when the only inspected content is an ``input`` list of plain
+    strings, the /embeddings batch shape, which :func:`apply_redacted_messages_back`
+    rewrites element-wise."""
+    if "messages" in data:
+        return False
+    input_value: Final = data.get("input")
+    return isinstance(input_value, list) and bool(input_value) and all(isinstance(item, str) for item in input_value)
+
+
+def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: Sequence[object]) -> bool:
     """Write redacted messages back to whichever field(s) the caller used.
 
     Mask/anonymize paths take a synthesised messages list (from
@@ -205,17 +231,39 @@ def apply_redacted_messages_back(data: dict[str, Any], redacted_messages: list[d
     only to ``data["messages"]`` leaves the Responses-API ``data["input"]``
     field untouched, so the unredacted text still reaches the LLM.
 
-    This helper updates both fields when both are present.
+    This helper updates both fields when both are present. A string batch
+    (``/embeddings`` ``input`` list) is rewritten element-wise: the n-th
+    redacted message replaces the n-th non-empty element, because
+    :func:`build_inspection_messages` emits one message per non-empty string.
+
+    Returns False, leaving ``data`` untouched, when a batch response does not
+    carry exactly one message per inspected element: a partial rewrite would
+    forward the remaining originals unredacted. Callers must block on False.
     """
+    if is_string_batch_input(data):
+        batch: Final = data["input"]
+        inspected_indices: Final = tuple(idx for idx, item in enumerate(batch) if item)
+        if len(redacted_messages) != len(inspected_indices):
+            return False
+        if any(not isinstance(message, Mapping) or message.get("content") is None for message in redacted_messages):
+            return False
+        redacted_texts: Final = tuple(
+            "\n".join(_iter_text_parts_in_content(message["content"])) for message in redacted_messages
+        )
+        for idx, text in zip(inspected_indices, redacted_texts):
+            batch[idx] = text
+        return True
     if "messages" in data:
         data["messages"] = redacted_messages
-    if isinstance(data.get("input"), str):
+    input_value: Final = data.get("input")
+    if isinstance(input_value, str):
         text_parts: Final[list[str]] = []
         for msg in redacted_messages:
             if not isinstance(msg, dict):
                 continue
             text_parts.extend(_iter_text_parts_in_content(msg.get("content")))
         data["input"] = "\n".join(text_parts)
+    return True
 
 
 def has_non_string_content(data: Mapping[str, object]) -> bool:

litellm/proxy/guardrails/guardrail_hooks/aim/__init__.py

@@ -18,6 +18,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         guardrail_name=guardrail.get("guardrail_name", ""),
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
+        inspect_embeddings=litellm_params.inspect_embeddings,
     )
     litellm.logging_callback_manager.add_litellm_callback(_aim_callback)
 

litellm/proxy/guardrails/guardrail_hooks/aim/aim.py

@@ -10,7 +10,7 @@ import os
 from collections.abc import AsyncGenerator, AsyncIterator, Mapping, Sequence
 from typing import TYPE_CHECKING, Final, TypeAlias
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter, ValidationError
 from typing_extensions import NotRequired, ReadOnly, TypedDict
 from websockets.asyncio.client import ClientConnection, connect
 
@@ -27,6 +27,8 @@ from litellm.proxy.guardrails._content_utils import (
     apply_redacted_messages_back,
     build_inspection_messages,
     has_non_string_content,
+    is_non_conversational_call_type,
+    is_string_batch_input,
 )
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import (
@@ -71,6 +73,9 @@ class AimRedactedChat(TypedDict):
     all_redacted_messages: ReadOnly[Sequence[AimRedactedMessage]]
 
 
+_REDACTED_CHAT_ADAPTER: Final = TypeAdapter(AimRedactedChat)
+
+
 class AimAnalyzeResponse(TypedDict):
     """Body returned by Aim's ``POST /fw/v1/analyze``."""
 
@@ -106,8 +111,15 @@ class AimGuardrail(CustomGuardrail):
             GuardrailEventHooks.post_call,
         ]
 
-    def __init__(self, api_key: str | None = None, api_base: str | None = None, **kwargs):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        inspect_embeddings: bool | None = None,
+        **kwargs,
+    ):
         kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
+        self.inspect_embeddings: Final = inspect_embeddings is True
         ssl_verify: Final = kwargs.pop("ssl_verify", None)
         self.async_handler = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.GuardrailCallback,
@@ -134,6 +146,12 @@ class AimGuardrail(CustomGuardrail):
         call_type: CallTypesLiteral,
     ) -> Exception | str | dict | None:
         verbose_proxy_logger.debug("Inside AIM Pre-Call Hook")
+        # /embeddings carries ``input`` — documents being indexed, not a prompt — which
+        # the flatten lifts into synthetic chat messages. A verdict on that text then
+        # blocks or silently rewrites a request that was never a conversation.
+        if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
+            verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
+            return data
         return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias)
 
     async def async_moderation_hook(
@@ -143,6 +161,9 @@ class AimGuardrail(CustomGuardrail):
         call_type: CallTypesLiteral,
     ) -> Exception | str | dict | None:
         verbose_proxy_logger.debug("Inside AIM Moderation Hook")
+        if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
+            verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
+            return data
 
         await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias)
         return data
@@ -215,24 +236,36 @@ class AimGuardrail(CustomGuardrail):
         # ``data["messages"]`` with that would silently strip image/audio
         # parts from a multimodal request — degrade to block so the
         # multimodal payload is never silently rewritten.
-        if has_non_string_content(data):
+        if has_non_string_content(data) and not is_string_batch_input(data):
             raise self._rejection(
                 "Aim: anonymize action requested for multimodal input "
                 "but mask-in-place would drop non-text parts. Send the "
                 "request with plain string content to use anonymize, "
                 "or rely on block-mode policies."
             )
-        redacted_messages: Final = [
-            {
-                "role": message["role"],
-                "content": message["content"],
-            }
-            for message in redacted_chat["all_redacted_messages"]
-        ]
+        try:
+            redacted_chat_model: Final = _REDACTED_CHAT_ADAPTER.validate_python(redacted_chat)
+        except ValidationError:
+            raise self._rejection(
+                "Aim: anonymize action returned malformed redacted messages, "
+                "so the request cannot be rewritten without forwarding unredacted text."
+            ) from None
+        redacted_messages: Final = list(redacted_chat_model["all_redacted_messages"])
+        if len(redacted_messages) != len(build_inspection_messages(data)):
+            raise self._rejection(
+                "Aim: anonymize action returned a redacted batch of a different "
+                "size than the inspected input, so the request cannot be "
+                "rewritten without forwarding unredacted text."
+            )
         # Write back to ``messages`` AND ``input``. The Responses-API
         # backend reads ``input``; writing only to ``messages`` would let
         # unredacted text reach the LLM for ``/v1/responses`` calls.
-        apply_redacted_messages_back(data, redacted_messages)
+        if not apply_redacted_messages_back(data, redacted_messages):
+            raise self._rejection(
+                "Aim: anonymize action returned a redacted batch of a different "
+                "size than the inspected input, so the request cannot be "
+                "rewritten without forwarding unredacted text."
+            )
         return data
 
     async def call_aim_guardrail_on_output(
@@ -261,9 +294,29 @@ class AimGuardrail(CustomGuardrail):
             return self._handle_block_action_on_output(res["analysis_result"], required_action)
         redacted_chat: Final = res.get("redacted_chat", None)
 
-        if action_type and action_type == "anonymize_action" and redacted_chat:
-            return {"redacted_output": redacted_chat["all_redacted_messages"][-1]["content"]}
-        return {"redacted_output": output}
+        if action_type != "anonymize_action":
+            return {"redacted_output": output}
+        try:
+            redacted_chat_model: Final = _REDACTED_CHAT_ADAPTER.validate_python(redacted_chat)
+        except ValidationError:
+            raise self._rejection(
+                "Aim: anonymize action returned malformed redacted output, "
+                "so the response cannot be rewritten without forwarding unredacted text."
+            ) from None
+        redacted_messages: Final = redacted_chat_model["all_redacted_messages"]
+        inspected_messages: Final = self._build_aim_inspection_messages(request_data)
+        if len(redacted_messages) != len(inspected_messages) + 1:
+            raise self._rejection(
+                "Aim: anonymize action returned an invalid redacted output count, "
+                "so the response cannot be rewritten without forwarding unredacted text."
+            )
+        redacted_output: Final = redacted_messages[-1]["content"]
+        if not redacted_output:
+            raise self._rejection(
+                "Aim: anonymize action returned empty redacted output, "
+                "so the response cannot be rewritten without forwarding unredacted text."
+            )
+        return {"redacted_output": redacted_output}
 
     def _handle_block_action_on_output(
         self, analysis_result: AimAnalysisResult, required_action: AimRequiredAction

litellm/proxy/guardrails/guardrail_hooks/cato_networks/__init__.py

@@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
         guardrail_name=guardrail.get("guardrail_name", ""),
         event_hook=litellm_params.mode,
         default_on=litellm_params.default_on,
+        inspect_embeddings=litellm_params.inspect_embeddings,
         ssl_verify=getattr(litellm_params, "ssl_verify", None),
     )
     litellm.logging_callback_manager.add_litellm_callback(_cato_callback)

litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py

@@ -32,6 +32,8 @@ from litellm.proxy._types import UserAPIKeyAuth
 from litellm.proxy.guardrails._content_utils import (
     apply_redacted_messages_back,
     build_inspection_messages,
+    is_non_conversational_call_type,
+    is_string_batch_input,
 )
 from litellm.types.guardrails import GuardrailEventHooks
 from litellm.types.utils import (
@@ -99,8 +101,15 @@ class CatoNetworksGuardrail(CustomGuardrail):
             GuardrailEventHooks.post_call,
         ]
 
-    def __init__(self, api_key: str | None = None, api_base: str | None = None, **kwargs):
+    def __init__(
+        self,
+        api_key: str | None = None,
+        api_base: str | None = None,
+        inspect_embeddings: bool | None = None,
+        **kwargs,
+    ):
         kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
+        self.inspect_embeddings: Final = inspect_embeddings is True
         ssl_verify: Final = kwargs.pop("ssl_verify", None)
         self.async_handler = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.GuardrailCallback,
@@ -154,6 +163,10 @@ class CatoNetworksGuardrail(CustomGuardrail):
         call_type: CallTypesLiteral,
     ) -> Exception | str | dict | None:
         verbose_proxy_logger.debug("Inside Cato Pre-Call Hook")
+        # /embeddings carries documents being indexed, not a conversation to inspect.
+        if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
+            verbose_proxy_logger.debug("Cato: skipping non-conversational call type %s", call_type)
+            return data
         return await self.call_cato_guardrail(
             data,
             hook="pre_call",
@@ -168,6 +181,9 @@ class CatoNetworksGuardrail(CustomGuardrail):
         call_type: CallTypesLiteral,
     ) -> Exception | str | dict | None:
         verbose_proxy_logger.debug("Inside Cato Moderation Hook")
+        if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
+            verbose_proxy_logger.debug("Cato: skipping non-conversational call type %s", call_type)
+            return data
         return await self.call_cato_guardrail(
             data,
             hook="moderation",
@@ -327,6 +343,16 @@ class CatoNetworksGuardrail(CustomGuardrail):
             return data
         redacted_messages: Final = redacted_chat.get("all_redacted_messages") or []
         original_messages: Final = data.get("messages")
+        sources: Final = self._extra_inspection_sources(data)
+        if is_string_batch_input(data) and len(redacted_messages) != sum(len(messages) for _, messages in sources):
+            raise HTTPException(
+                status_code=400,
+                detail=(
+                    "Cato: anonymize action returned a redacted batch of a different "
+                    "size than the inspected input, so the request cannot be rewritten "
+                    "without forwarding unredacted text."
+                ),
+            )
         offset = 0
         if original_messages:
             data["messages"] = [
@@ -338,26 +364,40 @@ class CatoNetworksGuardrail(CustomGuardrail):
                 for idx, original in enumerate(original_messages)
             ]
             offset = len(original_messages)
-        for field, messages in self._extra_inspection_sources(data):
+        for field, messages in sources:
             redacted_slice = redacted_messages[offset : offset + len(messages)]
             offset += len(messages)
-            if redacted_slice:
-                self._apply_extra_redaction(data, field, redacted_slice)
+            if not self._apply_extra_redaction(data, field, redacted_slice):
+                raise HTTPException(
+                    status_code=400,
+                    detail=(
+                        "Cato: anonymize action returned a redacted batch of a different "
+                        "size than the inspected input, so the request cannot be rewritten "
+                        "without forwarding unredacted text."
+                    ),
+                )
         return data
 
     @classmethod
-    def _apply_extra_redaction(cls, data: dict, field: str, redacted: list) -> None:
+    def _apply_extra_redaction(cls, data: dict, field: str, redacted: list) -> bool:
         if field == "input":
             input_only: Final = {"input": data["input"]}
-            apply_redacted_messages_back(input_only, redacted)
+            if not redacted:
+                return not is_string_batch_input(input_only)
+            if not apply_redacted_messages_back(input_only, redacted):
+                return False
             data["input"] = input_only["input"]
-        elif field == "instructions":
+            return True
+        if not redacted:
+            return True
+        if field == "instructions":
             if redacted[0].get("content") is not None:
                 data["instructions"] = redacted[0]["content"]
         elif field == "prompt":
             cls._apply_prompt_redaction(data, redacted)
         elif field == "schema_strings":
             cls._apply_schema_string_redaction(data, redacted)
+        return True
 
     @classmethod
     def _apply_schema_string_redaction(cls, data: dict, redacted: list) -> None:

litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py

@@ -36,6 +36,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) ->
         default_on=litellm_params.default_on or False,
         unreachable_fallback=litellm_params.unreachable_fallback,
         timeout=litellm_params.timeout,
+        ccr_retrieval=litellm_params.ccr_retrieval,
     )
     litellm.logging_callback_manager.add_litellm_callback(  # pyright: ignore[reportUnknownMemberType]
         _callback

litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py

@@ -6,6 +6,7 @@ import re
 import time
 import uuid
 from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard
 
 import httpx
@@ -60,7 +61,7 @@ _STREAM_CONVERTIBLE_CALL_TYPES: Final = frozenset(
 # stalled service holds the caller's request and a pooled connection for 600s or more.
 _COMPRESS_TIMEOUT_SECONDS: Final = 60.0
 HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve"
-_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})")
+_HASH_PATTERN: Final = re.compile(r"[a-f0-9]{12,24}")
 _HASH_CACHE_TTL_SECONDS: Final = 15 * 60
 # Narrows the base class's bare-dict ``request_data`` at the boundary so its
 # untranslated messages can be read with concrete types (values pass through by
@@ -239,15 +240,25 @@ def _protected_indices(
     tool exchanges the way ``compress()`` expands it, so a protected assistant
     tool call cannot end up answered by a marker standing in for the result the
     model just asked for.
+
+    Every assistant row is then withheld without expanding its tool exchange:
+    the service protects assistant text blocks but has no gate for assistant
+    strings, and the Anthropic adapter hands assistant blocks over as strings,
+    so the model's own earlier tables came back rewritten and it imitated the
+    shape. The tool results those turns asked for stay compressible.
     """
     protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices(
         messages, extra_retrieve_call_ids
     )
-    return protected | frozenset(
-        index
-        for group in group_tool_exchanges(messages)
-        if any(member in protected for member in group)
-        for index in group
+    return (
+        protected
+        | frozenset(
+            index
+            for group in group_tool_exchanges(messages)
+            if any(member in protected for member in group)
+            for index in group
+        )
+        | frozenset(index for index, message in enumerate(messages) if message.get("role") == "assistant")
     )
 
 
@@ -290,19 +301,23 @@ def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, obj
     return {"status_code": status_code, "body": body}
 
 
-def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]:
-    hashes: Final[list[str]] = []
-    for msg in messages:
-        content = msg.get("content")
-        if isinstance(content, str):
-            hashes.extend(_HASH_PATTERN.findall(content))
-        elif isinstance(content, list):
-            for block in content:
-                if isinstance(block, dict):
-                    text = block.get("text")
-                    if isinstance(text, str):
-                        hashes.extend(_HASH_PATTERN.findall(text))
-    return hashes
+def _read_ccr_hashes(body: Mapping[str, object]) -> frozenset[str]:
+    ccr_hashes: Final = body.get("ccr_hashes")
+    if not isinstance(ccr_hashes, list):
+        return frozenset()
+    return frozenset(
+        hash_value.lower()
+        for hash_value in ccr_hashes
+        if isinstance(hash_value, str) and _HASH_PATTERN.fullmatch(hash_value.lower())
+    )
+
+
+@dataclass(frozen=True, slots=True)
+class _CompressResult:
+    messages: list[dict[str, object]]
+    succeeded: bool
+    stats: dict[str, object]
+    ccr_hashes: frozenset[str] = frozenset()
 
 
 def _build_headroom_retrieve_tool() -> dict[str, object]:
@@ -319,7 +334,7 @@ def _build_headroom_retrieve_tool() -> dict[str, object]:
                 "properties": {
                     "hash": {
                         "type": "string",
-                        "description": "The 24-character hex hash from the compression marker.",
+                        "description": "The hex hash from the compression marker.",
                     },
                     "query": {
                         "type": "string",
@@ -479,6 +494,7 @@ class HeadroomGuardrail(CustomGuardrail):
         default_on: bool = False,
         unreachable_fallback: str | None = None,
         timeout: float | None = None,
+        ccr_retrieval: bool = True,
     ):
         self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/")
         if not self.headroom_api_base:
@@ -492,6 +508,7 @@ class HeadroomGuardrail(CustomGuardrail):
             "fail_open" if unreachable_fallback == "fail_open" else "fail_closed"
         )
         self.timeout: httpx.Timeout = self._resolve_timeout(timeout)
+        self.ccr_retrieval = ccr_retrieval
         self.async_handler = get_async_httpx_client(
             llm_provider=httpxSpecialProvider.GuardrailCallback,
         )
@@ -569,7 +586,7 @@ class HeadroomGuardrail(CustomGuardrail):
         self,
         messages: list[dict[str, object]],
         model: str | None,
-    ) -> tuple[list[dict[str, object]], bool, dict[str, object]]:
+    ) -> _CompressResult:
         payload: Final[dict[str, object]] = {"messages": messages}
         if model:
             payload["model"] = model
@@ -582,7 +599,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 timeout=self.timeout,
             )
         except httpx.HTTPStatusError as e:
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned an error",
@@ -592,7 +609,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 {},
             )
         except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e:
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service unreachable",
@@ -604,7 +621,7 @@ class HeadroomGuardrail(CustomGuardrail):
         response: Final[HttpxResponse] = raw_response
 
         if response.status_code != 200:
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned an error",
@@ -617,7 +634,7 @@ class HeadroomGuardrail(CustomGuardrail):
         try:
             body: Final[object] = response.json()
         except ValueError:
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned non-JSON response",
@@ -627,7 +644,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 {},
             )
         if not _is_str_object_dict(body):
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned unexpected response shape",
@@ -639,7 +656,7 @@ class HeadroomGuardrail(CustomGuardrail):
 
         compressed_messages: Final = body.get("messages")
         if not _is_object_list(compressed_messages):
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service response missing 'messages'",
@@ -651,7 +668,7 @@ class HeadroomGuardrail(CustomGuardrail):
 
         filtered: Final = [item for item in compressed_messages if _is_str_object_dict(item)]
         if not filtered:
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service returned empty message list",
@@ -664,7 +681,7 @@ class HeadroomGuardrail(CustomGuardrail):
         if len(filtered) != len(messages):
             # Rows are matched positionally when the never-compressed messages
             # are put back, so a reshaped conversation cannot be applied at all.
-            return (
+            return _CompressResult(
                 self._handle_compress_failure(
                     messages,
                     "Headroom compression service changed the message count",
@@ -705,7 +722,7 @@ class HeadroomGuardrail(CustomGuardrail):
             # tokens_saved, which the live compression service omits; derive it
             # so savings are counted, but let a service-sent value win.
             stats["tokens_saved"] = tokens_before - tokens_after
-        return filtered, True, stats
+        return _CompressResult(filtered, True, stats, _read_ccr_hashes(body))
 
     async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str:
         params: Final[dict[str, str]] = {}
@@ -793,7 +810,7 @@ class HeadroomGuardrail(CustomGuardrail):
 
         model: Final = self.headroom_model or request_data.get("model")
         start_time: Final = time.time()
-        returned, compression_succeeded, stats = await self._call_compress(
+        result: Final = await self._call_compress(
             messages=_flatten_messages_for_compression(compressible),
             model=model if isinstance(model, str) else None,
         )
@@ -803,7 +820,7 @@ class HeadroomGuardrail(CustomGuardrail):
             add_guardrail_to_applied_guardrails_header,
         )
 
-        if not compression_succeeded:
+        if not result.succeeded:
             self.add_standard_logging_guardrail_information_to_request_data(
                 guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"},
                 request_data=request_data,
@@ -822,12 +839,12 @@ class HeadroomGuardrail(CustomGuardrail):
 
         compressed: Final = _restore_protected_messages(
             messages=messages,
-            compressed=_restore_content_shapes(originals=compressible, returned=returned),
+            compressed=_restore_content_shapes(originals=compressible, returned=result.messages),
             protected_indices=protected_indices,
         )
 
         self.add_standard_logging_guardrail_information_to_request_data(
-            guardrail_json_response=stats,
+            guardrail_json_response=result.stats,
             request_data=request_data,
             guardrail_status="success",
             guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER,
@@ -837,7 +854,7 @@ class HeadroomGuardrail(CustomGuardrail):
         )
         add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
 
-        hashes: Final = extract_hashes_from_messages(compressed)
+        hashes: Final = result.ccr_hashes if self.ccr_retrieval else frozenset()
         if not hashes:
             return {**inputs, "structured_messages": compressed}  # pyright: ignore[reportReturnType]
 
@@ -918,14 +935,15 @@ class HeadroomGuardrail(CustomGuardrail):
         retrieved: Final[list[tuple[dict[str, object], str]]] = []
         for tc in tool_calls:
             arguments = tc.get("arguments", {})
-            hash_value = arguments.get("hash", "") if isinstance(arguments, dict) else ""
+            raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else ""
+            hash_value = str(raw_hash).lower()
             query = arguments.get("query") if isinstance(arguments, dict) else None
             # A hash is only honored if it was issued by *this request's own*
             # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by
             # message text alone is forgeable -- an attacker can plant a
             # hash-shaped string in their own prompt, and a hash issued for one
             # request would validate for any other request that echoes it back.
-            if str(hash_value) not in valid_hashes:
+            if hash_value not in valid_hashes:
                 verbose_proxy_logger.warning(
                     "Headroom CCR: rejecting hash=%s not produced by current request compression",
                     hash_value,
@@ -933,7 +951,7 @@ class HeadroomGuardrail(CustomGuardrail):
                 content = f"[Headroom: hash={hash_value} was not produced by the current request]"
             else:
                 content = await self._call_retrieve(
-                    hash_value=str(hash_value),
+                    hash_value=hash_value,
                     query=str(query) if query else None,
                 )
             verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content))

litellm/proxy/management_endpoints/mcp_management_endpoints.py

@@ -193,6 +193,7 @@ if MCP_AVAILABLE:
         UpdateMCPServerRequest,
         UserAPIKeyAuth,
         UserMCPManagementMode,
+        is_per_server_oauth_discovery_eligible,
     )
     from litellm.proxy.auth.user_api_key_auth import (
         _user_api_key_auth_builder,
@@ -2714,6 +2715,27 @@ if MCP_AVAILABLE:
             old_server_record = None
             old_server_record_read_failed = True
 
+        if payload.per_server_oauth_discovery and (old_server_record is not None or old_server_record_read_failed):
+            relay_eligible: Final = old_server_record is not None and is_per_server_oauth_discovery_eligible(
+                payload.auth_type if "auth_type" in payload_fields_set else old_server_record.auth_type,
+                payload.oauth2_flow if "oauth2_flow" in payload_fields_set else old_server_record.oauth2_flow,
+                (
+                    payload.delegate_auth_to_upstream
+                    if "delegate_auth_to_upstream" in payload_fields_set
+                    else old_server_record.delegate_auth_to_upstream
+                ),
+            )
+            if not relay_eligible:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail={  # mutable-ok: FastAPI HTTPException detail requires a plain dict
+                        "error": (
+                            "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow "
+                            "authorization_code and without delegate_auth_to_upstream."
+                        )
+                    },
+                )
+
         if (
             payload.dcr_bridge
             and payload.auth_type is None

litellm/proxy/schema.prisma

@@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable {
   delegate_auth_to_upstream Boolean @default(false)
   oauth_passthrough Boolean @default(false)
   dcr_bridge Boolean?
+  per_server_oauth_discovery Boolean @default(false)
   is_byok               Boolean  @default(false)
   byok_description      String[] @default([])
   byok_api_key_help_url String?

litellm/router.py

@@ -9073,10 +9073,12 @@ class Router:
             if prefs_raw is not None:
                 model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw)
 
-            # `input_cost_per_token` is a LiteLLM_Params field per types/router.py.
+            # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set.
             lp = d.get("litellm_params") if isinstance(d, dict) else d.litellm_params
             lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
             cost = lp_dict.get("input_cost_per_token")
+            if cost is None:
+                cost = mi_dict.get("input_cost_per_token")
             if cost is not None:
                 model_to_cost[name] = float(cost)
 

litellm/router_strategy/adaptive_router/adaptive_router.py

@@ -123,7 +123,12 @@ class AdaptiveRouter:
                 self._cells[(rt, model)] = initial_cell(prefs, rt)
 
     async def load_state_from_db(self, prisma_client: Any) -> None:
-        """Override cold-start cells with persisted state. Called once at startup."""
+        """Add each row's persisted delta to a freshly computed cold-start prior.
+
+        A row holds an accumulated delta, not a full posterior, and can be one-sided
+        (e.g. beta=0) - assigning it straight into the cell would zero out a Beta shape
+        parameter and crash thompson_sample() on every later draw for that cell.
+        """
         if prisma_client is None:
             return
         try:
@@ -139,7 +144,12 @@ class AdaptiveRouter:
                     continue
                 if row.model_name not in self.config.available_models:
                     continue
-                self._cells[(rt, row.model_name)] = BanditCell(alpha=row.alpha, beta=row.beta)
+                prefs = self.model_to_prefs.get(row.model_name) or _default_prefs()
+                prior = initial_cell(prefs, rt)
+                self._cells[(rt, row.model_name)] = BanditCell(
+                    alpha=prior.alpha + row.alpha,
+                    beta=prior.beta + row.beta,
+                )
                 loaded += 1
             verbose_router_logger.info(
                 "AdaptiveRouter[%s]: loaded %d cells from DB",

litellm/router_strategy/auto_router/auto_router.py

@@ -2,6 +2,7 @@
 Auto-Routing Strategy that works with a Semantic Router Config
 """
 
+import asyncio
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Final, Optional
 
@@ -75,6 +76,8 @@ class AutoRouter(CustomLogger):
         self.auto_sync_value = self.DEFAULT_AUTO_SYNC_VALUE
         self.loaded_routes: list[Route] = self._load_semantic_routing_routes()
         self.routelayer: SemanticRouter | None = None
+        self._routelayer_lock = asyncio.Lock()
+        self._routelayer_build_task: asyncio.Task[SemanticRouter] | None = None
         self.default_model = default_model
         self.embedding_model: str = embedding_model
         self.max_input_chars: int = max_input_chars
@@ -115,6 +118,45 @@ class AutoRouter(CustomLogger):
             )
         return auto_router_routes
 
+    def _build_routelayer(self) -> "SemanticRouter":
+        """Synchronous (embeds every route's utterances); run only via `_ensure_routelayer`."""
+        if self.routelayer is not None:
+            return self.routelayer
+
+        from semantic_router.routers import SemanticRouter
+
+        routelayer: Final = SemanticRouter(
+            routes=self.loaded_routes,
+            encoder=self.encoder,
+            auto_sync=self.auto_sync_value,
+        )
+        self.routelayer = routelayer
+        return routelayer
+
+    def _clear_build_task_on_failure(self, build_task: "asyncio.Task[SemanticRouter]") -> None:
+        """Runs even with no caller left awaiting, so a failure never stays cached forever."""
+        if build_task is self._routelayer_build_task and not build_task.cancelled() and build_task.exception():
+            self._routelayer_build_task = None
+
+    async def _ensure_routelayer(self) -> "SemanticRouter":
+        """Build the route layer once, off the event loop, shared across concurrent callers.
+
+        A shared task (not a bare `asyncio.to_thread` awaited under the lock) survives one
+        caller's cancellation, so `cancel_on_disconnect` can't free a second caller into
+        starting a duplicate build.
+        """
+        if self.routelayer is not None:
+            return self.routelayer
+        async with self._routelayer_lock:
+            if self.routelayer is not None:
+                return self.routelayer
+            build_task = self._routelayer_build_task
+            if build_task is None:
+                build_task = asyncio.ensure_future(asyncio.to_thread(self._build_routelayer))
+                build_task.add_done_callback(self._clear_build_task_on_failure)
+                self._routelayer_build_task = build_task
+        return await asyncio.shield(build_task)
+
     @staticmethod
     def _extract_text_from_messages(messages: list[dict[str, Any]]) -> str:
         """
@@ -151,8 +193,6 @@ class AutoRouter(CustomLogger):
 
         Used for the litellm auto-router to modify the request before the routing decision is made.
         """
-        from semantic_router.routers import SemanticRouter
-
         from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
         from litellm.types.router import PreRoutingHookResponse
 
@@ -164,17 +204,7 @@ class AutoRouter(CustomLogger):
         if resolved_messages is None:
             return None
 
-        routelayer = self.routelayer
-        if routelayer is None:
-            #######################
-            # Create the route layer
-            #######################
-            routelayer = SemanticRouter(
-                routes=self.loaded_routes,
-                encoder=self.encoder,
-                auto_sync=self.auto_sync_value,
-            )
-            self.routelayer = routelayer
+        routelayer = await self._ensure_routelayer()
 
         message_content: Final = self._extract_text_from_messages(resolved_messages)
         route_name: Final = await self._matched_route_name(routelayer, message_content, request_kwargs)

litellm/router_strategy/complexity_router/complexity_router.py

@@ -2174,9 +2174,12 @@ class ComplexityRouter(CustomLogger):
             else:
                 model_to_prefs[name] = AdaptiveRouterPreferences(quality_tier=2, strengths=[])
 
+            # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set.
             lp = deployment.get("litellm_params") if isinstance(deployment, dict) else deployment.litellm_params
             lp_dict: dict[str, Any] = lp if isinstance(lp, dict) else (lp.model_dump() if lp else {})
             cost = lp_dict.get("input_cost_per_token")
+            if cost is None:
+                cost = mi_dict.get("input_cost_per_token")
             model_to_cost[name] = float(cost) if cost is not None else 0.0
 
         self.adaptive_router = AdaptiveRouter(

litellm/types/guardrails.py

@@ -832,6 +832,15 @@ class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch up
         ),
     )
 
+    inspect_embeddings: bool | None = Field(
+        default=None,
+        description=(
+            "When True, the Aim and Cato Networks guardrails send /embeddings `input` to the vendor as "
+            "user messages. Off by default because embedding input is documents being indexed, not a "
+            "conversation."
+        ),
+    )
+
     # Lakera specific params
     category_thresholds: LakeraCategoryThresholds | None = Field(
         default=None,

litellm/types/integrations/azure_sentinel.py

@@ -1,5 +1,9 @@
+from typing import Final
+
 from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams
 
+AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES: Final = 1_000_000
+
 
 class AzureSentinelInitParams(StandardCustomLoggerInitParams):
     """

litellm/types/mcp_server/mcp_server_manager.py

@@ -158,6 +158,7 @@ class MCPServer(BaseModel):
     # be set explicitly to avoid regressing servers that did not opt in.
     oauth_passthrough: bool = False
     dcr_bridge: bool | None = None
+    per_server_oauth_discovery: bool = False
     is_byok: bool = False
     byok_description: list[str] = []
     byok_api_key_help_url: str | None = None
@@ -241,6 +242,16 @@ class MCPServer(BaseModel):
         so they are excluded by construction."""
         return self.auth_type == MCPAuth.oauth2 and not self.delegate_auth_to_upstream
 
+    @property
+    def uses_per_server_oauth_relay(self) -> bool:
+        """Whether named discovery should advertise the configured per-server OAuth relay."""
+        return self.per_server_oauth_discovery and self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials
+
+    @property
+    def advertises_gateway_authorization_server(self) -> bool:
+        """Whether named discovery should advertise the aggregate gateway authorization server."""
+        return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay
+
     @property
     def is_true_passthrough(self) -> bool:
         """True for the transparent-proxy mode: LiteLLM performs no admission auth and forwards the

litellm/types/proxy/guardrails/guardrail_hooks/aim.py

@@ -1,5 +1,7 @@
 from pydantic import Field
 
+from litellm.types.guardrails import GuardrailParamUITypes
+
 from .base import GuardrailConfigModel
 
 
@@ -12,6 +14,14 @@ class AimGuardrailConfigModel(GuardrailConfigModel):
         default=None,
         description="The API base for the Aim guardrail. Default is https://api.aim.security. Also checks if the `AIM_API_BASE` environment variable is set.",
     )
+    inspect_embeddings: bool | None = Field(
+        default=False,
+        description=(
+            "Send /embeddings `input` to Aim as user messages. Off by default because embedding input is "
+            "documents being indexed, not a conversation."
+        ),
+        json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},  # mutable-ok: pydantic accepts only a dict here
+    )
 
     @staticmethod
     def ui_friendly_name() -> str:

litellm/types/proxy/guardrails/guardrail_hooks/cato_networks.py

@@ -1,5 +1,7 @@
 from pydantic import Field
 
+from litellm.types.guardrails import GuardrailParamUITypes
+
 from .base import GuardrailConfigModel
 
 
@@ -12,6 +14,14 @@ class CatoNetworksGuardrailConfigModel(GuardrailConfigModel):
         default=None,
         description="The API base for the Cato Networks guardrail. Default is https://api.aisec.catonetworks.com. Also checks if the `CATO_API_BASE` environment variable is set.",
     )
+    inspect_embeddings: bool | None = Field(
+        default=False,
+        description=(
+            "Send /embeddings `input` to Cato Networks as user messages. Off by default because embedding "
+            "input is documents being indexed, not a conversation."
+        ),
+        json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},  # mutable-ok: pydantic accepts only a dict here
+    )
 
     @staticmethod
     def ui_friendly_name() -> str:

litellm/types/proxy/guardrails/guardrail_hooks/headroom.py

@@ -26,6 +26,10 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]):
             "forwards the request uncompressed instead of blocking it."
         ),
     )
+    ccr_retrieval: bool = Field(
+        default=True,
+        description="Inject the Headroom retrieval tool for hashes declared by the compression service.",
+    )
 
     @staticmethod
     def ui_friendly_name() -> str:

model_prices_and_context_window.json

@@ -12238,6 +12238,48 @@
         "output_cost_per_token": 2.65e-06,
         "supports_pdf_input": true
     },
+    "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0": {
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 8172,
+        "max_tokens": 8172,
+        "mode": "embedding",
+        "input_cost_per_token": 1.62e-07,
+        "input_cost_per_image": 7.2e-05,
+        "input_cost_per_video_per_second": 0.00084,
+        "input_cost_per_audio_per_second": 0.000168,
+        "output_cost_per_token": 0.0,
+        "output_vector_size": 3072,
+        "supports_embedding_image_input": true,
+        "supports_image_input": true,
+        "supports_video_input": true,
+        "supports_audio_input": true
+    },
+    "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 300000,
+        "max_output_tokens": 10000,
+        "max_tokens": 10000,
+        "mode": "chat",
+        "output_cost_per_token": 2.88e-07,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_response_schema": true,
+        "supports_vision": true
+    },
+    "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": {
+        "input_cost_per_token": 4.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 10000,
+        "max_tokens": 10000,
+        "mode": "chat",
+        "output_cost_per_token": 1.68e-07,
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_response_schema": true
+    },
     "bedrock/us-gov-west-1/amazon.nova-pro-v1:0": {
         "input_cost_per_token": 9.6e-07,
         "litellm_provider": "bedrock",
@@ -43648,6 +43690,23 @@
         "input_cost_per_token_batches": 1.65e-06,
         "output_cost_per_token_batches": 8.25e-06
     },
+    "us-gov.anthropic.claude-3-haiku-20240307-v1:0": {
+        "deprecation_date": "2026-09-10",
+        "input_cost_per_token": 3e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 200000,
+        "max_output_tokens": 4096,
+        "max_tokens": 4096,
+        "mode": "chat",
+        "output_cost_per_token": 1.5e-06,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "cache_read_input_token_cost": 3e-08,
+        "cache_creation_input_token_cost": 3.75e-07
+    },
     "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0": {
         "cache_creation_input_token_cost": 4.5e-06,
         "cache_creation_input_token_cost_above_1hr": 7.2e-06,
@@ -43742,6 +43801,160 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "us-gov.anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "us-gov.anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
+    "us-gov.nvidia.nemotron-nano-3-30b": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 262144,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.88e-07,
+        "source": "https://aws.amazon.com/bedrock/pricing/",
+        "supports_function_calling": true,
+        "supports_native_structured_output": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.nvidia.nemotron-nano-12b-v2": {
+        "input_cost_per_token": 2.4e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 7.2e-07,
+        "supports_system_messages": true,
+        "supports_vision": true
+    },
+    "us-gov.nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
+    "us-gov.nvidia.nemotron-super-3-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "output_cost_per_token": 7.8e-07,
+        "source": "https://aws.amazon.com/bedrock/pricing/",
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_system_messages": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.openai.gpt-oss-20b-1:0": {
+        "input_cost_per_token": 8.4e-08,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3.6e-07,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.openai.gpt-oss-120b-1:0": {
+        "input_cost_per_token": 1.8e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 7.2e-07,
+        "supports_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "us-gov.xai.grok-4.6": {
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_converse",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_prompt_caching": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
     "au.anthropic.claude-haiku-4-5-20251001-v1:0": {
         "cache_creation_input_token_cost": 1.375e-06,
         "cache_creation_input_token_cost_above_1hr": 2.2e-06,
@@ -59547,6 +59760,16 @@
         "supports_system_messages": true,
         "supports_vision": true
     },
+    "bedrock/us-gov-west-1/nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
     "bedrock/us-gov-west-1/nvidia.nemotron-super-3-120b": {
         "input_cost_per_token": 1.8e-07,
         "litellm_provider": "bedrock",
@@ -59651,6 +59874,70 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "bedrock/us-gov-west-1/anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
     "bedrock/us-gov-east-1/nvidia.nemotron-nano-3-30b": {
         "input_cost_per_token": 7.2e-08,
         "litellm_provider": "bedrock",
@@ -59676,6 +59963,16 @@
         "supports_system_messages": true,
         "supports_vision": true
     },
+    "bedrock/us-gov-east-1/nvidia.nemotron-nano-9b-v2": {
+        "input_cost_per_token": 7.2e-08,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 8192,
+        "max_tokens": 8192,
+        "mode": "chat",
+        "output_cost_per_token": 2.76e-07,
+        "supports_system_messages": true
+    },
     "bedrock/us-gov-east-1/nvidia.nemotron-super-3-120b": {
         "input_cost_per_token": 1.8e-07,
         "litellm_provider": "bedrock",
@@ -59780,6 +60077,70 @@
         "supports_vision": true,
         "supports_xhigh_reasoning_effort": true
     },
+    "bedrock/us-gov-east-1/anthropic.claude-opus-5": {
+        "bedrock_converse_supports_strict_tools": false,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+        "input_cost_per_token": 6e-06,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 3e-05,
+        "prompt_cache_min_tokens": 512,
+        "supports_adaptive_thinking": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_function_calling": true,
+        "supports_max_reasoning_effort": true,
+        "supports_mid_conversation_system": true,
+        "supports_native_structured_output": false,
+        "supports_output_config": true,
+        "supports_parallel_tool_use_config": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true
+    },
+    "bedrock/us-gov-east-1/anthropic.claude-fable-5-1": {
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+        "input_cost_per_token": 1.2e-05,
+        "litellm_provider": "bedrock",
+        "max_input_tokens": 1000000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "output_cost_per_token": 6e-05,
+        "supports_adaptive_thinking": true,
+        "thinking_always_on": true,
+        "supports_mid_conversation_system": true,
+        "supports_assistant_prefill": false,
+        "supports_computer_use": true,
+        "supports_forced_tool_use": false,
+        "supports_function_calling": true,
+        "supports_pdf_input": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_sampling_params": false,
+        "supports_tool_choice": true,
+        "supports_vision": true,
+        "supports_xhigh_reasoning_effort": true,
+        "supports_native_structured_output": false,
+        "supports_max_reasoning_effort": true,
+        "supports_output_config": true,
+        "bedrock_output_config_effort_ceiling": "xhigh",
+        "supports_parallel_tool_use_config": true,
+        "prompt_cache_min_tokens": 512
+    },
     "bedrock_mantle/us-gov-west-1/openai.gpt-5.6-terra": {
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
@@ -59894,6 +60255,120 @@
         "output_cost_per_token": 3e-06,
         "cache_read_input_token_cost": 2.4e-07
     },
+    "bedrock_mantle/us-gov-west-1/xai.grok-4.6": {
+        "use_openai_responses_path": true,
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": {
+        "input_cost_per_token": 4.8e-08,
+        "output_cost_per_token": 9.6e-08,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 128000,
+        "max_output_tokens": 128000,
+        "max_tokens": 128000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": {
+        "input_cost_per_token": 1.56e-07,
+        "output_cost_per_token": 4.8e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": {
+        "input_cost_per_token": 1.68e-07,
+        "output_cost_per_token": 4.8e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 256000,
+        "max_output_tokens": 256000,
+        "max_tokens": 256000,
+        "mode": "chat",
+        "use_openai_responses_path": true,
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": false,
+        "supports_reasoning": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": {
+        "input_cost_per_token": 8.4e-08,
+        "output_cost_per_token": 3.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "output_cost_per_token": 7.2e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
     "bedrock_mantle/us-gov-east-1/openai.gpt-5.4": {
         "litellm_provider": "bedrock_mantle",
         "max_input_tokens": 1050000,
@@ -59921,6 +60396,63 @@
         "cache_read_input_token_cost": 3.3e-07,
         "output_cost_per_token": 1.98e-05
     },
+    "bedrock_mantle/us-gov-east-1/xai.grok-4.6": {
+        "use_openai_responses_path": true,
+        "input_cost_per_token": 2.64e-06,
+        "output_cost_per_token": 7.92e-06,
+        "cache_read_input_token_cost": 6.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 500000,
+        "max_output_tokens": 500000,
+        "max_tokens": 500000,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_prompt_caching": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true,
+        "supports_vision": true
+    },
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": {
+        "input_cost_per_token": 8.4e-08,
+        "output_cost_per_token": 3.6e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": {
+        "input_cost_per_token": 1.8e-07,
+        "output_cost_per_token": 7.2e-07,
+        "litellm_provider": "bedrock_mantle",
+        "max_input_tokens": 131072,
+        "max_output_tokens": 32768,
+        "max_tokens": 32768,
+        "mode": "chat",
+        "supported_endpoints": [
+            "/v1/chat/completions",
+            "/v1/responses"
+        ],
+        "supports_function_calling": true,
+        "supports_parallel_function_calling": true,
+        "supports_reasoning": true,
+        "supports_response_schema": true,
+        "supports_tool_choice": true
+    },
     "azure/us-gov/gpt-5.1": {
         "cache_read_input_token_cost": 1.71875e-07,
         "default_reasoning_effort": "none",

schema.prisma

@@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable {
   delegate_auth_to_upstream Boolean @default(false)
   oauth_passthrough Boolean @default(false)
   dcr_bridge Boolean?
+  per_server_oauth_discovery Boolean @default(false)
   is_byok               Boolean  @default(false)
   byok_description      String[] @default([])
   byok_api_key_help_url String?

tests/test_litellm/integrations/datadog/test_datadog_logger_batching.py

@@ -4,6 +4,7 @@ from unittest.mock import AsyncMock, Mock, patch
 import httpx
 import pytest
 from httpx import Request, Response
+from pydantic import BaseModel, computed_field
 
 from litellm.integrations.datadog.datadog import DataDogLogger
 from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
@@ -31,9 +32,7 @@ def _payloads(n, message=None):
 def _raised_413():
     request = Request("POST", "https://example.com")
     response = Response(413, request=request, text="Payload Too Large")
-    return MaskedHTTPStatusError(
-        httpx.HTTPStatusError("413", request=request, response=response)
-    )
+    return MaskedHTTPStatusError(httpx.HTTPStatusError("413", request=request, response=response))
 
 
 def _make_send(max_ok, delivered, *, raise_413=True):
@@ -85,9 +84,7 @@ async def test_async_send_batch_keeps_events_appended_during_send(datadog_env):
                 status="info",
             )
         )
-        return Response(
-            202, request=Request("POST", "https://example.com"), text="Accepted"
-        )
+        return Response(202, request=Request("POST", "https://example.com"), text="Accepted")
 
     logger.async_send_compressed_data = AsyncMock(side_effect=_mock_send)
 
@@ -172,9 +169,7 @@ async def test_413_returned_response_also_splits(datadog_env):
 
     logger.log_queue = _payloads(4)
     delivered: list = []
-    logger.async_send_compressed_data = AsyncMock(
-        side_effect=_make_send(1, delivered, raise_413=False)
-    )
+    logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(1, delivered, raise_413=False))
 
     await logger.async_send_batch()
 
@@ -186,9 +181,7 @@ def _make_recording_send(sent_batches, delivered):
     async def _send(data):
         sent_batches.append(list(data))
         delivered.extend(data)
-        return Response(
-            202, request=Request("POST", "https://example.com"), text="Accepted"
-        )
+        return Response(202, request=Request("POST", "https://example.com"), text="Accepted")
 
     return _send
 
@@ -206,18 +199,13 @@ async def test_oversized_payload_splits_before_any_send(datadog_env):
     logger.log_queue = list(events)
     sent_batches: list = []
     delivered: list = []
-    logger.async_send_compressed_data = AsyncMock(
-        side_effect=_make_recording_send(sent_batches, delivered)
-    )
+    logger.async_send_compressed_data = AsyncMock(side_effect=_make_recording_send(sent_batches, delivered))
 
     await logger.async_send_batch()
 
     assert delivered == events
     assert len(sent_batches) == 3
-    assert all(
-        len(safe_dumps(batch).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES
-        for batch in sent_batches
-    )
+    assert all(len(safe_dumps(batch).encode("utf-8")) <= DD_MAX_PAYLOAD_SIZE_BYTES for batch in sent_batches)
     assert logger.log_queue == []
 
 
@@ -232,9 +220,7 @@ async def test_batch_over_max_event_count_splits_before_any_send(datadog_env):
     logger.log_queue = list(events)
     sent_batches: list = []
     delivered: list = []
-    logger.async_send_compressed_data = AsyncMock(
-        side_effect=_make_recording_send(sent_batches, delivered)
-    )
+    logger.async_send_compressed_data = AsyncMock(side_effect=_make_recording_send(sent_batches, delivered))
 
     await logger.async_send_batch()
 
@@ -281,9 +267,7 @@ async def test_partial_delivery_then_transient_error_requeues_only_undelivered(
         if messages == ['{"event": 2}', '{"event": 3}']:
             raise RuntimeError("transient network error")
         delivered.extend(messages)
-        return Response(
-            202, request=Request("POST", "https://example.com"), text="Accepted"
-        )
+        return Response(202, request=Request("POST", "https://example.com"), text="Accepted")
 
     logger.async_send_compressed_data = AsyncMock(side_effect=_send)
 
@@ -304,9 +288,7 @@ async def test_unexpected_non_202_status_requeues(datadog_env):
 
     logger.log_queue = _payloads(2)
     logger.async_send_compressed_data = AsyncMock(
-        return_value=Response(
-            200, request=Request("POST", "https://example.com"), text="OK"
-        )
+        return_value=Response(200, request=Request("POST", "https://example.com"), text="OK")
     )
 
     await logger.async_send_batch()
@@ -502,3 +484,77 @@ async def test_flush_queue_returns_without_lock(datadog_env):
     await logger.flush_queue()
 
     logger.async_send_batch.assert_not_awaited()
+
+
+class _RaisesWhileDumping(BaseModel):
+    @computed_field
+    @property
+    def rendered(self) -> str:
+        raise RuntimeError("this field cannot be rendered")
+
+
+@pytest.mark.asyncio
+async def test_event_whose_serialization_raises_is_dropped_alone(datadog_env):
+    """safe_dumps hands pydantic models to model_dump, so serialization can raise any exception
+    class. The intake-limit probe has to isolate that one event and drop it, not fail the whole
+    batch back onto the queue where it would poison every later flush."""
+    with patch("asyncio.create_task"):
+        logger = DataDogLogger()
+
+    logger.log_queue = _payloads(4)
+    logger.log_queue[1]["message"] = _RaisesWhileDumping()
+    delivered: list = []
+    logger.async_send_compressed_data = AsyncMock(side_effect=_make_send(DD_MAX_BATCH_SIZE, delivered))
+
+    await logger.async_send_batch()
+
+    assert delivered == ['{"event": 0}', '{"event": 2}', '{"event": 3}']
+    assert logger.log_queue == []
+
+
+@pytest.mark.asyncio
+async def test_cancellation_mid_split_requeues_only_the_undelivered_events(datadog_env):
+    """A cancelled split must keep the pieces Datadog never accepted, without resending the piece
+    it did, and must surface as a plain CancelledError so asyncio.wait_for still reads it as a
+    timeout on Python 3.12."""
+    with patch("asyncio.create_task"):
+        logger = DataDogLogger()
+
+    logger.log_queue = _payloads(4)
+    attempts: list = []
+
+    async def _send(data):
+        if len(data) > 2:
+            raise _raised_413()
+        attempts.append([event["message"] for event in data])
+        if len(attempts) > 1:
+            raise asyncio.CancelledError
+        return Response(202, request=Request("POST", "https://example.com"), text="Accepted")
+
+    logger.async_send_compressed_data = AsyncMock(side_effect=_send)
+
+    with pytest.raises(asyncio.CancelledError) as excinfo:
+        await logger.async_send_batch()
+
+    assert type(excinfo.value) is asyncio.CancelledError
+    assert attempts == [['{"event": 0}', '{"event": 1}'], ['{"event": 2}', '{"event": 3}']]
+    assert [event["message"] for event in logger.log_queue] == ['{"event": 2}', '{"event": 3}']
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status_code", [400, 403, 429, 500, 503])
+async def test_raised_intake_error_preserves_datadog_requeue_behavior(datadog_env, status_code):
+    """Datadog requeues every non-413 HTTP failure so a corrected key or endpoint can recover telemetry."""
+    with patch("asyncio.create_task"):
+        logger = DataDogLogger()
+
+    logger.log_queue = _payloads(2)
+    request = Request("POST", "https://example.com")
+    response = Response(status_code, request=request, text="rejected")
+    logger.async_send_compressed_data = AsyncMock(
+        side_effect=MaskedHTTPStatusError(httpx.HTTPStatusError(str(status_code), request=request, response=response))
+    )
+
+    await logger.async_send_batch()
+
+    assert [event["message"] for event in logger.log_queue] == ['{"event": 0}', '{"event": 1}']

tests/test_litellm/integrations/test_azure_sentinel.py

@@ -2,18 +2,23 @@
 Test Azure Sentinel logging integration
 """
 
+import asyncio
 import json
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+from httpx import Request, Response
+from pydantic import BaseModel, computed_field
 
 from litellm.integrations.azure_sentinel.azure_sentinel import AzureSentinelLogger
+from litellm.llms.custom_httpx.http_handler import MaskedHTTPStatusError
+from litellm.types.integrations.azure_sentinel import AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES
 from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload
 
 
 def _close_periodic_flush_task(coro):
     coro.close()
-    return None
 
 
 @pytest.mark.asyncio
@@ -414,3 +419,837 @@ def test_azure_sentinel_authority_host_argument_outranks_the_scoped_env_var(_no_
 
     assert logger.authority_host == "https://login.microsoftonline.com"
     assert logger.oauth_scope == "https://monitor.azure.com/.default"
+
+
+def _standard_payloads(count, filler_bytes=0):
+    return [
+        StandardLoggingPayload(
+            id=f"standard-{i}",
+            call_type="completion",
+            model="gpt-3.5-turbo",
+            status="success",
+            messages=[{"role": "user", "content": "x" * filler_bytes}],
+            response={"choices": [{"message": {"content": "Hi"}}]},
+        )
+        for i in range(count)
+    ]
+
+
+def _audit_payloads(count, filler_bytes=0):
+    return [
+        StandardAuditLogPayload(
+            id=f"audit-{i}",
+            updated_at="2026-05-06T04:39:00+00:00",
+            changed_by="user-1",
+            changed_by_api_key="sk-test",
+            action="created",
+            table_name="LiteLLM_TeamTable",
+            object_id="team-1",
+            before_value=None,
+            updated_values=json.dumps({"team_alias": "x" * filler_bytes}),
+        )
+        for i in range(count)
+    ]
+
+
+QUEUE_CASES = [
+    pytest.param("log_queue", "async_send_batch", _standard_payloads, id="standard"),
+    pytest.param("audit_log_queue", "async_send_audit_batch", _audit_payloads, id="audit"),
+]
+
+
+def _token_response():
+    response = MagicMock()
+    response.status_code = 200
+    response.json = MagicMock(return_value={"access_token": "test-bearer-token", "expires_in": 3600})
+    response.text = "Success"
+    return response
+
+
+def _install_ingestion(logger, on_ingest):
+    """Route the OAuth call to a canned token and every ingestion call to `on_ingest(body_bytes)`."""
+
+    async def _post(*args, **kwargs):
+        if "oauth2/v2.0/token" in kwargs.get("url", ""):
+            return _token_response()
+        return await on_ingest(kwargs["data"])
+
+    logger.async_httpx_client.post = AsyncMock(side_effect=_post)
+
+
+def _accepted():
+    return Response(204, request=Request("POST", "https://example.com"), text="")
+
+
+def _too_large(*, raised):
+    request = Request("POST", "https://example.com")
+    response = Response(413, request=request, text="Payload Too Large")
+    if raised:
+        raise MaskedHTTPStatusError(httpx.HTTPStatusError("413", request=request, response=response))
+    return response
+
+
+def _rejected(status_code, *, raised):
+    """litellm's http handler calls raise_for_status, so a real rejection arrives raised, not returned."""
+    request = Request("POST", "https://example.com")
+    response = Response(status_code, request=request, text=f"rejected with {status_code}")
+    if raised:
+        raise MaskedHTTPStatusError(httpx.HTTPStatusError(str(status_code), request=request, response=response))
+    return response
+
+
+def _awaiting_retry(logger, queue_attr):
+    return getattr(logger, "logs_awaiting_retry" if queue_attr == "log_queue" else "audit_logs_awaiting_retry")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_splits_a_batch_that_would_exceed_the_ingestion_cap(
+    queue_attr, send_method, build_payloads
+):
+    """Azure Monitor rejects a body over 1MB uncompressed, so an oversize batch has to be split
+    before it is sent instead of being posted whole and lost."""
+    logger = _build_logger()
+    records = build_payloads(4, filler_bytes=400_000)
+    setattr(logger, queue_attr, list(records))
+
+    sent_bodies = []
+
+    async def _on_ingest(data):
+        sent_bodies.append(data)
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert len(sent_bodies) > 1
+    assert all(len(body) <= AZURE_SENTINEL_MAX_PAYLOAD_SIZE_BYTES for body in sent_bodies)
+    delivered = [record["id"] for body in sent_bodies for record in json.loads(body.decode("utf-8"))]
+    assert delivered == [record["id"] for record in records]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raised", [True, False], ids=["raised", "returned"])
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_halves_the_batch_on_413(queue_attr, send_method, build_payloads, raised):
+    """A 413 the size estimate did not predict must halve the batch and retry, not drop it.
+
+    litellm's http handler raises MaskedHTTPStatusError on a 4xx, so the raised path is the one
+    a real Azure Monitor 413 takes, and both are covered here.
+    """
+    logger = _build_logger()
+    records = build_payloads(4)
+    setattr(logger, queue_attr, list(records))
+
+    delivered = []
+
+    async def _on_ingest(data):
+        body = json.loads(data.decode("utf-8"))
+        if len(body) > 1:
+            return _too_large(raised=raised)
+        delivered.extend(record["id"] for record in body)
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert delivered == [record["id"] for record in records]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_drops_only_the_lone_record_that_still_413s(queue_attr, send_method, build_payloads):
+    """One undeliverable record must not take its siblings down with it or wedge the queue."""
+    logger = _build_logger()
+    records = build_payloads(4)
+    poison = records[2]["id"]
+    setattr(logger, queue_attr, list(records))
+
+    delivered = []
+
+    async def _on_ingest(data):
+        body = json.loads(data.decode("utf-8"))
+        if any(record["id"] == poison for record in body):
+            return _too_large(raised=True)
+        delivered.extend(record["id"] for record in body)
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await asyncio.wait_for(getattr(logger, send_method)(), timeout=10)
+
+    assert delivered == [record["id"] for record in records if record["id"] != poison]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_requeues_only_what_a_transient_failure_left_undelivered(
+    queue_attr, send_method, build_payloads
+):
+    """Records Azure Monitor already accepted must not be sent twice, and the rest must survive
+    for the next flush instead of being cleared."""
+    logger = _build_logger()
+    records = build_payloads(4)
+    setattr(logger, queue_attr, list(records))
+
+    delivered = []
+
+    async def _on_ingest(data):
+        body = json.loads(data.decode("utf-8"))
+        if len(body) > 2:
+            return _too_large(raised=True)
+        if any(record["id"] == records[2]["id"] for record in body):
+            raise httpx.ConnectError("connection reset")
+        delivered.extend(record["id"] for record in body)
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert delivered == [records[0]["id"], records[1]["id"]]
+    assert getattr(logger, queue_attr) == records[2:]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_requeues_the_batch_on_a_non_success_status(queue_attr, send_method, build_payloads):
+    """A 500 from ingestion is retryable, so the batch has to stay queued."""
+    logger = _build_logger()
+    records = build_payloads(3)
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        return Response(500, request=Request("POST", "https://example.com"), text="Internal Server Error")
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == records
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_requeues_the_batch_when_the_oauth_token_call_fails(
+    queue_attr, send_method, build_payloads
+):
+    """Losing the token is transient, so the batch must not be dropped on the way to the wire."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    setattr(logger, queue_attr, list(records))
+
+    ingestion_calls = []
+
+    async def _post(*args, **kwargs):
+        if "oauth2/v2.0/token" in kwargs.get("url", ""):
+            failed = MagicMock()
+            failed.status_code = 401
+            failed.text = "Unauthorized"
+            return failed
+        ingestion_calls.append(kwargs["url"])
+        return _accepted()
+
+    logger.async_httpx_client.post = AsyncMock(side_effect=_post)
+
+    await getattr(logger, send_method)()
+
+    assert ingestion_calls == []
+    assert getattr(logger, queue_attr) == records
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_caps_the_retry_queue_at_max_queue_size(queue_attr, send_method, build_payloads):
+    """Retrying forever against an unreachable workspace must not grow the queue without bound,
+    so the oldest records go once the queue is over its limit."""
+    logger = _build_logger(max_queue_size=3)
+    records = build_payloads(4)
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        raise httpx.ConnectError("connection reset")
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == records[1:]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_keeps_records_queued_during_a_send(queue_attr, send_method, build_payloads):
+    """The queue is detached before sending, so a record logged mid-flush is kept and lands behind
+    anything the failed send hands back."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    late_record = build_payloads(1)[0]
+    late_record["id"] = "logged-during-send"
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        getattr(logger, queue_attr).append(late_record)
+        raise httpx.ConnectError("connection reset")
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == [*records, late_record]
+
+
+def _poison(record):
+    """A mixed-type set makes safe_dumps raise TypeError while sorting it, so the record can never be serialized."""
+    field = "messages" if "messages" in record else "updated_values"
+    record[field] = {1, "a"}
+    return record
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_drops_only_the_record_that_cannot_be_serialized(queue_attr, send_method, build_payloads):
+    """A record that raises during serialization used to escape the send, which killed the periodic
+    flush task for good and lost the already-detached batch with it. It has to be isolated and
+    dropped alone, with the flush completing normally."""
+    logger = _build_logger()
+    records = build_payloads(4)
+    poison = _poison(records[2])["id"]
+    setattr(logger, queue_attr, list(records))
+
+    delivered = []
+
+    async def _on_ingest(data):
+        delivered.extend(record["id"] for record in json.loads(data.decode("utf-8")))
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await asyncio.wait_for(logger.flush_queue(), timeout=10)
+
+    assert delivered == [record["id"] for record in records if record["id"] != poison]
+    assert getattr(logger, queue_attr) == []
+
+
+async def _log(logger, queue_attr, record):
+    if queue_attr == "log_queue":
+        await logger.async_log_success_event({"standard_logging_object": record}, None, None, None)
+        return
+    await logger.async_log_audit_log_event(record)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_retries_on_the_flush_timer_not_on_every_record_while_the_destination_is_down(
+    queue_attr, send_method, build_payloads
+):
+    """Requeued records keep the queue at or over batch_size, so without a guard every new record
+    re-sent the whole growing queue. While a retry is pending only the periodic flush may send, and
+    a successful flush hands the trigger back to the batch size."""
+    logger = _build_logger(batch_size=3)
+    records = build_payloads(11)
+
+    attempts = []
+    destination_down = True
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        if destination_down:
+            raise httpx.ConnectError("connection reset")
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    for record in records[:8]:
+        await _log(logger, queue_attr, record)
+
+    assert attempts == [[record["id"] for record in records[:3]]]
+    assert getattr(logger, queue_attr) == records[:8]
+
+    destination_down = False
+    await logger.flush_queue()
+    for record in records[8:]:
+        await _log(logger, queue_attr, record)
+
+    assert [record_id for attempt in attempts[1:-1] for record_id in attempt] == [record["id"] for record in records[:8]]
+    assert all(len(attempt) <= 3 for attempt in attempts[1:-1])
+    assert attempts[-1] == [record["id"] for record in records[8:]]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_threshold_send_waits_for_an_in_flight_timer_flush(
+    queue_attr, send_method, build_payloads
+):
+    """A batch-size send that overlapped the periodic flush could finish after it and requeue its
+    newer records in front of the older ones, so the max_queue_size trim would then drop the
+    newest records instead of the oldest. Both paths have to take the flush lock, and a waiter
+    that gets the lock after a failed flush stands down instead of resending the whole queue."""
+    logger = _build_logger(batch_size=2)
+    records = build_payloads(4)
+    setattr(logger, queue_attr, list(records[:2]))
+
+    attempts = []
+    timer_send_started = asyncio.Event()
+    release_timer_send = asyncio.Event()
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        if len(attempts) == 1:
+            timer_send_started.set()
+            await release_timer_send.wait()
+        raise httpx.ConnectError("connection reset")
+
+    _install_ingestion(logger, _on_ingest)
+
+    timer_flush = asyncio.create_task(logger.flush_queue())
+    await asyncio.wait_for(timer_send_started.wait(), timeout=10)
+    await _log(logger, queue_attr, records[2])
+    threshold_send = asyncio.create_task(_log(logger, queue_attr, records[3]))
+    await asyncio.sleep(0)
+    release_timer_send.set()
+    await asyncio.wait_for(timer_flush, timeout=10)
+    await asyncio.wait_for(threshold_send, timeout=10)
+
+    assert attempts == [[record["id"] for record in records[:2]]]
+    assert getattr(logger, queue_attr) == records
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_concurrent_threshold_sends_collapse_into_one_attempt_while_the_destination_is_down(
+    queue_attr, send_method, build_payloads
+):
+    """Records logged while a threshold send is blocked on the wire all see the retry flag still
+    unset and queue up on the flush lock. Each waiter has to recheck under the lock, or every one
+    of them resends the growing queue as soon as the first attempt fails."""
+    logger = _build_logger(batch_size=2)
+    records = build_payloads(6)
+
+    attempts = []
+    first_send_started = asyncio.Event()
+    release_first_send = asyncio.Event()
+    destination_down = True
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        if len(attempts) == 1:
+            first_send_started.set()
+            await release_first_send.wait()
+        if destination_down:
+            raise httpx.ConnectError("connection reset")
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await _log(logger, queue_attr, records[0])
+    first_send = asyncio.create_task(_log(logger, queue_attr, records[1]))
+    await asyncio.wait_for(first_send_started.wait(), timeout=10)
+    waiters = [asyncio.create_task(_log(logger, queue_attr, record)) for record in records[2:]]
+    await asyncio.sleep(0)
+    release_first_send.set()
+    await asyncio.wait_for(asyncio.gather(first_send, *waiters), timeout=10)
+
+    assert attempts == [[record["id"] for record in records[:2]]]
+    assert getattr(logger, queue_attr) == records
+
+    destination_down = False
+    await logger.flush_queue()
+
+    assert [record_id for attempt in attempts[1:] for record_id in attempt] == [record["id"] for record in records]
+    assert all(len(attempt) <= 2 for attempt in attempts[1:])
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_requeues_a_cancelled_send(
+    queue_attr, send_method, build_payloads
+):
+    """Cancellation after detaching a batch must preserve the detached records for a later flush."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        raise asyncio.CancelledError
+
+    _install_ingestion(logger, _on_ingest)
+
+    with pytest.raises(asyncio.CancelledError) as excinfo:
+        await getattr(logger, send_method)()
+
+    assert type(excinfo.value) is asyncio.CancelledError
+    assert getattr(logger, queue_attr) == records
+    assert _awaiting_retry(logger, queue_attr)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_requeues_a_send_cancelled_before_it_reached_the_wire(
+    queue_attr, send_method, build_payloads
+):
+    """Cancellation can land on the token call, before any record was sent, and the detached batch
+    has to survive that too."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    setattr(logger, queue_attr, list(records))
+    logger.async_httpx_client.post = AsyncMock(side_effect=asyncio.CancelledError)
+
+    with pytest.raises(asyncio.CancelledError) as excinfo:
+        await getattr(logger, send_method)()
+
+    assert type(excinfo.value) is asyncio.CancelledError
+    assert getattr(logger, queue_attr) == records
+    assert _awaiting_retry(logger, queue_attr)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_does_not_resend_the_half_delivered_before_a_cancelled_split(
+    queue_attr, send_method, build_payloads
+):
+    """A batch over the size cap goes out in pieces, so a cancellation partway through must requeue
+    only the pieces the destination never accepted, or the accepted ones land in Sentinel twice."""
+    logger = _build_logger()
+    records = build_payloads(8, filler_bytes=400_000)
+    setattr(logger, queue_attr, list(records))
+
+    attempts = []
+    cancel_after_the_first_piece = True
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        if cancel_after_the_first_piece and len(attempts) > 1:
+            raise asyncio.CancelledError
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    with pytest.raises(asyncio.CancelledError) as excinfo:
+        await getattr(logger, send_method)()
+
+    assert type(excinfo.value) is asyncio.CancelledError
+    assert attempts == [[record["id"] for record in records[:2]], [record["id"] for record in records[2:4]]]
+    assert getattr(logger, queue_attr) == records[2:]
+
+    cancel_after_the_first_piece = False
+    await logger.flush_queue()
+
+    assert [record_id for attempt in attempts[2:] for record_id in attempt] == [record["id"] for record in records[2:]]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_threshold_waiter_does_not_send_a_sub_batch_after_success(
+    queue_attr, send_method, build_payloads
+):
+    """A successful threshold send can leave one record behind, so a waiter must not send it
+    before the next record completes a batch."""
+    logger = _build_logger(batch_size=2)
+    records = build_payloads(3)
+
+    attempts = []
+    first_send_started = asyncio.Event()
+    release_first_send = asyncio.Event()
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        if len(attempts) == 1:
+            first_send_started.set()
+            await release_first_send.wait()
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await _log(logger, queue_attr, records[0])
+    first_send = asyncio.create_task(_log(logger, queue_attr, records[1]))
+    await asyncio.wait_for(first_send_started.wait(), timeout=10)
+    waiter = asyncio.create_task(_log(logger, queue_attr, records[2]))
+    await asyncio.sleep(0)
+    release_first_send.set()
+    await asyncio.wait_for(asyncio.gather(first_send, waiter), timeout=10)
+
+    assert attempts == [[record["id"] for record in records[:2]]]
+    assert getattr(logger, queue_attr) == [records[2]]
+
+
+@pytest.mark.asyncio
+async def test_azure_sentinel_threshold_send_only_sends_the_queue_that_crossed_the_threshold():
+    """The standard and audit queues retry independently: crossing the audit threshold must not
+    resend standard records that are waiting for the periodic flush."""
+    logger = _build_logger(batch_size=2)
+    standard_records = _standard_payloads(2)
+    audit_records = _audit_payloads(2)
+    logger.log_queue = list(standard_records)
+    logger.logs_awaiting_retry = True
+
+    attempts = []
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    for record in audit_records:
+        await logger.async_log_audit_log_event(record)
+
+    assert attempts == [[record["id"] for record in audit_records]]
+    assert logger.audit_log_queue == []
+    assert logger.log_queue == standard_records
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("status_code", [408, 429, 500, 503])
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_keeps_the_batch_when_ingestion_raises_a_retryable_status(
+    queue_attr, send_method, build_payloads, status_code
+):
+    """A 5xx, a timeout or a throttle can clear on the next flush, so the whole batch stays queued
+    and the awaiting-retry flag hands the send back to the timer."""
+    logger = _build_logger()
+    records = build_payloads(3)
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        return _rejected(status_code, raised=True)
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == records
+    assert _awaiting_retry(logger, queue_attr)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("raised", [True, False], ids=["raised", "returned"])
+@pytest.mark.parametrize("status_code", [400, 403, 404])
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_drops_the_batch_when_ingestion_rejects_it_for_good(
+    queue_attr, send_method, build_payloads, status_code, raised
+):
+    """A permanent 4xx is dropped, the flag is cleared and the next records go out on their own."""
+    logger = _build_logger(batch_size=2)
+    rejected_records = build_payloads(2)
+    later_records = build_payloads(4)[2:]
+    setattr(logger, queue_attr, list(rejected_records))
+
+    delivered = []
+    destination_rejects = True
+
+    async def _on_ingest(data):
+        if destination_rejects:
+            return _rejected(status_code, raised=raised)
+        delivered.extend(record["id"] for record in json.loads(data.decode("utf-8")))
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == []
+    assert not _awaiting_retry(logger, queue_attr)
+
+    destination_rejects = False
+    for record in later_records:
+        await _log(logger, queue_attr, record)
+
+    assert delivered == [record["id"] for record in later_records]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_keeps_the_whole_batch_when_the_first_piece_of_a_split_fails(
+    queue_attr, send_method, build_payloads
+):
+    """When the first half of a split hits a retryable error the untried second half must be kept
+    too, in the original order, instead of being sent ahead of records that are still pending."""
+    logger = _build_logger()
+    records = build_payloads(4)
+    setattr(logger, queue_attr, list(records))
+
+    attempts = []
+
+    async def _on_ingest(data):
+        body = json.loads(data.decode("utf-8"))
+        attempts.append([record["id"] for record in body])
+        if len(body) > 2:
+            return _too_large(raised=True)
+        return _rejected(503, raised=True)
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert attempts == [[record["id"] for record in records], [record["id"] for record in records[:2]]]
+    assert getattr(logger, queue_attr) == records
+    assert _awaiting_retry(logger, queue_attr)
+
+
+class _RaisesWhileDumping(BaseModel):
+    @computed_field
+    @property
+    def rendered(self) -> str:
+        raise RuntimeError("this field cannot be rendered")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_drops_only_the_record_whose_serialization_raises_an_unexpected_error(
+    queue_attr, send_method, build_payloads
+):
+    """Serialization can fail with any exception class, not just TypeError or ValueError, because
+    safe_dumps hands pydantic models to model_dump. A record that raises anything has to be isolated
+    and dropped alone, or the flush dies with the whole batch."""
+    logger = _build_logger()
+    records = build_payloads(4)
+    poison = records[1]
+    poison["messages" if "messages" in poison else "updated_values"] = _RaisesWhileDumping()
+    setattr(logger, queue_attr, list(records))
+
+    delivered = []
+
+    async def _on_ingest(data):
+        delivered.extend(record["id"] for record in json.loads(data.decode("utf-8")))
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await asyncio.wait_for(logger.flush_queue(), timeout=10)
+
+    assert delivered == [record["id"] for record in records if record is not poison]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_send_cancelled_by_a_timeout_surfaces_as_a_timeout(
+    queue_attr, send_method, build_payloads
+):
+    """The logging worker bounds each flush with asyncio.wait_for, which on Python 3.12 only turns
+    an exact CancelledError into TimeoutError. A subclass carrying the undelivered records would
+    escape the worker as an unhandled error, so the send must re-raise the plain class."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    setattr(logger, queue_attr, list(records))
+
+    async def _on_ingest(data):
+        await asyncio.sleep(60)
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(getattr(logger, send_method)(), timeout=0.05)
+
+    assert getattr(logger, queue_attr) == records
+    assert _awaiting_retry(logger, queue_attr)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_never_sends_more_than_batch_size_records_in_one_request(
+    queue_attr, send_method, build_payloads
+):
+    """A recovery flush can find far more than batch_size records queued. Splitting on the count
+    first keeps each request at the configured size and bounds how much of the queue is serialized
+    just to measure it."""
+    logger = _build_logger(batch_size=2)
+    records = build_payloads(5)
+    setattr(logger, queue_attr, list(records))
+
+    attempts = []
+
+    async def _on_ingest(data):
+        attempts.append([record["id"] for record in json.loads(data.decode("utf-8"))])
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert attempts == [
+        [records[0]["id"], records[1]["id"]],
+        [records[2]["id"]],
+        [records[3]["id"], records[4]["id"]],
+    ]
+    assert getattr(logger, queue_attr) == []
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "status_code, expected_queue",
+    [pytest.param(503, "kept", id="503-kept"), pytest.param(401, "dropped", id="401-dropped")],
+)
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_oauth_rejection_follows_the_same_retry_rule_as_ingestion(
+    queue_attr, send_method, build_payloads, status_code, expected_queue
+):
+    """The token endpoint raises through the same http handler as ingestion. A 5xx there is
+    transient and keeps the batch, a 401 means the client secret is wrong and would fail every
+    retry, so the batch is dropped instead of wedging the queue."""
+    logger = _build_logger()
+    records = build_payloads(2)
+    setattr(logger, queue_attr, list(records))
+
+    ingestion_calls = []
+
+    async def _post(*args, **kwargs):
+        if "oauth2/v2.0/token" in kwargs.get("url", ""):
+            return _rejected(status_code, raised=True)
+        ingestion_calls.append(kwargs["url"])
+        return _accepted()
+
+    logger.async_httpx_client.post = AsyncMock(side_effect=_post)
+
+    await getattr(logger, send_method)()
+
+    assert ingestion_calls == []
+    assert getattr(logger, queue_attr) == (records if expected_queue == "kept" else [])
+    assert _awaiting_retry(logger, queue_attr) is (expected_queue == "kept")
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("queue_attr, send_method, build_payloads", QUEUE_CASES)
+async def test_azure_sentinel_does_not_stay_in_retry_mode_when_the_queue_cap_trims_everything(
+    queue_attr, send_method, build_payloads
+):
+    """With max_queue_size at 0 the cap drops every requeued record, so there is nothing for the
+    timer to retry. The flag must follow the retained queue, or every later threshold send is
+    skipped until the timer happens to fire."""
+    logger = _build_logger(batch_size=2, max_queue_size=0)
+    lost_records = build_payloads(2)
+    later_records = build_payloads(4)[2:]
+    setattr(logger, queue_attr, list(lost_records))
+
+    delivered = []
+    destination_down = True
+
+    async def _on_ingest(data):
+        if destination_down:
+            raise httpx.ConnectError("connection reset")
+        delivered.extend(record["id"] for record in json.loads(data.decode("utf-8")))
+        return _accepted()
+
+    _install_ingestion(logger, _on_ingest)
+
+    await getattr(logger, send_method)()
+
+    assert getattr(logger, queue_attr) == []
+    assert not _awaiting_retry(logger, queue_attr)
+
+    destination_down = False
+    for record in later_records:
+        await _log(logger, queue_attr, record)
+
+    assert delivered == [record["id"] for record in later_records]

tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py

@@ -6,6 +6,7 @@ import pytest
 import litellm
 
 
+
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     TOOL_RESULT_IMAGE_PLACEHOLDER,
 )
@@ -52,7 +53,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
                 tool_calls=[
                     ChatCompletionDeltaToolCall(
                         id="call_d581d130-e234-4315-94e8-27e7ff7c4e55",
-                        function=Function(arguments='{"location": "Boston"}', name="get_weather"),
+                        function=Function(
+                            arguments='{"location": "Boston"}', name="get_weather"
+                        ),
                         type="function",
                         index=0,
                     )
@@ -66,7 +69,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block():
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     print(content_block_start)
 
@@ -96,7 +101,9 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
                 tool_calls=[
                     ChatCompletionDeltaToolCall(
                         id=combined,
-                        function=Function(arguments='{"a": 17, "b": 25}', name="add_numbers"),
+                        function=Function(
+                            arguments='{"a": 17, "b": 25}', name="add_numbers"
+                        ),
                         type="function",
                         index=0,
                     )
@@ -110,7 +117,9 @@ def test_translate_streaming_openai_chunk_strips_gemini_thought_from_tool_call_i
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "tool_use"
     assert content_block_start["id"] == base
@@ -155,7 +164,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_content_block():
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "thinking"
     assert content_block_start == {
@@ -191,7 +202,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_only_co
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "thinking"
     assert content_block_start == {
@@ -237,7 +250,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_signature_block(
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "thinking"
     assert content_block_start == {
@@ -290,7 +305,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_content_block_thinking_an
     (
         block_type,
         content_block_start,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "thinking"
 
@@ -333,7 +350,10 @@ def test_translate_anthropic_messages_to_openai_thinking_blocks():
     assert "thinking_blocks" in result[1]
     assert len(result[1]["thinking_blocks"]) == 2
     assert result[1]["thinking_blocks"][0]["type"] == "thinking"
-    assert result[1]["thinking_blocks"][0]["thinking"] == "I will call the get_weather tool."
+    assert (
+        result[1]["thinking_blocks"][0]["thinking"]
+        == "I will call the get_weather tool."
+    )
     assert result[1]["thinking_blocks"][0]["signature"] == "sigsig"
     assert result[1]["thinking_blocks"][1]["type"] == "redacted_thinking"
     assert result[1]["thinking_blocks"][1]["data"] == "REDACTED"
@@ -436,7 +456,9 @@ def test_translate_anthropic_messages_to_openai_tool_message_placement():
 
     assert tool_message_idx is not None, "Tool message not found"
     assert user_message_idx is not None, "User message not found"
-    assert tool_message_idx < user_message_idx, "Tool message should be placed before user message"
+    assert (
+        tool_message_idx < user_message_idx
+    ), "Tool message should be placed before user message"
 
 
 @pytest.mark.parametrize(
@@ -711,7 +733,9 @@ def test_translate_anthropic_to_openai_skips_prompt_cache_key_when_provider_lack
 
 
 def test_translate_anthropic_to_openai_skips_prompt_cache_key_for_chained_litellm_proxy():
-    assert "prompt_cache_key" in litellm.get_supported_openai_params(model="xai", custom_llm_provider="litellm_proxy")
+    assert "prompt_cache_key" in litellm.get_supported_openai_params(
+        model="xai", custom_llm_provider="litellm_proxy"
+    )
     openai_request = _translate_with_metadata("litellm_proxy/xai", {"user_id": "session-abc"}, "litellm_proxy")
     assert openai_request["user"] == "session-abc"
     assert "prompt_cache_key" not in openai_request
@@ -756,8 +780,7 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
                         id="call_empty_args",
                         type="function",
                         function=Function(
-                            name="test_function",
-                            arguments="",  # empty arguments string
+                            name="test_function", arguments=""  # empty arguments string
                         ),
                     )
                 ],
@@ -772,7 +795,10 @@ def test_translate_openai_content_to_anthropic_empty_function_arguments():
     assert result[0]["type"] == "tool_use"
     assert result[0]["id"] == "call_empty_args"
     assert result[0]["name"] == "test_function"
-    assert result[0]["input"] == {}, "Empty function arguments should result in empty dict"
+    assert (
+        result[0]["input"] == {}
+    ), "Empty function arguments should result in empty dict"
+    assert "provider_specific_fields" not in result[0]
 
 
 def test_translate_openai_content_to_anthropic_text_and_tool_calls():
@@ -818,6 +844,11 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c
     base = "call_3e9417b7925e49aca9a71dc1885e"
     sig = "CiIBDDnWx+/a=="
     combined = f"{base}{THOUGHT_SIGNATURE_SEPARATOR}{sig}"
+    function = Function(
+        name="get_weather",
+        arguments='{"location": "Boston"}',
+    )
+    function.provider_specific_fields = {"thought_signature": sig}
     openai_choices = [
         Choices(
             message=Message(
@@ -827,10 +858,7 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c
                     ChatCompletionAssistantToolCall(
                         id=combined,
                         type="function",
-                        function=Function(
-                            name="get_weather",
-                            arguments='{"location": "Boston"}',
-                        ),
+                        function=function,
                     )
                 ],
             )
@@ -846,6 +874,7 @@ def test_translate_openai_content_to_anthropic_strips_gemini_thought_from_tool_c
     assert THOUGHT_SIGNATURE_SEPARATOR not in result[0]["id"]
     assert result[0]["name"] == "get_weather"
     assert result[0]["input"] == {"location": "Boston"}
+    assert result[0]["provider_specific_fields"] == {"signature": sig}
 
 
 def test_translate_openai_content_to_anthropic_sanitizes_colon_dot_tool_call_ids():
@@ -892,7 +921,9 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
                         ChatCompletionAssistantToolCall(
                             id="call_tool_combo",
                             type="function",
-                            function=Function(name="get_weather", arguments='{"location": "Paris"}'),
+                            function=Function(
+                                name="get_weather", arguments='{"location": "Paris"}'
+                            ),
                         )
                     ],
                 ),
@@ -902,7 +933,9 @@ def test_translate_openai_response_to_anthropic_text_and_tool_calls():
     )
 
     adapter = LiteLLMAnthropicMessagesAdapter()
-    anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
+    anthropic_response = adapter.translate_openai_response_to_anthropic(
+        response=openai_response
+    )
 
     anthropic_content = anthropic_response.get("content")
     assert anthropic_content is not None
@@ -943,7 +976,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_partial_json():
     (
         type_of_content,
         content_block_delta,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
+        choices=choices
+    )
 
     print("Type of content:", type_of_content)
     print("Content block delta:", content_block_delta)
@@ -1052,7 +1087,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_thinking_delta():
     (
         type_of_content,
         content_block_delta,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
+        choices=choices
+    )
 
     assert type_of_content == "thinking_delta"
     assert content_block_delta["type"] == "thinking_delta"
@@ -1095,7 +1132,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_with_thinking():
     (
         type_of_content,
         content_block_delta,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
+        choices=choices
+    )
 
     assert type_of_content == "signature_delta"
     assert content_block_delta["type"] == "signature_delta"
@@ -1159,7 +1198,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_emits_signature_when_thin
     (
         block_type,
         content_block_start,
-    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "thinking"
 
@@ -1199,7 +1240,9 @@ def test_translate_anthropic_messages_to_openai_user_message_with_base64_image()
     # Check image content
     assert result[0]["content"][1]["type"] == "image_url"
     assert "image_url" in result[0]["content"][1]
-    assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
+    assert result[0]["content"][1]["image_url"]["url"].startswith(
+        "data:image/png;base64,"
+    )
     assert (
         "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
         in result[0]["content"][1]["image_url"]["url"]
@@ -1237,14 +1280,18 @@ def test_translate_anthropic_messages_to_openai_user_message_with_url_image():
     # Check image content
     assert result[0]["content"][1]["type"] == "image_url"
     assert "image_url" in result[0]["content"][1]
-    assert result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
+    assert (
+        result[0]["content"][1]["image_url"]["url"] == "https://example.com/forest.jpg"
+    )
 
 
 def test_translate_anthropic_messages_to_openai_tool_result_with_base64_image():
     """Test that base64 images in tool results are correctly translated to OpenAI format."""
 
     anthropic_messages = [
-        AnthropicMessagesUserMessageParam(role="user", content=[{"type": "text", "text": "Take a screenshot"}]),
+        AnthropicMessagesUserMessageParam(
+            role="user", content=[{"type": "text", "text": "Take a screenshot"}]
+        ),
         AnthopicMessagesAssistantMessageParam(
             role="assistant",
             content=[
@@ -1396,7 +1443,9 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
 
     # Check first image (base64)
     assert result[0]["content"][1]["type"] == "image_url"
-    assert result[0]["content"][1]["image_url"]["url"].startswith("data:image/png;base64,")
+    assert result[0]["content"][1]["image_url"]["url"].startswith(
+        "data:image/png;base64,"
+    )
 
     # Check middle text
     assert result[0]["content"][2]["type"] == "text"
@@ -1404,7 +1453,9 @@ def test_translate_anthropic_messages_to_openai_mixed_content_with_image():
 
     # Check second image (URL)
     assert result[0]["content"][3]["type"] == "image_url"
-    assert result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
+    assert (
+        result[0]["content"][3]["image_url"]["url"] == "https://example.com/image2.jpg"
+    )
 
     # Check final text
     assert result[0]["content"][4]["type"] == "text"
@@ -1450,7 +1501,10 @@ def test_translate_anthropic_messages_to_openai_tool_use_with_signature():
     assert tool_call["id"] == "call_386f67af31f9415781bc35071405"
     assert "function" in tool_call
     assert "provider_specific_fields" in tool_call["function"]
-    assert tool_call["function"]["provider_specific_fields"]["thought_signature"] == test_signature
+    assert (
+        tool_call["function"]["provider_specific_fields"]["thought_signature"]
+        == test_signature
+    )
 
 
 def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_content_items():
@@ -1508,7 +1562,9 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
     result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
 
     # Count how many tool messages have the same tool_call_id
-    tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
+    tool_messages = [
+        msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
+    ]
     tool_call_ids = [msg.get("tool_call_id") for msg in tool_messages]
 
     # The critical assertion: each tool_call_id should appear only ONCE
@@ -1524,8 +1580,12 @@ def test_translate_anthropic_messages_to_openai_tool_result_with_multiple_conten
     # The content should be a list with all items combined
     tool_message = tool_messages[0]
     assert tool_message["tool_call_id"] == "toolu_016hYHBkTf4JDF3p22UoYk5C"
-    assert isinstance(tool_message["content"], list), "Multiple content items should be combined into a list"
-    assert len(tool_message["content"]) == 3, f"Expected 3 content items, got {len(tool_message['content'])}"
+    assert isinstance(
+        tool_message["content"], list
+    ), "Multiple content items should be combined into a list"
+    assert (
+        len(tool_message["content"]) == 3
+    ), f"Expected 3 content items, got {len(tool_message['content'])}"
 
     # Verify content types
     assert tool_message["content"][0]["type"] == "text"
@@ -1574,14 +1634,17 @@ def test_translate_anthropic_messages_to_openai_tool_result_single_item_backward
     adapter = LiteLLMAnthropicMessagesAdapter()
     result = adapter.translate_anthropic_messages_to_openai(messages=anthropic_messages)
 
-    tool_messages = [msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"]
+    tool_messages = [
+        msg for msg in result if isinstance(msg, dict) and msg.get("role") == "tool"
+    ]
 
     assert len(tool_messages) == 1
     tool_message = tool_messages[0]
 
     # Single item should be a string for backward compatibility
     assert isinstance(tool_message["content"], str), (
-        f"Single content item should be a string for backward compatibility, got {type(tool_message['content'])}"
+        f"Single content item should be a string for backward compatibility, "
+        f"got {type(tool_message['content'])}"
     )
     assert tool_message["content"] == "72°F and sunny"
 
@@ -1630,7 +1693,9 @@ def test_streaming_chunk_with_both_text_and_tool_calls_issue_18238():
     (
         block_type,
         content_block_start,
-    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "tool_use"
     assert content_block_start["name"] == "Bash"
@@ -1674,7 +1739,9 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
     (
         block_type,
         content_block_start,
-    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(choices=choices)
+    ) = adapter._translate_streaming_openai_chunk_to_anthropic_content_block(
+        choices=choices
+    )
 
     assert block_type == "text"
     assert content_block_start == {"type": "text", "text": ""}
@@ -1685,12 +1752,15 @@ def test_streaming_chunk_with_text_and_empty_tool_calls_returns_text_delta():
 # ============================================================================
 
 # Model constant for cache control tests
-CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
+CACHE_CONTROL_BEDROCK_CONVERSE_MODEL = (
+    "bedrock/converse/global.anthropic.claude-opus-4-5-20251101-v1:0"
+)
 CACHE_CONTROL_NON_ANTHROPIC_MODEL = "gpt-4"
 # Bedrock Application Inference Profile ARN: the string contains neither
 # "anthropic" nor "claude", so the model can only be recognized via its ARN shape
 CACHE_CONTROL_BEDROCK_ARN_MODEL = (
-    "bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456"
+    "bedrock/converse/arn:aws:bedrock:us-east-1:123456789012:"
+    "application-inference-profile/abcdef123456"
 )
 
 
@@ -1706,7 +1776,9 @@ def test_should_add_cache_control_for_anthropic_model():
         "vertex_ai/claude-3-sonnet@20240229",
     ]:
         target = {}
-        adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
+        adapter._add_cache_control_if_applicable(
+            {"cache_control": cache_control}, target, model
+        )
         assert "cache_control" in target
         assert target["cache_control"] == cache_control
 
@@ -1722,7 +1794,9 @@ def test_should_not_add_cache_control_for_non_anthropic_model():
         "gemini-pro",
     ]:
         target = {}
-        adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
+        adapter._add_cache_control_if_applicable(
+            {"cache_control": cache_control}, target, model
+        )
         assert "cache_control" not in target
 
 
@@ -1737,7 +1811,9 @@ def test_should_not_add_cache_control_when_none():
         {},
     ]:
         target = {}
-        adapter._add_cache_control_if_applicable(source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL)
+        adapter._add_cache_control_if_applicable(
+            source, target, CACHE_CONTROL_BEDROCK_CONVERSE_MODEL
+        )
         assert "cache_control" not in target
 
 
@@ -1748,7 +1824,9 @@ def test_should_not_add_cache_control_when_model_none():
 
     for model in [None, ""]:
         target = {}
-        adapter._add_cache_control_if_applicable({"cache_control": cache_control}, target, model)
+        adapter._add_cache_control_if_applicable(
+            {"cache_control": cache_control}, target, model
+        )
         assert "cache_control" not in target
 
 
@@ -1854,7 +1932,12 @@ def test_cache_control_fix_does_not_broaden_claude_detection():
     make is_anthropic_claude_model treat ARN profiles as Claude, which would route
     thinking params through unmodified and break non-Claude Bedrock profiles.
     """
-    assert LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(CACHE_CONTROL_BEDROCK_ARN_MODEL) is False
+    assert (
+        LiteLLMAnthropicMessagesAdapter.is_anthropic_claude_model(
+            CACHE_CONTROL_BEDROCK_ARN_MODEL
+        )
+        is False
+    )
 
 
 def test_thinking_preserved_for_bedrock_arn_inference_profile():
@@ -2452,7 +2535,9 @@ def test_translate_streaming_openai_chunk_to_anthropic_reasoning_content_without
     (
         type_of_content,
         content_block_delta,
-    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(choices=choices)
+    ) = LiteLLMAnthropicMessagesAdapter()._translate_streaming_openai_chunk_to_anthropic(
+        choices=choices
+    )
 
     assert type_of_content == "thinking_delta"
     assert content_block_delta["type"] == "thinking_delta"
@@ -2484,7 +2569,9 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
     )
 
     adapter = LiteLLMAnthropicMessagesAdapter()
-    anthropic_response = adapter.translate_openai_response_to_anthropic(response=openai_response)
+    anthropic_response = adapter.translate_openai_response_to_anthropic(
+        response=openai_response
+    )
 
     anthropic_content = anthropic_response.get("content")
     assert anthropic_content is not None
@@ -2497,7 +2584,9 @@ def test_translate_openai_response_to_anthropic_with_reasoning_content_only():
 
     # Second block should be text
     assert anthropic_content[1]["type"] == "text"
-    assert anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
+    assert (
+        anthropic_content[1]["text"] == 'There are **3** "r"s in the word strawberry.'
+    )
 
     assert anthropic_response.get("stop_reason") == "end_turn"
 
@@ -2549,7 +2638,9 @@ def test_truncate_tool_name_deterministic():
 def test_truncate_tool_name_avoids_collisions():
     """Similar long names should produce different truncated names."""
     name1 = "process_user_data_with_validation_and_error_handling_for_production_environment"
-    name2 = "process_user_data_with_validation_and_error_handling_for_staging_environment"
+    name2 = (
+        "process_user_data_with_validation_and_error_handling_for_staging_environment"
+    )
 
     result1 = truncate_tool_name(name1)
     result2 = truncate_tool_name(name2)
@@ -2569,7 +2660,9 @@ def test_create_tool_name_mapping_no_long_names():
 
 def test_create_tool_name_mapping_with_long_names():
     """Mapping should contain entries for truncated names."""
-    long_name = "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
+    long_name = (
+        "a_very_long_tool_name_that_exceeds_the_64_character_limit_imposed_by_openai"
+    )
     tools = [
         {"name": "short_name"},
         {"name": long_name},
@@ -2594,7 +2687,9 @@ def test_translate_anthropic_tools_with_long_names():
     ]
 
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
+    result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
+        tools=tools, model="gpt-4"
+    )
 
     assert len(result) == 1
     # The tool name should be truncated
@@ -2616,7 +2711,9 @@ def test_translate_anthropic_tools_mixed_names():
     ]
 
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(tools=tools, model="gpt-4")
+    result, tool_name_mapping = adapter.translate_anthropic_tools_to_openai(
+        tools=tools, model="gpt-4"
+    )
 
     assert len(result) == 2
     # Short name unchanged
@@ -2630,7 +2727,9 @@ def test_translate_anthropic_tools_mixed_names():
 
 def test_translate_openai_response_restores_tool_names():
     """Tool names in responses should be restored to original."""
-    original_name = "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
+    original_name = (
+        "a_very_long_tool_name_that_needs_truncation_for_openai_api_compatibility"
+    )
     truncated_name = truncate_tool_name(original_name)
     tool_name_mapping = {truncated_name: original_name}
 
@@ -2662,7 +2761,9 @@ def test_translate_openai_response_restores_tool_names():
     )
 
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result = adapter.translate_openai_response_to_anthropic(response=response, tool_name_mapping=tool_name_mapping)
+    result = adapter.translate_openai_response_to_anthropic(
+        response=response, tool_name_mapping=tool_name_mapping
+    )
 
     # Find the tool_use block in the response
     tool_use_blocks = [c for c in result["content"] if c.get("type") == "tool_use"]
@@ -2828,7 +2929,9 @@ def test_translate_openai_usage_to_anthropic_cache_tokens_from_dict_details_with
         "cache_write_tokens": 20.0,
     }
 
-    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
 
     assert anthropic_usage["input_tokens"] == 70
     assert anthropic_usage["output_tokens"] == 50
@@ -2847,7 +2950,9 @@ def test_translate_openai_usage_to_anthropic_ignores_fractional_cache_tokens():
         "cache_creation_tokens": 20.25,
     }
 
-    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
 
     assert anthropic_usage["input_tokens"] == 120
     assert anthropic_usage["output_tokens"] == 50
@@ -2864,7 +2969,9 @@ def test_translate_openai_usage_to_anthropic_ignores_bool_cache_tokens():
     usage.cache_read_input_tokens = True
     usage.cache_creation_input_tokens = True
 
-    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(usage)
+    anthropic_usage = LiteLLMAnthropicMessagesAdapter._translate_openai_usage_to_anthropic_usage_delta(
+        usage
+    )
 
     assert anthropic_usage["input_tokens"] == 120
     assert anthropic_usage["output_tokens"] == 50
@@ -3083,7 +3190,9 @@ def test_translate_streaming_openai_response_to_anthropic_cache_tokens_with_appl
     assert message_delta["usage"]["output_tokens"] == 50
     assert message_delta["usage"]["cache_read_input_tokens"] == 30
     assert message_delta["usage"]["cache_creation_input_tokens"] == 20
-    assert message_delta["context_management"]["applied_edits"][0]["type"] == ("compact_20260112")
+    assert message_delta["context_management"]["applied_edits"][0]["type"] == (
+        "compact_20260112"
+    )
 
 
 # =====================================================================
@@ -3258,8 +3367,15 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
         assert schema["required"] == ["user"]
         assert schema["properties"]["user"]["additionalProperties"] is False
         assert schema["properties"]["user"]["required"] == ["name", "address"]
-        assert schema["properties"]["user"]["properties"]["address"]["additionalProperties"] is False
-        assert schema["properties"]["user"]["properties"]["address"]["required"] == ["city"]
+        assert (
+            schema["properties"]["user"]["properties"]["address"][
+                "additionalProperties"
+            ]
+            is False
+        )
+        assert schema["properties"]["user"]["properties"]["address"]["required"] == [
+            "city"
+        ]
 
     def test_array_items_object_adds_additional_properties_false(self):
         output_format = {
@@ -3334,9 +3450,19 @@ class TestTranslateAnthropicOutputFormatToOpenAI:
         assert sorted(schema["required"]) == ["age", "email", "name"]
 
     def test_invalid_output_format_returns_none(self):
-        assert self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
-        assert self.adapter.translate_anthropic_output_format_to_openai({"type": "text"}) is None
-        assert self.adapter.translate_anthropic_output_format_to_openai({"type": "json_schema"}) is None
+        assert (
+            self.adapter.translate_anthropic_output_format_to_openai("invalid") is None
+        )
+        assert (
+            self.adapter.translate_anthropic_output_format_to_openai({"type": "text"})
+            is None
+        )
+        assert (
+            self.adapter.translate_anthropic_output_format_to_openai(
+                {"type": "json_schema"}
+            )
+            is None
+        )
 
 
 class TestAnthropicStreamWrapperToolArgs:
@@ -3540,7 +3666,9 @@ def test_translate_openai_response_to_anthropic_with_polyfill_compaction_block()
     )
     response = _make_simple_openai_response(text="Hello after compaction.")
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
+    result = adapter.translate_openai_response_to_anthropic(
+        response=response, polyfill_result=polyfill
+    )
 
     content = result.get("content")
     assert content is not None
@@ -3572,7 +3700,9 @@ def test_translate_openai_response_to_anthropic_with_polyfill_iterations_usage()
     )
     response = _make_simple_openai_response(prompt_tokens=100, completion_tokens=30)
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
+    result = adapter.translate_openai_response_to_anthropic(
+        response=response, polyfill_result=polyfill
+    )
 
     usage = result.get("usage")
     assert usage is not None
@@ -3627,9 +3757,13 @@ def test_translate_openai_response_to_anthropic_with_polyfill_both_compaction_an
             {"type": "compaction", "input_tokens": 300, "output_tokens": 75},
         ],
     )
-    response = _make_simple_openai_response(text="After compaction.", prompt_tokens=120, completion_tokens=40)
+    response = _make_simple_openai_response(
+        text="After compaction.", prompt_tokens=120, completion_tokens=40
+    )
     adapter = LiteLLMAnthropicMessagesAdapter()
-    result = adapter.translate_openai_response_to_anthropic(response=response, polyfill_result=polyfill)
+    result = adapter.translate_openai_response_to_anthropic(
+        response=response, polyfill_result=polyfill
+    )
 
     # compaction block must come first
     content = result.get("content")
@@ -3731,9 +3865,7 @@ def test_translate_anthropic_tools_to_openai_omits_unset_strict():
     assert function["parameters"]["required"] == ["query"]
 
 
-TOOL_RESULT_IMAGE_B64 = (
-    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
-)
+TOOL_RESULT_IMAGE_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
 TOOL_RESULT_IMAGE_URL = "https://example.com/screenshot.png"
 
 
@@ -3741,7 +3873,8 @@ def _anthropic_tool_use_turn(*tool_use_ids):
     return AnthopicMessagesAssistantMessageParam(
         role="assistant",
         content=[
-            {"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}} for tid in tool_use_ids
+            {"type": "tool_use", "id": tid, "name": "read_file", "input": {"path": "img.png"}}
+            for tid in tool_use_ids
         ],
     )
 
@@ -3861,7 +3994,9 @@ def test_tool_result_parallel_tool_calls_keep_tool_message_adjacency():
     result = _run_chat_completions_pipeline(
         [
             _anthropic_tool_use_turn("toolu_01", "toolu_02"),
-            _anthropic_tool_result_turn({"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}),
+            _anthropic_tool_result_turn(
+                {"toolu_01": [_base64_image_block()], "toolu_02": [_url_image_block()]}
+            ),
         ]
     )
 
@@ -4025,9 +4160,7 @@ def test_translate_anthropic_to_openai_without_prompt_cache_breakpoint_adds_noth
 def test_translate_anthropic_messages_to_openai_carries_midturn_system_prompt_cache_breakpoint():
     explicit = {"mode": "explicit"}
     result = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(
-        messages=[
-            {"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}
-        ],
+        messages=[{"role": "system", "content": [{"type": "text", "text": "fix", "prompt_cache_breakpoint": explicit}]}],
         model="gpt-5.6",
     )
     assert result == [

tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py

@@ -1072,6 +1072,54 @@ async def test_messages_strips_provider_prefix_exactly_once(requested_model, exp
     assert captured["url"] == expected_url
 
 
+@pytest.mark.asyncio
+async def test_native_messages_strips_replayed_provider_specific_fields_from_wire():
+    captured = {}
+
+    async def fake_send(self, request, **kwargs):
+        captured["body"] = json.loads(request.content)
+        raise httpx.ConnectError("cut at the wire", request=request)
+
+    messages = [
+        {
+            "role": "assistant",
+            "content": [
+                {
+                    "type": "tool_use",
+                    "id": "toolu_01",
+                    "name": "get_weather",
+                    "input": {"city": "Paris"},
+                    "provider_specific_fields": {"signature": "sig_abc"},
+                }
+            ],
+        },
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "tool_result",
+                    "tool_use_id": "toolu_01",
+                    "content": "Sunny",
+                }
+            ],
+        },
+    ]
+
+    with (
+        patch.object(httpx.AsyncClient, "send", fake_send),
+        pytest.raises(litellm.exceptions.InternalServerError),
+    ):
+        await litellm.anthropic.messages.acreate(
+            max_tokens=100,
+            messages=messages,
+            model="anthropic/claude-haiku-4-5-20251001",
+            api_key="test-api-key",
+        )
+
+    assert "provider_specific_fields" in messages[0]["content"][0]
+    assert "provider_specific_fields" not in captured["body"]["messages"][0]["content"][0]
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     "requested_model, expected_reported_model",

tests/test_litellm/llms/anthropic/experimental_pass_through/responses_adapters/test_responses_adapters_transformation.py

@@ -1265,6 +1265,7 @@ class TestTranslateResponse:
         assert block["id"] == "call_99"
         assert block["name"] == "get_weather"
         assert block["input"] == {"city": "NYC"}
+        assert "provider_specific_fields" not in block
 
     def test_function_call_sets_stop_reason_tool_use(self):
         """Presence of a function_call sets stop_reason to 'tool_use'."""
@@ -1447,6 +1448,7 @@ class TestTranslateResponse:
         assert result["content"][0]["type"] == "tool_use"
         assert result["content"][0]["name"] == "search"
         assert result["content"][0]["input"] == {"query": "cats"}
+        assert "provider_specific_fields" not in result["content"][0]
         assert result["stop_reason"] == "tool_use"
 
     def test_mixed_reasoning_text_and_tool_use(self):

tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py

@@ -7184,6 +7184,7 @@ class TestAggregateGatewayDcrChallenge:
 
         cases = [
             (_server(MCPAuth.oauth2), "srv"),
+            (_server(MCPAuth.oauth2, per_server_oauth_discovery=True), None),
             (_server(MCPAuth.oauth2, oauth2_flow="client_credentials"), "srv"),
             (_server(MCPAuth.oauth2, delegate_auth_to_upstream=True), None),
             (_server(MCPAuth.oauth2_token_exchange), None),

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

@@ -1362,3 +1362,75 @@ async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer
 
     assert result is not None
     assert captured["url"] == "https://idp.example.com/token"
+
+
+def test_prepare_mcp_server_data_carries_per_server_oauth_discovery():
+    request = NewMCPServerRequest(
+        server_name="relay_create",
+        url="https://upstream.example.com/mcp",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2,
+        oauth2_flow="authorization_code",
+        per_server_oauth_discovery=True,
+    )
+
+    data = _prepare_mcp_server_data(request)
+
+    assert data["per_server_oauth_discovery"] is True
+
+
+def test_prepare_mcp_server_data_update_carries_per_server_oauth_discovery():
+    request = UpdateMCPServerRequest(
+        server_id="relay-update",
+        url="https://upstream.example.com/mcp",
+        transport=MCPTransport.http,
+        auth_type=MCPAuth.oauth2,
+        oauth2_flow="authorization_code",
+        per_server_oauth_discovery=True,
+    )
+
+    data = _prepare_mcp_server_data(request, exclude_unset=True)
+
+    assert data["per_server_oauth_discovery"] is True
+
+
+@pytest.mark.parametrize(
+    "request_cls, extra, overrides",
+    [
+        (NewMCPServerRequest, {"server_name": "relay_create"}, {"auth_type": MCPAuth.oauth_delegate}),
+        (NewMCPServerRequest, {"server_name": "relay_create"}, {"oauth2_flow": "client_credentials"}),
+        (UpdateMCPServerRequest, {"server_id": "relay-update"}, {"delegate_auth_to_upstream": True}),
+    ],
+)
+def test_request_models_reject_unsupported_per_server_oauth_discovery(request_cls, extra, overrides):
+    payload = {
+        "url": "https://upstream.example.com/mcp",
+        "transport": MCPTransport.http,
+        "auth_type": MCPAuth.oauth2,
+        "oauth2_flow": "authorization_code",
+        "per_server_oauth_discovery": True,
+        **extra,
+        **overrides,
+    }
+
+    with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"):
+        request_cls(**payload)
+
+
+@pytest.mark.parametrize(
+    "partial_payload",
+    [
+        {"oauth2_flow": "client_credentials"},
+        {"delegate_auth_to_upstream": True},
+        {"auth_type": MCPAuth.api_key},
+    ],
+)
+def test_partial_update_rejects_ineligible_field_alongside_per_server_oauth_discovery(partial_payload):
+    with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"):
+        UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True, **partial_payload)
+
+
+def test_partial_update_defers_omitted_eligibility_fields_to_the_stored_row():
+    request = UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True)
+
+    assert request.per_server_oauth_discovery is True

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

@@ -3342,12 +3342,13 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa
     mock_request.headers = {}
 
     interactive = _oauth2_server("github_mcp")
+    relay = _oauth2_server("relay_mcp", per_server_oauth_discovery=True)
     m2m = _oauth2_server("m2m_mcp", oauth2_flow="client_credentials", client_id="cid", client_secret="cs")
     delegated = _oauth2_server("delegated_mcp", delegate_auth_to_upstream=True)
 
     global_mcp_server_manager.registry.clear()
     try:
-        for server in (interactive, m2m, delegated):
+        for server in (interactive, relay, m2m, delegated):
             global_mcp_server_manager.registry[server.server_id] = server
 
         for name in ("github_mcp", "m2m_mcp"):
@@ -3363,6 +3364,15 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa
             assert legacy["authorization_servers"] == ["https://litellm.example.com/mcp"], name
             assert legacy["resource"] == f"https://litellm.example.com/{name}/mcp"
 
+        relay_response = await _build_oauth_protected_resource_response(
+            request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=True
+        )
+        assert relay_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"]
+        relay_legacy_response = await _build_oauth_protected_resource_response(
+            request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=False
+        )
+        assert relay_legacy_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"]
+
         delegated_response = await _build_oauth_protected_resource_response(
             request=mock_request, mcp_server_name="delegated_mcp", use_standard_pattern=True
         )

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

@@ -95,7 +95,7 @@ def test_interpolate_headers_returns_independent_copy():
 def test_build_env_var_setup_url_includes_server_id(monkeypatch):
     monkeypatch.delenv("PROXY_BASE_URL", raising=False)
     url = _u("build_env_var_setup_url")("abc-123")
-    assert url.startswith("/ui/?page=mcp-servers")
+    assert url.startswith("/ui/mcp-servers?")
     assert "fill_env_vars=abc-123" in url
 
 
@@ -123,7 +123,7 @@ def test_missing_user_env_vars_error_message_is_friendly():
             server_id="abc-123",
             server_name="CorporateDB",
             missing=["CORP_USERNAME", "CORP_PASSWORD"],
-            setup_url="https://proxy.example.com/ui/?page=mcp-servers&fill_env_vars=abc-123",
+            setup_url="https://proxy.example.com/ui/mcp-servers?fill_env_vars=abc-123",
         )
     err = exc_info.value
     text = str(err)
@@ -1694,7 +1694,7 @@ async def test_missing_user_env_vars_error_renders_in_mcp_call_tool():
         server_id="srv-99",
         server_name="CorporateDB",
         missing=["CORP_USERNAME"],
-        setup_url="/ui/?page=mcp-servers&fill_env_vars=srv-99",
+        setup_url="/ui/mcp-servers?fill_env_vars=srv-99",
     )
     # We don't want to spin up the full MCP server framework — just
     # mimic the except-clause behavior the @server.call_tool handler uses.

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

@@ -1229,6 +1229,50 @@ class TestMCPServerManager:
         base.update(overrides)
         return {"bridgeserver": base}
 
+    @pytest.mark.asyncio
+    async def test_load_servers_from_config_accepts_per_server_oauth_discovery_for_oauth2(self):
+        manager = MCPServerManager()
+
+        with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)):
+            await manager.load_servers_from_config(
+                self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery=True)
+            )
+
+        server = next(iter(manager.config_mcp_servers.values()))
+        assert server.per_server_oauth_discovery is True
+        assert server.uses_per_server_oauth_relay is True
+        assert server.advertises_gateway_authorization_server is False
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "config",
+        [
+            {"auth_type": MCPAuth.oauth_delegate},
+            {"oauth2_flow": "client_credentials"},
+            {"oauth2_flow": "authorization_code", "delegate_auth_to_upstream": True},
+        ],
+    )
+    async def test_load_servers_from_config_rejects_unsupported_per_server_oauth_discovery(self, config):
+        manager = MCPServerManager()
+
+        with (
+            patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)),
+            pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"),
+        ):
+            await manager.load_servers_from_config(self._oauth2_config(per_server_oauth_discovery=True, **config))
+
+    @pytest.mark.asyncio
+    async def test_load_servers_from_config_rejects_non_boolean_per_server_oauth_discovery(self):
+        manager = MCPServerManager()
+
+        with (
+            patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)),
+            pytest.raises(ValueError, match="per_server_oauth_discovery.*must be a boolean"),
+        ):
+            await manager.load_servers_from_config(
+                self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery="yes")
+            )
+
     @pytest.mark.asyncio
     async def test_load_servers_from_config_rejects_dcr_bridge_on_gateway_managed_auth_type(self):
         manager = MCPServerManager()

tests/test_litellm/proxy/client/cli/test_agents.py

@@ -54,6 +54,15 @@ class _Recorder:
         return self.returns
 
 
+class _FakeJsonResponse:
+    def __init__(self, status_code, payload=None):
+        self.status_code = status_code
+        self._payload = payload
+
+    def json(self):
+        return self._payload
+
+
 class TestAgentProfile:
     def test_claude_is_anthropic(self):
         name, profiles = agent_profile("claude")
@@ -69,6 +78,9 @@ class TestAgentProfile:
         assert agent_profile("codex") == ("Codex", frozenset({"openai"}))
         assert agent_profile("opencode") == ("OpenCode", frozenset({"openai"}))
 
+    def test_pi_is_litellm(self):
+        assert agent_profile("pi") == ("pi", frozenset({"litellm"}))
+
     def test_unknown_command_gets_both_profiles(self):
         name, profiles = agent_profile("mytool")
         assert name == "mytool"
@@ -134,6 +146,15 @@ class TestBuildAgentEnv:
         assert env["OPENAI_API_KEY"] == "sk-key"
         assert env["ENABLE_TOOL_SEARCH"] == "true"
 
+    def test_litellm_profile_exports_only_the_proxy_key(self):
+        env = build_agent_env(
+            {}, "http://localhost:4000/", "sk-key", frozenset({"litellm"})
+        )
+        assert env["LITELLM_PROXY_API_KEY"] == "sk-key"
+        assert "ANTHROPIC_BASE_URL" not in env
+        assert "OPENAI_BASE_URL" not in env
+        assert "OPENAI_API_KEY" not in env
+
     def test_preserves_unrelated_env_and_does_not_mutate_input(self):
         base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"}
         env = build_agent_env(
@@ -166,6 +187,9 @@ class TestAgentLaunchArgs:
             agent_launch_args("codex", "http://localhost:4000")
         )
 
+    def test_pi_gets_no_static_args(self):
+        assert agent_launch_args("pi", "http://localhost:4000") == []
+
 
 class TestVerifyProxyKey:
     def test_ok_status_passes_and_uses_models_endpoint(self):
@@ -520,6 +544,126 @@ class TestRunAgent:
         # overrides must precede the codex subcommand so codex parses them
         assert args.index('model_provider="litellm"') < args.index("exec")
 
+    def test_pi_preparer_runs_after_verify_and_before_launch(self):
+        order = []
+        captured = {}
+
+        def fake_prepare(base_url, api_key, base_env):
+            order.append("prepare")
+            captured["args"] = (base_url, api_key, dict(base_env))
+            return []
+
+        run_agent(
+            "http://localhost:4000",
+            "sk-key",
+            ["pi"],
+            base_env={"HOME": "/home/u"},
+            which=lambda name: "/usr/local/bin/pi",
+            verify=lambda *a: order.append("verify"),
+            launcher=lambda *a: order.append("launch"),
+            preparers={"pi": fake_prepare},
+        )
+        assert order == ["verify", "prepare", "launch"]
+        assert captured["args"] == (
+            "http://localhost:4000",
+            "sk-key",
+            {"HOME": "/home/u"},
+        )
+
+    def test_pi_prepared_args_precede_user_args_and_env_has_proxy_key(self):
+        calls = {}
+        run_agent(
+            "http://localhost:4000",
+            "sk-key",
+            ["pi", "-p", "hello"],
+            base_env={},
+            which=lambda name: "/usr/local/bin/pi",
+            verify=lambda *a: None,
+            launcher=lambda p, a, e: calls.update(args=tuple(a), env=dict(e)),
+            preparers={"pi": lambda *a: ["--model", "litellm/m-1"]},
+        )
+        assert calls["args"] == ("pi", "--model", "litellm/m-1", "-p", "hello")
+        assert calls["env"]["LITELLM_PROXY_API_KEY"] == "sk-key"
+        assert "OPENAI_API_KEY" not in calls["env"]
+        assert "ANTHROPIC_BASE_URL" not in calls["env"]
+
+    def test_failed_preparer_aborts_before_launch(self):
+        launched = []
+
+        def boom(*a):
+            raise AgentRunError("sync failed")
+
+        with pytest.raises(AgentRunError, match="sync failed"):
+            run_agent(
+                "http://localhost:4000",
+                "sk-key",
+                ["pi"],
+                base_env={},
+                which=lambda name: "/usr/local/bin/pi",
+                verify=lambda *a: None,
+                launcher=lambda *a: launched.append(a),
+                preparers={"pi": boom},
+            )
+        assert launched == []
+
+    def test_prepare_pi_syncs_models_json_and_pins_first_model(self, tmp_path):
+        from litellm.proxy.client.cli.commands.agents import prepare_pi
+
+        def fake_get(url, headers, timeout):
+            if url.endswith("/model_group/info"):
+                return _FakeJsonResponse(
+                    200,
+                    {"data": [{"model_group": "m-first", "max_input_tokens": 131072, "max_output_tokens": 8192}]},
+                )
+            return _FakeJsonResponse(200, {"data": [{"id": "m-first"}, {"id": "m-second"}]})
+
+        pin = prepare_pi(
+            "http://localhost:4000",
+            "sk-key",
+            {"PI_CODING_AGENT_DIR": str(tmp_path)},
+            get=fake_get,
+        )
+
+        assert pin == ("--model", "litellm/m-first")
+        import json
+
+        written = json.loads((tmp_path / "models.json").read_text())
+        assert written["providers"]["litellm"]["apiKey"] == "$LITELLM_PROXY_API_KEY"
+        assert written["providers"]["litellm"]["models"] == [
+            {"id": "m-first", "contextWindow": 131072, "maxTokens": 8192},
+            {"id": "m-second"},
+        ]
+
+    def test_prepare_pi_surfaces_fetch_failure_as_agent_error(self, tmp_path):
+        from litellm.proxy.client.cli.commands.agents import prepare_pi
+
+        with pytest.raises(AgentRunError, match="HTTP 500"):
+            prepare_pi(
+                "http://localhost:4000",
+                "sk-key",
+                {"PI_CODING_AGENT_DIR": str(tmp_path)},
+                get=lambda *a, **k: _FakeJsonResponse(500),
+            )
+
+    def test_claude_has_no_preparer(self):
+        prepared = []
+
+        def fake_prepare(*a):
+            prepared.append(a)
+            return []
+
+        run_agent(
+            "http://localhost:4000",
+            "sk-key",
+            ["claude"],
+            base_env={},
+            which=lambda name: "/usr/local/bin/claude",
+            verify=lambda *a: None,
+            launcher=lambda *a: None,
+            preparers={"pi": fake_prepare},
+        )
+        assert prepared == []
+
     def test_claude_launches_without_injected_args(self):
         calls = {}
         run_agent(
@@ -877,7 +1021,11 @@ class TestAgentCommands:
         self.runner = CliRunner()
 
     def test_one_command_per_known_agent(self):
-        assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode"}
+        assert {c.name for c in agent_commands()} == {"claude", "codex", "opencode", "pi"}
+
+    def test_pi_is_hidden_from_help_but_still_registered(self):
+        hidden_by_name = {c.name: c.hidden for c in agent_commands()}
+        assert hidden_by_name == {"claude": False, "codex": False, "opencode": False, "pi": True}
 
     def test_claude_launches_with_stored_key_and_forwards_args(self):
         captured = {}

tests/test_litellm/proxy/client/cli/test_pi.py

@@ -0,0 +1,236 @@
+import json
+import os
+import stat
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import requests
+
+from litellm.proxy.client.cli.commands.pi import (
+    ModelLimits,
+    PiSyncError,
+    fetch_model_ids,
+    fetch_model_limits,
+    models_json_path,
+    provider_block,
+    sync_models_json,
+)
+
+
+class _FakeResponse:
+    def __init__(self, status_code, payload=None):
+        self.status_code = status_code
+        self._payload = payload
+
+    def json(self):
+        if self._payload is None:
+            raise ValueError("not json")
+        return self._payload
+
+
+class TestFetchModelIds:
+    def test_returns_ids_in_proxy_order_deduped(self):
+        captured = {}
+
+        def fake_get(url, headers, timeout):
+            captured["url"] = url
+            captured["headers"] = headers
+            return _FakeResponse(
+                200,
+                {"data": [{"id": "m-b"}, {"id": "m-a"}, {"id": "m-b"}]},
+            )
+
+        assert fetch_model_ids("http://localhost:4000/", "sk-key", get=fake_get) == ("m-b", "m-a")
+        assert captured["url"] == "http://localhost:4000/v1/models"
+        assert captured["headers"] == {"Authorization": "Bearer sk-key"}
+
+    def test_network_error_is_a_value(self):
+        def boom(*a, **k):
+            raise requests.ConnectionError("refused")
+
+        result = fetch_model_ids("http://localhost:4000", "sk-key", get=boom)
+        assert isinstance(result, PiSyncError)
+        assert "Could not list models" in result.message
+
+    def test_non_200_is_a_value(self):
+        result = fetch_model_ids(
+            "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(500)
+        )
+        assert isinstance(result, PiSyncError)
+        assert "HTTP 500" in result.message
+
+    def test_malformed_body_is_a_value(self):
+        result = fetch_model_ids(
+            "http://localhost:4000",
+            "sk-key",
+            get=lambda *a, **k: _FakeResponse(200, {"data": "nope"}),
+        )
+        assert isinstance(result, PiSyncError)
+
+    def test_empty_model_list_is_a_value(self):
+        result = fetch_model_ids(
+            "http://localhost:4000",
+            "sk-key",
+            get=lambda *a, **k: _FakeResponse(200, {"data": []}),
+        )
+        assert isinstance(result, PiSyncError)
+        assert "no models" in result.message
+
+
+class TestFetchModelLimits:
+    def test_maps_group_limits_and_hits_model_group_info(self):
+        captured = {}
+
+        def fake_get(url, headers, timeout):
+            captured["url"] = url
+            return _FakeResponse(
+                200,
+                {
+                    "data": [
+                        {"model_group": "m-a", "max_input_tokens": 131072, "max_output_tokens": 8192},
+                        {"model_group": "m-b", "max_input_tokens": None, "max_output_tokens": None},
+                    ]
+                },
+            )
+
+        limits = fetch_model_limits("http://localhost:4000/", "sk-key", get=fake_get)
+        assert captured["url"] == "http://localhost:4000/model_group/info"
+        assert limits["m-a"] == ModelLimits(context_window=131072, max_tokens=8192)
+        assert limits["m-b"] == ModelLimits(context_window=None, max_tokens=None)
+
+    def test_non_200_degrades_to_no_limits(self):
+        assert fetch_model_limits("http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(403)) == {}
+
+    def test_network_error_degrades_to_no_limits(self):
+        def boom(*a, **k):
+            raise requests.ConnectionError("refused")
+
+        assert fetch_model_limits("http://localhost:4000", "sk-key", get=boom) == {}
+
+    def test_malformed_body_degrades_to_no_limits(self):
+        assert (
+            fetch_model_limits(
+                "http://localhost:4000", "sk-key", get=lambda *a, **k: _FakeResponse(200, {"data": "nope"})
+            )
+            == {}
+        )
+
+
+class TestModelsJsonPath:
+    def test_env_override_wins(self):
+        assert models_json_path({"PI_CODING_AGENT_DIR": "/custom/dir"}) == Path("/custom/dir/models.json")
+
+    def test_defaults_to_home_pi_agent(self):
+        assert models_json_path({}) == Path.home() / ".pi" / "agent" / "models.json"
+
+
+class TestProviderBlock:
+    def test_points_pi_at_proxy_with_env_interpolated_key(self):
+        block = provider_block("http://localhost:4000/", ("m-1", "m-2"))
+        assert block == {
+            "baseUrl": "http://localhost:4000/v1",
+            "api": "openai-completions",
+            "apiKey": "$LITELLM_PROXY_API_KEY",
+            "models": [{"id": "m-1"}, {"id": "m-2"}],
+        }
+
+    def test_known_limits_become_context_window_and_max_tokens(self):
+        block = provider_block(
+            "http://localhost:4000",
+            ("m-1", "m-2"),
+            {
+                "m-1": ModelLimits(context_window=131072, max_tokens=8192),
+                "m-2": ModelLimits(context_window=None, max_tokens=None),
+            },
+        )
+        assert block["models"] == [
+            {"id": "m-1", "contextWindow": 131072, "maxTokens": 8192},
+            {"id": "m-2"},
+        ]
+
+
+class TestSyncModelsJson:
+    def test_creates_file_and_parent_dirs(self, tmp_path):
+        path = tmp_path / "agent" / "models.json"
+        assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
+        written = json.loads(path.read_text())
+        assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1"
+        assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}]
+
+    def test_preserves_other_providers_and_top_level_keys(self, tmp_path):
+        path = tmp_path / "models.json"
+        path.write_text(
+            json.dumps(
+                {
+                    "somethingElse": True,
+                    "providers": {
+                        "ollama": {"baseUrl": "http://localhost:11434/v1"},
+                        "litellm": {"baseUrl": "http://stale:1234/v1", "models": []},
+                    },
+                }
+            )
+        )
+        assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
+        written = json.loads(path.read_text())
+        assert written["somethingElse"] is True
+        assert written["providers"]["ollama"] == {"baseUrl": "http://localhost:11434/v1"}
+        assert written["providers"]["litellm"]["baseUrl"] == "http://localhost:4000/v1"
+        assert written["providers"]["litellm"]["models"] == [{"id": "m-1"}]
+
+    def test_write_leaves_no_staging_file_behind(self, tmp_path):
+        path = tmp_path / "models.json"
+        assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
+        assert [p.name for p in tmp_path.iterdir()] == ["models.json"]
+
+    def test_written_file_is_private(self, tmp_path):
+        path = tmp_path / "models.json"
+        assert sync_models_json(path, "http://localhost:4000", ("m-1",)) is None
+        if os.name != "nt":
+            assert stat.S_IMODE(path.stat().st_mode) == 0o600
+
+        path.write_text(json.dumps({"providers": {"other": {"apiKey": "literal-secret"}}}))
+        path.chmod(0o644)
+        assert sync_models_json(path, "http://localhost:4000", ("m-2",)) is None
+        if os.name != "nt":
+            assert stat.S_IMODE(path.stat().st_mode) == 0o600
+
+    def test_concurrent_syncs_do_not_collide(self, tmp_path):
+        path = tmp_path / "models.json"
+        model_lists = (("m-a",), ("m-b",))
+
+        def sync(model_ids):
+            return sync_models_json(path, "http://localhost:4000", model_ids)
+
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            results = [result for _ in range(30) for result in executor.map(sync, model_lists)]
+
+        assert results == [None] * 60
+        written = json.loads(path.read_text())
+        assert written["providers"]["litellm"]["models"] in ([{"id": "m-a"}], [{"id": "m-b"}])
+        assert list(tmp_path.glob("models.json.*.tmp")) == []
+
+    def test_invalid_json_is_a_value_and_file_untouched(self, tmp_path):
+        path = tmp_path / "models.json"
+        path.write_text("{not json")
+        result = sync_models_json(path, "http://localhost:4000", ("m-1",))
+        assert isinstance(result, PiSyncError)
+        assert path.read_text() == "{not json"
+
+    def test_non_object_providers_is_a_value(self, tmp_path):
+        path = tmp_path / "models.json"
+        path.write_text(json.dumps({"providers": ["nope"]}))
+        result = sync_models_json(path, "http://localhost:4000", ("m-1",))
+        assert isinstance(result, PiSyncError)
+
+    def test_top_level_non_object_is_a_value(self, tmp_path):
+        path = tmp_path / "models.json"
+        path.write_text(json.dumps(["nope"]))
+        result = sync_models_json(path, "http://localhost:4000", ("m-1",))
+        assert isinstance(result, PiSyncError)
+
+    def test_unwritable_path_is_a_value(self, tmp_path):
+        blocker = tmp_path / "agent"
+        blocker.write_text("i am a file, not a directory")
+        result = sync_models_json(blocker / "models.json", "http://localhost:4000", ("m-1",))
+        assert isinstance(result, PiSyncError)
+        assert "Could not" in result.message

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_aim.py

@@ -1,6 +1,15 @@
 """Tests for the AIM guardrail's inspection-payload construction."""
 
+from copy import deepcopy
+from unittest.mock import AsyncMock, patch
+
+import pytest
+from httpx import Request, Response
+
+from litellm import DualCache
+from litellm.proxy._types import ProxyException, UserAPIKeyAuth
 from litellm.proxy.guardrails.guardrail_hooks.aim.aim import AimGuardrail
+from litellm.types.utils import ModelResponse
 
 
 def test_aim_inspection_messages_coerces_chat_completions_tool_role_to_user():
@@ -86,3 +95,354 @@ def test_aim_inspection_messages_preserves_safe_roles():
         {"role": "user", "content": "hi"},
         {"role": "assistant", "content": "hello"},
     ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("hook", ["pre_call", "moderation"])
+@pytest.mark.parametrize("call_type", ["embedding", "aembedding"])
+async def test_aim_skips_embeddings_without_calling_the_guardrail(hook: str, call_type: str):
+    """/embeddings is not a conversation, so neither hook should reach AIM."""
+    guardrail = AimGuardrail(api_key="hs-aim-key", guardrail_name="aim", event_hook="pre_call")
+    data = {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+    with patch(  # test-quality-ok: transport is litellm's aiohttp-backed handler; respx cannot intercept it
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new_callable=AsyncMock,
+    ) as mock_post:
+        if hook == "pre_call":
+            result = await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+        else:
+            result = await guardrail.async_moderation_hook(
+                data=data,
+                user_api_key_dict=UserAPIKeyAuth(),
+                call_type=call_type,
+            )
+
+    mock_post.assert_not_called()
+    assert result == {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+
+@pytest.mark.parametrize(
+    ("configured", "expected"),
+    [
+        ({}, False),
+        ({"inspect_embeddings": True}, True),
+        ({"inspect_embeddings": "true"}, True),
+        ({"inspect_embeddings": "false"}, False),
+    ],
+)
+def test_aim_config_plumbs_inspect_embeddings(
+    configured: dict[str, object], expected: bool, monkeypatch: pytest.MonkeyPatch
+):
+    import litellm
+    from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
+
+    monkeypatch.setattr(litellm, "guardrail_name_config_map", {})
+    monkeypatch.setattr(litellm, "callbacks", [])
+
+    init_guardrails_v2(
+        all_guardrails=[
+            {
+                "guardrail_name": "aim-guard",
+                "litellm_params": {
+                    "guardrail": "aim",
+                    "mode": "pre_call",
+                    "api_key": "hs-aim-key",
+                    **configured,
+                },
+            },
+        ],
+        config_file_path="",
+    )
+
+    aim_guardrails = [callback for callback in litellm.callbacks if isinstance(callback, AimGuardrail)]
+    assert len(aim_guardrails) == 1
+    assert aim_guardrails[0].inspect_embeddings is expected
+
+
+@pytest.mark.asyncio
+async def test_aim_anonymize_action_redacts_batched_embeddings():
+    """A batched ``input`` list of plain strings is redactable: AIM returns one
+    redacted message per string, so the list is rewritten element-wise instead
+    of being hard-blocked as non-text content."""
+    guardrail = AimGuardrail(
+        api_key="hs-aim-key",
+        guardrail_name="aim",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+    response = Response(
+        json={
+            "required_action": {"action_type": "anonymize_action"},
+            "analysis_result": {"policy_drill_down": {}},
+            "redacted_chat": {
+                "all_redacted_messages": [
+                    {"role": "user", "content": "first [REDACTED]"},
+                    {"role": "user", "content": "second [REDACTED]"},
+                ]
+            },
+        },
+        status_code=200,
+        request=Request(method="POST", url="http://aim"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", return_value=response):
+        result = await guardrail.async_pre_call_hook(
+            user_api_key_dict=UserAPIKeyAuth(),
+            cache=DualCache(),
+            data=data,
+            call_type="aembedding",
+        )
+
+    assert result is not None
+    assert result["input"] == ["first [REDACTED]", "second [REDACTED]"]
+
+
+@pytest.mark.asyncio
+async def test_aim_anonymize_action_blocks_when_batch_redaction_count_differs():
+    """AIM returning fewer redacted messages than the batch carries cannot be
+    applied element-wise. Blocking is the only safe answer: a partial rewrite
+    would forward the unmatched elements to the provider unredacted."""
+    guardrail = AimGuardrail(
+        api_key="hs-aim-key",
+        guardrail_name="aim",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"model": "text-embedding-3-small", "input": ["first SSN", "second SSN", "third SSN"]}
+    response = Response(
+        json={
+            "required_action": {"action_type": "anonymize_action"},
+            "analysis_result": {"policy_drill_down": {}},
+            "redacted_chat": {"all_redacted_messages": [{"role": "user", "content": "first [REDACTED]"}]},
+        },
+        status_code=200,
+        request=Request(method="POST", url="http://aim"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", return_value=response):
+        with pytest.raises(ProxyException) as exc_info:
+            await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type="aembedding",
+            )
+
+    assert exc_info.value.code == "400"
+    assert data["input"] == ["first SSN", "second SSN", "third SSN"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "all_redacted_messages",
+    [
+        pytest.param([{"role": "user"}], id="content-missing"),
+        pytest.param([{"role": "user", "content": None}], id="content-null"),
+        pytest.param(["first [REDACTED]"], id="not-a-mapping"),
+        pytest.param([], id="empty-list"),
+        pytest.param("invalid", id="missing-collection"),
+    ],
+)
+@pytest.mark.parametrize(
+    ("request_body", "call_type"),
+    [
+        pytest.param({"model": "text-embedding-3-small", "input": ["first SSN"]}, "aembedding", id="batch-input"),
+        pytest.param(
+            {"model": "gpt-4o", "messages": [{"role": "user", "content": "first SSN"}]},
+            "acompletion",
+            id="chat-messages",
+        ),
+    ],
+)
+async def test_aim_anonymize_action_blocks_malformed_redacted_messages(
+    all_redacted_messages: object, request_body: dict, call_type: str
+):
+    """Malformed AIM redactions return a controlled 400 without changing the request."""
+    guardrail = AimGuardrail(
+        api_key="hs-aim-key",
+        guardrail_name="aim",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = deepcopy(request_body)
+    response = Response(
+        json={
+            "required_action": {"action_type": "anonymize_action"},
+            "analysis_result": {"policy_drill_down": {}},
+            "redacted_chat": {"all_redacted_messages": all_redacted_messages},
+        },
+        status_code=200,
+        request=Request(method="POST", url="http://aim"),
+    )
+
+    with patch.object(guardrail.async_handler, "post", return_value=response):
+        with pytest.raises(ProxyException) as exc_info:
+            await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+
+    assert exc_info.value.code == "400"
+    assert data == request_body
+
+
+_OUTPUT_REQUEST = {
+    "messages": [
+        {"role": "system", "content": "be terse"},
+        {"role": "user", "content": "repeat my SSN"},
+    ]
+}
+_OUTPUT_ECHO = [
+    {"role": "system", "content": "be terse"},
+    {"role": "user", "content": "repeat my SSN"},
+]
+
+
+def _completion(content: str) -> ModelResponse:
+    return ModelResponse(
+        choices=[{"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": content}}]
+    )
+
+
+def _anonymize_response(all_redacted_messages: object) -> Response:
+    return Response(
+        json={
+            "required_action": {"action_type": "anonymize_action"},
+            "analysis_result": {"policy_drill_down": {}},
+            "redacted_chat": {"all_redacted_messages": all_redacted_messages},
+        },
+        status_code=200,
+        request=Request(method="POST", url="http://aim"),
+    )
+
+
+@pytest.mark.asyncio
+async def test_aim_output_anonymize_takes_the_assistant_entry_after_the_echoed_request():
+    """AIM echoes every inspected request message before the assistant turn, so the
+    redacted completion is the final entry of a batch one longer than the request."""
+    guardrail = AimGuardrail(api_key="hs-aim-key", guardrail_name="aim", event_hook="post_call")
+    response = _completion("your SSN is 123-45-6789")
+
+    with patch.object(
+        guardrail.async_handler,
+        "post",
+        return_value=_anonymize_response([*_OUTPUT_ECHO, {"role": "assistant", "content": "your SSN is [REDACTED]"}]),
+    ):
+        result = await guardrail.async_post_call_success_hook(
+            data=deepcopy(_OUTPUT_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response
+        )
+
+    assert result["choices"][0]["message"]["content"] == "your SSN is [REDACTED]"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "all_redacted_messages",
+    [
+        pytest.param(_OUTPUT_ECHO, id="assistant-entry-missing"),
+        pytest.param([{"role": "assistant", "content": "your SSN is [REDACTED]"}], id="request-echo-missing"),
+        pytest.param([*_OUTPUT_ECHO, {"role": "assistant", "content": ""}], id="assistant-content-empty"),
+        pytest.param([*_OUTPUT_ECHO, {"role": "assistant", "content": None}], id="assistant-content-null"),
+        pytest.param([*_OUTPUT_ECHO, "your SSN is [REDACTED]"], id="not-a-mapping"),
+        pytest.param([], id="empty-list"),
+        pytest.param("invalid", id="missing-collection"),
+    ],
+)
+async def test_aim_output_anonymize_blocks_malformed_redactions(all_redacted_messages: object):
+    """A redaction AIM cannot be aligned to the completion is a 400, never the
+    unredacted completion and never a 500."""
+    guardrail = AimGuardrail(api_key="hs-aim-key", guardrail_name="aim", event_hook="post_call")
+    response = _completion("your SSN is 123-45-6789")
+
+    with patch.object(guardrail.async_handler, "post", return_value=_anonymize_response(all_redacted_messages)):
+        with pytest.raises(ProxyException) as exc_info:
+            await guardrail.async_post_call_success_hook(
+                data=deepcopy(_OUTPUT_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response
+            )
+
+    assert exc_info.value.code == "400"
+    assert response["choices"][0]["message"]["content"] == "your SSN is 123-45-6789"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("hook", ["pre_call", "moderation"])
+@pytest.mark.parametrize("call_type", ["embedding", "aembedding"])
+async def test_aim_inspects_embeddings_when_enabled(hook: str, call_type: str):
+    guardrail = AimGuardrail(
+        api_key="hs-aim-key",
+        guardrail_name="aim",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+    with patch.object(
+        guardrail.async_handler,
+        "post",
+        return_value=Response(
+            json={"required_action": None, "analysis_result": {"policy_drill_down": {}}},
+            status_code=200,
+            request=Request(method="POST", url="http://aim"),
+        ),
+    ) as mock_post:
+        if hook == "pre_call":
+            result = await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+        else:
+            result = await guardrail.async_moderation_hook(
+                data=data,
+                user_api_key_dict=UserAPIKeyAuth(),
+                call_type=call_type,
+            )
+
+    mock_post.assert_called_once()
+    assert result == data
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "call_type",
+    ["completion", "acompletion", "responses", "aresponses", "anthropic_messages", "call_mcp_tool"],
+)
+async def test_aim_still_inspects_every_conversational_call_type(call_type: str):
+    """Deny-list, not allow-list: ``TEXT_CONTENT_CALL_TYPES`` omits these, so gating
+    on it would silently stop inspecting real chat traffic."""
+    guardrail = AimGuardrail(api_key="hs-aim-key", guardrail_name="aim", event_hook="pre_call")
+    data = {"messages": [{"role": "user", "content": "Hi my name is Brian"}]}
+
+    with patch(  # test-quality-ok: transport is litellm's aiohttp-backed handler; respx cannot intercept it
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=Response(
+            json={
+                "analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}},
+                "required_action": {
+                    "action_type": "block_action",
+                    "detection_message": "PII detected",
+                },
+            },
+            status_code=200,
+            request=Request(method="POST", url="http://aim"),
+        ),
+    ) as mock_post:
+        with pytest.raises(ProxyException, match="PII detected"):
+            await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+
+    mock_post.assert_called_once()

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_cato_networks.py

@@ -8,19 +8,19 @@ from fastapi.exceptions import HTTPException
 from httpx import Request, Response
 from websockets.exceptions import ConnectionClosed
 
+import litellm
 from litellm import DualCache
 from litellm.proxy.guardrails.guardrail_hooks.cato_networks.cato_networks import (
     CatoNetworksGuardrail,
     CatoNetworksGuardrailMissingSecrets,
 )
+from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
 from litellm.proxy.proxy_server import UserAPIKeyAuth
 from litellm.types.utils import ModelResponse, ResponsesAPIResponse
 
-import litellm
-from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
 
-
-def test_cato_guard_config():
+def test_cato_guard_config(monkeypatch):
+    monkeypatch.setattr(litellm, "callbacks", [])
     litellm.guardrail_name_config_map = {}
 
     init_guardrails_v2(
@@ -32,11 +32,15 @@ def test_cato_guard_config():
                     "guard_name": "gibberish_guard",
                     "mode": "pre_call",
                     "api_key": "hs-cato-key",
+                    "inspect_embeddings": True,
                 },
             },
         ],
         config_file_path="",
     )
+    cato_guardrails = [callback for callback in litellm.callbacks if isinstance(callback, CatoNetworksGuardrail)]
+    assert len(cato_guardrails) == 1
+    assert cato_guardrails[0].inspect_embeddings is True
 
 
 def test_cato_guard_config_no_api_key(monkeypatch):
@@ -218,7 +222,7 @@ async def test_post_call__with_anonymized_entities__it_doesnt_deanonymize_output
             elif request_body["messages"][-1]["role"] == "assistant":
                 return response_without_detections
             else:
-                raise ValueError("Unexpected request: {}".format(request_body))
+                raise ValueError(f"Unexpected request: {request_body}")
 
         mock_post.side_effect = mock_post_detect_side_effect
 
@@ -772,6 +776,92 @@ async def test_call_cato_guardrail_on_output_flattens_multimodal_context():
     assert sent[-1] == {"role": "assistant", "content": "the answer"}
 
 
+@pytest.mark.asyncio
+async def test_anonymize_action_redacts_batched_embeddings_input():
+    """A batched ``input`` list of plain strings is redactable, so the redacted
+    text is written back element-wise instead of the request going out with the
+    original strings intact."""
+    guard = CatoNetworksGuardrail(
+        api_key="hs-cato-key",
+        guardrail_name="cato",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"input": ["first SSN", "second SSN"]}
+    response = _make_response(
+        {
+            "analysis_result": {"policy_drill_down": {}},
+            "required_action": {"action_type": "anonymize_action"},
+            "redacted_chat": {
+                "all_redacted_messages": [
+                    {"role": "user", "content": "first [REDACTED]"},
+                    {"role": "user", "content": "second [REDACTED]"},
+                ]
+            },
+        }
+    )
+
+    with patch.object(guard.async_handler, "post", return_value=response):
+        result = await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
+
+    assert result["input"] == ["first [REDACTED]", "second [REDACTED]"]
+
+
+@pytest.mark.asyncio
+async def test_anonymize_action_blocks_when_batch_redaction_count_differs():
+    """Cato returning fewer redacted messages than the batch carries cannot be
+    applied element-wise. Blocking is the only safe answer: a partial rewrite
+    would forward the unmatched elements to the provider unredacted."""
+    guard = CatoNetworksGuardrail(
+        api_key="hs-cato-key",
+        guardrail_name="cato",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"input": ["first SSN", "second SSN", "third SSN"]}
+    response = _make_response(
+        {
+            "analysis_result": {"policy_drill_down": {}},
+            "required_action": {"action_type": "anonymize_action"},
+            "redacted_chat": {"all_redacted_messages": [{"role": "user", "content": "first [REDACTED]"}]},
+        }
+    )
+
+    with patch.object(guard.async_handler, "post", return_value=response):
+        with pytest.raises(HTTPException) as exc_info:
+            await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
+
+    assert exc_info.value.status_code == 400
+    assert data["input"] == ["first SSN", "second SSN", "third SSN"]
+
+
+@pytest.mark.asyncio
+async def test_anonymize_action_blocks_when_batch_redaction_is_empty():
+    """An anonymize verdict with no redacted messages at all is the extreme case
+    of the same mismatch, and must not silently forward the raw batch."""
+    guard = CatoNetworksGuardrail(
+        api_key="hs-cato-key",
+        guardrail_name="cato",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"input": ["first SSN", "second SSN"]}
+    response = _make_response(
+        {
+            "analysis_result": {"policy_drill_down": {}},
+            "required_action": {"action_type": "anonymize_action"},
+            "redacted_chat": {"all_redacted_messages": []},
+        }
+    )
+
+    with patch.object(guard.async_handler, "post", return_value=response):
+        with pytest.raises(HTTPException) as exc_info:
+            await guard.call_cato_guardrail(data, hook="pre_call", key_alias=None)
+
+    assert exc_info.value.status_code == 400
+    assert data["input"] == ["first SSN", "second SSN"]
+
+
 @pytest.mark.asyncio
 async def test_anonymize_action_redacts_responses_api_input():
     """Anonymized text must be written back to ``input`` for Responses-API requests."""
@@ -2590,3 +2680,108 @@ async def test_forward_the_stream_to_cato_serializes_chunks():
     assert sent[2] == "raw-sse-chunk"
     assert sent[3] == json.dumps([1, 2, 3])
     assert json.loads(sent[-1]) == {"done": True}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("hook", ["pre_call", "moderation"])
+@pytest.mark.parametrize("call_type", ["embedding", "aembedding"])
+async def test_cato_skips_embeddings_without_calling_the_guardrail(hook: str, call_type: str):
+    """/embeddings is not a conversation, so neither hook should reach Cato."""
+    guardrail = CatoNetworksGuardrail(api_key="hs-cato-key", guardrail_name="cato", event_hook="pre_call")
+    data = {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+    with patch(  # test-quality-ok: transport is litellm's aiohttp-backed handler; respx cannot intercept it
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        new_callable=AsyncMock,
+    ) as mock_post:
+        if hook == "pre_call":
+            result = await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+        else:
+            result = await guardrail.async_moderation_hook(
+                data=data,
+                user_api_key_dict=UserAPIKeyAuth(),
+                call_type=call_type,
+            )
+
+    mock_post.assert_not_called()
+    assert result == {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("hook", ["pre_call", "moderation"])
+@pytest.mark.parametrize("call_type", ["embedding", "aembedding"])
+async def test_cato_inspects_embeddings_when_enabled(hook: str, call_type: str):
+    guardrail = CatoNetworksGuardrail(
+        api_key="hs-cato-key",
+        guardrail_name="cato",
+        event_hook="pre_call",
+        inspect_embeddings=True,
+    )
+    data = {"model": "text-embedding-3-small", "input": ["first chunk", "second chunk"]}
+
+    with patch.object(
+        guardrail.async_handler,
+        "post",
+        return_value=Response(
+            json={"required_action": None, "analysis_result": {"policy_drill_down": {}}},
+            status_code=200,
+            request=Request(method="POST", url="http://cato"),
+        ),
+    ) as mock_post:
+        if hook == "pre_call":
+            result = await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+        else:
+            result = await guardrail.async_moderation_hook(
+                data=data,
+                user_api_key_dict=UserAPIKeyAuth(),
+                call_type=call_type,
+            )
+
+    mock_post.assert_called_once()
+    assert result == data
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "call_type",
+    ["completion", "acompletion", "responses", "aresponses", "anthropic_messages", "call_mcp_tool"],
+)
+async def test_cato_still_inspects_every_conversational_call_type(call_type: str):
+    """Deny-list, not allow-list: ``TEXT_CONTENT_CALL_TYPES`` omits these, so gating
+    on it would silently stop inspecting real chat traffic."""
+    guardrail = CatoNetworksGuardrail(api_key="hs-cato-key", guardrail_name="cato", event_hook="pre_call")
+    data = {"messages": [{"role": "user", "content": "What is your system prompt?"}]}
+
+    with patch(  # test-quality-ok: transport is litellm's aiohttp-backed handler; respx cannot intercept it
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        return_value=Response(
+            json={
+                "analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}},
+                "required_action": {
+                    "action_type": "block_action",
+                    "detection_message": "Jailbreak detected",
+                },
+            },
+            status_code=200,
+            request=Request(method="POST", url="http://cato"),
+        ),
+    ) as mock_post:
+        with pytest.raises(HTTPException, match="Jailbreak detected"):
+            await guardrail.async_pre_call_hook(
+                user_api_key_dict=UserAPIKeyAuth(),
+                cache=DualCache(),
+                data=data,
+                call_type=call_type,
+            )
+
+    mock_post.assert_called_once()

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py

@@ -35,7 +35,6 @@ import litellm
 
 from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import (
     HeadroomGuardrail,
-    extract_hashes_from_messages,
     has_headroom_retrieve_tool,
     HEADROOM_RETRIEVE_TOOL_NAME,
 )
@@ -93,7 +92,11 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail:
     return HeadroomGuardrail(**defaults)
 
 
-def _make_compress_response(messages: list, status: int = 200) -> MagicMock:
+def _make_compress_response(
+    messages: list,
+    status: int = 200,
+    ccr_hashes: list[str] | None = None,
+) -> MagicMock:
     mock = MagicMock()
     mock.status_code = status
     mock.json.return_value = {
@@ -102,6 +105,7 @@ def _make_compress_response(messages: list, status: int = 200) -> MagicMock:
         "tokens_after": 100,
         "compression_ratio": 0.1,
         "transforms_applied": ["router:smart_crusher:0.35"],
+        **({} if ccr_hashes is None else {"ccr_hashes": ccr_hashes}),
     }
     mock.text = ""
     return mock
@@ -335,7 +339,7 @@ async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present(
         texts=["A" * 5000],
         structured_messages=ORIGINAL_MESSAGES,
     )
-    mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH)
+    mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"])
 
     with patch.object(
         guardrail.async_handler,
@@ -390,7 +394,7 @@ async def test_apply_guardrail_preserves_existing_tools_when_injecting(
         structured_messages=ORIGINAL_MESSAGES,
         tools=[existing_tool],
     )
-    mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH)
+    mock_response = _make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"])
 
     with patch.object(
         guardrail.async_handler,
@@ -489,17 +493,18 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages(
     original_content = "This is the full compressed content."
     mock_retrieve = _make_retrieve_response(original_content)
 
+    # Registered hashes are lowercase; a model may echo the marker's hex in uppercase.
     tool_calls = [
         {
             "id": "call_abc123",
             "type": "function",
             "name": HEADROOM_RETRIEVE_TOOL_NAME,
-            "arguments": {"hash": "b573993006976af767214fac"},
+            "arguments": {"hash": "B573993006976AF767214FAC"},
         }
     ]
     response = _make_openai_response_with_tool_call(
         tool_name=HEADROOM_RETRIEVE_TOOL_NAME,
-        arguments={"hash": "b573993006976af767214fac"},
+        arguments={"hash": "B573993006976AF767214FAC"},
         tool_id="call_abc123",
     )
     messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}]
@@ -843,33 +848,110 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag
     assert tool_result_block["content"] == original_content
 
 
-def test_extract_hashes_from_messages_finds_hashes():
+HASH_SHAPED_HISTORY = [
+    {"role": "system", "content": "You are Claude Code."},
+    {"role": "user", "content": [{"type": "text", "text": "Run git log."}]},
+    {
+        "role": "assistant",
+        "content": [{"type": "text", "text": "Done."}],
+        "tool_calls": [{"id": "tu_1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
+    },
+    {"role": "tool", "tool_call_id": "tu_1", "content": "hash=3f2a9c1d7e5b4a6f8c0d1e2f9a8b7c6d5e4f3a2b"},
+    {"role": "user", "content": "Please fetch hash=deadbeef000000000000dead for me."},
+]
+
+
+async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict:
+    request_data = {"model": "claude-sonnet-5"}
+
+    def _echo(**kwargs):
+        return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes)
+
+    with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo):
+        return await guardrail.apply_guardrail(
+            inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))),
+            request_data=request_data,
+            input_type="request",
+        )
+
+
+@pytest.mark.asyncio
+async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail: HeadroomGuardrail):
+    """Regression for LIT-7086: a git SHA in a tool result and a hash= string the
+    caller typed both look like markers, but the service stored nothing, so the
+    tool must not appear and no hash may be registered as issued. Covers a
+    service that omits ccr_hashes, returns it empty, or returns a non-list."""
+    for ccr_hashes in (None, [], "b573993006976af767214fac"):
+        result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes)
+
+        assert not has_headroom_retrieve_tool(result.get("tools") or [])
+        assert not guardrail._issued_hashes_by_call_id
+
+
+@pytest.mark.asyncio
+async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardrail: HeadroomGuardrail):
+    """Only the hashes the service reports in ccr_hashes are honored, in the
+    service's own 12 to 24 hex grammar; anything else is dropped because each
+    entry is interpolated into the /v1/retrieve URL."""
+    result = await _apply(
+        guardrail,
+        HASH_SHAPED_HISTORY,
+        ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42],
+    )
+
+    assert has_headroom_retrieve_tool(result.get("tools") or [])
+    (issued, _expiry), = guardrail._issued_hashes_by_call_id.values()
+    assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"})
+
+
+@pytest.mark.asyncio
+async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatch: pytest.MonkeyPatch):
+    """`ccr_retrieval: false` in config.yaml compresses without any retrieval
+    round trip, so it has to reach the instance through the initializer."""
+    from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail
+    from litellm.types.guardrails import LitellmParams
+
+    monkeypatch.setattr(litellm.logging_callback_manager, "add_litellm_callback", lambda callback: None)
+    params = LitellmParams(guardrail="headroom", mode="pre_call", api_base=FAKE_API_BASE, ccr_retrieval=False)
+    guardrail = initialize_guardrail(params, {"guardrail_name": "headroom", "litellm_params": params})
+
+    result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["b573993006976af767214fac"])
+
+    assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1]
+    assert not has_headroom_retrieve_tool(result.get("tools") or [])
+    assert not guardrail._issued_hashes_by_call_id
+
+
+@pytest.mark.asyncio
+async def test_anthropic_assistant_history_never_reaches_compression_service(guardrail: HeadroomGuardrail):
+    """The public Anthropic handler translates assistant content blocks to a
+    string before Headroom sees them, so model-authored rows must be excluded
+    from the compression payload rather than protected by their content shape."""
+    from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
+
+    table = "| Guardrail | Model |\n|---|---|\n" + "\n".join(f"| gr-{i} | model-{i} |" for i in range(40))
     messages = [
-        {"role": "user", "content": "Retrieve more: hash=b573993006976af767214fac"},
-        {"role": "assistant", "content": "Also: hash=aabbccdd001122334455aabb"},
+        {"role": "user", "content": [{"type": "text", "text": "List the guardrails."}]},
+        {"role": "assistant", "content": [{"type": "text", "text": table}]},
+        {"role": "user", "content": [{"type": "text", "text": "Earlier follow-up. " + "B" * 5000}]},
+        {"role": "assistant", "content": [{"type": "text", "text": "Noted."}]},
+        {"role": "user", "content": "Re-print the table."},
     ]
-    hashes = extract_hashes_from_messages(messages)
-    assert "b573993006976af767214fac" in hashes
-    assert "aabbccdd001122334455aabb" in hashes
+    sent: dict = {}
 
+    def _echo(**kwargs):
+        sent["messages"] = kwargs["json"]["messages"]
+        return _make_compress_response(json.loads(json.dumps(sent["messages"])))
 
-def test_extract_hashes_from_messages_ignores_short_hashes():
-    messages = [{"role": "user", "content": "hash=tooshort"}]
-    hashes = extract_hashes_from_messages(messages)
-    assert not hashes
+    data = {"model": "claude-sonnet-5", "messages": json.loads(json.dumps(messages))}
+    with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo):
+        result = await AnthropicMessagesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
 
+    assert [row["role"] for row in sent["messages"]] == ["user", "user"]
+    assert sent["messages"][1]["content"] == "Earlier follow-up. " + "B" * 5000
+    assert table not in json.dumps(sent["messages"])
+    assert result["messages"][1]["content"] == [{"type": "text", "text": table}]
 
-def test_extract_hashes_from_list_content_blocks():
-    messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": "hash=b573993006976af767214fac found here"},
-            ],
-        }
-    ]
-    hashes = extract_hashes_from_messages(messages)
-    assert "b573993006976af767214fac" in hashes
 
 
 def test_has_headroom_retrieve_tool_recognizes_anthropic_native_shape():
@@ -1001,7 +1083,7 @@ async def test_responses_request_sends_compressed_input_and_retrieve_tool_upstre
         guardrail.async_handler,
         "post",
         new_callable=AsyncMock,
-        return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
+        return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]),
     ):
         result = await OpenAIResponsesHandler().process_input_messages(data=data, guardrail_to_apply=guardrail)
 
@@ -1793,7 +1875,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row(
     )
     compressed = _echo_wire_view()
     compressed[0]["content"] = "compressed history. Retrieve more: hash=b573993006976af767214fac"
-    mock_response = _make_compress_response(compressed)
+    mock_response = _make_compress_response(compressed, ccr_hashes=["b573993006976af767214fac"])
 
     with patch.object(
         guardrail.async_handler,
@@ -1819,7 +1901,7 @@ async def test_apply_guardrail_restores_rewritten_all_text_row(
     assert history_content[0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"}
     # Mixed row passes through byte-identical.
     assert messages[2]["content"] == PARTS_MESSAGES[2]["content"]
-    # Hashes inside restored parts still drive retrieve-tool injection.
+    # The service-declared hash still drives retrieve-tool injection on a restored row.
     assert has_headroom_retrieve_tool(result.get("tools") or [])
 
 
@@ -2159,7 +2241,7 @@ async def test_pre_call_deployment_hook_converts_stream_after_deployment_level_c
         guardrail.async_handler,
         "post",
         new_callable=AsyncMock,
-        return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH),
+        return_value=_make_compress_response(COMPRESSED_MESSAGES_WITH_HASH, ccr_hashes=["b573993006976af767214fac"]),
     ):
         result = await guardrail.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.acompletion)
 
@@ -2345,7 +2427,11 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end(
 AGENTIC_MESSAGES = [
     {"role": "system", "content": "You are Claude Code. " + "S" * 5000},
     {"role": "user", "content": "H" * 5000},
-    {"role": "assistant", "content": "Older answer. " + "O" * 5000},
+    {
+        "role": "assistant",
+        "content": "Older answer. " + "O" * 5000,
+        "tool_calls": [{"id": "old_1", "type": "function", "function": {"name": "Read", "arguments": "{}"}}],
+    },
     {"role": "tool", "tool_call_id": "old_1", "content": "older tool output " + "T" * 5000},
     {
         "role": "assistant",
@@ -2420,23 +2506,21 @@ async def test_history_is_still_compressed(guardrail: HeadroomGuardrail):
     """Negative control: protection must not turn compression into a no-op."""
     compressed_history = [
         {"role": "user", "content": "hist. hash=b573993006976af767214fac"},
-        {"role": "assistant", "content": "older. hash=a73993006976af767214fac1"},
         {"role": "tool", "tool_call_id": "old_1", "content": "older tool. hash=c73993006976af767214fac2"},
     ]
     wire, result = await _wire_and_result(guardrail, AGENTIC_MESSAGES, returned=compressed_history)
 
-    # Exactly the three history rows go to the service, in order.
-    assert [row["role"] for row in wire] == ["user", "assistant", "tool"]
+    # Older user and tool rows go to the service, in order. Every assistant row
+    # stays out, but the tool results those turns asked for remain compressible.
+    assert [row["role"] for row in wire] == ["user", "tool"]
     assert wire[0]["content"] == "H" * 5000
-    assert wire[2]["tool_call_id"] == "old_1"
+    assert wire[1]["tool_call_id"] == "old_1"
 
     messages = result["structured_messages"]
     assert len(messages) == len(AGENTIC_MESSAGES)
     assert messages[1] == compressed_history[0]
-    assert messages[2] == compressed_history[1]
-    assert messages[3] == compressed_history[2]
-    # Hashes in the compressed history still drive retrieve-tool injection.
-    assert has_headroom_retrieve_tool(result.get("tools") or [])
+    assert messages[2] == AGENTIC_MESSAGES[2]
+    assert messages[3] == compressed_history[1]
 
 
 # ---------------------------------------------------------------------------

tests/test_litellm/proxy/guardrails/test_content_utils.py

@@ -4,6 +4,8 @@ from litellm.proxy.guardrails._content_utils import (
     apply_redacted_messages_back,
     build_inspection_messages,
     has_non_string_content,
+    is_non_conversational_call_type,
+    is_string_batch_input,
     iter_message_text,
     walk_user_text,
 )
@@ -580,6 +582,95 @@ def test_apply_redacted_messages_back_skips_input_when_not_string():
     assert data["input"] == [{"type": "text", "text": "leak"}]
 
 
+def test_apply_redacted_messages_back_rewrites_string_batches():
+    """An /embeddings batch is a list of plain strings; each is rewritten in place
+    from the matching redacted message so no element reaches the LLM unredacted."""
+    data = {"input": ["first SSN", "second SSN"]}
+    apply_redacted_messages_back(
+        data,
+        [
+            {"role": "user", "content": "first [REDACTED]"},
+            {"role": "user", "content": "second [REDACTED]"},
+        ],
+    )
+    assert data["input"] == ["first [REDACTED]", "second [REDACTED]"]
+
+
+def test_apply_redacted_messages_back_keeps_batch_elements_aligned():
+    """A guardrail that redacts a whole element away returns it as empty text.
+    Each element still has to take its own redaction, never the next one's."""
+    data = {"input": ["all secret", "second doc", "third doc"]}
+    apply_redacted_messages_back(
+        data,
+        [
+            {"role": "user", "content": ""},
+            {"role": "user", "content": "second doc"},
+            {"role": "user", "content": "third doc"},
+        ],
+    )
+    assert data["input"] == ["", "second doc", "third doc"]
+
+
+def test_apply_redacted_messages_back_skips_empty_batch_elements():
+    """Empty elements are never sent to the guardrail, so the redactions line up
+    with the elements that were."""
+    data = {"input": ["", "secret doc"]}
+    assert apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED] doc"}]) is True
+    assert data["input"] == ["", "[REDACTED] doc"]
+
+
+def test_apply_redacted_messages_back_rejects_short_batch_response():
+    """A guardrail that returns fewer messages than were inspected cannot be
+    applied element-wise: writing the prefix would forward the rest of the batch
+    unredacted, so nothing is written and the caller has to block."""
+    data = {"input": ["first SSN", "second SSN", "third SSN"]}
+    assert apply_redacted_messages_back(data, [{"role": "user", "content": "first [REDACTED]"}]) is False
+    assert data["input"] == ["first SSN", "second SSN", "third SSN"]
+
+
+def test_apply_redacted_messages_back_rejects_long_batch_response():
+    """More redactions than inspected elements means the alignment is unknown."""
+    data = {"input": ["only SSN"]}
+    assert (
+        apply_redacted_messages_back(
+            data,
+            [
+                {"role": "user", "content": "only [REDACTED]"},
+                {"role": "user", "content": "spurious"},
+            ],
+        )
+        is False
+    )
+    assert data["input"] == ["only SSN"]
+
+
+def test_apply_redacted_messages_back_rejects_batch_content_missing():
+    """A message without content cannot safely replace the original batch element."""
+    data = {"input": ["secret doc"]}
+    assert apply_redacted_messages_back(data, [{"role": "user"}]) is False
+    assert data["input"] == ["secret doc"]
+
+
+def test_apply_redacted_messages_back_returns_true_for_non_batch_shapes():
+    data = {"messages": [{"role": "user", "content": "secret"}]}
+    assert apply_redacted_messages_back(data, [{"role": "user", "content": "[REDACTED]"}]) is True
+
+
+# ── is_string_batch_input ─────────────────────────────────────────────────────
+
+
+def test_is_string_batch_input_embeddings_batch():
+    assert is_string_batch_input({"input": ["a", "b"]}) is True
+
+
+def test_is_string_batch_input_rejects_other_shapes():
+    assert is_string_batch_input({"input": "a"}) is False
+    assert is_string_batch_input({"input": []}) is False
+    assert is_string_batch_input({"input": [1, 2]}) is False
+    assert is_string_batch_input({"input": ["a", {"type": "text", "text": "b"}]}) is False
+    assert is_string_batch_input({"messages": [], "input": ["a"]}) is False
+
+
 # -------------------------------------------------------------------
 # LIT-4302: custom_tool_call_output walking
 # -------------------------------------------------------------------
@@ -617,3 +708,36 @@ def test_build_inspection_messages_custom_tool_call_output():
     }
     msgs = build_inspection_messages(data)
     assert any("custom-tool-leak" in m["content"] for m in msgs)
+
+
+# ── is_non_conversational_call_type ──────────────────────────────────────────────
+
+
+def test_is_non_conversational_call_type_flags_embeddings():
+    """An /embeddings body carries documents being indexed, not a prompt."""
+    assert is_non_conversational_call_type("embedding") is True
+    assert is_non_conversational_call_type("aembedding") is True
+
+
+def test_is_non_conversational_call_type_passes_every_conversational_call_type():
+    """Deliberately a deny-list: ``anthropic_messages``, ``responses`` and
+    ``call_mcp_tool`` carry conversations but are absent from
+    ``TEXT_CONTENT_CALL_TYPES``, so a guardrail gating on that allow-list would
+    stop inspecting them."""
+    for call_type in (
+        "completion",
+        "acompletion",
+        "text_completion",
+        "responses",
+        "aresponses",
+        "anthropic_messages",
+        "aanthropic_messages",
+        "call_mcp_tool",
+    ):
+        assert is_non_conversational_call_type(call_type) is False
+
+
+def test_is_non_conversational_call_type_defaults_to_inspecting_unknown_call_types():
+    """A call type this module has never heard of must still be inspected —
+    failing closed is the point of the deny-list."""
+    assert is_non_conversational_call_type("some_future_call_type") is False

tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py

@@ -670,6 +670,18 @@ def test_get_provider_specific_params():
     )  # Literal type should be select
 
 
+@pytest.mark.asyncio
+async def test_provider_specific_params_includes_embedding_toggle():
+    from litellm.proxy.guardrails.guardrail_endpoints import get_provider_specific_params
+
+    provider_params = await get_provider_specific_params()
+
+    for provider in ("aim", "cato_networks"):
+        field = provider_params[provider]["inspect_embeddings"]
+        assert field["type"] == "bool"
+        assert field["default_value"] is False
+
+
 @pytest.mark.asyncio
 async def test_provider_specific_params_includes_hide_secrets():
     """hide-secrets lives in the enterprise package so it is not in

tests/test_litellm/router_strategy/adaptive_router/test_adaptive_router.py

@@ -269,7 +269,9 @@ async def test_record_turn_bounds_feedback_contexts_and_evicts_least_recent_sess
 
 
 @pytest.mark.asyncio
-async def test_load_state_from_db_overrides_cold_start():
+async def test_load_state_from_db_adds_the_persisted_delta_to_the_cold_start_prior():
+    """A row holds an accumulated delta, not a full posterior; loading must add it to the
+    cold-start prior, not replace the cell outright."""
     r = _make_router()
     cold = r._cells[(RequestType.GENERAL, "fast")]
 
@@ -284,14 +286,38 @@ async def test_load_state_from_db_overrides_cold_start():
     await r.load_state_from_db(prisma)
 
     new_cell = r._cells[(RequestType.GENERAL, "fast")]
-    assert (new_cell.alpha, new_cell.beta) == (42.0, 13.0)
-    assert (new_cell.alpha, new_cell.beta) != (cold.alpha, cold.beta)
+    assert (new_cell.alpha, new_cell.beta) == (cold.alpha + 42.0, cold.beta + 13.0)
+
+
+@pytest.mark.asyncio
+async def test_load_state_from_db_keeps_a_one_sided_delta_row_sampleable():
+    """A cell whose only DB activity is one signal type persists a one-sided row (e.g.
+    beta=0.0); loading it must not zero out a Beta shape parameter and crash thompson_sample()."""
+    from litellm.router_strategy.adaptive_router.bandit import thompson_sample
+
+    r = _make_router()
+
+    one_sided_row = MagicMock()
+    one_sided_row.request_type = "general"
+    one_sided_row.model_name = "fast"
+    one_sided_row.alpha = 1.0
+    one_sided_row.beta = 0.0
+
+    prisma = MagicMock()
+    prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[one_sided_row])
+    await r.load_state_from_db(prisma)
+
+    loaded_cell = r._cells[(RequestType.GENERAL, "fast")]
+    assert loaded_cell.alpha > 0.0
+    assert loaded_cell.beta > 0.0
+    thompson_sample(loaded_cell)  # must not raise
 
 
 @pytest.mark.asyncio
 async def test_load_state_from_db_handles_unknown_request_type():
     r = _make_router()
-    cold = r._cells[(RequestType.GENERAL, "fast")]
+    cold_general = r._cells[(RequestType.GENERAL, "fast")]
+    cold_writing = r._cells[(RequestType.WRITING, "fast")]
 
     bad_row = MagicMock()
     bad_row.request_type = "nonexistent_type_v999"
@@ -309,10 +335,11 @@ async def test_load_state_from_db_handles_unknown_request_type():
     prisma.db.litellm_adaptiverouterstate.find_many = AsyncMock(return_value=[bad_row, good_row])
     await r.load_state_from_db(prisma)
 
-    # Unknown skipped; good applied.
-    assert r._cells[(RequestType.GENERAL, "fast")].alpha == 7.0
-    # Other request types kept their cold-start values.
-    assert r._cells[(RequestType.WRITING, "fast")] == cold or True
+    # Unknown skipped; good added to the cold-start prior.
+    new_general = r._cells[(RequestType.GENERAL, "fast")]
+    assert new_general.alpha == cold_general.alpha + 7.0
+    # Other request types kept their own cold-start values.
+    assert r._cells[(RequestType.WRITING, "fast")] == cold_writing
 
 
 # ---- Session state eviction ---------------------------------------------

tests/test_litellm/router_strategy/adaptive_router/test_e2e_adaptive_router.py

@@ -186,8 +186,12 @@ async def test_failure_signal_increments_beta_after_flush():
 
 
 @pytest.mark.asyncio
-async def test_load_state_from_db_overrides_cold_start():
+async def test_load_state_from_db_adds_persisted_delta_to_cold_start():
+    """A row holds an accumulated delta, not a full posterior; loading must add it to the
+    cold-start prior, not replace the cell outright."""
     router = _make_router()
+    cold = router._cells[(RequestType.GENERAL, "gpt-4o")]
+
     fake_row = MagicMock()
     fake_row.request_type = RequestType.GENERAL.value
     fake_row.model_name = "gpt-4o"
@@ -200,8 +204,8 @@ async def test_load_state_from_db_overrides_cold_start():
     await router.load_state_from_db(prisma)
 
     cell = router._cells[(RequestType.GENERAL, "gpt-4o")]
-    assert cell.alpha == 90.0
-    assert cell.beta == 10.0
+    assert cell.alpha == cold.alpha + 90.0
+    assert cell.beta == cold.beta + 10.0
 
 
 @pytest.mark.asyncio

tests/test_litellm/router_strategy/adaptive_router/test_router_dispatch.py

@@ -133,6 +133,102 @@ def test_init_adaptive_router_reads_cost_from_litellm_params():
     }
 
 
+def test_init_adaptive_router_falls_back_to_model_info_cost():
+    """Custom pricing declared under model_info (the conventional location everywhere else in
+    LiteLLM: cost_calculator.py, add_deployment's litellm.model_cost registration) must still
+    feed cost-weighted routing, not silently zero it out."""
+    r = Router(
+        model_list=[
+            {
+                "model_name": "smart-cheap-router",
+                "litellm_params": {
+                    "model": "auto_router/adaptive_router",
+                    "adaptive_router_config": {
+                        "available_models": ["fast", "smart"],
+                    },
+                },
+            },
+            {
+                "model_name": "fast",
+                "litellm_params": {"model": "openai/gpt-4o-mini"},
+                "model_info": {"input_cost_per_token": 0.00000015},
+            },
+            {
+                "model_name": "smart",
+                "litellm_params": {"model": "openai/gpt-4o"},
+                "model_info": {"input_cost_per_token": 0.0000050},
+            },
+        ]
+    )
+    assert _adaptive(r, "smart-cheap-router").model_to_cost == {
+        "fast": 0.00000015,
+        "smart": 0.0000050,
+    }
+
+
+@pytest.mark.asyncio
+async def test_pick_model_favors_the_cheaper_model_info_priced_deployment():
+    """Same fix, exercised through pick_model's actual scoring rather than the model_to_cost
+    dict alone: with cost as the only weight and equal quality priors, the cheaper deployment
+    must win every draw. `smart` (expensive) is listed first deliberately: before the fix both
+    models silently cost 0.0, tying every score, and pick_best's insertion-order tie-break would
+    hand every request to the first-listed (expensive) model instead."""
+    r = Router(
+        model_list=[
+            {
+                "model_name": "smart-cheap-router",
+                "litellm_params": {
+                    "model": "auto_router/adaptive_router",
+                    "adaptive_router_config": {
+                        "available_models": ["smart", "fast"],
+                        "weights": {"quality": 0.0, "cost": 1.0},
+                    },
+                },
+            },
+            {
+                "model_name": "smart",
+                "litellm_params": {"model": "openai/gpt-4o"},
+                "model_info": {"input_cost_per_token": 0.0000050},
+            },
+            {
+                "model_name": "fast",
+                "litellm_params": {"model": "openai/gpt-4o-mini"},
+                "model_info": {"input_cost_per_token": 0.00000015},
+            },
+        ]
+    )
+    adaptive = _adaptive(r, "smart-cheap-router")
+
+    picks = [await adaptive.pick_model(RequestType.GENERAL) for _ in range(10)]
+
+    assert picks == ["fast"] * 10
+
+
+def test_init_adaptive_router_prefers_litellm_params_cost_over_model_info():
+    r = Router(
+        model_list=[
+            {
+                "model_name": "smart-cheap-router",
+                "litellm_params": {
+                    "model": "auto_router/adaptive_router",
+                    "adaptive_router_config": {
+                        "available_models": ["fast"],
+                    },
+                },
+            },
+            {
+                "model_name": "fast",
+                "litellm_params": {
+                    "model": "openai/gpt-4o-mini",
+                    "input_cost_per_token": 0.00000015,
+                },
+                "model_info": {"input_cost_per_token": 0.0000050},
+            },
+        ]
+    )
+    assert _adaptive(r, "smart-cheap-router").model_to_cost == {"fast": 0.00000015}
+
+
 # ---- Fix 4: pre-routing dispatch ---------------------------------------
 
 

tests/test_litellm/router_strategy/test_auto_router.py

@@ -316,6 +316,11 @@ class TestAutoRouter:
 
 semantic_router = pytest.importorskip("semantic_router", reason="auto-router needs the semantic-router extra")
 
+# SemanticRouter(auto_sync="local") calls the encoder's sync embedding path twice per build:
+# once to probe the encoder's output dimension, once to embed ROUTER_CONFIG's one route's
+# utterances.
+_EMBEDDING_CALLS_PER_ROUTELAYER_BUILD: Final = 2
+
 ROUTER_CONFIG: Final = json.dumps(
     {
         "routes": [
@@ -604,3 +609,164 @@ class TestAutoRouterAttributesItsEmbeddingSpend:
         assert router.aembedding_kwargs["proxy_server_request"] == {
             "body": {"model": "text-embedding-3-small", "input": ["fix this stack trace"]}
         }
+
+
+class ThreadTrackingEmbeddingRouter(StubEmbeddingRouter):
+    """Records which OS thread and how many times `embedding()` was called during a build."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.embedding_call_threads: list[int] = []
+
+    def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any:
+        import threading
+
+        self.embedding_call_threads.append(threading.get_ident())
+        return super().embedding(input, model, **kwargs)
+
+
+class TestAutoRouterColdStartDoesNotBlockTheEventLoop:
+    """The first request through a fresh alias builds the route layer off the event loop thread,
+    and concurrent first requests build it exactly once."""
+
+    @pytest.mark.asyncio
+    async def test_should_build_the_routelayer_on_a_worker_thread_not_the_event_loop_thread(self):
+        import threading
+
+        embedding_router: Final = ThreadTrackingEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router)
+        event_loop_thread: Final = threading.get_ident()
+
+        result: Final = await auto_router.async_pre_routing_hook(
+            model="my-auto-router",
+            request_kwargs={},
+            messages=[{"role": "user", "content": "fix this stack trace"}],
+        )
+
+        assert result is not None
+        assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD
+        assert set(embedding_router.embedding_call_threads) == {embedding_router.embedding_call_threads[0]}
+        assert embedding_router.embedding_call_threads[0] != event_loop_thread
+
+    @pytest.mark.asyncio
+    async def test_should_build_the_routelayer_exactly_once_under_concurrent_cold_start_requests(self):
+        embedding_router: Final = ThreadTrackingEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router)
+
+        results: Final = await asyncio.gather(
+            *(
+                auto_router.async_pre_routing_hook(
+                    model="my-auto-router",
+                    request_kwargs={},
+                    messages=[{"role": "user", "content": "fix this stack trace"}],
+                )
+                for _ in range(10)
+            )
+        )
+
+        assert all(result is not None for result in results)
+        assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD
+
+    @pytest.mark.asyncio
+    async def test_should_not_duplicate_the_build_when_a_caller_is_cancelled_mid_build(self):
+        """A caller arriving while the first is cancelled mid-build must reuse it, not duplicate it."""
+        import threading
+
+        class BlockingEmbeddingRouter(ThreadTrackingEmbeddingRouter):
+            def __init__(self) -> None:
+                super().__init__()
+                self.started = threading.Event()
+                self.release = threading.Event()
+
+            def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any:
+                self.started.set()
+                self.release.wait(timeout=5)
+                return super().embedding(input, model, **kwargs)
+
+        embedding_router: Final = BlockingEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router)
+
+        first_call: Final = asyncio.ensure_future(
+            auto_router.async_pre_routing_hook(
+                model="my-auto-router",
+                request_kwargs={},
+                messages=[{"role": "user", "content": "fix this stack trace"}],
+            )
+        )
+        while not embedding_router.started.is_set():
+            await asyncio.sleep(0.01)
+
+        first_call.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await first_call
+
+        second_call: Final = asyncio.ensure_future(
+            auto_router.async_pre_routing_hook(
+                model="my-auto-router",
+                request_kwargs={},
+                messages=[{"role": "user", "content": "fix this stack trace"}],
+            )
+        )
+        await asyncio.sleep(0.01)  # let the second call observe the still-running build
+        embedding_router.release.set()
+        result: Final = await second_call
+
+        assert result is not None
+        assert len(embedding_router.embedding_call_threads) == _EMBEDDING_CALLS_PER_ROUTELAYER_BUILD
+
+    @pytest.mark.asyncio
+    async def test_should_clear_a_failed_build_even_with_no_caller_left_to_observe_it(self):
+        """A build failing after its only caller was cancelled must still clear, not stay cached."""
+        import threading
+
+        class FailsOnFirstAttemptEmbeddingRouter(ThreadTrackingEmbeddingRouter):
+            def __init__(self) -> None:
+                super().__init__()
+                self.started = threading.Event()
+                self.release = threading.Event()
+                self.attempts = 0
+
+            def embedding(self, input: list[str], model: str, **kwargs: Any) -> Any:
+                self.attempts += 1
+                attempt = self.attempts
+                self.started.set()
+                self.release.wait(timeout=5)
+                if attempt == 1:
+                    raise ValueError("boom")
+                return super().embedding(input, model, **kwargs)
+
+        embedding_router: Final = FailsOnFirstAttemptEmbeddingRouter()
+        auto_router: Final = _auto_router(None, litellm_router_instance=embedding_router)
+
+        first_call: Final = asyncio.ensure_future(
+            auto_router.async_pre_routing_hook(
+                model="my-auto-router",
+                request_kwargs={},
+                messages=[{"role": "user", "content": "fix this stack trace"}],
+            )
+        )
+        while not embedding_router.started.is_set():
+            await asyncio.sleep(0.01)
+        first_call.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await first_call
+
+        # Nobody awaits the build now. Let the first attempt fail on its own.
+        embedding_router.release.set()
+        build_task = auto_router._routelayer_build_task
+        assert build_task is not None
+        while not build_task.done():
+            await asyncio.sleep(0.01)
+        await asyncio.sleep(0.01)  # let the done-callback (scheduled via call_soon) run
+
+        assert auto_router._routelayer_build_task is None
+
+        embedding_router.started.clear()
+        embedding_router.release.clear()
+        result: Final = await auto_router.async_pre_routing_hook(
+            model="my-auto-router",
+            request_kwargs={},
+            messages=[{"role": "user", "content": "fix this stack trace"}],
+        )
+
+        assert result is not None

tests/test_litellm/router_strategy/test_complexity_router.py

@@ -1512,6 +1512,110 @@ class TestRouterComplexityDeploymentMethods:
         assert adaptive.model_to_prefs["cheap"].quality_tier == 1
         assert adaptive.model_to_prefs["premium"].quality_tier == 3
 
+    def test_hybrid_adaptive_router_falls_back_to_model_info_cost(self):
+        """Custom pricing declared under model_info (the conventional location everywhere else
+        in LiteLLM) must still feed the hybrid adaptive router's cost-weighted scoring, not
+        silently cost the deployment at 0.0."""
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "hybrid",
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_default_model": "cheap",
+                        "complexity_router_config": {
+                            "adaptive": True,
+                            "tiers": {"SIMPLE": ["cheap"], "MEDIUM": ["cheap", "premium"]},
+                        },
+                    },
+                },
+                {
+                    "model_name": "cheap",
+                    "litellm_params": {"model": "openai/gpt-4o-mini"},
+                    "model_info": {"input_cost_per_token": 0.00000015},
+                },
+                {
+                    "model_name": "premium",
+                    "litellm_params": {"model": "openai/gpt-4o"},
+                    "model_info": {"input_cost_per_token": 0.000005},
+                },
+            ]
+        )
+
+        adaptive = router.adaptive_routers["hybrid"][0].strategy
+        assert adaptive.model_to_cost == {
+            "cheap": pytest.approx(0.00000015),
+            "premium": pytest.approx(0.000005),
+        }
+
+    @pytest.mark.asyncio
+    async def test_hybrid_adaptive_router_pick_model_favors_the_cheaper_model_info_priced_deployment(self):
+        """Same fix, exercised through pick_model's actual scoring rather than the model_to_cost
+        dict alone. `premium` is listed first (SIMPLE tier) deliberately: before the fix both
+        models silently cost 0.0, tying every score, and pick_best's insertion-order tie-break
+        would hand every request to the first-listed (expensive) model instead."""
+        from litellm.types.router import RequestType
+
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "hybrid",
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_default_model": "cheap",
+                        "complexity_router_config": {
+                            "adaptive": True,
+                            "adaptive_weights": {"quality": 0.0, "cost": 1.0},
+                            "tiers": {"SIMPLE": ["premium"], "MEDIUM": ["premium", "cheap"]},
+                        },
+                    },
+                },
+                {
+                    "model_name": "premium",
+                    "litellm_params": {"model": "openai/gpt-4o"},
+                    "model_info": {"input_cost_per_token": 0.000005},
+                },
+                {
+                    "model_name": "cheap",
+                    "litellm_params": {"model": "openai/gpt-4o-mini"},
+                    "model_info": {"input_cost_per_token": 0.00000015},
+                },
+            ]
+        )
+        adaptive = router.adaptive_routers["hybrid"][0].strategy
+
+        picks = [await adaptive.pick_model(RequestType.GENERAL) for _ in range(10)]
+
+        assert picks == ["cheap"] * 10
+
+    def test_hybrid_adaptive_router_prefers_litellm_params_cost_over_model_info(self):
+        router = Router(
+            model_list=[
+                {
+                    "model_name": "hybrid",
+                    "litellm_params": {
+                        "model": "auto_router/complexity_router",
+                        "complexity_router_default_model": "cheap",
+                        "complexity_router_config": {
+                            "adaptive": True,
+                            "tiers": {"SIMPLE": ["cheap"]},
+                        },
+                    },
+                },
+                {
+                    "model_name": "cheap",
+                    "litellm_params": {
+                        "model": "openai/gpt-4o-mini",
+                        "input_cost_per_token": 0.00000015,
+                    },
+                    "model_info": {"input_cost_per_token": 0.000005},
+                },
+            ]
+        )
+
+        adaptive = router.adaptive_routers["hybrid"][0].strategy
+        assert adaptive.model_to_cost == {"cheap": pytest.approx(0.00000015)}
+
 
 class TestComplexityRouterTagBasedRouting:
     """Regression tests for https://github.com/BerriAI/litellm/issues/33655.

tests/test_litellm/test_bedrock_usgov_pricing.py

@@ -132,6 +132,20 @@ CLAUDE_GOV_EXPECTED = {
         "cache_creation_input_token_cost_above_1hr": 1.2e-05,
         "cache_read_input_token_cost": 6e-07,
     },
+    "anthropic.claude-opus-5": {
+        "input_cost_per_token": 6e-06,
+        "output_cost_per_token": 3e-05,
+        "cache_creation_input_token_cost": 7.5e-06,
+        "cache_creation_input_token_cost_above_1hr": 1.2e-05,
+        "cache_read_input_token_cost": 6e-07,
+    },
+    "anthropic.claude-fable-5-1": {
+        "input_cost_per_token": 1.2e-05,
+        "output_cost_per_token": 6e-05,
+        "cache_creation_input_token_cost": 1.5e-05,
+        "cache_creation_input_token_cost_above_1hr": 2.4e-05,
+        "cache_read_input_token_cost": 3e-07,
+    },
 }
 
 
@@ -144,15 +158,18 @@ USGOV_CLAUDE_KEY_TEMPLATES = {
 
 @pytest.mark.parametrize("base_key", CLAUDE_GOV_EXPECTED)
 @pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
-def test_usgov_claude_sonnet5_opus48_pricing(model_data, key_template, expected_provider, base_key):
-    """Sonnet 5 and Opus 4.8 gov entries, both in-region keys and the us-gov.
-    geo inference profile the model cards list for GovCloud, must match the
-    rates AWS publishes on the Bedrock pricing page (1.2x global).
+def test_usgov_claude_pricing(model_data, key_template, expected_provider, base_key):
+    """Sonnet 5, Opus 4.8, Opus 5, and Fable 5.1 gov entries, both in-region keys
+    and the us-gov. geo inference profile the model cards list for GovCloud, must
+    carry the 1.2x GovCloud premium over the global anthropic.* rates. No public
+    AWS source (offer files, pricing page) lists Claude GovCloud rows; the premium
+    is the one AWS quotes for Opus 4.8 in GovCloud ($6/$30 per million).
     """
     gov_key = key_template.format(base_key=base_key)
     assert gov_key in model_data, f"Missing model entry: {gov_key}"
     info = model_data[gov_key]
     assert info["litellm_provider"] == expected_provider
+    assert "search_context_cost_per_query" not in info
     for field, expected in CLAUDE_GOV_EXPECTED[base_key].items():
         assert info[field] == expected, f"{gov_key}: {field} should be {expected} (got {info[field]})"
         ratio = info[field] / model_data[base_key][field]
@@ -161,6 +178,7 @@ def test_usgov_claude_sonnet5_opus48_pricing(model_data, key_template, expected_
 
 CONVERSE_GOV_EXPECTED = {
     "nvidia.nemotron-nano-3-30b": (7.2e-08, 2.88e-07),
+    "nvidia.nemotron-nano-9b-v2": (7.2e-08, 2.76e-07),
     "nvidia.nemotron-nano-12b-v2": (2.4e-07, 7.2e-07),
     "nvidia.nemotron-super-3-120b": (1.8e-07, 7.8e-07),
     "openai.gpt-oss-20b-1:0": (8.4e-08, 3.6e-07),
@@ -169,18 +187,19 @@ CONVERSE_GOV_EXPECTED = {
 
 
 @pytest.mark.parametrize("base_key", CONVERSE_GOV_EXPECTED)
-@pytest.mark.parametrize("region", ["us-gov-east-1", "us-gov-west-1"])
-def test_usgov_converse_model_pricing(model_data, region, base_key):
-    """Nemotron and gpt-oss gov entries must match the AWS Bedrock offer file,
-    which prices both GovCloud regions identically at 1.2x commercial.
+@pytest.mark.parametrize("key_template,expected_provider", USGOV_CLAUDE_KEY_TEMPLATES.items())
+def test_usgov_converse_model_pricing(model_data, key_template, expected_provider, base_key):
+    """Nemotron and gpt-oss gov entries, in-region and the us-gov. geo inference
+    profile both GovCloud regions list as ACTIVE, must match the AWS Bedrock
+    offer file, which prices both regions identically at 1.2x commercial.
     """
-    gov_key = f"bedrock/{region}/{base_key}"
+    gov_key = key_template.format(base_key=base_key)
     assert gov_key in model_data, f"Missing model entry: {gov_key}"
     info = model_data[gov_key]
     expected_input, expected_output = CONVERSE_GOV_EXPECTED[base_key]
     assert info["input_cost_per_token"] == expected_input
     assert info["output_cost_per_token"] == expected_output
-    assert info["litellm_provider"] == "bedrock"
+    assert info["litellm_provider"] == expected_provider
     base = model_data[base_key]
     assert abs(info["input_cost_per_token"] / base["input_cost_per_token"] - 1.2) < 1e-9
     assert abs(info["output_cost_per_token"] / base["output_cost_per_token"] - 1.2) < 1e-9
@@ -259,6 +278,147 @@ def test_usgov_mantle_grok_4_3_west_only(model_data):
     assert "bedrock_mantle/us-gov-east-1/xai.grok-4.3" not in model_data
 
 
+def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data):
+    """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile
+    only, so the profile row must bill exactly like the in-region gov row.
+    """
+    profile = model_data["us-gov.anthropic.claude-3-haiku-20240307-v1:0"]
+    in_region = model_data["bedrock/us-gov-east-1/anthropic.claude-3-haiku-20240307-v1:0"]
+    assert profile["litellm_provider"] == "bedrock_converse"
+    assert {k: v for k, v in profile.items() if k != "litellm_provider"} == {
+        k: v for k, v in in_region.items() if k != "litellm_provider"
+    }
+
+
+GROK_4_6_GOV_KEYS = {
+    "us-gov.xai.grok-4.6": ("us.xai.grok-4.6", "bedrock_converse"),
+    "bedrock_mantle/us-gov-west-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
+    "bedrock_mantle/us-gov-east-1/xai.grok-4.6": ("bedrock_mantle/xai.grok-4.6", "bedrock_mantle"),
+}
+
+
+@pytest.mark.parametrize("gov_key", GROK_4_6_GOV_KEYS)
+def test_usgov_grok_4_6_pricing(model_data, gov_key):
+    """Both GovCloud regions serve grok-4.6 through the us-gov. profile only, and
+    both offer files price its standard SKU at 1.2x the commercial US rate.
+    """
+    base_key, expected_provider = GROK_4_6_GOV_KEYS[gov_key]
+    assert gov_key in model_data, f"Missing model entry: {gov_key}"
+    info = model_data[gov_key]
+    assert info["litellm_provider"] == expected_provider
+    assert info["input_cost_per_token"] == 2.64e-06
+    assert info["output_cost_per_token"] == 7.92e-06
+    assert info["cache_read_input_token_cost"] == 6.6e-07
+    for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost"):
+        assert abs(info[field] / model_data[base_key][field] - 1.2) < 1e-9
+
+
+NOVA_GOV_WEST_EXPECTED = {
+    "amazon.nova-lite-v1:0": (7.2e-08, 2.88e-07),
+    "amazon.nova-micro-v1:0": (4.2e-08, 1.68e-07),
+}
+
+
+@pytest.mark.parametrize("base_key", NOVA_GOV_WEST_EXPECTED)
+def test_usgov_west_nova_lite_micro_pricing(model_data, base_key):
+    """Nova Lite and Micro are on-demand in us-gov-west-1 only; the offer file
+    prices them at 1.2x commercial, like the Nova Pro row that was already there.
+    """
+    gov_key = f"bedrock/us-gov-west-1/{base_key}"
+    assert gov_key in model_data, f"Missing model entry: {gov_key}"
+    info = model_data[gov_key]
+    expected_input, expected_output = NOVA_GOV_WEST_EXPECTED[base_key]
+    assert info["litellm_provider"] == "bedrock"
+    assert info["input_cost_per_token"] == expected_input
+    assert info["output_cost_per_token"] == expected_output
+    assert abs(info["input_cost_per_token"] / model_data[base_key]["input_cost_per_token"] - 1.2) < 1e-9
+    assert abs(info["output_cost_per_token"] / model_data[base_key]["output_cost_per_token"] - 1.2) < 1e-9
+    assert f"bedrock/us-gov-east-1/{base_key}" not in model_data
+
+
+def test_usgov_west_nova_2_multimodal_embeddings_pricing(model_data):
+    """Every meter of the multimodal embedding model (tokens, images, audio and
+    video seconds) carries the 1.2x uplift the us-gov-west-1 offer file lists.
+    """
+    gov_key = "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0"
+    assert gov_key in model_data, f"Missing model entry: {gov_key}"
+    info = model_data[gov_key]
+    assert info["litellm_provider"] == "bedrock"
+    assert info["mode"] == "embedding"
+    assert info["input_cost_per_token"] == 1.62e-07
+    assert info["input_cost_per_image"] == 7.2e-05
+    assert info["input_cost_per_audio_per_second"] == 0.000168
+    assert info["input_cost_per_video_per_second"] == 0.00084
+    assert "bedrock/us-gov-east-1/amazon.nova-2-multimodal-embeddings-v1:0" not in model_data
+
+
+MANTLE_GOV_FLAT_EXPECTED = {
+    "google.gemma-4-e2b": (4.8e-08, 9.6e-08, ("us-gov-west-1",)),
+    "google.gemma-4-26b-a4b": (1.56e-07, 4.8e-07, ("us-gov-west-1",)),
+    "google.gemma-4-31b": (1.68e-07, 4.8e-07, ("us-gov-west-1",)),
+    "openai.gpt-oss-20b": (8.4e-08, 3.6e-07, ("us-gov-west-1", "us-gov-east-1")),
+    "openai.gpt-oss-120b": (1.8e-07, 7.2e-07, ("us-gov-west-1", "us-gov-east-1")),
+}
+
+
+@pytest.mark.parametrize("model", MANTLE_GOV_FLAT_EXPECTED)
+def test_usgov_mantle_gemma_and_gpt_oss_pricing(model_data, model):
+    """Gemma 4 is priced in the us-gov-west-1 offer file only and gpt-oss in both;
+    each Mantle gov row carries the offer file's standard SKU, and no row exists
+    for a region whose offer file has no SKU.
+    """
+    expected_input, expected_output, regions = MANTLE_GOV_FLAT_EXPECTED[model]
+    for region in ("us-gov-west-1", "us-gov-east-1"):
+        gov_key = f"bedrock_mantle/{region}/{model}"
+        if region not in regions:
+            assert gov_key not in model_data
+            continue
+        assert gov_key in model_data, f"Missing model entry: {gov_key}"
+        info = model_data[gov_key]
+        assert info["litellm_provider"] == "bedrock_mantle"
+        assert info["input_cost_per_token"] == expected_input
+        assert info["output_cost_per_token"] == expected_output
+
+
+GOV_ROW_SOURCES = {
+    "us-gov.anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
+    "bedrock/us-gov-west-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
+    "bedrock/us-gov-east-1/anthropic.claude-fable-5-1": "anthropic.claude-fable-5-1",
+    "us-gov.nvidia.nemotron-nano-9b-v2": "nvidia.nemotron-nano-9b-v2",
+    "bedrock/us-gov-west-1/nvidia.nemotron-nano-9b-v2": "nvidia.nemotron-nano-9b-v2",
+    "bedrock/us-gov-east-1/nvidia.nemotron-nano-9b-v2": "nvidia.nemotron-nano-9b-v2",
+    "us-gov.xai.grok-4.6": "us.xai.grok-4.6",
+    "bedrock_mantle/us-gov-west-1/xai.grok-4.6": "bedrock_mantle/xai.grok-4.6",
+    "bedrock_mantle/us-gov-east-1/xai.grok-4.6": "bedrock_mantle/xai.grok-4.6",
+    "bedrock/us-gov-west-1/amazon.nova-2-multimodal-embeddings-v1:0": "amazon.nova-2-multimodal-embeddings-v1:0",
+    "bedrock/us-gov-west-1/amazon.nova-lite-v1:0": "amazon.nova-lite-v1:0",
+    "bedrock/us-gov-west-1/amazon.nova-micro-v1:0": "amazon.nova-micro-v1:0",
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-e2b": "bedrock_mantle/google.gemma-4-e2b",
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-26b-a4b": "bedrock_mantle/google.gemma-4-26b-a4b",
+    "bedrock_mantle/us-gov-west-1/google.gemma-4-31b": "bedrock_mantle/google.gemma-4-31b",
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-20b": "bedrock_mantle/openai.gpt-oss-20b",
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-20b": "bedrock_mantle/openai.gpt-oss-20b",
+    "bedrock_mantle/us-gov-west-1/openai.gpt-oss-120b": "bedrock_mantle/openai.gpt-oss-120b",
+    "bedrock_mantle/us-gov-east-1/openai.gpt-oss-120b": "bedrock_mantle/openai.gpt-oss-120b",
+}
+
+
+def _non_pricing_fields(info):
+    return {k: v for k, v in info.items() if "cost" not in k and k not in ("litellm_provider", "source")}
+
+
+@pytest.mark.parametrize("gov_key", GOV_ROW_SOURCES)
+def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key):
+    """A gov row differs from the commercial row it mirrors only in price and
+    provider: context limits, mode, and capability flags stay identical, so a
+    hand-copied row cannot silently drop tool calling or shrink the context window.
+    """
+    gov = model_data[gov_key]
+    assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]])
+    assert "search_context_cost_per_query" not in gov
+    assert "source" not in gov
+
+
 AZURE_GOV_EXPECTED = {
     "azure/us-gov/gpt-5.1": {
         "input_cost_per_token": 1.71875e-06,

ui/litellm-dashboard/eslint-budgets.json

@@ -3,9 +3,9 @@
   "no-console": { "max": 12, "target": 0 },
   "complexity": { "max": 140, "target": 80 },
   "max-depth": { "max": 70, "target": 30 },
-  "local/no-large-inline-object-arg": { "max": 554, "target": 300 },
-  "local/no-long-condition-chain": { "max": 265, "target": 120 },
+  "local/no-large-inline-object-arg": { "max": 551, "target": 300 },
+  "local/no-long-condition-chain": { "max": 196, "target": 120 },
   "testing-library/no-container": { "max": 133, "target": 50 },
-  "testing-library/no-node-access": { "max": 716, "target": 500 },
+  "testing-library/no-node-access": { "max": 707, "target": 500 },
   "testing-library/prefer-screen-queries": { "max": 18, "target": 18 }
 }

ui/litellm-dashboard/eslint-suppressions.json

@@ -1619,14 +1619,6 @@
       "count": 1
     }
   },
-  "src/components/common_components/fetch_teams.tsx": {
-    "local/filename-pascal-case": {
-      "count": 1
-    },
-    "max-params": {
-      "count": 1
-    }
-  },
   "src/components/common_components/simple_table.tsx": {
     "local/filename-pascal-case": {
       "count": 1
@@ -1823,7 +1815,7 @@
       "count": 5
     },
     "no-restricted-syntax": {
-      "count": 152
+      "count": 150
     },
     "prefer-const": {
       "count": 32
@@ -1871,9 +1863,6 @@
   "src/components/per_user_usage.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "react-hooks/set-state-in-effect": {
-      "count": 1
     }
   },
   "src/components/permissions/MCPServerPermissions.tsx": {
@@ -2303,17 +2292,6 @@
       "count": 1
     }
   },
-  "src/components/user_dashboard.tsx": {
-    "local/filename-pascal-case": {
-      "count": 1
-    },
-    "prefer-const": {
-      "count": 1
-    },
-    "react-hooks/set-state-in-effect": {
-      "count": 1
-    }
-  },
   "src/components/vector_store_management/types.tsx": {
     "local/filename-pascal-case": {
       "count": 1

ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.test.tsx

@@ -1,59 +1,90 @@
-import { render } from "@testing-library/react";
-import { describe, it, expect, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
-const { userDashboardSpy } = vi.hoisted(() => ({
-  userDashboardSpy: vi.fn((_props: Record<string, unknown>) => null),
+const { teamListCall, authorizedSession } = vi.hoisted(() => ({
+  teamListCall: vi.fn(() => new Promise(() => {})),
+  authorizedSession: vi.fn(),
 }));
 
-vi.mock("@/components/user_dashboard", () => ({
-  default: (props: Record<string, unknown>) => userDashboardSpy(props),
-}));
+const session = (overrides: { userRole?: string; isViewOnly?: boolean } = {}) => ({
+  isLoading: false,
+  isAuthorized: true,
+  token: "jwt",
+  accessToken: "sk-access",
+  userId: "u-123",
+  userEmail: "admin@example.com",
+  userRole: "Admin",
+  isViewOnly: false,
+  premiumUser: false,
+  disabledPersonalKeyCreation: false,
+  showSSOBanner: false,
+  ...overrides,
+});
 
-// AuthContext is still hydrating: userID has not been populated yet (the regression).
-vi.mock("@/contexts/AuthContext", () => ({
-  useAuth: () => ({
-    userID: null,
-    userRole: "",
-    userEmail: null,
-    accessToken: null,
-    premiumUser: false,
-    setUserRole: vi.fn(),
-    setUserEmail: vi.fn(),
-  }),
-}));
-
-// useAuthorized decodes the cookie synchronously, so identity is already available.
 vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({
-  default: () => ({
-    isLoading: false,
-    isAuthorized: true,
-    token: "jwt",
-    accessToken: "sk-access",
-    userId: "u-123",
-    userEmail: "admin@example.com",
-    userRole: "Admin",
-    premiumUser: false,
-    disabledPersonalKeyCreation: false,
-    showSSOBanner: false,
-  }),
+  default: () => authorizedSession(),
 }));
 
 vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
-  teamListCall: vi.fn(() => new Promise(() => {})),
+  teamListCall,
 }));
 
 vi.mock("next/navigation", () => ({
   useSearchParams: () => new URLSearchParams(""),
 }));
 
+vi.mock("@/components/VirtualKeysPage/VirtualKeysTable", () => ({
+  VirtualKeysTable: ({ headerActions }: { headerActions?: React.ReactNode }) => (
+    <div>
+      {headerActions}
+      <table aria-label="Virtual Keys" />
+    </div>
+  ),
+}));
+
+vi.mock("@/components/organisms/create_key_button", () => ({
+  default: () => <button type="button">Create Key</button>,
+}));
+
 import ApiKeysDashboard from "./ApiKeysDashboard";
 
-describe("ApiKeysDashboard identity source", () => {
-  it("passes the useAuthorized userID through even while AuthContext.userID is still null", () => {
+describe("ApiKeysDashboard", () => {
+  beforeEach(() => {
+    teamListCall.mockClear();
+    authorizedSession.mockReturnValue(session());
+    sessionStorage.clear();
+  });
+
+  it("scopes the team list to the signed-in user for non-admin roles", () => {
+    authorizedSession.mockReturnValue(session({ userRole: "Internal User" }));
     render(<ApiKeysDashboard />);
 
-    expect(userDashboardSpy).toHaveBeenCalled();
-    const props = userDashboardSpy.mock.calls[0][0];
-    expect(props.userID).toBe("u-123");
+    expect(teamListCall).toHaveBeenCalledWith("sk-access", 1, 100, { userID: "u-123" });
+  });
+
+  it("renders the keys table with a Create Key action for roles that can write", () => {
+    render(<ApiKeysDashboard />);
+
+    expect(screen.getByRole("table", { name: "Virtual Keys" })).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: "Create Key" })).toBeInTheDocument();
+  });
+
+  it("hides Create Key for view-only roles", () => {
+    authorizedSession.mockReturnValue(session({ isViewOnly: true }));
+    render(<ApiKeysDashboard />);
+
+    expect(screen.getByRole("table", { name: "Virtual Keys" })).toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: "Create Key" })).not.toBeInTheDocument();
+  });
+
+  it("leaves other pages' session state intact when the tab reloads", () => {
+    sessionStorage.setItem("chatHistory", '[{"role":"user","content":"hi"}]');
+    sessionStorage.setItem("selectedModel", "gpt-5.5");
+    render(<ApiKeysDashboard />);
+
+    window.dispatchEvent(new Event("beforeunload"));
+
+    expect(sessionStorage.getItem("chatHistory")).toBe('[{"role":"user","content":"hi"}]');
+    expect(sessionStorage.getItem("selectedModel")).toBe("gpt-5.5");
   });
 });

ui/litellm-dashboard/src/app/(dashboard)/api-keys/ApiKeysDashboard.tsx

@@ -3,22 +3,17 @@
 import { teamListCall as v2TeamListCall } from "@/app/(dashboard)/hooks/teams/useTeams";
 import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { KeyResponse, Team } from "@/components/key_team_helpers/key_list";
-import { CreateKeyPrefillData } from "@/components/organisms/create_key_button";
-import UserDashboard from "@/components/user_dashboard";
-import { useAuth } from "@/contexts/AuthContext";
+import CreateKey, { CreateKeyPrefillData } from "@/components/organisms/create_key_button";
+import { VirtualKeysTable } from "@/components/VirtualKeysPage/VirtualKeysTable";
 import { useSearchParams } from "next/navigation";
 import { useEffect, useMemo, useState } from "react";
 
 export default function ApiKeysDashboard() {
-  // Identity comes from useAuthorized (synchronous cookie decode) so userID is set whenever the
-  // route is authorized; useAuth only supplies the backfill setters UserDashboard still expects.
-  const { userId: userID, userRole, userEmail, accessToken, premiumUser } = useAuthorized();
-  const { setUserRole, setUserEmail } = useAuth();
+  const { userId: userID, userRole, accessToken, isViewOnly } = useAuthorized();
   const searchParams = useSearchParams()!;
 
   const [teams, setTeams] = useState<Team[] | null>(null);
   const [keys, setKeys] = useState<KeyResponse[] | null>([]);
-  const [createClicked, setCreateClicked] = useState<boolean>(false);
 
   const autoOpenCreate = searchParams.get("create") === "true";
   const prefillData: CreateKeyPrefillData | undefined = useMemo(() => {
@@ -63,7 +58,6 @@ export default function ApiKeysDashboard() {
 
   const addKey = (data: KeyResponse) => {
     setKeys((prevData) => (prevData ? [...prevData, data] : [data]));
-    setCreateClicked((prev) => !prev);
   };
 
   useEffect(() => {
@@ -77,21 +71,21 @@ export default function ApiKeysDashboard() {
   }, [accessToken, userID, userRole]);
 
   return (
-    <UserDashboard
-      userID={userID}
-      userRole={userRole}
-      premiumUser={premiumUser ?? false}
-      teams={teams}
-      keys={keys}
-      setUserRole={setUserRole}
-      userEmail={userEmail}
-      setUserEmail={setUserEmail}
-      setTeams={setTeams}
-      setKeys={setKeys}
-      addKey={addKey}
-      createClicked={createClicked}
-      autoOpenCreate={autoOpenCreate}
-      prefillData={prefillData}
-    />
+    <main className="flex h-full flex-col p-8">
+      <VirtualKeysTable
+        headerActions={
+          isViewOnly ? undefined : (
+            <CreateKey
+              team={null}
+              teams={teams}
+              data={keys}
+              addKey={addKey}
+              autoOpenCreate={autoOpenCreate}
+              prefillData={prefillData}
+            />
+          )
+        }
+      />
+    </main>
   );
 }

ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx

@@ -6,18 +6,11 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useEffect, useState } from "react";
 
 interface SidebarProviderProps {
-  setPage: (page: string) => void;
-  defaultSelectedKey: string;
   sidebarCollapsed: boolean;
   onToggleCollapsed?: () => void;
 }
 
-const SidebarProvider = ({
-  setPage,
-  defaultSelectedKey,
-  sidebarCollapsed,
-  onToggleCollapsed,
-}: SidebarProviderProps) => {
+const SidebarProvider = ({ sidebarCollapsed, onToggleCollapsed }: SidebarProviderProps) => {
   const { accessToken } = useAuthorized();
   const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState<string[] | null>(null);
   const [enableProjectsUI, setEnableProjectsUI] = useState<boolean>(false);
@@ -70,8 +63,6 @@ const SidebarProvider = ({
 
   return (
     <Sidebar
-      setPage={setPage}
-      defaultSelectedKey={defaultSelectedKey}
       collapsed={sidebarCollapsed}
       onToggleCollapsed={onToggleCollapsed}
       enabledPagesInternalUsers={enabledPagesInternalUsers}

ui/litellm-dashboard/src/app/(dashboard)/hooks/proxySettings/useProxySettings.ts

@@ -8,6 +8,8 @@ export interface ProxySettings {
   PROXY_BASE_URL: string;
   PROXY_LOGOUT_URL: string;
   LITELLM_UI_API_DOC_BASE_URL?: string | null;
+  DISABLE_EXPENSIVE_DB_QUERIES?: boolean;
+  NUM_SPEND_LOGS_ROWS?: number;
 }
 
 const EMPTY_PROXY_SETTINGS: ProxySettings = {

ui/litellm-dashboard/src/app/(dashboard)/layout.tsx

@@ -7,12 +7,12 @@ import LoadingScreen from "@/components/common_components/LoadingScreen";
 import { ThemeProvider } from "@/contexts/ThemeContext";
 import { useAuth } from "@/contexts/AuthContext";
 import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider";
-import { useRouter, useSearchParams, usePathname } from "next/navigation";
+import { useRouter, useSearchParams } from "next/navigation";
 import { DebugWarningBanner } from "@/components/DebugWarningBanner";
 import { NoRedisWarningBanner } from "@/components/NoRedisWarningBanner";
 import { LicenseExpiryBanner } from "@/components/LicenseExpiryBanner";
 import { UserBanner } from "@/components/UserBanner";
-import { MIGRATED_PAGES, migratedHref, legacyPageHref, legacyKeyForPathname } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 import { PluginModeProvider, usePluginMode } from "@/contexts/PluginModeContext";
 import { createApiClient } from "@/lib/http/client";
 import { getProxyBaseUrl } from "@/components/networking";
@@ -97,21 +97,12 @@ export function AgentControlPlaneView() {
 }
 
 function DashboardShell({ children }: { children: React.ReactNode }) {
-  const router = useRouter();
-  const searchParams = useSearchParams();
-  const pathname = usePathname();
   const { accessToken } = useAuth();
   const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
   const { mode } = usePluginMode();
 
-  const page = legacyKeyForPathname(pathname) || searchParams.get("page") || "api-keys";
   const isGateway = mode === "ai-gateway";
 
-  const navigateToPage = (newPage: string) => {
-    const migratedRoute = MIGRATED_PAGES[newPage];
-    router.push(migratedRoute ? migratedHref(migratedRoute) : legacyPageHref(newPage));
-  };
-
   // Non-gateway (agent control plane) mode keeps the original full-width Navbar,
   // which carries the account menu; the redesigned sidebar + header shell is
   // scoped to the ai-gateway dashboard. Chat and the public model hub are
@@ -136,14 +127,9 @@ function DashboardShell({ children }: { children: React.ReactNode }) {
   // so the page can't be dragged past the end of the nav.
   return (
     <div className="flex h-screen overflow-hidden bg-background">
-      <SidebarProvider
-        setPage={navigateToPage}
-        defaultSelectedKey={page}
-        sidebarCollapsed={sidebarCollapsed}
-        onToggleCollapsed={() => setSidebarCollapsed((v) => !v)}
-      />
+      <SidebarProvider sidebarCollapsed={sidebarCollapsed} onToggleCollapsed={() => setSidebarCollapsed((v) => !v)} />
       <div className="flex min-w-0 flex-1 flex-col overflow-hidden">
-        <DashboardHeader page={page} />
+        <DashboardHeader />
         <DebugWarningBanner accessToken={accessToken} />
         <NoRedisWarningBanner accessToken={accessToken} />
         <LicenseExpiryBanner accessToken={accessToken} />
@@ -161,10 +147,10 @@ function LayoutContent({ children }: { children: React.ReactNode }) {
   const isInvitationFlow = Boolean(searchParams.get("invitation_id"));
 
   // Legacy invitation links point at /ui/?invitation_id=; the onboarding form now lives at its own
-  // /onboarding route. Redirect once ui-config has loaded so migratedHref resolves the SERVER_ROOT_PATH base.
+  // /onboarding route. Redirect once ui-config has loaded so uiHref resolves the SERVER_ROOT_PATH base.
   useEffect(() => {
     if (!authLoading && isInvitationFlow) {
-      router.replace(`${migratedHref("onboarding")}?${searchParams.toString()}`);
+      router.replace(`${uiHref("onboarding")}?${searchParams.toString()}`);
     }
   }, [authLoading, isInvitationFlow, router, searchParams]);
 

ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts

@@ -0,0 +1,47 @@
+import { describe, expect, it } from "vitest";
+import { menuGroups } from "@/components/leftnav";
+import { legacyPageRedirectHref } from "./legacyPageRoutes";
+
+const redirect = (query: string) => legacyPageRedirectHref(new URLSearchParams(query));
+
+describe("legacyPageRedirectHref", () => {
+  it("sends an old ?page= bookmark to the path route that replaced it", () => {
+    expect(redirect("page=logs")).toBe("/ui/logs");
+    expect(redirect("page=models")).toBe("/ui/models-and-endpoints");
+    expect(redirect("page=llm-playground")).toBe("/ui/playground");
+    expect(redirect("page=new_usage")).toBe("/ui/usage");
+    expect(redirect("page=usage")).toBe("/ui/old-usage");
+  });
+
+  it("keeps the older aliases for renamed pages", () => {
+    expect(redirect("page=api_ref")).toBe("/ui/api-reference");
+    expect(redirect("page=api-reference")).toBe("/ui/api-reference");
+    expect(redirect("page=claude-code-plugins")).toBe("/ui/skills");
+  });
+
+  it("forwards the remaining query params so the MCP env-var setup link still opens its form", () => {
+    expect(redirect("page=mcp-servers&fill_env_vars=srv-1")).toBe("/ui/mcp-servers?fill_env_vars=srv-1");
+    expect(redirect("fill_env_vars=srv-1&page=mcp-servers")).toBe("/ui/mcp-servers?fill_env_vars=srv-1");
+  });
+
+  it("keeps forwarded values encoded", () => {
+    expect(redirect("page=mcp-servers&fill_env_vars=a%26b%3Dc")).toBe("/ui/mcp-servers?fill_env_vars=a%26b%3Dc");
+  });
+
+  it("returns null when there is no page param or the id is unknown", () => {
+    expect(redirect("")).toBeNull();
+    expect(redirect("login=success")).toBeNull();
+    expect(redirect("page=does-not-exist")).toBeNull();
+    expect(redirect("page=constructor")).toBeNull();
+  });
+
+  it("covers every sidebar page id with the route the sidebar itself links to", () => {
+    const leaves = menuGroups
+      .flatMap((group) => group.items.flatMap((item) => item.children ?? [item]))
+      .filter((item) => !item.external_url);
+    expect(leaves.length).toBeGreaterThan(30);
+    for (const leaf of leaves) {
+      expect(redirect(`page=${leaf.page}`), leaf.page).toBe(`/ui/${leaf.route ?? leaf.page}`);
+    }
+  });
+});

ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts

@@ -0,0 +1,54 @@
+import { uiHref } from "@/utils/uiHref";
+
+const LEGACY_PAGE_ROUTES: ReadonlyMap<string, string> = new Map(
+  Object.entries({
+    "api-keys": "api-keys",
+    models: "models-and-endpoints",
+    api_ref: "api-reference",
+    "api-reference": "api-reference",
+    "llm-playground": "playground",
+    projects: "projects",
+    chat: "chat",
+    "access-groups": "access-groups",
+    budgets: "budgets",
+    workflows: "workflows",
+    "guardrails-monitor": "guardrails-monitor",
+    "mcp-servers": "mcp-servers",
+    "search-tools": "search-tools",
+    "tag-management": "tag-management",
+    "vector-stores": "vector-stores",
+    memory: "memory",
+    policies: "policies",
+    guardrails: "guardrails",
+    prompts: "prompts",
+    "tool-policies": "tool-policies",
+    skills: "skills",
+    "claude-code-plugins": "skills",
+    caching: "caching",
+    "cost-tracking": "cost-tracking",
+    "transform-request": "transform-request",
+    "ui-theme": "ui-theme",
+    logs: "logs",
+    "admin-panel": "admin-panel",
+    "logging-and-alerts": "logging-and-alerts",
+    "model-hub-table": "model-hub-table",
+    new_usage: "usage",
+    usage: "old-usage",
+    "cost-optimization": "cost-optimization",
+    agents: "agents",
+    "router-settings": "router-settings",
+    users: "users",
+    teams: "teams",
+    organizations: "organizations",
+  }),
+);
+
+export function legacyPageRedirectHref(searchParams: URLSearchParams): string | null {
+  const page = searchParams.get("page");
+  const route = page === null ? undefined : LEGACY_PAGE_ROUTES.get(page);
+  if (route === undefined) return null;
+  const rest = new URLSearchParams(searchParams);
+  rest.delete("page");
+  const query = rest.toString();
+  return query ? `${uiHref(route)}?${query}` : uiHref(route);
+}

ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx

@@ -7,7 +7,7 @@ import DeleteResourceModal from "@/components/common_components/DeleteResourceMo
 import ModelSettingsModal from "@/components/model_dashboard/ModelSettingsModal/ModelSettingsModal";
 import { ModelData } from "@/components/model_dashboard/types";
 import { toast } from "@/lib/toast";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 import { modelDeleteCall, modelPatchUpdateCall } from "@/components/networking";
 import { useQueryClient } from "@tanstack/react-query";
 import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
@@ -294,7 +294,7 @@ const AllModelsTab = ({
             {selectedTeamValue === PERSONAL_TEAM_VALUE ? (
               <span>
                 To access these models, create a Virtual Key without selecting a team on the{" "}
-                <a href={migratedHref("api-keys")} className="font-medium text-info hover:underline">
+                <a href={uiHref("api-keys")} className="font-medium text-info hover:underline">
                   Virtual Keys page
                 </a>
                 .
@@ -302,7 +302,7 @@ const AllModelsTab = ({
             ) : (
               <span>
                 To access these models, create a Virtual Key and select Team as &quot;{teamAccessLabel}&quot; on the{" "}
-                <a href={migratedHref("api-keys")} className="font-medium text-info hover:underline">
+                <a href={uiHref("api-keys")} className="font-medium text-info hover:underline">
                   Virtual Keys page
                 </a>
                 .

ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx

@@ -1,7 +1,7 @@
 import React, { useState, useEffect } from "react";
 
 import ViewUserSpend from "@/components/view_user_spend";
-import { ProxySettings } from "@/components/user_dashboard";
+import { ProxySettings } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings";
 import AdvancedDatePicker from "@/components/shared/advanced_date_picker";
 import { Button } from "@/components/ui/button";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";

ui/litellm-dashboard/src/app/(dashboard)/page.test.tsx

@@ -6,9 +6,9 @@ interface KeyRow {
   token: string;
 }
 
-const { mockReplace, mockUseKeys, mockMigratedHref, state } = vi.hoisted(() => {
+const { mockReplace, mockUseKeys, mockUiHref, state } = vi.hoisted(() => {
   const state = {
-    login: "success" as string | null,
+    search: "login=success",
     userRole: "Internal User",
     keys: [] as KeyRow[],
     returnUrl: null as string | null,
@@ -16,7 +16,7 @@ const { mockReplace, mockUseKeys, mockMigratedHref, state } = vi.hoisted(() => {
   return {
     state,
     mockReplace: vi.fn(),
-    mockMigratedHref: vi.fn((segment: string) => `/mocked-ui/${segment}`),
+    mockUiHref: vi.fn((segment: string) => `/mocked-ui/${segment}`),
     mockUseKeys: vi.fn(() => ({
       data: { keys: state.keys, total_count: state.keys.length },
       isLoading: false,
@@ -26,7 +26,7 @@ const { mockReplace, mockUseKeys, mockMigratedHref, state } = vi.hoisted(() => {
 
 vi.mock("next/navigation", () => ({
   useRouter: () => ({ replace: mockReplace }),
-  useSearchParams: () => ({ get: (key: string) => (key === "login" ? state.login : null) }),
+  useSearchParams: () => new URLSearchParams(state.search),
 }));
 vi.mock("@/contexts/AuthContext", () => ({
   useAuth: () => ({
@@ -44,7 +44,7 @@ vi.mock("@/components/common_components/LoadingScreen", () => ({
   default: () => <div data-testid="loading-screen" />,
 }));
 vi.mock("@/components/networking", () => ({ proxyBaseUrl: "" }));
-vi.mock("@/utils/migratedPages", () => ({ MIGRATED_PAGES: {}, migratedHref: mockMigratedHref }));
+vi.mock("@/utils/uiHref", () => ({ uiHref: mockUiHref }));
 vi.mock("@/utils/returnUrlUtils", () => ({
   buildLoginUrlWithReturn: (u: string) => u,
   consumeReturnUrl: () => state.returnUrl,
@@ -71,13 +71,13 @@ describe("dashboard landing", () => {
 
   afterEach(() => {
     Object.defineProperty(window, "location", { configurable: true, value: realLocation });
-    state.login = "success";
+    state.search = "login=success";
     state.userRole = "Internal User";
     state.keys = [];
     state.returnUrl = null;
     mockReplace.mockClear();
     mockUseKeys.mockClear();
-    mockMigratedHref.mockClear();
+    mockUiHref.mockClear();
     mockLocationReplace.mockClear();
   });
 
@@ -89,7 +89,7 @@ describe("dashboard landing", () => {
       expect(screen.getByTestId("api-keys-dashboard")).toBeInTheDocument();
       expect(screen.queryByTestId("loading-screen")).not.toBeInTheDocument();
       expect(mockReplace).not.toHaveBeenCalled();
-      expect(mockMigratedHref).not.toHaveBeenCalledWith("connect");
+      expect(mockUiHref).not.toHaveBeenCalledWith("connect");
     },
   );
 
@@ -105,6 +105,20 @@ describe("dashboard landing", () => {
     expect(mockUseKeys).not.toHaveBeenCalled();
   });
 
+  it("redirects an old ?page= bookmark to its path route without rendering the keys dashboard", () => {
+    state.search = "page=logs";
+    render(<CreateKeyPage />);
+    expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/logs");
+    expect(screen.getByTestId("loading-screen")).toBeInTheDocument();
+    expect(screen.queryByTestId("api-keys-dashboard")).not.toBeInTheDocument();
+  });
+
+  it("carries the MCP env-var deep link's other params through the legacy redirect", () => {
+    state.search = "page=mcp-servers&fill_env_vars=srv-1";
+    render(<CreateKeyPage />);
+    expect(mockReplace).toHaveBeenCalledWith("/mocked-ui/mcp-servers?fill_env_vars=srv-1");
+  });
+
   it("still sends the user to an explicit stored return URL", () => {
     state.returnUrl = "/ui/models-and-endpoints";
     render(<CreateKeyPage />);

ui/litellm-dashboard/src/app/(dashboard)/page.tsx

@@ -12,7 +12,7 @@ import {
   normalizeUrlForCompare,
   storeReturnUrl,
 } from "@/utils/returnUrlUtils";
-import { MIGRATED_PAGES, migratedHref } from "@/utils/migratedPages";
+import { legacyPageRedirectHref } from "@/app/(dashboard)/legacyPageRoutes";
 import { useRouter, useSearchParams } from "next/navigation";
 import { Suspense, useEffect, useRef } from "react";
 
@@ -22,8 +22,6 @@ function CreateKeyPageContent() {
   const router = useRouter();
   const searchParams = useSearchParams()!;
 
-  const explicitPage = searchParams.get("page");
-
   // Track if we've already attempted a return URL redirect to prevent race conditions
   const hasAttemptedReturnRedirectRef = useRef(false);
 
@@ -41,13 +39,12 @@ function CreateKeyPageContent() {
     }
   }, [redirectToLogin]);
 
-  // Redirect legacy ?page= deep links (old bookmarks) to their path-based routes.
-  const isLegacyRedirect = explicitPage !== null && explicitPage in MIGRATED_PAGES;
+  const legacyRedirectHref = legacyPageRedirectHref(searchParams);
   useEffect(() => {
-    if (!authLoading && isLegacyRedirect) {
-      router.replace(migratedHref(MIGRATED_PAGES[explicitPage]));
+    if (!authLoading && legacyRedirectHref !== null) {
+      router.replace(legacyRedirectHref);
     }
-  }, [authLoading, isLegacyRedirect, explicitPage, router]);
+  }, [authLoading, legacyRedirectHref, router]);
 
   // Check for a stored return URL after successful authentication
   // This handles the case where user comes back from SSO and we need to redirect to the original URL
@@ -86,7 +83,7 @@ function CreateKeyPageContent() {
     }
   }, [token]);
 
-  const isRedirecting = redirectToLogin || isLegacyRedirect;
+  const isRedirecting = redirectToLogin || legacyRedirectHref !== null;
 
   if (authLoading || isRedirecting) {
     return <LoadingScreen />;

ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/ChatUI.tsx

@@ -77,6 +77,7 @@ import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover
 import { Select as ShadcnSelect, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
+import { uiHref } from "@/utils/uiHref";
 import {
   AUDIO_ACCEPT,
   IMAGE_EDIT_ACCEPT,
@@ -1650,7 +1651,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                       </TooltipTrigger>
                       <TooltipContent className="max-w-xs">
                         Select vector store(s) to use for this LLM API call. You can set up your vector store{" "}
-                        <a href="?page=vector-stores" className="text-info underline">
+                        <a href={uiHref("vector-stores")} className="text-info underline">
                           here
                         </a>
                         .
@@ -1674,7 +1675,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                       </TooltipTrigger>
                       <TooltipContent className="max-w-xs">
                         Select guardrail(s) to use for this LLM API call. You can set up your guardrails{" "}
-                        <a href="?page=guardrails" className="text-info underline">
+                        <a href={uiHref("guardrails")} className="text-info underline">
                           here
                         </a>
                         .
@@ -1700,7 +1701,7 @@ const ChatUI: React.FC<ChatUIProps> = ({
                         <TooltipContent className="max-w-xs">
                           Select policy/policies to apply to this LLM API call. Policies define which guardrails are
                           applied based on conditions. You can set up your policies{" "}
-                          <a href="?page=policies" className="text-info underline">
+                          <a href={uiHref("policies")} className="text-info underline">
                             here
                           </a>
                           .

ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.test.tsx

@@ -48,7 +48,6 @@ vi.mock("next/navigation", () => ({
   useSearchParams: () => new URLSearchParams(window.location.search),
 }));
 
-// entityLinks -> migratedPages imports serverRootPath from the same module, so the mock must export it too.
 vi.mock("@/components/networking", () => {
   return {
     serverRootPath: "/",

ui/litellm-dashboard/src/app/chat/layout.test.tsx

@@ -2,7 +2,7 @@ import { afterEach, describe, expect, it, vi } from "vitest";
 import { render, screen } from "@testing-library/react";
 import ChatLayout from "./layout";
 
-const { mockUseAuthorized, mockUseUISettings, mockReplace, mockMigratedHref, state } = vi.hoisted(() => {
+const { mockUseAuthorized, mockUseUISettings, mockReplace, mockUiHref, state } = vi.hoisted(() => {
   const state = {
     enableChatUI: false,
     isUISettingsLoading: false,
@@ -10,7 +10,7 @@ const { mockUseAuthorized, mockUseUISettings, mockReplace, mockMigratedHref, sta
   return {
     state,
     mockReplace: vi.fn(),
-    mockMigratedHref: vi.fn((segment: string) => `/mocked-ui/${segment}`),
+    mockUiHref: vi.fn((segment: string) => `/mocked-ui/${segment}`),
     mockUseAuthorized: vi.fn(() => ({
       accessToken: "token-123",
       userRole: "Internal User",
@@ -30,7 +30,7 @@ vi.mock("next/navigation", () => ({
 }));
 vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: mockUseAuthorized }));
 vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings }));
-vi.mock("@/utils/migratedPages", () => ({ migratedHref: mockMigratedHref }));
+vi.mock("@/utils/uiHref", () => ({ uiHref: mockUiHref }));
 vi.mock("@/components/navbar", () => ({ default: () => <div data-testid="navbar" /> }));
 vi.mock("@/contexts/ThemeContext", () => ({
   ThemeProvider: ({ children }: { children: React.ReactNode }) => <>{children}</>,
@@ -47,7 +47,7 @@ describe("ChatLayout", () => {
     state.enableChatUI = false;
     state.isUISettingsLoading = false;
     mockReplace.mockClear();
-    mockMigratedHref.mockClear();
+    mockUiHref.mockClear();
   });
 
   it("renders the chat shell when enable_chat_ui is on", () => {

ui/litellm-dashboard/src/app/chat/layout.tsx

@@ -8,7 +8,7 @@ import Navbar from "@/components/navbar";
 import { ThemeProvider } from "@/contexts/ThemeContext";
 import { ChatShellProvider } from "@/contexts/ChatShellContext";
 import ChatShell from "@/components/chat/ChatShell";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 
 // ChatShellProvider uses useSearchParams(), which requires a Suspense boundary for static export.
 function ChatLayoutContent({ children }: { children: React.ReactNode }) {
@@ -20,7 +20,7 @@ function ChatLayoutContent({ children }: { children: React.ReactNode }) {
   const blocked = !isUISettingsLoading && !chatEnabled;
 
   useEffect(() => {
-    if (blocked) router.replace(migratedHref(""));
+    if (blocked) router.replace(uiHref(""));
   }, [blocked, router]);
 
   if (isUISettingsLoading || blocked) return null;

ui/litellm-dashboard/src/components/DashboardHeader.test.tsx

@@ -7,6 +7,7 @@ const { mockUsePluginMode, mockUseUISettings, state } = vi.hoisted(() => {
   const state = {
     plugins: [] as { name: string; display_name: string; url: string }[],
     enableChatUI: false,
+    pathname: "/ui/logs",
   };
   return {
     state,
@@ -17,8 +18,7 @@ const { mockUsePluginMode, mockUseUISettings, state } = vi.hoisted(() => {
 
 vi.mock("@/contexts/PluginModeContext", () => ({ usePluginMode: mockUsePluginMode }));
 vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings }));
-vi.mock("next/navigation", () => ({ usePathname: () => "/ui/" }));
-vi.mock("@/utils/migratedPages", () => ({ migratedHref: (seg: string) => `/ui/${seg}` }));
+vi.mock("next/navigation", () => ({ usePathname: () => state.pathname }));
 vi.mock("@/hooks/useWorker", () => ({ useWorker: () => ({ isControlPlane: false, selectedWorker: null }) }));
 vi.mock("@/app/(dashboard)/hooks/useDisableShowPrompts", () => ({ useDisableShowPrompts: () => false }));
 vi.mock("@/components/Navbar/BlogDropdown/BlogDropdown", () => ({ BlogDropdown: () => null }));
@@ -32,11 +32,26 @@ describe("DashboardHeader breadcrumb", () => {
   afterEach(() => {
     state.plugins = [];
     state.enableChatUI = false;
+    state.pathname = "/ui/logs";
+  });
+
+  it("titles the breadcrumb from the current route, not from a sidebar page id", () => {
+    state.pathname = "/ui/models-and-endpoints";
+    render(<DashboardHeader />);
+
+    expect(screen.getByText("Models + Endpoints")).toBeInTheDocument();
+  });
+
+  it("titles the dashboard root as Virtual Keys", () => {
+    state.pathname = "/ui/";
+    render(<DashboardHeader />);
+
+    expect(screen.getByText("Virtual Keys")).toBeInTheDocument();
   });
 
   it("roots the breadcrumb in the AI Gateway selector (with a Chat option) and drops the static section crumb when the selector is available", async () => {
     state.enableChatUI = true;
-    render(<DashboardHeader page="logs" />);
+    render(<DashboardHeader />);
 
     expect(screen.getByText("Logs")).toBeInTheDocument();
     expect(screen.queryByText("Observability")).not.toBeInTheDocument();
@@ -49,7 +64,7 @@ describe("DashboardHeader breadcrumb", () => {
   });
 
   it("keeps the AI Gateway selector at the root even when there is nothing to switch to (discovery)", () => {
-    render(<DashboardHeader page="logs" />);
+    render(<DashboardHeader />);
 
     expect(screen.getByRole("button", { name: /AI Gateway/i })).toBeInTheDocument();
     expect(screen.getByText("Logs")).toBeInTheDocument();
@@ -57,7 +72,7 @@ describe("DashboardHeader breadcrumb", () => {
   });
 
   it("styles Docs with the shared product-link class instead of a muted toolbar button", () => {
-    render(<DashboardHeader page="logs" />);
+    render(<DashboardHeader />);
 
     const docs = screen.getByRole("link", { name: "Docs" });
     for (const cls of NAV_PRODUCT_LINK_CLASS.trim().split(/\s+/)) {
@@ -67,7 +82,7 @@ describe("DashboardHeader breadcrumb", () => {
   });
 
   it("renders the tools divider centered rather than stretched to the top of the row", () => {
-    const { container } = render(<DashboardHeader page="logs" />);
+    const { container } = render(<DashboardHeader />);
 
     const separators = container.querySelectorAll('[data-slot="separator"][data-orientation="vertical"]');
     expect(separators).toHaveLength(1);

ui/litellm-dashboard/src/components/DashboardHeader.tsx

@@ -20,15 +20,12 @@ import { useWorker } from "@/hooks/useWorker";
 import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils";
-
-interface DashboardHeaderProps {
-  page: string;
-}
+import { usePathname } from "next/navigation";
 
 // Top bar for the dashboard shell. Sits only over the content column (the brand
 // lives in the sidebar header); mirrors the design's breadcrumb-left / tools-right layout.
-export function DashboardHeader({ page }: DashboardHeaderProps) {
-  const { title } = getBreadcrumb(page);
+export function DashboardHeader() {
+  const { title } = getBreadcrumb(usePathname());
   const { isControlPlane, selectedWorker } = useWorker();
   const showWorkerSwitch = isControlPlane && selectedWorker !== null;
   const hideCommunityLinks = useDisableShowPrompts();

ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.test.tsx

@@ -28,7 +28,7 @@ vi.mock("@/contexts/PluginModeContext", () => ({ usePluginMode: mockUsePluginMod
 vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: mockUseUISettings }));
 vi.mock("next/navigation", () => ({ usePathname: mockUsePathname }));
 // Deterministic hrefs so navigation assertions don't depend on server_root_path.
-vi.mock("@/utils/migratedPages", () => ({ migratedHref: (seg: string) => `/ui/${seg}` }));
+vi.mock("@/utils/uiHref", () => ({ uiHref: (seg: string) => `/ui/${seg}` }));
 
 describe("ViewSwitcher", () => {
   let assignSpy: ReturnType<typeof vi.fn>;

ui/litellm-dashboard/src/components/Navbar/ViewSwitcher.tsx

@@ -9,7 +9,7 @@ import {
 import { Check, ChevronsUpDown, LayoutGrid } from "lucide-react";
 import { usePluginMode } from "@/contexts/PluginModeContext";
 import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 
 const GATEWAY = "ai-gateway";
 const CHAT = "chat";
@@ -28,7 +28,7 @@ export default function ViewSwitcher() {
 
   const chatEnabled = Boolean(uiSettings?.values?.enable_chat_ui);
 
-  const chatHref = migratedHref(CHAT);
+  const chatHref = uiHref(CHAT);
   const normalizedPathname = (pathname ?? "").replace(/\/+$/, "");
   const isChatRoute = chatEnabled && (normalizedPathname === chatHref || normalizedPathname.startsWith(`${chatHref}/`));
 
@@ -44,7 +44,7 @@ export default function ViewSwitcher() {
     // The chat route lives outside the dashboard SPA shell that reacts to `mode`,
     // so switching modes from there needs a real navigation, not just state.
     if (isChatRoute) {
-      window.location.assign(migratedHref(""));
+      window.location.assign(uiHref(""));
     }
   };
 
@@ -57,7 +57,7 @@ export default function ViewSwitcher() {
             {isChatRoute && <Check className="size-4 text-info" />}
           </div>
         ),
-        onClick: () => window.location.assign(migratedHref(CHAT)),
+        onClick: () => window.location.assign(uiHref(CHAT)),
       }
     : {
         key: CHAT,

ui/litellm-dashboard/src/components/chat/ChatShell.test.tsx

@@ -18,7 +18,7 @@ vi.mock("next/navigation", () => ({
   usePathname: mockUsePathname,
 }));
 // Deterministic hrefs so navigation/active-state assertions don't depend on server_root_path.
-vi.mock("@/utils/migratedPages", () => ({ migratedHref: (seg: string) => `/ui/${seg}`.replace(/\/$/, "") || "/ui" }));
+vi.mock("@/utils/uiHref", () => ({ uiHref: (seg: string) => `/ui/${seg}`.replace(/\/$/, "") || "/ui" }));
 vi.mock("@/contexts/ChatShellContext", () => ({ useChatShell: mockUseChatShell }));
 vi.mock("./ConversationList", () => ({ default: () => <div data-testid="conversation-list" /> }));
 

ui/litellm-dashboard/src/components/chat/ChatShell.tsx

@@ -5,12 +5,12 @@ import { usePathname, useRouter } from "next/navigation";
 import { Plus, MessageSquare, LayoutGrid, KeyRound, Lock, BarChart3, ScrollText } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Separator } from "@/components/ui/separator";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 import { useChatShell } from "@/contexts/ChatShellContext";
 import ConversationList from "./ConversationList";
 
 export function getChatRoutes() {
-  const base = migratedHref("chat");
+  const base = uiHref("chat");
   return {
     chats: base,
     integrations: `${base}/integrations`,

ui/litellm-dashboard/src/components/common_components/fetch_teams.tsx

@@ -1,18 +0,0 @@
-import { teamListCall, Organization } from "../networking";
-
-export const fetchTeams = async (
-  accessToken: string,
-  userID: string | null,
-  userRole: string | null,
-  currentOrg: Organization | null,
-  setTeams: (teams: any[]) => void,
-) => {
-  let givenTeams;
-  if (userRole != "Admin" && userRole != "Admin Viewer") {
-    givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || null, userID);
-  } else {
-    givenTeams = await teamListCall(accessToken, currentOrg?.organization_id || null);
-  }
-
-  setTeams(givenTeams);
-};

ui/litellm-dashboard/src/components/leftnav.test.tsx

@@ -17,6 +17,12 @@ vi.mock("../utils/roles", async (importOriginal) => {
   };
 });
 
+const navState = vi.hoisted(() => ({ pathname: "/ui/api-keys" }));
+
+vi.mock("next/navigation", () => ({
+  usePathname: () => navState.pathname,
+}));
+
 const { mockUseAuthorized, mockUseOrganizations } = vi.hoisted(() => {
   const mockUseAuthorized = vi.fn(() => ({
     userId: "test-user-id",
@@ -98,8 +104,6 @@ const placementsOf = (page: string): string[] =>
 
 describe("Sidebar (leftnav)", () => {
   const defaultProps = {
-    setPage: vi.fn(),
-    defaultSelectedKey: "api-keys",
     collapsed: false,
   };
 
@@ -107,6 +111,7 @@ describe("Sidebar (leftnav)", () => {
     mockUseAuthorized.mockReset();
     mockUseOrganizations.mockReset();
     mockUseThemeImpl = unbrandedTheme;
+    navState.pathname = "/ui/api-keys";
   });
 
   it("should link the logo to the UI home route rather than the proxy origin", () => {
@@ -509,12 +514,52 @@ describe("Sidebar (leftnav)", () => {
     expect(screen.getByText("Organizations")).toBeInTheDocument();
   });
 
-  it("marks the selected page's nav item active", () => {
-    renderWithProviders(<Sidebar {...defaultProps} defaultSelectedKey="logs" />);
-    const logs = screen.getByText("Logs").closest("a");
-    expect(logs).toHaveAttribute("data-active", "true");
-    // A different item must not be active.
-    expect(screen.getByText("Virtual Keys").closest("a")).not.toHaveAttribute("data-active");
+  it("marks the nav item for the current route active", () => {
+    navState.pathname = "/ui/logs";
+    renderWithProviders(<Sidebar {...defaultProps} />);
+    expect(screen.getByRole("link", { name: "Logs" })).toHaveAttribute("data-active", "true");
+    expect(screen.getByRole("link", { name: "Virtual Keys" })).not.toHaveAttribute("data-active");
+  });
+
+  it("marks Virtual Keys active at the dashboard root", () => {
+    navState.pathname = "/ui/";
+    renderWithProviders(<Sidebar {...defaultProps} />);
+    expect(screen.getByRole("link", { name: "Virtual Keys" })).toHaveAttribute("data-active", "true");
+  });
+
+  it("expands the parent group of the current nested route and marks the child active", () => {
+    navState.pathname = "/ui/search-tools";
+    renderWithProviders(<Sidebar {...defaultProps} />);
+    expect(screen.getByRole("link", { name: "Search Tools" })).toHaveAttribute("data-active", "true");
+    expect(screen.getByRole("button", { name: "Tools" })).toHaveAttribute("aria-expanded", "true");
+  });
+
+  it("links every leaf to its path route, including the ids that differ from their route", () => {
+    renderWithProviders(<Sidebar {...defaultProps} />);
+    act(() => {
+      fireEvent.click(screen.getByText("Experimental"));
+    });
+
+    const expectHref = (label: string, href: string) =>
+      expect(screen.getByRole("link", { name: label })).toHaveAttribute("href", href);
+    expectHref("Virtual Keys", "/ui/api-keys");
+    expectHref("Playground", "/ui/playground");
+    expectHref("Models + Endpoints", "/ui/models-and-endpoints");
+    expectHref("Usage", "/ui/usage");
+    expectHref("API Reference", "/ui/api-reference");
+    expectHref("Old Usage", "/ui/old-usage");
+  });
+
+  it("never links a leaf to the legacy ?page= switch", () => {
+    renderWithProviders(<Sidebar {...defaultProps} enableProjectsUI />);
+    for (const group of ["Agentic", "Tools", "Experimental", "Settings"]) {
+      act(() => {
+        fireEvent.click(screen.getByText(group));
+      });
+    }
+    const hrefs = screen.getAllByRole("link").map((link) => link.getAttribute("href") ?? "");
+    expect(hrefs.filter((href) => href.includes("page="))).toHaveLength(0);
+    expect(hrefs.filter((href) => href.startsWith("/ui/")).length).toBeGreaterThan(30);
   });
 
   it("hides labels but keeps items reachable (icon + link) when collapsed to the rail", () => {
@@ -550,20 +595,30 @@ describe("Sidebar (leftnav)", () => {
 });
 
 describe("getBreadcrumb", () => {
-  it("resolves a top-level page to its section + title", () => {
-    expect(getBreadcrumb("api-keys")).toEqual({ section: "AI Gateway", title: "Virtual Keys" });
-    expect(getBreadcrumb("logs")).toEqual({ section: "Observability", title: "Logs" });
+  it("resolves a top-level route to its section + title", () => {
+    expect(getBreadcrumb("/ui/api-keys")).toEqual({ section: "AI Gateway", title: "Virtual Keys" });
+    expect(getBreadcrumb("/ui/logs")).toEqual({ section: "Observability", title: "Logs" });
   });
 
-  it("resolves a nested child page to its parent section", () => {
-    expect(getBreadcrumb("search-tools")).toEqual({ section: "AI Gateway", title: "Search Tools" });
+  it("resolves routes whose segment differs from the sidebar page id", () => {
+    expect(getBreadcrumb("/ui/models-and-endpoints")).toEqual({ section: "AI Gateway", title: "Models + Endpoints" });
+    expect(getBreadcrumb("/ui/usage")).toEqual({ section: "Observability", title: "Usage" });
+    expect(getBreadcrumb("/ui/old-usage")).toEqual({ section: "Developer Tools", title: "Old Usage" });
+  });
+
+  it("titles the dashboard root as Virtual Keys", () => {
+    expect(getBreadcrumb("/ui/")).toEqual({ section: "AI Gateway", title: "Virtual Keys" });
+  });
+
+  it("resolves a nested child route to its parent section", () => {
+    expect(getBreadcrumb("/ui/search-tools/")).toEqual({ section: "AI Gateway", title: "Search Tools" });
   });
 
   it("resolves router-settings under the Settings section", () => {
-    expect(getBreadcrumb("router-settings")).toEqual({ section: "Settings", title: "Router Settings" });
+    expect(getBreadcrumb("/ui/router-settings")).toEqual({ section: "Settings", title: "Router Settings" });
   });
 
-  it("falls back to a prettified title with no section for unknown pages", () => {
-    expect(getBreadcrumb("some-unknown-page")).toEqual({ section: null, title: "Some Unknown Page" });
+  it("falls back to a prettified title with no section for unknown routes", () => {
+    expect(getBreadcrumb("/ui/some-unknown-page")).toEqual({ section: null, title: "Some Unknown Page" });
   });
 });

ui/litellm-dashboard/src/components/leftnav.tsx

@@ -62,6 +62,7 @@ import {
   Workflow,
 } from "lucide-react";
 import Link from "next/link";
+import { usePathname } from "next/navigation";
 import { useMemo, useState } from "react";
 import { cn } from "@/lib/cva.config";
 import { rolesWithCapability } from "../utils/capabilities";
@@ -76,15 +77,13 @@ import {
 import BetaBadge from "./BetaBadge";
 import SidebarAccountMenu from "./SidebarAccountMenu/SidebarAccountMenu";
 import SidebarUsageCard from "./SidebarUsageCard";
-import { MIGRATED_PAGES, migratedHref, legacyPageHref } from "@/utils/migratedPages";
+import { routeSegmentForPathname, uiHref } from "@/utils/uiHref";
 
 const ICON = { strokeWidth: 1.75 } as const;
 
 const LOGO_CLASS_NAME = "h-7 w-auto max-w-[150px] object-contain group-data-[collapsed=true]/sidebar:w-7";
 
 interface SidebarProps {
-  setPage: (page: string) => void;
-  defaultSelectedKey: string;
   collapsed?: boolean;
   onToggleCollapsed?: () => void;
   enabledPagesInternalUsers?: string[] | null;
@@ -98,6 +97,7 @@ interface SidebarProps {
 interface MenuItem {
   key: string;
   page: string;
+  route?: string;
   label: string | React.ReactNode;
   roles?: string[];
   children?: MenuItem[];
@@ -122,6 +122,7 @@ const menuGroups: MenuGroup[] = [
       {
         key: "llm-playground",
         page: "llm-playground",
+        route: "playground",
         label: "Playground",
         icon: <PlayCircle {...ICON} />,
         roles: rolesWithWriteAccess,
@@ -129,6 +130,7 @@ const menuGroups: MenuGroup[] = [
       {
         key: "models",
         page: "models",
+        route: "models-and-endpoints",
         label: "Models + Endpoints",
         icon: <Network {...ICON} />,
         roles: rolesAllowedToViewWriteScopedPages,
@@ -197,6 +199,7 @@ const menuGroups: MenuGroup[] = [
       {
         key: "new_usage",
         page: "new_usage",
+        route: "usage",
         icon: <BarChart3 {...ICON} />,
         roles: [...all_admin_roles, ...internalUserRoles],
         label: "Usage",
@@ -258,7 +261,7 @@ const menuGroups: MenuGroup[] = [
   {
     groupLabel: "DEVELOPER TOOLS",
     items: [
-      { key: "api_ref", page: "api_ref", label: "API Reference", icon: <Code2 {...ICON} /> },
+      { key: "api_ref", page: "api_ref", route: "api-reference", label: "API Reference", icon: <Code2 {...ICON} /> },
       { key: "model-hub-table", page: "model-hub-table", label: "AI Hub", icon: <LayoutGrid {...ICON} /> },
       {
         key: "learning-resources",
@@ -304,6 +307,7 @@ const menuGroups: MenuGroup[] = [
           {
             key: "4",
             page: "usage",
+            route: "old-usage",
             label: "Old Usage",
             icon: <BarChart3 {...ICON} />,
             roles: rolesWithCapability("viewGlobalSpend"),
@@ -358,24 +362,30 @@ const menuGroups: MenuGroup[] = [
   },
 ];
 
-const findParentKey = (page: string): string | null => {
+const HOME_ROUTE = "api-keys";
+
+const routeOf = (item: MenuItem): string => item.route ?? item.page;
+
+const routeForPathname = (pathname: string): string => routeSegmentForPathname(pathname) || HOME_ROUTE;
+
+const findParentKey = (route: string): string | null => {
   for (const group of menuGroups) {
     for (const item of group.items) {
-      if (item.children?.some((c) => c.page === page || c.key === page)) return item.key;
+      if (item.children?.some((c) => routeOf(c) === route)) return item.key;
     }
   }
   return null;
 };
 
-const findMenuItemKey = (page: string): string => {
+const findMenuItemKey = (route: string): string => {
   for (const group of menuGroups) {
     for (const item of group.items) {
-      if (item.page === page) return item.key;
-      const child = item.children?.find((c) => c.page === page);
+      if (routeOf(item) === route) return item.key;
+      const child = item.children?.find((c) => routeOf(c) === route);
       if (child) return child.key;
     }
   }
-  return "api-keys";
+  return HOME_ROUTE;
 };
 
 const SECTION_DISPLAY: Record<string, string> = {
@@ -395,22 +405,20 @@ const prettify = (key: string): string =>
 const labelText = (item: MenuItem): string => (typeof item.label === "string" ? item.label : prettify(item.key));
 
 // Breadcrumb ("Section" / "Page") for the top bar, derived from the same nav config.
-export const getBreadcrumb = (page: string): { section: string | null; title: string } => {
+export const getBreadcrumb = (pathname: string): { section: string | null; title: string } => {
+  const route = routeForPathname(pathname);
   for (const group of menuGroups) {
     for (const item of group.items) {
       const section = SECTION_DISPLAY[group.groupLabel] ?? group.groupLabel;
-      if (item.page === page)
-        return { section, title: typeof item.label === "string" ? item.label : prettify(item.key) };
-      const child = item.children?.find((c) => c.page === page);
-      if (child) return { section, title: typeof child.label === "string" ? child.label : prettify(child.key) };
+      if (routeOf(item) === route) return { section, title: labelText(item) };
+      const child = item.children?.find((c) => routeOf(c) === route);
+      if (child) return { section, title: labelText(child) };
     }
   }
-  return { section: null, title: prettify(page) };
+  return { section: null, title: prettify(route) };
 };
 
 const Sidebar_: React.FC<SidebarProps> = ({
-  setPage,
-  defaultSelectedKey,
   collapsed = false,
   onToggleCollapsed,
   enabledPagesInternalUsers,
@@ -430,20 +438,21 @@ const Sidebar_: React.FC<SidebarProps> = ({
 
   const baseUrl = getProxyBaseUrl();
   const version = healthData?.litellm_version;
-  const selectedKey = findMenuItemKey(defaultSelectedKey);
+  const currentRoute = routeForPathname(usePathname());
+  const selectedKey = findMenuItemKey(currentRoute);
 
   const [openGroups, setOpenGroups] = useState<Set<string>>(() => {
-    const parent = findParentKey(defaultSelectedKey);
+    const parent = findParentKey(currentRoute);
     return new Set(parent ? [parent] : []);
   });
 
   // Keep the active page's parent group expanded as the user navigates, using the
   // "adjust state during render" pattern rather than an effect (avoids a
   // setState-in-effect render cascade).
-  const [prevSelectedKey, setPrevSelectedKey] = useState(defaultSelectedKey);
-  if (defaultSelectedKey !== prevSelectedKey) {
-    setPrevSelectedKey(defaultSelectedKey);
-    const parent = findParentKey(defaultSelectedKey);
+  const [prevRoute, setPrevRoute] = useState(currentRoute);
+  if (currentRoute !== prevRoute) {
+    setPrevRoute(currentRoute);
+    const parent = findParentKey(currentRoute);
     if (parent && !openGroups.has(parent)) {
       setOpenGroups((prev) => new Set(prev).add(parent));
     }
@@ -512,13 +521,6 @@ const Sidebar_: React.FC<SidebarProps> = ({
     });
   };
 
-  const handleLeafClick = (e: React.MouseEvent, item: MenuItem) => {
-    if (item.external_url) return;
-    if (e.metaKey || e.ctrlKey || e.shiftKey || e.button === 1) return;
-    e.preventDefault();
-    setPage(item.page);
-  };
-
   const renderLeaf = (item: MenuItem, isChild: boolean) => {
     const active = selectedKey === item.key;
     const size = isChild ? "sub" : "default";
@@ -542,19 +544,17 @@ const Sidebar_: React.FC<SidebarProps> = ({
       );
     }
 
-    const href = MIGRATED_PAGES[item.page] ? migratedHref(MIGRATED_PAGES[item.page]) : legacyPageHref(item.page);
     return (
-      <a
+      <Link
         key={item.key}
-        href={href}
-        onClick={(e) => handleLeafClick(e, item)}
+        href={uiHref(routeOf(item))}
         title={collapsed ? labelText(item) : undefined}
         data-active={active || undefined}
         className={cn(sidebarMenuButtonVariants({ isActive: active, size }))}
       >
         {item.icon}
         {label}
-      </a>
+      </Link>
     );
   };
 
@@ -603,7 +603,7 @@ const Sidebar_: React.FC<SidebarProps> = ({
       <SidebarHeader className="h-14 border-b border-border group-data-[collapsed=true]/sidebar:h-auto">
         <div className="flex items-center justify-between gap-2 group-data-[collapsed=true]/sidebar:flex-col">
           <div className="flex min-w-0 items-center gap-2">
-            <Link href={migratedHref("")} className="flex min-w-0 items-center" aria-label="LiteLLM home">
+            <Link href={uiHref("")} className="flex min-w-0 items-center" aria-label="LiteLLM home">
               <img src={logoSrc} alt="LiteLLM" className={cn(LOGO_CLASS_NAME, "dark:hidden")} />
               <img
                 src={darkLogoSrc}

ui/litellm-dashboard/src/components/navbar.tsx

@@ -3,7 +3,7 @@ import { useDisableBouncingIcon } from "@/app/(dashboard)/hooks/useDisableBounci
 import { useDisableShowPrompts } from "@/app/(dashboard)/hooks/useDisableShowPrompts";
 import { useWorker } from "@/hooks/useWorker";
 import { getProxyBaseUrl } from "@/components/networking";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 import { useTheme } from "@/contexts/ThemeContext";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import { clearStoredReturnUrl, getLoginUrl } from "@/utils/returnUrlUtils";
@@ -87,7 +87,7 @@ const Navbar: React.FC<NavbarProps> = ({
             )}
 
             <div className="flex items-center gap-2">
-              <Link href={migratedHref("")} className="flex items-center">
+              <Link href={uiHref("")} className="flex items-center">
                 <div className="relative">
                   <div className="flex h-10 max-w-48 items-center justify-center overflow-hidden">
                     <img src={imageUrl} alt="LiteLLM Brand" className={cn(NAV_LOGO_CLASS_NAME, "dark:hidden")} />

ui/litellm-dashboard/src/components/networking.test.ts

@@ -1,7 +1,7 @@
 import { describe, it, expect, beforeEach, afterEach, vi } from "vitest";
 import { clearTokenCookies } from "@/utils/cookieUtils";
 import * as Networking from "./networking";
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 
 vi.mock("@/utils/cookieUtils", () => ({
   clearTokenCookies: vi.fn(),
@@ -392,7 +392,7 @@ describe("UI config and public endpoints", () => {
     await Networking.getUiConfig();
 
     expect(Networking.serverRootPath).toBe("/litellm");
-    expect(migratedHref("api-reference")).toBe("/litellm/ui/api-reference");
+    expect(uiHref("api-reference")).toBe("/litellm/ui/api-reference");
   });
 });
 

ui/litellm-dashboard/src/components/networking.tsx

@@ -140,7 +140,7 @@ const resolveDefaultBase = (fallback: string | null): string | null =>
 const defaultProxyBaseUrl = resolveDefaultBase(null);
 const WORKER_URL_KEY = "litellm_worker_url";
 // If a worker URL is in localStorage, use it as the initial proxyBaseUrl.
-// This survives page navigation and the sessionStorage.clear() in user_dashboard.
+// This survives page navigation.
 const _rawWorkerUrl = typeof window !== "undefined" ? window.localStorage.getItem(WORKER_URL_KEY) : null;
 // Validate stored worker URL — reject non-HTTP schemes to prevent exfiltration
 const _initialWorkerUrl = (() => {
@@ -195,10 +195,9 @@ export const getProxyBaseUrl = (): string => {
 
 /**
  * Switch API calls to point at a worker (or back to the control plane).
- * Persists to localStorage so it survives page navigation and the
- * sessionStorage.clear() in user_dashboard. Also updates the module-level
- * proxyBaseUrl so in-flight code in this JS execution sees the new value
- * immediately.
+ * Persists to localStorage so it survives page navigation. Also updates the
+ * module-level proxyBaseUrl so in-flight code in this JS execution sees the
+ * new value immediately.
  */
 function isValidHttpUrl(url: string): boolean {
   try {

ui/litellm-dashboard/src/components/organization/organization_view.test.tsx

@@ -12,8 +12,7 @@ vi.mock("next/navigation", () => ({
   useSearchParams: () => new URLSearchParams(window.location.search),
 }));
 
-// Mock networking calls used by the component's mutation handlers. entityLinks -> migratedPages
-// imports serverRootPath from the same module, so the mock must export it too.
+// Mock networking calls used by the component's mutation handlers.
 vi.mock("../networking", () => {
   return {
     __esModule: true,

ui/litellm-dashboard/src/components/user_dashboard.test.tsx

@@ -1,133 +0,0 @@
-import { vi, describe, it, expect, beforeEach, afterEach } from "vitest";
-import { cleanup } from "@testing-library/react";
-import React from "react";
-import { renderWithProviders } from "../../tests/test-utils";
-
-// Track addEventListener/removeEventListener calls for "beforeunload"
-const addEventListenerSpy = vi.spyOn(window, "addEventListener");
-const removeEventListenerSpy = vi.spyOn(window, "removeEventListener");
-
-// Mock next/navigation
-vi.mock("next/navigation", () => ({
-  useSearchParams: () => new URLSearchParams(),
-}));
-
-// Mock networking with importOriginal so all exports are available
-vi.mock("./networking", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("./networking")>();
-  return {
-    ...actual,
-    getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"),
-    getProxyUISettings: vi.fn().mockResolvedValue({}),
-    keyInfoCall: vi.fn().mockResolvedValue({}),
-    modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }),
-    userGetInfoV2: vi.fn().mockResolvedValue({
-      user_id: "user-1",
-      user_email: "test@example.com",
-      spend: 0,
-      max_budget: null,
-      models: [],
-      teams: [],
-    }),
-  };
-});
-
-// Mock jwt-decode to return a valid token structure
-vi.mock("jwt-decode", () => ({
-  jwtDecode: vi.fn().mockReturnValue({
-    key: "test-access-token",
-    user_role: "proxy_admin",
-    user_email: "test@example.com",
-    exp: Math.floor(Date.now() / 1000) + 3600,
-  }),
-}));
-
-// Mock cookie utility
-vi.mock("@/utils/cookieUtils", () => ({
-  clearTokenCookies: vi.fn(),
-  getCookie: vi.fn().mockReturnValue("fake-jwt-token"),
-}));
-
-// Mock fetchTeams
-vi.mock("./common_components/fetch_teams", () => ({
-  fetchTeams: vi.fn(),
-}));
-
-// Mock heavy child components to isolate UserDashboard behavior
-vi.mock("./organisms/create_key_button", () => ({
-  default: () => <div data-testid="create-key-mock" />,
-}));
-
-vi.mock("./VirtualKeysPage/VirtualKeysTable", () => ({
-  VirtualKeysTable: () => <div data-testid="virtual-keys-table-mock" />,
-}));
-
-vi.mock("../app/onboarding/page", () => ({
-  default: () => <div data-testid="onboarding-mock" />,
-}));
-
-// Provide a token cookie so the component doesn't redirect to login
-Object.defineProperty(document, "cookie", {
-  writable: true,
-  value: "token=fake-jwt-token",
-});
-
-import UserDashboard from "./user_dashboard";
-
-const defaultProps = {
-  userID: "user-1",
-  userRole: "Admin",
-  userEmail: "test@example.com",
-  teams: [] as any[],
-  keys: [] as any[],
-  setUserRole: vi.fn(),
-  setUserEmail: vi.fn(),
-  setTeams: vi.fn(),
-  setKeys: vi.fn(),
-  premiumUser: false,
-  addKey: vi.fn(),
-  createClicked: false,
-};
-
-function renderDashboard(props = {}) {
-  return renderWithProviders(<UserDashboard {...defaultProps} {...props} />);
-}
-
-describe("UserDashboard beforeunload listener", () => {
-  beforeEach(() => {
-    addEventListenerSpy.mockClear();
-    removeEventListenerSpy.mockClear();
-  });
-
-  afterEach(() => {
-    cleanup();
-  });
-
-  it("registers exactly one beforeunload listener on mount", () => {
-    renderDashboard();
-
-    const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload");
-    expect(beforeUnloadCalls).toHaveLength(1);
-  });
-
-  it("does not add duplicate listeners on re-render", () => {
-    const { rerender } = renderWithProviders(<UserDashboard {...defaultProps} />);
-
-    addEventListenerSpy.mockClear();
-
-    // Re-render with different props to trigger a render cycle
-    rerender(<UserDashboard {...defaultProps} userEmail="updated@example.com" />);
-
-    const beforeUnloadCalls = addEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload");
-    expect(beforeUnloadCalls).toHaveLength(0);
-  });
-
-  it("removes the beforeunload listener on unmount", () => {
-    const { unmount } = renderDashboard();
-
-    unmount();
-
-    const removeCalls = removeEventListenerSpy.mock.calls.filter(([event]) => event === "beforeunload");
-    expect(removeCalls).toHaveLength(1);
-  });
-});

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -1,239 +0,0 @@
-"use client";
-import { clearTokenCookies, getCookie } from "@/utils/cookieUtils";
-import { jwtDecode } from "jwt-decode";
-import React, { useEffect, useState } from "react";
-import { fetchTeams } from "./common_components/fetch_teams";
-import { KeyResponse, Team } from "./key_team_helpers/key_list";
-import { effectiveSessionRole } from "@/utils/roles";
-import { getProxyBaseUrl, keyInfoCall, modelAvailableCall, Organization, userGetInfoV2 } from "./networking";
-import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button";
-import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable";
-
-export interface ProxySettings {
-  PROXY_BASE_URL: string | null;
-  PROXY_LOGOUT_URL: string | null;
-  LITELLM_UI_API_DOC_BASE_URL?: string | null;
-  DEFAULT_TEAM_DISABLED: boolean;
-  SSO_ENABLED: boolean;
-  DISABLE_EXPENSIVE_DB_QUERIES: boolean;
-  NUM_SPEND_LOGS_ROWS: number;
-}
-
-export type UserInfo = {
-  models: string[];
-  max_budget?: number | null;
-  spend: number;
-};
-
-interface UserDashboardProps {
-  userID: string | null;
-  userRole: string | null;
-  userEmail: string | null;
-  teams: Team[] | null;
-  keys: any[] | null;
-  setUserRole: React.Dispatch<React.SetStateAction<string>>;
-  setUserEmail: React.Dispatch<React.SetStateAction<string | null>>;
-  setTeams: React.Dispatch<React.SetStateAction<Team[] | null>>;
-  setKeys: (keys: KeyResponse[]) => void;
-  premiumUser: boolean;
-  addKey: (data: any) => void;
-  createClicked: boolean;
-  autoOpenCreate?: boolean;
-  prefillData?: CreateKeyPrefillData;
-}
-
-const UserDashboard: React.FC<UserDashboardProps> = ({
-  userID,
-  userRole,
-  teams,
-  keys,
-  setUserRole,
-  userEmail,
-  setUserEmail,
-  setTeams,
-  setKeys,
-  premiumUser,
-  addKey,
-  createClicked,
-  autoOpenCreate,
-  prefillData,
-}) => {
-  const [userSpendData, setUserSpendData] = useState<UserInfo | null>(null);
-  const [currentOrg] = useState<Organization | null>(null);
-
-  const token = getCookie("token");
-
-  const [accessToken, setAccessToken] = useState<string | null>(null);
-  const [selectedTeam] = useState<any | null>(null);
-
-  // Clear session storage on page unload so next load fetches fresh data.
-  // Note: MCP auth tokens are persistent and should not be cleared on page refresh
-  // They are only cleared on logout
-  useEffect(() => {
-    const handleBeforeUnload = () => {
-      const token = sessionStorage.getItem("token");
-      sessionStorage.clear();
-      if (token) {
-        sessionStorage.setItem("token", token);
-      }
-    };
-    window.addEventListener("beforeunload", handleBeforeUnload);
-    return () => window.removeEventListener("beforeunload", handleBeforeUnload);
-  }, []);
-
-  // console.log(`selectedTeam: ${Object.entries(selectedTeam)}`);
-  // Moved useEffect inside the component and used a condition to run fetch only if the params are available
-  useEffect(() => {
-    if (token) {
-      const decoded = jwtDecode(token) as { [key: string]: any };
-      if (decoded) {
-        // cast decoded to dictionary
-
-        // set accessToken
-        setAccessToken(decoded.key);
-
-        // check if userRole is defined
-        if (decoded.user_role) {
-          setUserRole(effectiveSessionRole(decoded.user_role));
-        } else {
-        }
-
-        if (decoded.user_email) {
-          setUserEmail(decoded.user_email);
-        } else {
-        }
-      }
-    }
-    if (userID && accessToken && userRole && !userSpendData) {
-      const cachedUserModels = sessionStorage.getItem("userModels" + userID);
-      if (!cachedUserModels) {
-        const fetchData = async () => {
-          try {
-            const response = await userGetInfoV2(accessToken, userID);
-
-            setUserSpendData(response);
-
-            sessionStorage.setItem("userSpendData" + userID, JSON.stringify(response));
-
-            const model_available = await modelAvailableCall(accessToken, userID, userRole);
-            // loop through model_info["data"] and create an array of element.model_name
-            let available_model_names = model_available["data"].map((element: { id: string }) => element.id);
-
-            sessionStorage.setItem("userModels" + userID, JSON.stringify(available_model_names));
-          } catch (error: any) {
-            console.error("There was an error fetching the data", error);
-            if (error.message.includes("Invalid proxy server token passed")) {
-              gotoLogin();
-            }
-            // Optionally, update your UI to reflect the error state here as well
-          }
-        };
-        fetchData();
-        fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
-      }
-    }
-  }, [userID, token, accessToken, userRole]);
-
-  useEffect(() => {
-    // check key health - if it's invalid, redirect to login
-    if (accessToken) {
-      const fetchKeyInfo = async () => {
-        try {
-          await keyInfoCall(accessToken, [accessToken]);
-        } catch (error: any) {
-          if (error.message.includes("Invalid proxy server token passed")) {
-            gotoLogin();
-          }
-        }
-      };
-      fetchKeyInfo();
-    }
-  }, [accessToken]);
-
-  useEffect(() => {
-    if (accessToken) {
-      fetchTeams(accessToken, userID, userRole, currentOrg, setTeams);
-    }
-  }, [currentOrg]);
-
-  function gotoLogin() {
-    // Clear token cookies using the utility function
-    clearTokenCookies();
-
-    const baseUrl = getProxyBaseUrl();
-
-    const url = baseUrl ? `${baseUrl}/sso/key/generate` : `/sso/key/generate`;
-
-    window.location.href = url;
-
-    return null;
-  }
-
-  if (token == null) {
-    // user is not logged in as yet
-
-    // Clear token cookies using the utility function
-    gotoLogin();
-    return null;
-  } else {
-    // Check if token is expired
-    try {
-      const decoded = jwtDecode(token) as { [key: string]: any };
-      const expTime = decoded.exp;
-      const currentTime = Math.floor(Date.now() / 1000);
-
-      if (expTime && currentTime >= expTime) {
-        gotoLogin();
-
-        return null;
-      }
-    } catch (error) {
-      console.error("Error decoding token:", error);
-      // If there's an error decoding the token, consider it invalid
-      clearTokenCookies();
-
-      gotoLogin();
-
-      return null;
-    }
-
-    if (accessToken == null) {
-      return null;
-    }
-  }
-
-  if (userID == null) {
-    return <h1>User ID is not set</h1>;
-  }
-
-  if (userRole == null) {
-    setUserRole("App Owner");
-  }
-
-  // Admin Viewer can view keys read-only — gate "Create Key" but render the
-  // virtual-keys table the same as for Proxy Admin (read parity). Every
-  // other role keeps its existing ability to create keys.
-  const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer";
-
-  return (
-    <main className="flex h-full flex-col p-8">
-      <VirtualKeysTable
-        headerActions={
-          canCreateKey ? (
-            <CreateKey
-              key={selectedTeam ? selectedTeam.team_id : null}
-              team={selectedTeam as Team | null}
-              teams={teams as Team[]}
-              data={keys}
-              addKey={addKey}
-              autoOpenCreate={autoOpenCreate}
-              prefillData={prefillData}
-            />
-          ) : undefined
-        }
-      />
-    </main>
-  );
-};
-
-export default UserDashboard;

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -23733,6 +23733,11 @@ export interface components {
              * @description Name of the guardrail in guardrails.ai
              */
             guard_name?: string | null;
+            /**
+             * Inspect Embeddings
+             * @description When True, the Aim and Cato Networks guardrails send /embeddings `input` to the vendor as user messages. Off by default because embedding input is documents being indexed, not a conversation.
+             */
+            inspect_embeddings?: boolean | null;
             /**
              * Keyword Redaction Tag
              * @description Tag to use for keyword redaction
@@ -28904,6 +28909,11 @@ export interface components {
              * @default false
              */
             oauth_passthrough: boolean;
+            /**
+             * Per Server Oauth Discovery
+             * @default false
+             */
+            per_server_oauth_discovery: boolean;
             /** Registration Url */
             registration_url?: string | null;
             /** Review Notes */
@@ -30575,6 +30585,12 @@ export interface components {
             categories?: components["schemas"]["ContentFilterCategoryConfig"][] | null;
             /** @description Threshold configuration for Lakera guardrail categories */
             category_thresholds?: components["schemas"]["LakeraCategoryThresholds"] | null;
+            /**
+             * Ccr Retrieval
+             * @description Inject the Headroom retrieval tool for hashes declared by the compression service.
+             * @default true
+             */
+            ccr_retrieval: boolean;
             /** @description Inline safeguards for the resource-less InvokeGuardrailChecks API (contentFilter / promptAttack / sensitiveInformation). When set, the guardrail calls InvokeGuardrailChecks instead of ApplyGuardrail and no guardrailIdentifier is required. Mutually exclusive with guardrailIdentifier. */
             checks?: components["schemas"]["BedrockChecksConfigModel"] | null;
             /**
@@ -30755,6 +30771,11 @@ export interface components {
              * @default true
              */
             include_scanners: boolean | null;
+            /**
+             * Inspect Embeddings
+             * @description When True, the Aim and Cato Networks guardrails send /embeddings `input` to the vendor as user messages. Off by default because embedding input is documents being indexed, not a conversation.
+             */
+            inspect_embeddings?: boolean | null;
             /**
              * Is Detector Server
              * @description Boolean flag to determine if calling a detector server (True) or the FMS Orchestrator (False). Defaults to True.
@@ -32075,6 +32096,11 @@ export interface components {
              * @default false
              */
             oauth_passthrough: boolean;
+            /**
+             * Per Server Oauth Discovery
+             * @default false
+             */
+            per_server_oauth_discovery: boolean;
             /** Registration Url */
             registration_url?: string | null;
             /** Server Id */
@@ -37898,6 +37924,11 @@ export interface components {
              * @default false
              */
             oauth_passthrough: boolean;
+            /**
+             * Per Server Oauth Discovery
+             * @default false
+             */
+            per_server_oauth_discovery: boolean;
             /** Registration Url */
             registration_url?: string | null;
             /** Server Id */

ui/litellm-dashboard/src/utils/entityLinks.ts

@@ -1,4 +1,4 @@
-import { migratedHref } from "@/utils/migratedPages";
+import { uiHref } from "@/utils/uiHref";
 
 const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
   "all-proxy-models",
@@ -7,22 +7,22 @@ const MODEL_GRANT_SENTINELS: ReadonlySet<string> = new Set([
 ]);
 
 export function teamDetailHref(teamId: string): string {
-  return `${migratedHref("teams")}?team=${encodeURIComponent(teamId)}`;
+  return `${uiHref("teams")}?team=${encodeURIComponent(teamId)}`;
 }
 
 export function keyDetailHref(keyToken: string): string {
-  return `${migratedHref("api-keys")}?key=${encodeURIComponent(keyToken)}`;
+  return `${uiHref("api-keys")}?key=${encodeURIComponent(keyToken)}`;
 }
 
 export function userDetailHref(userId: string): string {
-  return `${migratedHref("users")}?user=${encodeURIComponent(userId)}`;
+  return `${uiHref("users")}?user=${encodeURIComponent(userId)}`;
 }
 
 export function orgDetailHref(orgId: string): string {
-  return `${migratedHref("organizations")}?org=${encodeURIComponent(orgId)}`;
+  return `${uiHref("organizations")}?org=${encodeURIComponent(orgId)}`;
 }
 
 export function modelGroupHref(modelGroup: string): string | undefined {
   if (MODEL_GRANT_SENTINELS.has(modelGroup)) return undefined;
-  return `${migratedHref("models-and-endpoints")}?model_group=${encodeURIComponent(modelGroup)}`;
+  return `${uiHref("models-and-endpoints")}?model_group=${encodeURIComponent(modelGroup)}`;
 }

ui/litellm-dashboard/src/utils/migratedPages.test.ts

@@ -1,239 +0,0 @@
-import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
-
-describe("migratedHref / legacyPageHref", () => {
-  beforeEach(() => {
-    vi.resetModules();
-    vi.stubEnv("NODE_ENV", "test");
-  });
-
-  afterEach(() => {
-    vi.unstubAllEnvs();
-  });
-
-  it("builds a /ui-rooted path when serverRootPath is /", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { migratedHref, legacyPageHref } = await import("./migratedPages");
-
-    expect(migratedHref("api-reference")).toBe("/ui/api-reference");
-    expect(legacyPageHref("models")).toBe("/ui/?page=models");
-  });
-
-  it("prefixes a non-root serverRootPath without duplicating slashes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" }));
-    const { migratedHref, legacyPageHref } = await import("./migratedPages");
-
-    expect(migratedHref("api-reference")).toBe("/team-x/ui/api-reference");
-    expect(legacyPageHref("models")).toBe("/team-x/ui/?page=models");
-  });
-
-  it("tolerates a leading slash in the route segment", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { migratedHref } = await import("./migratedPages");
-
-    expect(migratedHref("/api-reference")).toBe("/ui/api-reference");
-  });
-
-  it("maps both the api_ref id and the hyphenated alias to the api-reference route", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.api_ref).toBe("api-reference");
-    expect(MIGRATED_PAGES["api-reference"]).toBe("api-reference");
-  });
-
-  it("maps the api-keys landing id to its route and builds its redirect href", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES, migratedHref } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES["api-keys"]).toBe("api-keys");
-    expect(migratedHref(MIGRATED_PAGES["api-keys"])).toBe("/ui/api-keys");
-  });
-
-  it("maps the llm-playground sidebar id to the playground route", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES["llm-playground"]).toBe("playground");
-  });
-
-  it("maps the models sidebar id to the models-and-endpoints route and builds its redirect href", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES, migratedHref } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.models).toBe("models-and-endpoints");
-    expect(migratedHref(MIGRATED_PAGES.models)).toBe("/ui/models-and-endpoints");
-  });
-
-  it("maps the projects and access-groups sidebar ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.projects).toBe("projects");
-    expect(MIGRATED_PAGES["access-groups"]).toBe("access-groups");
-  });
-
-  it("maps the budgets, workflows, and guardrails-monitor sidebar ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.budgets).toBe("budgets");
-    expect(MIGRATED_PAGES.workflows).toBe("workflows");
-    expect(MIGRATED_PAGES["guardrails-monitor"]).toBe("guardrails-monitor");
-  });
-
-  it("maps the mcp-servers, search-tools, tag-management, vector-stores, and memory ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES["mcp-servers"]).toBe("mcp-servers");
-    expect(MIGRATED_PAGES["search-tools"]).toBe("search-tools");
-    expect(MIGRATED_PAGES["tag-management"]).toBe("tag-management");
-    expect(MIGRATED_PAGES["vector-stores"]).toBe("vector-stores");
-    expect(MIGRATED_PAGES.memory).toBe("memory");
-  });
-
-  it("maps the policies, guardrails, prompts, tool-policies, and skills ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.policies).toBe("policies");
-    expect(MIGRATED_PAGES.guardrails).toBe("guardrails");
-    expect(MIGRATED_PAGES.prompts).toBe("prompts");
-    expect(MIGRATED_PAGES["tool-policies"]).toBe("tool-policies");
-    expect(MIGRATED_PAGES.skills).toBe("skills");
-    // Old bookmarks used ?page=claude-code-plugins for the same panel.
-    expect(MIGRATED_PAGES["claude-code-plugins"]).toBe("skills");
-  });
-
-  it("maps the caching, cost-tracking, transform-request, ui-theme, and logs ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.caching).toBe("caching");
-    expect(MIGRATED_PAGES["cost-tracking"]).toBe("cost-tracking");
-    expect(MIGRATED_PAGES["transform-request"]).toBe("transform-request");
-    expect(MIGRATED_PAGES["ui-theme"]).toBe("ui-theme");
-    expect(MIGRATED_PAGES.logs).toBe("logs");
-  });
-
-  it("maps the admin-panel, logging-and-alerts, model-hub-table, and new_usage ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES["admin-panel"]).toBe("admin-panel");
-    expect(MIGRATED_PAGES["logging-and-alerts"]).toBe("logging-and-alerts");
-    expect(MIGRATED_PAGES["model-hub-table"]).toBe("model-hub-table");
-    // new_usage routes to /usage; the legacy ?page=usage report routes to /old-usage (asserted below).
-    expect(MIGRATED_PAGES.new_usage).toBe("usage");
-  });
-
-  it("maps the legacy usage report id to the old-usage route and builds its redirect href", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES, migratedHref } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.usage).toBe("old-usage");
-    expect(migratedHref(MIGRATED_PAGES.usage)).toBe("/ui/old-usage");
-  });
-
-  it("maps the agents and router-settings ids to their routes", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.agents).toBe("agents");
-    expect(MIGRATED_PAGES["router-settings"]).toBe("router-settings");
-  });
-
-  it("maps the users id to its route", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.users).toBe("users");
-  });
-
-  it("maps the teams id to its route", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.teams).toBe("teams");
-  });
-
-  it("maps the organizations id to its route", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { MIGRATED_PAGES } = await import("./migratedPages");
-
-    expect(MIGRATED_PAGES.organizations).toBe("organizations");
-  });
-});
-
-describe("dev server (NODE_ENV=development)", () => {
-  beforeEach(() => {
-    vi.resetModules();
-    vi.stubEnv("NODE_ENV", "development");
-  });
-
-  afterEach(() => {
-    vi.unstubAllEnvs();
-  });
-
-  it("builds root-relative hrefs because next dev serves the app at /, not /ui", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { migratedHref, legacyPageHref } = await import("./migratedPages");
-
-    expect(migratedHref("api-reference")).toBe("/api-reference");
-    expect(legacyPageHref("models")).toBe("/?page=models");
-  });
-
-  it("ignores serverRootPath, which only applies to proxy-mounted deployments", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" }));
-    const { migratedHref } = await import("./migratedPages");
-
-    expect(migratedHref("api-reference")).toBe("/api-reference");
-  });
-
-  it("maps a bare migrated path back to its legacy sidebar key", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { legacyKeyForPathname } = await import("./migratedPages");
-
-    expect(legacyKeyForPathname("/api-reference/")).toBe("api_ref");
-    expect(legacyKeyForPathname("/")).toBeNull();
-  });
-});
-
-describe("legacyKeyForPathname", () => {
-  beforeEach(() => {
-    vi.resetModules();
-    vi.stubEnv("NODE_ENV", "test");
-  });
-
-  afterEach(() => {
-    vi.unstubAllEnvs();
-  });
-
-  it("maps a migrated path back to its legacy sidebar key (including trailing slash)", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { legacyKeyForPathname } = await import("./migratedPages");
-
-    // Resolves to the sidebar key api_ref, not the hyphenated alias, so highlighting works.
-    expect(legacyKeyForPathname("/ui/api-reference")).toBe("api_ref");
-    expect(legacyKeyForPathname("/ui/api-reference/")).toBe("api_ref");
-    // Same for skills: the claude-code-plugins alias maps to the same segment,
-    // and first-match-wins iteration must keep returning the sidebar key.
-    expect(legacyKeyForPathname("/ui/skills")).toBe("skills");
-  });
-
-  it("returns null for a not-yet-migrated path", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/" }));
-    const { legacyKeyForPathname } = await import("./migratedPages");
-
-    expect(legacyKeyForPathname("/ui/")).toBeNull();
-    expect(legacyKeyForPathname("/ui/some-legacy-page")).toBeNull();
-  });
-
-  it("strips a non-root serverRootPath prefix before matching", async () => {
-    vi.doMock("@/components/networking", () => ({ serverRootPath: "/team-x/" }));
-    const { legacyKeyForPathname } = await import("./migratedPages");
-
-    expect(legacyKeyForPathname("/team-x/ui/api-reference")).toBe("api_ref");
-    expect(legacyKeyForPathname("/ui/api-reference")).toBeNull();
-  });
-});

ui/litellm-dashboard/src/utils/migratedPages.ts

@@ -1,83 +0,0 @@
-import { serverRootPath } from "@/components/networking";
-
-/**
- * Single source of truth for pages cut over from the legacy `?page=` switch in
- * app/page.tsx to path-based routes under app/(dashboard)/.
- *
- * Key = legacy page id emitted by the sidebar. Value = route segment under (dashboard)/.
- * Add an entry to route the sidebar and deep links to the new path and redirect the
- * legacy `?page=` URL; remove it to roll back.
- */
-export const MIGRATED_PAGES: Record<string, string> = {
-  "api-keys": "api-keys",
-  models: "models-and-endpoints",
-  api_ref: "api-reference",
-  // Legacy alias: older bookmarks used the hyphenated ?page=api-reference form.
-  "api-reference": "api-reference",
-  "llm-playground": "playground",
-  projects: "projects",
-  chat: "chat",
-  "access-groups": "access-groups",
-  budgets: "budgets",
-  workflows: "workflows",
-  "guardrails-monitor": "guardrails-monitor",
-  "mcp-servers": "mcp-servers",
-  "search-tools": "search-tools",
-  "tag-management": "tag-management",
-  "vector-stores": "vector-stores",
-  memory: "memory",
-  policies: "policies",
-  guardrails: "guardrails",
-  prompts: "prompts",
-  "tool-policies": "tool-policies",
-  skills: "skills",
-  // Legacy alias: the old switch matched ?page=claude-code-plugins for the same panel.
-  "claude-code-plugins": "skills",
-  caching: "caching",
-  "cost-tracking": "cost-tracking",
-  "transform-request": "transform-request",
-  "ui-theme": "ui-theme",
-  logs: "logs",
-  "admin-panel": "admin-panel",
-  "logging-and-alerts": "logging-and-alerts",
-  "model-hub-table": "model-hub-table",
-  // The modern usage dashboard; the legacy ?page=usage report routes to /old-usage.
-  new_usage: "usage",
-  usage: "old-usage",
-  "cost-optimization": "cost-optimization",
-  agents: "agents",
-  "router-settings": "router-settings",
-  users: "users",
-  teams: "teams",
-  organizations: "organizations",
-};
-
-function uiBase(): string {
-  // next dev serves the app at the root; only the proxy mounts the static export under /ui
-  // (and optionally under server_root_path). Inlined at build time, so production is unaffected.
-  if (process.env.NODE_ENV === "development") {
-    return "";
-  }
-  const root = serverRootPath && serverRootPath !== "/" ? `/${serverRootPath.replace(/^\/+|\/+$/g, "")}` : "";
-  return `${root}/ui`;
-}
-
-/** Absolute (same-origin) href for a migrated route segment, e.g. "api-reference" -> "/ui/api-reference". */
-export function migratedHref(routeSegment: string): string {
-  return `${uiBase()}/${routeSegment.replace(/^\/+/, "")}`;
-}
-
-/** Href for a not-yet-migrated page, served by the legacy `?page=` switch at the UI root. */
-export function legacyPageHref(pageKey: string): string {
-  return `${uiBase()}/?page=${pageKey}`;
-}
-
-/** Reverse-maps a path-routed location back to its legacy page id, e.g. "/ui/api-reference" -> "api_ref". */
-export function legacyKeyForPathname(pathname: string): string | null {
-  const base = uiBase();
-  const rel = (pathname.startsWith(base) ? pathname.slice(base.length) : pathname).replace(/^\/+|\/+$/g, "");
-  for (const [key, segment] of Object.entries(MIGRATED_PAGES)) {
-    if (rel === segment) return key;
-  }
-  return null;
-}