From c11ebbed27a7b2871b0741af67504aafd7407e8e Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Wed, 12 Aug 2026 23:45:36 +0000
Subject: [PATCH 01/31] fix(batches): stop uncostable batches from starving the
 cost poll page

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../proxy/common_utils/check_batch_cost.py    |  89 ++++++++-
 .../proxy_unit_tests/test_check_batch_cost.py | 180 ++++++++++++++++++
 2 files changed, 265 insertions(+), 4 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index dc8f17fb665..73d5fef08e0 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -3,7 +3,7 @@ Polls LiteLLM_ManagedObjectTable to check if the batch job is complete, and if t
 """
 
 from datetime import datetime, timedelta, timezone
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple
 
 from litellm._logging import verbose_proxy_logger
 from litellm._uuid import uuid
@@ -23,6 +23,15 @@ if TYPE_CHECKING:
 
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 
+TERMINAL_MANAGED_OBJECT_STATUSES: Final[Tuple[str, ...]] = (
+    "completed",
+    "complete",
+    "failed",
+    "expired",
+    "cancelled",
+    "stale_expired",
+)
+
 
 class CheckBatchCost:
     def __init__(
@@ -132,11 +141,11 @@ class CheckBatchCost:
         in non-terminal states as 'stale_expired'. These will never complete and
         should not be polled.
         """
-        cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
-        result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+        cutoff: Final = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+        result: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
             where={
                 "file_purpose": "batch",
-                "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
+                "status": {"not_in": list(TERMINAL_MANAGED_OBJECT_STATUSES)},
                 "created_at": {"lt": cutoff},
             },
             data={"status": "stale_expired"},
@@ -147,6 +156,26 @@ class CheckBatchCost:
                 f"(older than {MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days) as stale_expired"
             )
 
+        if not self._has_batch_processed_column:
+            return
+
+        # A row already in a terminal status is never rewritten by the sweep above, so
+        # without this it keeps a poll-page slot forever and starves newer batches.
+        retired: Final = await self.prisma_client.db.litellm_managedobjecttable.update_many(
+            where={
+                "file_purpose": "batch",
+                "batch_processed": False,
+                "status": {"in": ["complete", "completed"]},
+                "created_at": {"lt": cutoff},
+            },
+            data={"batch_processed": True},
+        )
+        if retired > 0:
+            verbose_proxy_logger.warning(
+                f"CheckBatchCost: gave up on {retired} completed managed objects older than "
+                f"{MANAGED_OBJECT_STALENESS_CUTOFF_DAYS} days that were never costed"
+            )
+
     async def _fallback_find_jobs(self) -> list:
         """Query batch jobs without the batch_processed filter (for older schemas)."""
         return await self.prisma_client.db.litellm_managedobjecttable.find_many(
@@ -167,6 +196,54 @@ class CheckBatchCost:
             order={"created_at": "asc"},
         )
 
+    async def _retire_job(self, job: "LiteLLM_ManagedObjectTable", reason: str) -> None:
+        """
+        Take a row that can never be costed out of the poll page. Leaving it selectable
+        would burn one of the MAX_OBJECTS_PER_POLL_CYCLE slots on every future cycle, and
+        once enough such rows accumulate no newer batch is ever reached. Older schemas
+        without batch_processed can only be excluded through the status filter.
+        """
+        data: Final = (
+            {"batch_processed": True}
+            if self._has_batch_processed_column
+            else {"status": "stale_expired"}
+        )
+        try:
+            await self.prisma_client.db.litellm_managedobjecttable.update(
+                where={"id": job.id},
+                data=data,
+            )
+        except Exception as db_err:
+            verbose_proxy_logger.error(
+                f"CheckBatchCost: failed to retire uncostable job {job.id} ({reason}): {db_err}"
+            )
+            return
+        verbose_proxy_logger.warning(
+            f"CheckBatchCost: job {job.id} can never be costed ({reason}), "
+            "so it will no longer be polled"
+        )
+
+    @staticmethod
+    def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
+        """A unified id that decodes but carries no model_id can never be routed."""
+        from litellm.proxy.openai_files_endpoints.common_utils import (
+            _is_base64_encoded_unified_file_id,
+            get_model_id_from_unified_batch_id,
+        )
+
+        decoded: Final = _is_base64_encoded_unified_file_id(job.unified_object_id)
+        return bool(decoded) and get_model_id_from_unified_batch_id(decoded) is None
+
+    @staticmethod
+    def _is_batch_gone_at_provider(error: Exception) -> bool:
+        """A 404 from the provider means it dropped its record of the batch, so no later
+        retrieve can ever succeed."""
+        import openai
+
+        from litellm.exceptions import NotFoundError
+
+        return isinstance(error, (NotFoundError, openai.NotFoundError))
+
     @staticmethod
     def _record_error(
         prom_logger: Optional["PrometheusLogger"], error_type: str
@@ -645,6 +722,8 @@ class CheckBatchCost:
         for job in jobs:
             routing = self._resolve_job_routing(job, prom_logger)
             if routing is None:
+                if self._has_unified_id_without_model(job):
+                    await self._retire_job(job, "unified object id has no model id")
                 continue
             model_id, batch_id = routing
 
@@ -667,6 +746,8 @@ class CheckBatchCost:
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
+                if self._is_batch_gone_at_provider(e):
+                    await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
                 continue
 
             ## RETRIEVE THE BATCH JOB OUTPUT FILE
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index 6d7ada17ec5..0418bff0cc3 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1791,3 +1791,183 @@ class TestBatchCostAttribution:
         metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
 
         assert metadata["user_api_key_alias"] == "prod-key"
+
+
+class TestPollPageStarvation:
+    """LIT-5462 regression: a row that can never be costed used to keep its slot in the
+    MAX_OBJECTS_PER_POLL_CYCLE page forever, so once enough of them accumulated no newer
+    batch was ever polled or costed."""
+
+    def _instance(self, prisma, llm_router):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import CheckBatchCost
+
+        proxy_logging_obj = MagicMock()
+        proxy_logging_obj.get_proxy_hook.return_value = None
+        return CheckBatchCost(
+            proxy_logging_obj=proxy_logging_obj,
+            prisma_client=prisma,
+            llm_router=llm_router,
+        )
+
+    def _prisma(self, jobs):
+        prisma = MagicMock()
+        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
+        prisma.db.litellm_managedobjecttable.update = AsyncMock()
+        prisma.db.litellm_managedobjecttable.find_many = AsyncMock(return_value=jobs)
+        prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+        return prisma
+
+    def _job(self, job_id, unified_object_id):
+        job = MagicMock()
+        job.id = job_id
+        job.unified_object_id = unified_object_id
+        job.created_by = "user-1"
+        return job
+
+    @staticmethod
+    def _encode(unified_id: str) -> str:
+        import base64
+
+        return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
+
+    @pytest.mark.asyncio
+    async def test_unified_id_without_model_id_is_retired(self):
+        """A unified id that decodes but carries no model_id is unroutable no matter what
+        the config says, so it must leave the poll page instead of being retried forever."""
+        prisma = self._prisma(
+            [self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
+        )
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock()
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        llm_router.aretrieve_batch.assert_not_awaited()
+        prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
+        call = prisma.db.litellm_managedobjecttable.update.call_args[1]
+        assert call["where"] == {"id": "job-no-model"}
+        assert call["data"] == {"batch_processed": True}
+
+    @pytest.mark.asyncio
+    async def test_provider_404_retires_job(self):
+        """The provider dropping its record of the batch is permanent: no later retrieve
+        can succeed, so the row must stop occupying a slot."""
+        import litellm
+
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-gone",
+                    self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock(
+            side_effect=litellm.NotFoundError(
+                message="No batch found with id 'batch_deadbeef'.",
+                model="model-123",
+                llm_provider="openai",
+            )
+        )
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_awaited_once()
+        assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
+            "batch_processed": True
+        }
+
+    @pytest.mark.asyncio
+    async def test_transient_provider_error_keeps_job_for_retry(self):
+        """A failure that may clear up (timeout, 5xx) must still leave the row unprocessed."""
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-flaky",
+                    self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_flaky"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock(side_effect=Exception("connection reset"))
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
+
+    @pytest.mark.asyncio
+    async def test_retirement_falls_back_to_status_without_batch_processed_column(self):
+        """Older schemas have no batch_processed column, so the only way to stop selecting
+        the row is the status filter the poll query already applies."""
+        prisma = self._prisma(
+            [self._job("job-legacy", self._encode("litellm_proxy;llm_batch_id:poison-no-model"))]
+        )
+        instance = self._instance(prisma, MagicMock())
+        instance._has_batch_processed_column = False
+
+        await instance.check_batch_cost()
+
+        assert prisma.db.litellm_managedobjecttable.update.call_args[1]["data"] == {
+            "status": "stale_expired"
+        }
+
+    @pytest.mark.asyncio
+    async def test_stale_cleanup_gives_up_on_never_costed_completed_rows(self):
+        """A row already in a terminal status is never rewritten by the staleness sweep, so
+        it needs its own bound or it starves newer batches indefinitely."""
+        prisma = self._prisma([])
+
+        await self._instance(prisma, MagicMock()).check_batch_cost()
+
+        calls = prisma.db.litellm_managedobjecttable.update_many.call_args_list
+        assert len(calls) == 2, "expected the staleness sweep plus the never-costed sweep"
+        where = calls[1][1]["where"]
+        assert where["file_purpose"] == "batch"
+        assert where["batch_processed"] is False
+        assert where["status"] == {"in": ["complete", "completed"]}
+        assert "created_at" in where
+        assert calls[1][1]["data"] == {"batch_processed": True}
+
+    @pytest.mark.asyncio
+    async def test_newer_batch_is_polled_once_dead_rows_are_retired(self):
+        """The end state the customer cares about: dead rows retire on the cycle they are
+        first seen, and the healthy batch behind them keeps getting polled."""
+        dead_rows = [
+            self._job("job-no-model", self._encode("litellm_proxy;llm_batch_id:poison-no-model")),
+            self._job(
+                "job-gone",
+                self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_deadbeef"),
+            ),
+        ]
+        live_row = self._job(
+            "job-live",
+            self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_live"),
+        )
+        prisma = self._prisma(dead_rows + [live_row])
+
+        import litellm
+
+        in_progress = MagicMock()
+        in_progress.status = "in_progress"
+
+        async def _retrieve(model, batch_id, litellm_metadata):
+            if batch_id == "batch_deadbeef":
+                raise litellm.NotFoundError(
+                    message="No batch found", model=model, llm_provider="openai"
+                )
+            return in_progress
+
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock(side_effect=_retrieve)
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        retired = [
+            call[1]["where"]["id"]
+            for call in prisma.db.litellm_managedobjecttable.update.call_args_list
+        ]
+        assert retired == ["job-no-model", "job-gone"]
+        assert (
+            llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live"
+        ), "the newer healthy batch must still be polled in the same cycle"

From da8414228831c36605bee568f3c90cc0d1c7b3f9 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Wed, 12 Aug 2026 21:19:37 -0700
Subject: [PATCH 02/31] fix(batches): only trust a 404 from the batch's own
 deployment

---
 .../proxy/common_utils/check_batch_cost.py    |  9 +++++-
 .../proxy_unit_tests/test_check_batch_cost.py | 29 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 73d5fef08e0..44371c8299a 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -244,6 +244,13 @@ class CheckBatchCost:
 
         return isinstance(error, (NotFoundError, openai.NotFoundError))
 
+    def _batch_deployment_exists(self, model_id: str) -> bool:
+        """A 404 only proves the batch is gone when it came from the batch's own
+        deployment. Once that deployment leaves the router, default fallbacks can
+        silently send the retrieve to a provider that never saw the batch, so its
+        404 must not retire the row; the staleness sweep bounds it instead."""
+        return self.llm_router.get_deployment(model_id=model_id) is not None
+
     @staticmethod
     def _record_error(
         prom_logger: Optional["PrometheusLogger"], error_type: str
@@ -746,7 +753,7 @@ class CheckBatchCost:
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
-                if self._is_batch_gone_at_provider(e):
+                if self._is_batch_gone_at_provider(e) and self._batch_deployment_exists(model_id):
                     await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
                 continue
 
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index 0418bff0cc3..fc3743a0490 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1878,6 +1878,35 @@ class TestPollPageStarvation:
             "batch_processed": True
         }
 
+    @pytest.mark.asyncio
+    async def test_provider_404_with_deployment_gone_keeps_job(self):
+        """With the batch's own deployment removed from the router, default fallbacks can
+        send the retrieve to a provider that never saw the batch. That 404 proves nothing,
+        so the row must stay unprocessed instead of losing its spend forever."""
+        import litellm
+
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-misrouted",
+                    self._encode("litellm_proxy;model_id:model-gone;llm_batch_id:batch_alive"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.get_deployment = MagicMock(return_value=None)
+        llm_router.aretrieve_batch = AsyncMock(
+            side_effect=litellm.NotFoundError(
+                message="No batch found with id 'batch_alive'.",
+                model="model-gone",
+                llm_provider="openai",
+            )
+        )
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_not_awaited()
+
     @pytest.mark.asyncio
     async def test_transient_provider_error_keeps_job_for_retry(self):
         """A failure that may clear up (timeout, 5xx) must still leave the row unprocessed."""

From 8947008fd2a9fbb461e18fbfbe2bd7c072c832f2 Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Thu, 13 Aug 2026 04:22:23 +0000
Subject: [PATCH 03/31] fix(batches): only retire on a 404 that names the batch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../proxy/common_utils/check_batch_cost.py    | 14 ++++++---
 .../proxy_unit_tests/test_check_batch_cost.py | 31 ++++++++++++++++++-
 2 files changed, 39 insertions(+), 6 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 44371c8299a..818c44d2039 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -235,14 +235,18 @@ class CheckBatchCost:
         return bool(decoded) and get_model_id_from_unified_batch_id(decoded) is None
 
     @staticmethod
-    def _is_batch_gone_at_provider(error: Exception) -> bool:
-        """A 404 from the provider means it dropped its record of the batch, so no later
-        retrieve can ever succeed."""
+    def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:
+        """
+        A 404 naming the batch means the provider dropped its record of it, so no later
+        retrieve can ever succeed. A 404 about anything else, a renamed Azure deployment
+        or a fallback deployment that never saw this batch, is still fixable in config, so
+        it keeps retrying.
+        """
         import openai
 
         from litellm.exceptions import NotFoundError
 
-        return isinstance(error, (NotFoundError, openai.NotFoundError))
+        return isinstance(error, (NotFoundError, openai.NotFoundError)) and batch_id in str(error)
 
     def _batch_deployment_exists(self, model_id: str) -> bool:
         """A 404 only proves the batch is gone when it came from the batch's own
@@ -753,7 +757,7 @@ class CheckBatchCost:
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
-                if self._is_batch_gone_at_provider(e) and self._batch_deployment_exists(model_id):
+                if self._is_batch_gone_at_provider(e, batch_id) and self._batch_deployment_exists(model_id):
                     await self._retire_job(job, f"batch {batch_id} no longer exists at the provider")
                 continue
 
diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py
index fc3743a0490..fa274324fd6 100644
--- a/tests/proxy_unit_tests/test_check_batch_cost.py
+++ b/tests/proxy_unit_tests/test_check_batch_cost.py
@@ -1983,7 +1983,9 @@ class TestPollPageStarvation:
         async def _retrieve(model, batch_id, litellm_metadata):
             if batch_id == "batch_deadbeef":
                 raise litellm.NotFoundError(
-                    message="No batch found", model=model, llm_provider="openai"
+                    message=f"No batch found with id '{batch_id}'.",
+                    model=model,
+                    llm_provider="openai",
                 )
             return in_progress
 
@@ -2000,3 +2002,30 @@ class TestPollPageStarvation:
         assert (
             llm_router.aretrieve_batch.await_args_list[-1][1]["batch_id"] == "batch_live"
         ), "the newer healthy batch must still be polled in the same cycle"
+
+    @pytest.mark.asyncio
+    async def test_404_that_does_not_name_the_batch_keeps_job_for_retry(self):
+        """A 404 about something other than the batch, e.g. a renamed Azure deployment, is
+        fixable in config, so the row must survive to be costed after the fix."""
+        import litellm
+
+        prisma = self._prisma(
+            [
+                self._job(
+                    "job-bad-deployment",
+                    self._encode("litellm_proxy;model_id:model-123;llm_batch_id:batch_real"),
+                )
+            ]
+        )
+        llm_router = MagicMock()
+        llm_router.aretrieve_batch = AsyncMock(
+            side_effect=litellm.NotFoundError(
+                message="Error code: 404 - DeploymentNotFound",
+                model="model-123",
+                llm_provider="azure",
+            )
+        )
+
+        await self._instance(prisma, llm_router).check_batch_cost()
+
+        prisma.db.litellm_managedobjecttable.update.assert_not_awaited()

From 63464974980e825ec226052d8ed2d786c8dde975 Mon Sep 17 00:00:00 2001
From: yassin <yassin@berri.ai>
Date: Thu, 13 Aug 2026 13:01:40 +0000
Subject: [PATCH 04/31] fix(ui): add nvidia riva to the model provider list

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../provider_create_fields.json               | 38 ++++++++++
 .../public_endpoints/test_public_endpoints.py | 33 +++++++++
 .../components/provider_info_helpers.test.tsx | 15 ++++
 .../src/components/provider_info_helpers.tsx  | 69 ++++++++-----------
 4 files changed, 113 insertions(+), 42 deletions(-)

diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json
index fcc6aac1c14..e24e5b21583 100644
--- a/litellm/proxy/public_endpoints/provider_create_fields.json
+++ b/litellm/proxy/public_endpoints/provider_create_fields.json
@@ -2062,6 +2062,44 @@
     ],
     "default_model_placeholder": "gpt-3.5-turbo"
   },
+  {
+    "provider": "NVIDIA_RIVA",
+    "provider_display_name": "Nvidia Riva",
+    "litellm_provider": "nvidia_riva",
+    "credential_fields": [
+      {
+        "key": "api_base",
+        "label": "API Base",
+        "placeholder": "grpc.nvcf.nvidia.com:443",
+        "tooltip": "host:port of the Riva gRPC endpoint. Use grpc.nvcf.nvidia.com:443 for NVCF-hosted Riva, or your own host (e.g. localhost:50051) when self-hosting. Riva has no public default, so this is required.",
+        "required": true,
+        "field_type": "text",
+        "options": null,
+        "default_value": null
+      },
+      {
+        "key": "api_key",
+        "label": "API Key",
+        "placeholder": "nvapi-...",
+        "tooltip": "Sent as gRPC authorization metadata. Required for NVCF-hosted Riva, optional for self-hosted deployments without auth.",
+        "required": false,
+        "field_type": "password",
+        "options": null,
+        "default_value": null
+      },
+      {
+        "key": "nvcf_function_id",
+        "label": "NVCF Function ID",
+        "placeholder": "1598d209-5e27-4d3c-8079-4751568b1081",
+        "tooltip": "NVCF function id of the hosted Riva model. Setting it turns on TLS and the function-id gRPC metadata. Leave empty for self-hosted Riva.",
+        "required": false,
+        "field_type": "text",
+        "options": null,
+        "default_value": null
+      }
+    ],
+    "default_model_placeholder": "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr"
+  },
   {
     "provider": "Ollama",
     "provider_display_name": "Ollama",
diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index 88dc07e741b..ead5f4ab5cc 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -243,6 +243,39 @@ def test_bedrock_mantle_provider_fields():
     assert fields_by_key["api_base"]["field_type"] == "text"
 
 
+def test_nvidia_riva_provider_fields():
+    """The Add Model provider dropdown is populated from /public/providers/fields, so a
+    missing entry meant Riva could not be added through the UI. Riva is gRPC only with no
+    public default endpoint, hence the required api_base, and NVCF hosted Riva cannot be
+    called without nvcf_function_id.
+    """
+    app_instance = FastAPI()
+    app_instance.include_router(router)
+    test_client = TestClient(app_instance)
+
+    response = test_client.get("/public/providers/fields")
+    assert response.status_code == 200
+    providers = response.json()
+
+    riva = next((p for p in providers if p["provider"] == "NVIDIA_RIVA"), None)
+    assert riva is not None, "NVIDIA Riva provider entry not found"
+
+    assert riva["provider_display_name"] == "Nvidia Riva"
+    assert riva["litellm_provider"] == LlmProviders.NVIDIA_RIVA.value
+    assert riva["default_model_placeholder"].startswith("nvidia_riva/")
+
+    fields_by_key = {f["key"]: f for f in riva["credential_fields"]}
+
+    assert fields_by_key["api_base"]["required"] is True
+    assert fields_by_key["api_base"]["field_type"] == "text"
+
+    assert fields_by_key["api_key"]["required"] is False
+    assert fields_by_key["api_key"]["field_type"] == "password"
+
+    assert "nvcf_function_id" in fields_by_key
+    assert fields_by_key["nvcf_function_id"]["required"] is False
+
+
 def test_google_ai_studio_provider_fields_expose_api_base():
     """The Google AI Studio (gemini) credential form must let admins set a custom
     api_base so they can point at a Gemini-compatible gateway (e.g. a self-hosted
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
index 777cdc62987..f43200570b0 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
@@ -94,6 +94,17 @@ describe("provider_info_helpers", () => {
       expect(result.displayName).toBe(Providers.ZAI);
     });
 
+    it("should resolve the nvidia_riva provider value to the Nvidia Riva display name and logo", () => {
+      // The backend registers nvidia_riva and it has a docs page, but the UI
+      // registry had no entry, so it could not be picked in Add Model and the
+      // slug rendered raw with no logo.
+      const result = getProviderLogoAndName("nvidia_riva");
+      expect(result.displayName).toBe(Providers.NVIDIA_RIVA);
+      expect(provider_map.NVIDIA_RIVA).toBe("nvidia_riva");
+      expect(result.logo).toBe(providerLogoMap[Providers.NVIDIA_RIVA]);
+      expect(result.logo).toBeTruthy();
+    });
+
     it("should return provider value as display name when no mapping exists", () => {
       const unknownProvider = "unknown_provider";
       const result = getProviderLogoAndName(unknownProvider);
@@ -225,6 +236,10 @@ describe("provider_info_helpers", () => {
       expect(getPlaceholder(Providers.ZAI)).toBe("zai/glm-4.5");
     });
 
+    it("should return the riva asr placeholder for NVIDIA_RIVA provider", () => {
+      expect(getPlaceholder(Providers.NVIDIA_RIVA)).toBe("nvidia_riva/nvidia/parakeet-ctc-1_1b-asr");
+    });
+
     it("should return default gpt-3.5-turbo placeholder for unknown provider", () => {
       expect(getPlaceholder("UnknownProvider" as any)).toBe("gpt-3.5-turbo");
     });
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
index fa6b3c79230..86d67867fcf 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
@@ -130,6 +130,7 @@ export enum Providers {
   NOVITA = "Novita",
   NSCALE = "Nscale",
   NVIDIA_NIM = "Nvidia Nim",
+  NVIDIA_RIVA = "Nvidia Riva",
   Ollama = "Ollama",
   OLLAMA_CHAT = "Ollama Chat",
   OOBABOOGA = "Oobabooga",
@@ -238,6 +239,7 @@ export const provider_map: Record<string, string> = {
   NOVITA: "novita",
   NSCALE: "nscale",
   NVIDIA_NIM: "nvidia_nim",
+  NVIDIA_RIVA: "nvidia_riva",
   Ollama: "ollama",
   OLLAMA_CHAT: "ollama_chat",
   OOBABOOGA: "oobabooga",
@@ -334,6 +336,7 @@ export const providerLogoMap: Partial<Record<Providers, string>> = {
   [Providers.NEBIUS]: nebiusLogo.src,
   [Providers.NOVITA]: novitaLogo.src,
   [Providers.NVIDIA_NIM]: nvidiaNimLogo.src,
+  [Providers.NVIDIA_RIVA]: nvidiaNimLogo.src,
   [Providers.Ollama]: ollamaLogo.src,
   [Providers.OLLAMA_CHAT]: ollamaLogo.src,
   [Providers.OOBABOOGA]: openaiSmallLogo.src,
@@ -400,50 +403,32 @@ export const getProviderLogoAndName = (providerValue: string): { logo: string; d
   return { logo, displayName };
 };
 
-export const getPlaceholder = (selectedProvider: string): string => {
-  if (selectedProvider === Providers.AIML) {
-    return "aiml/flux-pro/v1.1";
-  } else if (selectedProvider === Providers.Vertex_AI) {
-    return "gemini-pro";
-  } else if (selectedProvider == Providers.Anthropic) {
-    return "claude-3-opus";
-  } else if (selectedProvider == Providers.Bedrock) {
-    return "claude-3-opus";
-  } else if (selectedProvider == Providers.SageMaker) {
-    return "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b";
-  } else if (selectedProvider == Providers.Google_AI_Studio) {
-    return "gemini-pro";
-  } else if (selectedProvider == Providers.Azure_AI_Studio) {
-    return "azure_ai/command-r-plus";
-  } else if (selectedProvider == Providers.Azure) {
-    return "my-deployment";
-  } else if (selectedProvider == Providers.Oracle) {
-    return "oci/xai.grok-4";
-  } else if (selectedProvider == Providers.Snowflake) {
-    return "snowflake/mistral-7b";
-  } else if (selectedProvider == Providers.Voyage) {
-    return "voyage/";
-  } else if (selectedProvider == Providers.JinaAI) {
-    return "jina_ai/";
-  } else if (selectedProvider == Providers.VolcEngine) {
-    return "volcengine/<any-model-on-volcengine>";
-  } else if (selectedProvider == Providers.DeepInfra) {
-    return "deepinfra/<any-model-on-deepinfra>";
-  } else if (selectedProvider == Providers.FalAI) {
-    return "fal_ai/fal-ai/flux-pro/v1.1-ultra";
-  } else if (selectedProvider == Providers.RunwayML) {
-    return "runwayml/gen4_turbo";
-  } else if (selectedProvider === Providers.WATSONX) {
-    return "watsonx/ibm/granite-3-3-8b-instruct";
-  } else if (selectedProvider === Providers.Cursor) {
-    return "cursor/claude-4-sonnet";
-  } else if (selectedProvider === Providers.ZAI) {
-    return "zai/glm-4.5";
-  } else {
-    return "gpt-3.5-turbo";
-  }
+const providerPlaceholderMap: Partial<Record<Providers, string>> = {
+  [Providers.AIML]: "aiml/flux-pro/v1.1",
+  [Providers.Anthropic]: "claude-3-opus",
+  [Providers.Azure]: "my-deployment",
+  [Providers.Azure_AI_Studio]: "azure_ai/command-r-plus",
+  [Providers.Bedrock]: "claude-3-opus",
+  [Providers.Cursor]: "cursor/claude-4-sonnet",
+  [Providers.DeepInfra]: "deepinfra/<any-model-on-deepinfra>",
+  [Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra",
+  [Providers.Google_AI_Studio]: "gemini-pro",
+  [Providers.JinaAI]: "jina_ai/",
+  [Providers.NVIDIA_RIVA]: "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr",
+  [Providers.Oracle]: "oci/xai.grok-4",
+  [Providers.RunwayML]: "runwayml/gen4_turbo",
+  [Providers.SageMaker]: "sagemaker/jumpstart-dft-meta-textgeneration-llama-2-7b",
+  [Providers.Snowflake]: "snowflake/mistral-7b",
+  [Providers.Vertex_AI]: "gemini-pro",
+  [Providers.VolcEngine]: "volcengine/<any-model-on-volcengine>",
+  [Providers.Voyage]: "voyage/",
+  [Providers.WATSONX]: "watsonx/ibm/granite-3-3-8b-instruct",
+  [Providers.ZAI]: "zai/glm-4.5",
 };
 
+export const getPlaceholder = (selectedProvider: string): string =>
+  providerPlaceholderMap[selectedProvider as Providers] ?? "gpt-3.5-turbo";
+
 export const getProviderModels = (provider: Providers, modelMap: any): Array<string> => {
   let providerKey = provider;
   let custom_llm_provider = provider_map[providerKey];

From 3b09484344defaf64a475f97f9b5ad0de35bce8b Mon Sep 17 00:00:00 2001
From: mateo <mateo@berri.ai>
Date: Thu, 13 Aug 2026 19:48:16 +0000
Subject: [PATCH 05/31] refactor(batches): decode unified ids through the
 public helper

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
---
 .../proxy/common_utils/check_batch_cost.py               | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
index 818c44d2039..6fe37f0aacb 100644
--- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
+++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
@@ -227,12 +227,15 @@ class CheckBatchCost:
     def _has_unified_id_without_model(job: "LiteLLM_ManagedObjectTable") -> bool:
         """A unified id that decodes but carries no model_id can never be routed."""
         from litellm.proxy.openai_files_endpoints.common_utils import (
-            _is_base64_encoded_unified_file_id,
+            convert_b64_uid_to_unified_uid,
             get_model_id_from_unified_batch_id,
         )
 
-        decoded: Final = _is_base64_encoded_unified_file_id(job.unified_object_id)
-        return bool(decoded) and get_model_id_from_unified_batch_id(decoded) is None
+        decoded: Final = convert_b64_uid_to_unified_uid(job.unified_object_id)
+        return (
+            decoded != job.unified_object_id
+            and get_model_id_from_unified_batch_id(decoded) is None
+        )
 
     @staticmethod
     def _is_batch_gone_at_provider(error: Exception, batch_id: str) -> bool:

From 8ada34fe2cfa2ae948a75f8b6bd6a3395a0c1274 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 13:45:33 -0700
Subject: [PATCH 06/31] test(ui): characterise guardrails route behaviour
 before shadcn migration

Rewrite the PatternModal test off antd class selectors onto role and text
queries, and add characterisation tests for the guardrails components that
had none: the keyword modal, the content filter display and configuration,
the guardrail garden and the custom code modal.

Also covers two behaviours the migration must preserve: the action and
severity dropdowns in the content filter tables, and test playground state
surviving a tab switch away and back.

Every assertion here passes against the current antd and Tremor components
so the same file can prove the shadcn versions unedited.
---
 .../_components/GuardrailsPanel.test.tsx      |  32 +++-
 .../ContentFilterConfiguration.test.tsx       | 136 +++++++++++++++
 .../ContentFilterDisplay.test.tsx             |  89 ++++++++++
 .../ContentFilterManager.test.tsx             |  24 +--
 .../ContentFilterTables.test.tsx              |  73 ++++++++
 .../content_filter/KeywordModal.test.tsx      |  87 ++++++++++
 .../content_filter/PatternModal.test.tsx      | 130 ++++++++------
 .../custom_code/CustomCodeModal.test.tsx      | 162 ++++++++++++++++++
 .../_components/guardrail_garden.test.tsx     | 104 +++++++++++
 9 files changed, 763 insertions(+), 74 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
index 29655e5910a..88e52e0cdd1 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
@@ -35,10 +35,19 @@ vi.mock("./guardrail_info", () => ({
   default: () => <div>Mock Guardrail Info View</div>,
 }));
 
-vi.mock("./GuardrailTestPlayground", () => ({
-  __esModule: true,
-  default: () => <div>Mock Guardrail Test Playground</div>,
-}));
+vi.mock("./GuardrailTestPlayground", async () => {
+  const { useState } = await import("react");
+  const MockGuardrailTestPlayground = () => {
+    const [draft, setDraft] = useState("");
+    return (
+      <div>
+        <div>Mock Guardrail Test Playground</div>
+        <input aria-label="playground draft" value={draft} onChange={(e) => setDraft(e.target.value)} />
+      </div>
+    );
+  };
+  return { __esModule: true, default: MockGuardrailTestPlayground };
+});
 
 vi.mock("./TeamGuardrailsTab", () => ({
   TeamGuardrailsTab: () => <div>Mock Team Guardrails Tab</div>,
@@ -129,6 +138,21 @@ describe("GuardrailsPanel", () => {
     expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2);
   });
 
+  it("should keep test playground state when switching tabs away and back", async () => {
+    render(<GuardrailsPanel {...defaultProps} />);
+
+    fireEvent.click(screen.getByText("Test Playground"));
+
+    const draft = await screen.findByLabelText("playground draft");
+    fireEvent.change(draft, { target: { value: "keep me" } });
+    expect(draft).toHaveValue("keep me");
+
+    fireEvent.click(screen.getByText("Guardrails"));
+    fireEvent.click(screen.getByText("Test Playground"));
+
+    expect(await screen.findByLabelText("playground draft")).toHaveValue("keep me");
+  });
+
   it("should not delete anything when the modal is cancelled", async () => {
     render(<GuardrailsPanel {...defaultProps} />);
     fireEvent.click(screen.getByText("Guardrails"));
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.test.tsx
new file mode 100644
index 00000000000..37a3d40128e
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.test.tsx
@@ -0,0 +1,136 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { renderWithProviders, screen } from "@/../tests/test-utils";
+import userEvent from "@testing-library/user-event";
+import ContentFilterConfiguration from "./ContentFilterConfiguration";
+
+vi.mock("@/components/networking", () => ({
+  validateBlockedWordsFile: vi.fn(),
+  getCategoryYaml: vi.fn(),
+}));
+
+vi.mock("@/components/molecules/notifications_manager", () => ({
+  default: { success: vi.fn(), error: vi.fn(), fromBackend: vi.fn() },
+}));
+
+const PREBUILT = [
+  { name: "us_ssn", display_name: "US Social Security Number", category: "PII Patterns", description: "d" },
+];
+
+describe("ContentFilterConfiguration", () => {
+  const handlers = {
+    onPatternAdd: vi.fn(),
+    onPatternRemove: vi.fn(),
+    onPatternActionChange: vi.fn(),
+    onBlockedWordAdd: vi.fn(),
+    onBlockedWordRemove: vi.fn(),
+    onBlockedWordUpdate: vi.fn(),
+  };
+
+  const renderConfig = (overrides = {}) =>
+    renderWithProviders(
+      <ContentFilterConfiguration
+        prebuiltPatterns={PREBUILT}
+        categories={["PII Patterns"]}
+        selectedPatterns={[]}
+        blockedWords={[]}
+        accessToken="test-token"
+        {...handlers}
+        {...overrides}
+      />,
+    );
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should render the pattern and keyword sections", () => {
+    renderConfig();
+
+    expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
+    expect(
+      screen.getByText("Detect sensitive information using regex patterns (SSN, credit cards, API keys, etc.)"),
+    ).toBeInTheDocument();
+    expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
+    expect(screen.getByText("Block or mask specific sensitive terms and phrases")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /add prebuilt pattern/i })).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /add custom regex/i })).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /add keyword/i })).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /upload yaml file/i })).toBeInTheDocument();
+  });
+
+  it("should show the empty states for patterns and keywords", () => {
+    renderConfig();
+
+    expect(screen.getByText("No patterns added.")).toBeInTheDocument();
+    expect(screen.getByText("No keywords added.")).toBeInTheDocument();
+  });
+
+  it("should open the prebuilt pattern modal", async () => {
+    const user = userEvent.setup();
+    renderConfig();
+
+    expect(screen.queryByText("Pattern type")).not.toBeInTheDocument();
+
+    await user.click(screen.getByRole("button", { name: /add prebuilt pattern/i }));
+
+    expect(await screen.findByText("Pattern type")).toBeInTheDocument();
+  });
+
+  it("should open the custom regex modal", async () => {
+    const user = userEvent.setup();
+    renderConfig();
+
+    expect(screen.queryByText("Add custom regex pattern")).not.toBeInTheDocument();
+
+    await user.click(screen.getByRole("button", { name: /add custom regex/i }));
+
+    expect(await screen.findByText("Add custom regex pattern")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., ID-[0-9]{6}")).toBeInTheDocument();
+  });
+
+  it("should open the keyword modal", async () => {
+    const user = userEvent.setup();
+    renderConfig();
+
+    expect(screen.queryByText("Add blocked keyword")).not.toBeInTheDocument();
+
+    await user.click(screen.getByRole("button", { name: /add keyword/i }));
+
+    expect(await screen.findByText("Add blocked keyword")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("Enter sensitive keyword or phrase")).toBeInTheDocument();
+  });
+
+  it("should list already selected patterns and keywords", () => {
+    renderConfig({
+      selectedPatterns: [
+        {
+          id: "pattern-1",
+          type: "prebuilt" as const,
+          name: "us_ssn",
+          display_name: "US Social Security Number",
+          action: "BLOCK" as const,
+        },
+      ],
+      blockedWords: [{ id: "word-1", keyword: "secret", action: "MASK" as const, description: "Sensitive" }],
+    });
+
+    expect(screen.getByText("US Social Security Number")).toBeInTheDocument();
+    expect(screen.getByText("secret")).toBeInTheDocument();
+    expect(screen.queryByText("No patterns added.")).not.toBeInTheDocument();
+    expect(screen.queryByText("No keywords added.")).not.toBeInTheDocument();
+  });
+
+  it("should show only the keyword section when the keywords step is requested", () => {
+    renderConfig({ showStep: "keywords" });
+
+    expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
+    expect(screen.queryByText("Pattern Detection")).not.toBeInTheDocument();
+  });
+
+  it("should show only the pattern section when the patterns step is requested", () => {
+    renderConfig({ showStep: "patterns" });
+
+    expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
+    expect(screen.queryByText("Blocked Keywords")).not.toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.test.tsx
new file mode 100644
index 00000000000..8fb48791f76
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.test.tsx
@@ -0,0 +1,89 @@
+import { describe, it, expect } from "vitest";
+import { renderWithProviders, screen } from "@/../tests/test-utils";
+import ContentFilterDisplay from "./ContentFilterDisplay";
+
+const PATTERN = {
+  id: "pattern-1",
+  type: "prebuilt" as const,
+  name: "email",
+  display_name: "Email address",
+  action: "BLOCK" as const,
+};
+
+const KEYWORD = {
+  id: "word-1",
+  keyword: "secret",
+  action: "MASK" as const,
+  description: "Sensitive term",
+};
+
+const CATEGORY = {
+  id: "category-1",
+  category: "self_harm",
+  display_name: "Self Harm",
+  action: "BLOCK" as const,
+  severity_threshold: "high" as const,
+};
+
+describe("ContentFilterDisplay", () => {
+  it("should render nothing when there is no content filter data", () => {
+    const { container } = renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[]} categories={[]} />);
+
+    expect(container).toBeEmptyDOMElement();
+  });
+
+  it("should render the categories section with a configured count", () => {
+    renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[]} categories={[CATEGORY]} />);
+
+    expect(screen.getByText("Content Categories")).toBeInTheDocument();
+    expect(screen.getByText("1 categories configured")).toBeInTheDocument();
+    expect(screen.getByText("Self Harm")).toBeInTheDocument();
+    expect(screen.queryByText("Pattern Detection")).not.toBeInTheDocument();
+    expect(screen.queryByText("Blocked Keywords")).not.toBeInTheDocument();
+  });
+
+  it("should render the patterns section with a configured count", () => {
+    renderWithProviders(<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[]} categories={[]} />);
+
+    expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
+    expect(screen.getByText("1 patterns configured")).toBeInTheDocument();
+    expect(screen.getByText("Email address")).toBeInTheDocument();
+    expect(screen.queryByText("Content Categories")).not.toBeInTheDocument();
+  });
+
+  it("should render the keywords section with a configured count", () => {
+    renderWithProviders(<ContentFilterDisplay patterns={[]} blockedWords={[KEYWORD]} categories={[]} />);
+
+    expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
+    expect(screen.getByText("1 keywords configured")).toBeInTheDocument();
+    expect(screen.getByText("secret")).toBeInTheDocument();
+    expect(screen.getByText("Sensitive term")).toBeInTheDocument();
+  });
+
+  it("should render every section when all three kinds of data are present", () => {
+    renderWithProviders(<ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} />);
+
+    expect(screen.getByText("Content Categories")).toBeInTheDocument();
+    expect(screen.getByText("Pattern Detection")).toBeInTheDocument();
+    expect(screen.getByText("Blocked Keywords")).toBeInTheDocument();
+  });
+
+  it("should render category severity and action as static text in read-only mode", () => {
+    renderWithProviders(
+      <ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} readOnly={true} />,
+    );
+
+    expect(screen.getByText("HIGH")).toBeInTheDocument();
+    expect(screen.getByText("BLOCK")).toBeInTheDocument();
+    expect(screen.getAllByRole("button", { name: /delete/i })).toHaveLength(2);
+  });
+
+  it("should render category severity and action as editable controls when not read-only", () => {
+    renderWithProviders(
+      <ContentFilterDisplay patterns={[PATTERN]} blockedWords={[KEYWORD]} categories={[CATEGORY]} readOnly={false} />,
+    );
+
+    expect(screen.queryByText("HIGH")).not.toBeInTheDocument();
+    expect(screen.getAllByRole("button", { name: /delete/i })).toHaveLength(3);
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.test.tsx
index 66758ac72f1..c3fd2e3eb59 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.test.tsx
@@ -2,7 +2,6 @@ import { describe, it, expect, vi, beforeEach } from "vitest";
 import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import ContentFilterManager, { formatContentFilterDataForAPI } from "./ContentFilterManager";
-import React from "react";
 
 const CONTENT_FILTER_GUARDRAIL_DATA = {
   litellm_params: {
@@ -85,18 +84,7 @@ vi.mock("./ContentFilterDisplay", () => ({
   ),
 }));
 
-vi.mock("antd", async (importOriginal) => {
-  const actual = await importOriginal<typeof import("antd")>();
-  return {
-    ...actual,
-    Divider: ({ children }: { children: React.ReactNode }) => <div data-testid="divider">{children}</div>,
-    Alert: ({ message, type }: { message: React.ReactNode; type: string }) => (
-      <div data-testid="unsaved-alert" data-type={type}>
-        {message}
-      </div>
-    ),
-  };
-});
+const UNSAVED_CHANGES_TEXT = /You have unsaved changes to patterns or keywords/;
 
 describe("ContentFilterManager", () => {
   beforeEach(() => {
@@ -117,7 +105,7 @@ describe("ContentFilterManager", () => {
       expect(screen.getByTestId("content-filter-config")).toBeInTheDocument();
     });
 
-    expect(screen.getByTestId("divider")).toHaveTextContent("Content Filter Configuration");
+    expect(screen.getByText("Content Filter Configuration")).toBeInTheDocument();
   });
 
   it("should return null when guardrail is not litellm_content_filter", () => {
@@ -275,15 +263,15 @@ describe("ContentFilterManager", () => {
       expect(screen.getByTestId("content-filter-config")).toBeInTheDocument();
     });
 
-    expect(screen.queryByTestId("unsaved-alert")).not.toBeInTheDocument();
+    expect(screen.queryByText(UNSAVED_CHANGES_TEXT)).not.toBeInTheDocument();
 
     await user.click(screen.getByRole("button", { name: /add pattern/i }));
 
     await waitFor(() => {
-      expect(screen.getByTestId("unsaved-alert")).toBeInTheDocument();
+      expect(screen.getByText(UNSAVED_CHANGES_TEXT)).toBeInTheDocument();
     });
 
-    expect(screen.getByTestId("unsaved-alert")).toHaveTextContent(/unsaved changes.*Save Changes/i);
+    expect(screen.getByText(UNSAVED_CHANGES_TEXT)).toHaveTextContent(/Save Changes/i);
   });
 
   it("should call onDataChange when patterns or keywords change", async () => {
@@ -371,7 +359,7 @@ describe("ContentFilterManager", () => {
     );
 
     await waitFor(() => {
-      expect(screen.getByTestId("divider")).toBeInTheDocument();
+      expect(screen.getByText("Content Filter Configuration")).toBeInTheDocument();
     });
 
     expect(screen.queryByTestId("content-filter-config")).not.toBeInTheDocument();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterTables.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterTables.test.tsx
index cc27394f2c3..c56ee8ff528 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterTables.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterTables.test.tsx
@@ -130,4 +130,77 @@ describe("content filter tables", () => {
 
     expect(onCategoryRemove).toHaveBeenCalledWith("category-1");
   });
+
+  it("should report a pattern action change", async () => {
+    const onActionChange = vi.fn();
+    const user = userEvent.setup();
+
+    renderWithProviders(
+      <PatternTable
+        patterns={[{ id: "pattern-1", type: "prebuilt", name: "email", action: "BLOCK" }]}
+        onActionChange={onActionChange}
+        onRemove={vi.fn()}
+      />,
+    );
+
+    await user.click(screen.getByRole("combobox"));
+    const maskOptions = await screen.findAllByText("Mask");
+    await user.click(maskOptions[maskOptions.length - 1]);
+
+    expect(onActionChange).toHaveBeenCalledWith("pattern-1", "MASK");
+  });
+
+  it("should report a keyword action change", async () => {
+    const onActionChange = vi.fn();
+    const user = userEvent.setup();
+
+    renderWithProviders(
+      <KeywordTable
+        keywords={[{ id: "keyword-1", keyword: "secret", action: "BLOCK" }]}
+        onActionChange={onActionChange}
+        onRemove={vi.fn()}
+      />,
+    );
+
+    await user.click(screen.getByRole("combobox"));
+    const maskOptions = await screen.findAllByText("Mask");
+    await user.click(maskOptions[maskOptions.length - 1]);
+
+    expect(onActionChange).toHaveBeenCalledWith("keyword-1", "action", "MASK");
+  });
+
+  it("should report category severity and action changes", async () => {
+    const onSeverityChange = vi.fn();
+    const onActionChange = vi.fn();
+    const user = userEvent.setup();
+
+    renderWithProviders(
+      <CategoryTable
+        categories={[
+          {
+            id: "category-1",
+            category: "self_harm",
+            display_name: "Self Harm",
+            action: "BLOCK",
+            severity_threshold: "high",
+          },
+        ]}
+        onActionChange={onActionChange}
+        onSeverityChange={onSeverityChange}
+        onRemove={vi.fn()}
+      />,
+    );
+
+    await user.click(screen.getAllByRole("combobox")[0]);
+    const lowOptions = await screen.findAllByText("Low");
+    await user.click(lowOptions[lowOptions.length - 1]);
+
+    expect(onSeverityChange).toHaveBeenCalledWith("category-1", "low");
+
+    await user.click(screen.getAllByRole("combobox")[1]);
+    const maskOptions = await screen.findAllByText("Mask");
+    await user.click(maskOptions[maskOptions.length - 1]);
+
+    expect(onActionChange).toHaveBeenCalledWith("category-1", "MASK");
+  });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx
new file mode 100644
index 00000000000..dda479f1579
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.test.tsx
@@ -0,0 +1,87 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import KeywordModal from "./KeywordModal";
+
+describe("KeywordModal", () => {
+  const handlers = {
+    onKeywordChange: vi.fn(),
+    onActionChange: vi.fn(),
+    onDescriptionChange: vi.fn(),
+    onAdd: vi.fn(),
+    onCancel: vi.fn(),
+  };
+
+  const renderModal = (overrides: Partial<React.ComponentProps<typeof KeywordModal>> = {}) =>
+    render(<KeywordModal visible keyword="" action="BLOCK" description="" {...handlers} {...overrides} />);
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it("should render the keyword, action and description fields", async () => {
+    renderModal();
+
+    expect(await screen.findByText("Add blocked keyword")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("Enter sensitive keyword or phrase")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("Explain why this keyword is sensitive")).toBeInTheDocument();
+    expect(screen.getByText("Description (optional)")).toBeInTheDocument();
+    expect(
+      screen.getByText("Choose what action the guardrail should take when this keyword is detected"),
+    ).toBeInTheDocument();
+  });
+
+  it("should report keyword edits", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.type(await screen.findByPlaceholderText("Enter sensitive keyword or phrase"), "s");
+
+    expect(handlers.onKeywordChange).toHaveBeenCalledWith("s");
+  });
+
+  it("should report description edits", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.type(await screen.findByPlaceholderText("Explain why this keyword is sensitive"), "x");
+
+    expect(handlers.onDescriptionChange).toHaveBeenCalledWith("x");
+  });
+
+  it("should report the chosen action", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.click(await screen.findByRole("combobox"));
+    const maskOptions = await screen.findAllByText("Mask");
+    await user.click(maskOptions[maskOptions.length - 1]);
+
+    expect(handlers.onActionChange).toHaveBeenCalled();
+    expect(handlers.onActionChange.mock.calls[0][0]).toBe("MASK");
+  });
+
+  it("should show the current keyword and description values", async () => {
+    renderModal({ keyword: "secret", description: "sensitive term" });
+
+    expect(await screen.findByDisplayValue("secret")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("sensitive term")).toBeInTheDocument();
+  });
+
+  it("should add and cancel through the footer buttons", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.click(await screen.findByRole("button", { name: "Add" }));
+    expect(handlers.onAdd).toHaveBeenCalledTimes(1);
+
+    await user.click(screen.getByRole("button", { name: "Cancel" }));
+    expect(handlers.onCancel).toHaveBeenCalledTimes(1);
+  });
+
+  it("should not render its content when not visible", () => {
+    renderModal({ visible: false });
+
+    expect(screen.queryByText("Add blocked keyword")).not.toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
index 4e6fd0db809..4311f4e13aa 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
@@ -1,5 +1,5 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { render, screen, waitFor } from "@testing-library/react";
+import { render, screen } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import PatternModal from "./PatternModal";
 
@@ -10,25 +10,25 @@ describe("PatternModal", () => {
   const mockOnActionChange = vi.fn();
 
   const mockPrebuiltPatterns = [
-    { name: "us_ssn", category: "PII Patterns", description: "US Social Security Number" },
-    { name: "email", category: "PII Patterns", description: "Email addresses" },
-    { name: "visa", category: "Financial Patterns", description: "Visa credit card numbers" },
-    { name: "aws_access_key", category: "Credential Patterns", description: "AWS Access Keys" },
+    {
+      name: "us_ssn",
+      display_name: "US Social Security Number",
+      category: "PII Patterns",
+      description: "US Social Security Number",
+    },
+    { name: "email", display_name: "Email address", category: "PII Patterns", description: "Email addresses" },
+    { name: "visa", display_name: "Visa card", category: "Financial Patterns", description: "Visa credit cards" },
+    {
+      name: "aws_access_key",
+      display_name: "AWS access key",
+      category: "Credential Patterns",
+      description: "AWS Access Keys",
+    },
   ];
 
   const mockCategories = ["PII Patterns", "Financial Patterns", "Credential Patterns"];
 
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
-  it("should show dropdown with prebuilt pattern options grouped by category", async () => {
-    /**
-     * Tests that the modal displays a dropdown with prebuilt patterns
-     * organized by category. This verifies the pattern selection UI is working.
-     */
-    const user = userEvent.setup();
-
+  const renderModal = () =>
     render(
       <PatternModal
         visible={true}
@@ -43,49 +43,75 @@ describe("PatternModal", () => {
       />,
     );
 
-    // Wait for modal to be visible
-    await waitFor(() => {
-      expect(screen.getByText("Add prebuilt pattern")).toBeInTheDocument();
-    });
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
 
-    // Find the pattern type dropdown by looking for the first combobox input
-    const comboboxes = screen.getAllByRole("combobox");
-    const dropdown = comboboxes[0]; // First combobox is the pattern selector
-    expect(dropdown).toBeInTheDocument();
+  it("should show prebuilt pattern options grouped by category and report the picked pattern", async () => {
+    const user = userEvent.setup();
+    renderModal();
 
-    // Click to open the dropdown
-    await user.click(dropdown);
+    expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
 
-    // Verify that pattern options are available in the dropdown
-    // Ant Design renders Select options in a portal, so we need to query the whole document
-    await waitFor(() => {
-      const options = document.querySelectorAll(".ant-select-item-option");
-      expect(options.length).toBeGreaterThan(0);
-    });
+    await user.click(screen.getAllByRole("combobox")[0]);
 
-    // Verify categories are shown as group labels
-    await waitFor(() => {
-      expect(document.body).toHaveTextContent("PII Patterns");
-      expect(document.body).toHaveTextContent("Financial Patterns");
-      expect(document.body).toHaveTextContent("Credential Patterns");
-    });
+    expect(await screen.findByText("PII Patterns")).toBeInTheDocument();
+    expect(screen.getByText("Financial Patterns")).toBeInTheDocument();
+    expect(screen.getByText("Credential Patterns")).toBeInTheDocument();
 
-    // Verify pattern options are available
-    expect(document.body).toHaveTextContent("us_ssn");
-    expect(document.body).toHaveTextContent("email");
-    expect(document.body).toHaveTextContent("visa");
-    expect(document.body).toHaveTextContent("aws_access_key");
+    expect(screen.getByText("Email address")).toBeInTheDocument();
+    expect(screen.getByText("Visa card")).toBeInTheDocument();
+    expect(screen.getByText("AWS access key")).toBeInTheDocument();
 
-    // Select a pattern by clicking on its option element
-    const ssnOption = Array.from(document.querySelectorAll(".ant-select-item-option")).find(
-      (el) => el.textContent === "us_ssn",
-    ) as HTMLElement;
-    await user.click(ssnOption);
+    const ssnOptions = await screen.findAllByText("US Social Security Number");
+    await user.click(ssnOptions[ssnOptions.length - 1]);
 
-    // Verify the change handler was called with the pattern name
-    // Note: Ant Design Select calls onChange with (value, option), so we check if it was called
     expect(mockOnPatternNameChange).toHaveBeenCalled();
-    const callArgs = mockOnPatternNameChange.mock.calls[0];
-    expect(callArgs[0]).toBe("us_ssn");
+    expect(mockOnPatternNameChange.mock.calls[0][0]).toBe("us_ssn");
+  });
+
+  it("should report the chosen action", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
+
+    await user.click(screen.getAllByRole("combobox")[1]);
+    const maskOptions = await screen.findAllByText("Mask");
+    await user.click(maskOptions[maskOptions.length - 1]);
+
+    expect(mockOnActionChange).toHaveBeenCalled();
+    expect(mockOnActionChange.mock.calls[0][0]).toBe("MASK");
+  });
+
+  it("should add and cancel through the footer buttons", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
+
+    await user.click(screen.getByRole("button", { name: "Add" }));
+    expect(mockOnAdd).toHaveBeenCalledTimes(1);
+
+    await user.click(screen.getByRole("button", { name: "Cancel" }));
+    expect(mockOnCancel).toHaveBeenCalledTimes(1);
+  });
+
+  it("should not render its content when not visible", () => {
+    render(
+      <PatternModal
+        visible={false}
+        prebuiltPatterns={mockPrebuiltPatterns}
+        categories={mockCategories}
+        selectedPatternName=""
+        patternAction="BLOCK"
+        onPatternNameChange={mockOnPatternNameChange}
+        onActionChange={mockOnActionChange}
+        onAdd={mockOnAdd}
+        onCancel={mockOnCancel}
+      />,
+    );
+
+    expect(screen.queryByText("Add prebuilt pattern")).not.toBeInTheDocument();
   });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
new file mode 100644
index 00000000000..ff3432a0d82
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
@@ -0,0 +1,162 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import CustomCodeModal from "./CustomCodeModal";
+import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
+
+vi.mock("@/components/networking", () => ({
+  createGuardrailCall: vi.fn(),
+  updateGuardrailCall: vi.fn(),
+  testCustomCodeGuardrail: vi.fn(),
+}));
+
+vi.mock("@/components/molecules/notifications_manager", () => ({
+  default: { success: vi.fn(), error: vi.fn(), fromBackend: vi.fn() },
+}));
+
+const mockCreate = vi.mocked(createGuardrailCall);
+const mockUpdate = vi.mocked(updateGuardrailCall);
+const mockTest = vi.mocked(testCustomCodeGuardrail);
+
+describe("CustomCodeModal", () => {
+  const onClose = vi.fn();
+  const onSuccess = vi.fn();
+
+  const renderModal = (overrides = {}) =>
+    render(<CustomCodeModal visible onClose={onClose} onSuccess={onSuccess} accessToken="test-token" {...overrides} />);
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockCreate.mockResolvedValue({} as never);
+    mockUpdate.mockResolvedValue({} as never);
+  });
+
+  it("should render the create heading and the editor scaffolding", async () => {
+    renderModal();
+
+    expect(await screen.findByText("Create Custom Guardrail")).toBeInTheDocument();
+    expect(screen.getByText("Define custom logic using Python-like syntax")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., block-pii-custom")).toBeInTheDocument();
+    expect(screen.getByText("Guardrail Name")).toBeInTheDocument();
+    expect(screen.getByText("Mode (can select multiple)")).toBeInTheDocument();
+    expect(screen.getByText("Available Primitives")).toBeInTheDocument();
+    expect(screen.getByText("Python Logic")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /save guardrail/i })).toBeInTheDocument();
+  });
+
+  it("should seed the editor with the empty template", async () => {
+    renderModal();
+
+    const editor = await screen.findByDisplayValue(/async def apply_guardrail/);
+    expect(editor).toBeInTheDocument();
+  });
+
+  it("should not render its content when not visible", () => {
+    renderModal({ visible: false });
+
+    expect(screen.queryByText("Create Custom Guardrail")).not.toBeInTheDocument();
+  });
+
+  it("should render the edit heading and existing values in edit mode", async () => {
+    renderModal({
+      editData: {
+        guardrail_id: "g-1",
+        guardrail_name: "existing-guardrail",
+        litellm_params: { mode: "post_call", default_on: true, custom_code: "def apply_guardrail(): pass" },
+      },
+    });
+
+    expect(await screen.findByText("Edit Custom Guardrail")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("existing-guardrail")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("def apply_guardrail(): pass")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /update guardrail/i })).toBeInTheDocument();
+  });
+
+  it("should keep save disabled until a guardrail name is entered", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.click(await screen.findByRole("button", { name: /save guardrail/i }));
+    expect(mockCreate).not.toHaveBeenCalled();
+
+    await user.type(screen.getByPlaceholderText("e.g., block-pii-custom"), "my-guardrail");
+    await user.click(screen.getByRole("button", { name: /save guardrail/i }));
+
+    await waitFor(() => {
+      expect(mockCreate).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  it("should create the guardrail with the entered name, mode and code", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.type(await screen.findByPlaceholderText("e.g., block-pii-custom"), "block-pii");
+    await user.click(screen.getByRole("button", { name: /save guardrail/i }));
+
+    await waitFor(() => {
+      expect(mockCreate).toHaveBeenCalled();
+    });
+
+    const [token, payload] = mockCreate.mock.calls[0] as [string, Record<string, never>];
+    expect(token).toBe("test-token");
+    expect(payload).toMatchObject({
+      guardrail_name: "block-pii",
+      litellm_params: { guardrail: "custom_code", mode: ["pre_call"], default_on: false },
+    });
+    await waitFor(() => {
+      expect(onSuccess).toHaveBeenCalled();
+    });
+  });
+
+  it("should switch the editor contents when a template is chosen", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    expect(await screen.findByDisplayValue(/async def apply_guardrail/)).toBeInTheDocument();
+
+    const comboboxes = screen.getAllByRole("combobox");
+    await user.click(comboboxes[comboboxes.length - 1]);
+    const options = await screen.findAllByText("Block SSN");
+    await user.click(options[options.length - 1]);
+
+    expect(await screen.findByDisplayValue(/SSN detected/)).toBeInTheDocument();
+  });
+
+  it("should expand the test section and run a test against the backend", async () => {
+    const user = userEvent.setup();
+    mockTest.mockResolvedValue({ success: true, result: { action: "allow" } } as never);
+    renderModal();
+
+    await user.click(await screen.findByText("Test Your Guardrail"));
+
+    const runButton = await screen.findByRole("button", { name: /run test/i });
+    await user.click(runButton);
+
+    await waitFor(() => {
+      expect(mockTest).toHaveBeenCalled();
+    });
+    expect(await screen.findByText("Allowed")).toBeInTheDocument();
+  });
+
+  it("should surface a backend test error", async () => {
+    const user = userEvent.setup();
+    mockTest.mockResolvedValue({ success: false, error: "boom", error_type: "SyntaxError" } as never);
+    renderModal();
+
+    await user.click(await screen.findByText("Test Your Guardrail"));
+    await user.click(await screen.findByRole("button", { name: /run test/i }));
+
+    expect(await screen.findByText("boom")).toBeInTheDocument();
+    expect(screen.getByText("[SyntaxError]")).toBeInTheDocument();
+  });
+
+  it("should cancel through the footer button", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.click(await screen.findByRole("button", { name: "Cancel" }));
+
+    expect(onClose).toHaveBeenCalledTimes(1);
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx
new file mode 100644
index 00000000000..e20c8a6fa55
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.test.tsx
@@ -0,0 +1,104 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import GuardrailGarden from "./guardrail_garden";
+import { ALL_CARDS } from "./guardrail_garden_data";
+
+vi.mock("./guardrail_garden_detail", () => ({
+  __esModule: true,
+  default: ({ card, onBack }: { card: { name: string }; onBack: () => void }) => (
+    <div>
+      <span>Detail for {card.name}</span>
+      <button onClick={onBack}>Back to garden</button>
+    </div>
+  ),
+}));
+
+const LITELLM_CARDS = ALL_CARDS.filter((c) => c.category === "litellm");
+const PARTNER_CARDS = ALL_CARDS.filter((c) => c.category === "partner");
+
+describe("GuardrailGarden", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  const renderGarden = () => render(<GuardrailGarden accessToken="test-token" onGuardrailCreated={vi.fn()} />);
+
+  it("should render both sections with their descriptions", () => {
+    renderGarden();
+
+    expect(screen.getByText("LiteLLM Content Filter")).toBeInTheDocument();
+    expect(
+      screen.getByText(
+        "Built-in guardrails powered by LiteLLM. Zero latency, no external dependencies, no additional cost.",
+      ),
+    ).toBeInTheDocument();
+    expect(screen.getByText("Partner Guardrails")).toBeInTheDocument();
+    expect(
+      screen.getByText("Third-party guardrail integrations from leading AI security providers."),
+    ).toBeInTheDocument();
+  });
+
+  it("should show a capped set of litellm cards behind a show all toggle", async () => {
+    const user = userEvent.setup();
+    renderGarden();
+
+    expect(screen.getByText(`Show all (${LITELLM_CARDS.length})`)).toBeInTheDocument();
+    expect(screen.getByText(LITELLM_CARDS[0].name)).toBeInTheDocument();
+    expect(screen.queryByText(LITELLM_CARDS[LITELLM_CARDS.length - 1].name)).not.toBeInTheDocument();
+
+    await user.click(screen.getByText(`Show all (${LITELLM_CARDS.length})`));
+
+    expect(screen.getByText("Show less")).toBeInTheDocument();
+    expect(screen.getByText(LITELLM_CARDS[LITELLM_CARDS.length - 1].name)).toBeInTheDocument();
+  });
+
+  it("should always render every partner card", () => {
+    renderGarden();
+
+    PARTNER_CARDS.forEach((card) => {
+      expect(screen.getByText(card.name)).toBeInTheDocument();
+    });
+  });
+
+  it("should filter cards by the search query", async () => {
+    const user = userEvent.setup();
+    renderGarden();
+
+    const target = PARTNER_CARDS[0];
+    await user.type(screen.getByPlaceholderText("Search guardrails"), target.name);
+
+    expect(await screen.findByText(target.name)).toBeInTheDocument();
+    const otherPartner = PARTNER_CARDS.find((c) => c.name !== target.name);
+    if (otherPartner) {
+      expect(screen.queryByText(otherPartner.name)).not.toBeInTheDocument();
+    }
+  });
+
+  it("should show an empty result set for a query that matches nothing", async () => {
+    const user = userEvent.setup();
+    renderGarden();
+
+    await user.type(screen.getByPlaceholderText("Search guardrails"), "zzzzznotaguardrailzzzzz");
+
+    expect(screen.getByText("Show all (0)")).toBeInTheDocument();
+    PARTNER_CARDS.forEach((card) => {
+      expect(screen.queryByText(card.name)).not.toBeInTheDocument();
+    });
+  });
+
+  it("should open the detail view for a clicked card and return to the garden", async () => {
+    const user = userEvent.setup();
+    renderGarden();
+
+    const target = PARTNER_CARDS[0];
+    await user.click(screen.getByText(target.name));
+
+    expect(await screen.findByText(`Detail for ${target.name}`)).toBeInTheDocument();
+    expect(screen.queryByPlaceholderText("Search guardrails")).not.toBeInTheDocument();
+
+    await user.click(screen.getByRole("button", { name: "Back to garden" }));
+
+    expect(await screen.findByPlaceholderText("Search guardrails")).toBeInTheDocument();
+  });
+});

From 28032eb4c9acadaa10b10090cfb09ad94eb736e5 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 14:01:00 -0700
Subject: [PATCH 07/31] test(ui): decouple usage-route tests from antd and
 Tremor markup

Rewrites the usage route's library-coupled assertions to role and text
queries so they characterise behaviour rather than the widget library that
happens to render it. Every test here is green against the current antd and
Tremor components and is meant to survive a migration unedited.

Drops the wholesale vi.mock of antd, @ant-design/icons and @tremor/react in
UsagePageView, UsageViewSelect and activity_metrics, and drives the real
controls instead. Two of those mocks were hiding behaviour: activity_metrics
asserted panel order through a heading role that only existed because the
mock faked an h2, and the Tremor tab stubs flattened panel selection away
entirely.

Replaces DOM-shape lookups with anchors that do not move: a chart is found
from its own heading rather than a fixed wrapper depth, an active panel is
identified by the inactive markers both tab libraries set, and select options
are matched by text since antd's real options carry no option role.

Adds the missing characterisation test for team_multi_select, and covers the
mount contracts the route depends on: Tremor keeps every tab panel mounted,
and antd Collapse keeps a section mounted once it has been expanded, which is
what preserves the view-mode state a model section owns.
---
 .../EntityUsage/EntityUsage.test.tsx          |  66 +++--
 .../EntityUsage/TopModelView.test.tsx         |  53 ++--
 .../components/UsageAIChatPanel.test.tsx      |   5 +-
 .../components/UsagePageView.test.tsx         | 239 +++---------------
 .../UsageViewSelect/UsageViewSelect.test.tsx  | 156 ++++--------
 .../src/components/activity_metrics.test.tsx  |  70 +++--
 .../team_multi_select.test.tsx                | 124 +++++++++
 .../src/components/per_user_usage.test.tsx    |  44 +++-
 .../components/user_agent_activity.test.tsx   |  79 ++++--
 9 files changed, 397 insertions(+), 439 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/common_components/team_multi_select.test.tsx

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
index c85a9fb71f6..1f3dd5642ee 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.test.tsx
@@ -1,4 +1,4 @@
-import { act, cleanup, fireEvent, render, screen, waitFor, within } from "@testing-library/react";
+import { act, cleanup, fireEvent, render, screen, waitFor } from "@testing-library/react";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import * as networking from "@/components/networking";
 import EntityUsage from "./EntityUsage";
@@ -500,18 +500,33 @@ describe("EntityUsage", () => {
     expect(screen.getAllByText("Activity Metrics")[1]).toBeInTheDocument();
   });
 
-  const selectedPanels = (container: HTMLElement) =>
-    Array.from(container.querySelectorAll("div.tremor-TabPanel-root")).filter(
-      (panel) => panel.getAttribute("aria-selected") === "true",
-    );
+  // An inactive tab panel is marked aria-selected="false" by one tab library and hidden by the
+  // other, so treat either as "not on screen" and the assertion holds whichever one is rendering.
+  const isShowing = (element: HTMLElement): boolean => {
+    for (let node: HTMLElement | null = element; node; node = node.parentElement) {
+      if (node.hasAttribute("hidden")) return false;
+      if (node.getAttribute("aria-selected") === "false") return false;
+    }
+    return true;
+  };
 
-  it.each([
+  const showingCount = (marker: string): number => screen.queryAllByText(marker).filter(isShowing).length;
+
+  const showingText = (text: string): HTMLElement => {
+    const [element] = screen.getAllByText(text).filter(isShowing);
+    expect(element).toBeDefined();
+    return element;
+  };
+
+  const NON_TEAM_PANELS: [string, string][] = [
     ["Cost", "Tag Spend Overview"],
     ["Model Activity", "metrics-source:model_groups"],
     ["Key Activity", "metrics-source:api_keys"],
     ["Endpoint Activity", "Endpoint Usage Panel"],
-  ])("shows only the %s panel for a non-team entity type", async (tabLabel, marker) => {
-    const { container } = render(<EntityUsage {...defaultProps} />);
+  ];
+
+  it.each(NON_TEAM_PANELS)("shows only the %s panel for a non-team entity type", async (tabLabel, marker) => {
+    render(<EntityUsage {...defaultProps} />);
 
     await waitFor(() => {
       expect(mockTagDailyActivityCall).toHaveBeenCalled();
@@ -521,19 +536,23 @@ describe("EntityUsage", () => {
       fireEvent.click(screen.getByText(tabLabel));
     });
 
-    const selected = selectedPanels(container);
-    expect(selected).toHaveLength(1);
-    expect(selected[0].textContent).toContain(marker);
+    expect(showingCount(marker)).toBeGreaterThan(0);
+    for (const [otherLabel, otherMarker] of NON_TEAM_PANELS) {
+      if (otherLabel === tabLabel) continue;
+      expect(showingCount(otherMarker)).toBe(0);
+    }
   });
 
-  it.each([
+  const TEAM_PANELS: [string, string][] = [
     ["Cost", "Team Spend Overview"],
     ["Model Activity", "metrics-source:model_groups"],
     ["Agent Activity", "metrics-source:entities"],
     ["Key Activity", "metrics-source:api_keys"],
     ["Endpoint Activity", "Endpoint Usage Panel"],
-  ])("shows only the %s panel for the team entity type", async (tabLabel, marker) => {
-    const { container } = render(<EntityUsage {...defaultProps} entityType="team" />);
+  ];
+
+  it.each(TEAM_PANELS)("shows only the %s panel for the team entity type", async (tabLabel, marker) => {
+    render(<EntityUsage {...defaultProps} entityType="team" />);
 
     await waitFor(() => {
       expect(mockTeamDailyActivityCall).toHaveBeenCalled();
@@ -543,9 +562,11 @@ describe("EntityUsage", () => {
       fireEvent.click(screen.getByText(tabLabel));
     });
 
-    const selected = selectedPanels(container);
-    expect(selected).toHaveLength(1);
-    expect(selected[0].textContent).toContain(marker);
+    expect(showingCount(marker)).toBeGreaterThan(0);
+    for (const [otherLabel, otherMarker] of TEAM_PANELS) {
+      if (otherLabel === tabLabel) continue;
+      expect(showingCount(otherMarker)).toBe(0);
+    }
   });
 
   it("should handle empty data gracefully", async () => {
@@ -615,20 +636,19 @@ describe("EntityUsage", () => {
       fireEvent.click(screen.getByText("Model Activity"));
     });
 
-    const modelActivityPanel = () => selectedPanels(container)[0] as HTMLElement;
-    expect(modelActivityPanel().textContent).toContain("metrics-source:model_groups");
+    expect(showingCount("metrics-source:model_groups")).toBeGreaterThan(0);
 
     act(() => {
-      fireEvent.click(within(modelActivityPanel()).getByText("Litellm Model Name"));
+      fireEvent.click(showingText("Litellm Model Name"));
     });
 
-    expect(modelActivityPanel().textContent).toContain("metrics-source:models");
+    expect(showingCount("metrics-source:models")).toBeGreaterThan(0);
 
     act(() => {
-      fireEvent.click(within(modelActivityPanel()).getByText("Public Model Name"));
+      fireEvent.click(showingText("Public Model Name"));
     });
 
-    expect(modelActivityPanel().textContent).toContain("metrics-source:model_groups");
+    expect(showingCount("metrics-source:model_groups")).toBeGreaterThan(0);
   });
 
   it("should display Top Agents title for agent entity type", async () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.test.tsx
index 2770bac47d9..961b19ab27a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.test.tsx
@@ -10,6 +10,14 @@ describe("TopModelView", () => {
     mockSetTopModelsLimit.mockClear();
   });
 
+  // Which element a control library gives its label to is its own business, so drive the
+  // control by its visible text and judge the result by what the panel renders.
+  const clickControl = async (user: ReturnType<typeof userEvent.setup>, label: string) => {
+    await user.click(screen.getByText(label));
+  };
+
+  const showsChart = (container: HTMLElement) => container.querySelector(".recharts-wrapper") !== null;
+
   it("should render", () => {
     render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
     expect(screen.getByText("Table View")).toBeInTheDocument();
@@ -17,12 +25,12 @@ describe("TopModelView", () => {
 
   it("should display table view button", () => {
     render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
-    expect(screen.getByRole("button", { name: "Table View" })).toBeInTheDocument();
+    expect(screen.getByText("Table View")).toBeInTheDocument();
   });
 
   it("should display chart view button", () => {
     render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
-    expect(screen.getByRole("button", { name: "Chart View" })).toBeInTheDocument();
+    expect(screen.getByText("Chart View")).toBeInTheDocument();
   });
 
   it("should display all table column headers", () => {
@@ -60,27 +68,32 @@ describe("TopModelView", () => {
     expect(screen.getByText("50,000")).toBeInTheDocument();
   });
 
+  const oneModel = [{ key: "gpt-4", spend: 150.5, successful_requests: 100, failed_requests: 5, tokens: 50000 }];
+
   it("should switch to chart view when chart view button is clicked", async () => {
     const user = userEvent.setup();
-    render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
+    const { container } = render(
+      <TopModelView topModels={oneModel} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />,
+    );
 
-    const chartViewButton = screen.getByRole("button", { name: "Chart View" });
-    await user.click(chartViewButton);
+    expect(showsChart(container)).toBe(false);
+    await clickControl(user, "Chart View");
 
-    expect(chartViewButton).toHaveClass("bg-blue-100");
+    expect(showsChart(container)).toBe(true);
+    expect(screen.queryByText("Spend (USD)")).not.toBeInTheDocument();
   });
 
   it("should switch to table view when table view button is clicked", async () => {
     const user = userEvent.setup();
-    render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
+    const { container } = render(
+      <TopModelView topModels={oneModel} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />,
+    );
 
-    const chartViewButton = screen.getByRole("button", { name: "Chart View" });
-    const tableViewButton = screen.getByRole("button", { name: "Table View" });
+    await clickControl(user, "Chart View");
+    await clickControl(user, "Table View");
 
-    await user.click(chartViewButton);
-    await user.click(tableViewButton);
-
-    expect(tableViewButton).toHaveClass("bg-blue-100");
+    expect(showsChart(container)).toBe(false);
+    expect(screen.getByText("Spend (USD)")).toBeInTheDocument();
   });
 
   it("renders one cyan bar per model with model names on the axis in chart view", async () => {
@@ -108,7 +121,7 @@ describe("TopModelView", () => {
       />,
     );
 
-    await user.click(screen.getByRole("button", { name: "Chart View" }));
+    await clickControl(user, "Chart View");
 
     const bars = container.querySelectorAll("path.recharts-rectangle");
     expect(bars).toHaveLength(2);
@@ -118,19 +131,11 @@ describe("TopModelView", () => {
     expect(screen.getAllByText("claude-3").length).toBeGreaterThan(0);
   });
 
-  it("should call setTopModelsLimit when limit is changed via Segmented control", async () => {
+  it("should call setTopModelsLimit when the limit control is changed", async () => {
     const user = userEvent.setup();
     render(<TopModelView topModels={[]} topModelsLimit={5} setTopModelsLimit={mockSetTopModelsLimit} />);
 
-    const limit10Radio = screen.getByRole("radio", { name: "10" });
-    const limit10Label = limit10Radio.closest("label");
-    if (limit10Label) {
-      await user.click(limit10Label);
-    } else {
-      // Fallback: click the div with title="10"
-      const limit10Div = screen.getByTitle("10");
-      await user.click(limit10Div);
-    }
+    await clickControl(user, "10");
 
     expect(mockSetTopModelsLimit).toHaveBeenCalledWith(10);
   });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx
index 54bf4fc25ce..d971f5d0506 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.test.tsx
@@ -37,7 +37,10 @@ describe("UsageAIChatPanel", () => {
   it("should render model selector", () => {
     renderWithProviders(<UsageAIChatPanel {...defaultProps} />);
 
-    expect(screen.getByText("Select a model (optional, defaults to gpt-4o-mini)")).toBeInTheDocument();
+    // One library paints the prompt as its own text node and the other leaves it on the input's
+    // placeholder attribute, so either one means the user is being told what to pick.
+    const prompt = "Select a model (optional, defaults to gpt-4o-mini)";
+    expect(screen.queryAllByText(prompt).length + screen.queryAllByPlaceholderText(prompt).length).toBeGreaterThan(0);
   });
 
   it("should render empty state message when no conversation", () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
index 9085cf961a9..9cd1e53c9b4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx
@@ -4,6 +4,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized";
 import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser";
 import { useInfiniteUsers } from "@/app/(dashboard)/hooks/users/useUsers";
 import { act, fireEvent, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest";
 import { renderWithProviders } from "@/../tests/test-utils";
 import type { Organization } from "@/components/networking";
@@ -143,207 +144,6 @@ vi.mock("@/app/(dashboard)/hooks/users/useUsers", () => ({
   useInfiniteUsers: vi.fn(),
 }));
 
-vi.mock("antd", async (importOriginal) => {
-  const React = await import("react");
-  const actual = await importOriginal<typeof import("antd")>();
-
-  function Select(props: any) {
-    const { value, onChange, options, ...rest } = props;
-    return React.createElement(
-      "select",
-      {
-        ...rest,
-        value,
-        onChange: (e: any) => onChange?.(e.target.value),
-        role: "combobox",
-      },
-      options?.map((opt: any) => React.createElement("option", { key: opt.value, value: opt.value }, opt.label)),
-    );
-  }
-  (Select as any).displayName = "AntdSelect";
-
-  function Alert(props: any) {
-    const { message, description, type, closable, onClose, ...rest } = props;
-    return React.createElement(
-      "div",
-      { ...rest, "data-testid": "antd-alert", "data-type": type },
-      message && React.createElement("div", null, message),
-      description && React.createElement("div", null, description),
-      closable && React.createElement("button", { onClick: onClose, "aria-label": "Close" }, "×"),
-    );
-  }
-  (Alert as any).displayName = "AntdAlert";
-
-  function Badge(props: any) {
-    const { count, color, children, ...rest } = props;
-    return React.createElement(
-      "div",
-      { ...rest, "data-testid": "antd-badge", "data-color": color },
-      count && React.createElement("span", { "data-testid": "antd-badge-count" }, count),
-      children,
-    );
-  }
-  (Badge as any).displayName = "AntdBadge";
-
-  function Table({ columns, dataSource, ...rest }: any) {
-    return React.createElement(
-      "div",
-      { ...rest, "data-testid": "antd-table" },
-      columns?.map((col: any) =>
-        React.createElement("div", { key: col.key, "data-testid": `column-${col.key}` }, col.title),
-      ),
-      dataSource?.map((row: any) =>
-        React.createElement(
-          "div",
-          { key: row.key, "data-testid": `row-${row.key}` },
-          columns?.map((col: any) => {
-            const value = col.render ? col.render(row[col.dataIndex], row) : row[col.dataIndex];
-            return React.createElement("div", { key: col.key }, value);
-          }),
-        ),
-      ),
-    );
-  }
-  (Table as any).displayName = "Table";
-
-  function Segmented(props: any) {
-    const { value, onChange, options, ...rest } = props;
-    return React.createElement(
-      "div",
-      { ...rest, "data-testid": "antd-segmented" },
-      options?.map((opt: any) =>
-        React.createElement(
-          "button",
-          {
-            key: opt.value,
-            onClick: () => onChange?.(opt.value),
-            "data-selected": value === opt.value,
-          },
-          opt.label,
-        ),
-      ),
-    );
-  }
-  (Segmented as any).displayName = "AntdSegmented";
-
-  function Tooltip(props: any) {
-    const { title, children, ...rest } = props;
-    return React.createElement("div", { ...rest, "data-testid": "antd-tooltip", title }, children);
-  }
-  (Tooltip as any).displayName = "AntdTooltip";
-
-  return {
-    ...actual,
-    Select,
-    Alert,
-    Badge,
-    Table,
-    Segmented,
-    Tooltip,
-  };
-});
-
-vi.mock("@ant-design/icons", async () => {
-  const React = await import("react");
-
-  function Icon() {
-    return React.createElement("span");
-  }
-
-  function LoadingOutlined(props: any) {
-    return React.createElement("span", { "data-testid": "loading-icon", ...props });
-  }
-
-  return {
-    GlobalOutlined: Icon,
-    BankOutlined: Icon,
-    TeamOutlined: Icon,
-    ShoppingCartOutlined: Icon,
-    TagsOutlined: Icon,
-    RobotOutlined: Icon,
-    LineChartOutlined: Icon,
-    BarChartOutlined: Icon,
-    ClockCircleOutlined: Icon,
-    CalendarOutlined: Icon,
-    InfoCircleOutlined: Icon,
-    UserOutlined: Icon,
-    DownOutlined: Icon,
-    RightOutlined: Icon,
-    ExportOutlined: Icon,
-    LoadingOutlined,
-  };
-});
-
-// Mock Tremor components
-vi.mock("@tremor/react", async () => {
-  const React = await import("react");
-  const actual = await import("@tremor/react");
-
-  function TabGroup({ children }: any) {
-    return React.createElement("div", { "data-testid": "tremor-tab-group" }, children);
-  }
-
-  function TabList({ children }: any) {
-    return React.createElement("div", { "data-testid": "tremor-tab-list" }, children);
-  }
-
-  function Tab({ children, ...props }: any) {
-    return React.createElement("button", { ...props, "data-testid": "tremor-tab" }, children);
-  }
-
-  function TabPanels({ children }: any) {
-    return React.createElement("div", { "data-testid": "tremor-tab-panels" }, children);
-  }
-
-  function TabPanel({ children }: any) {
-    return React.createElement("div", { "data-testid": "tremor-tab-panel" }, children);
-  }
-
-  function Card({ children, ...props }: any) {
-    return React.createElement("div", { ...props, "data-testid": "tremor-card" }, children);
-  }
-
-  function Grid({ children, numItems, ...props }: any) {
-    return React.createElement("div", { ...props, "data-testid": "tremor-grid" }, children);
-  }
-
-  function Col({ children, numColSpan, ...props }: any) {
-    return React.createElement("div", { ...props, "data-testid": "tremor-col" }, children);
-  }
-
-  function Title({ children, ...props }: any) {
-    return React.createElement("h2", { ...props, "data-testid": "tremor-title" }, children);
-  }
-
-  function Text({ children, ...props }: any) {
-    return React.createElement("p", { ...props, "data-testid": "tremor-text" }, children);
-  }
-
-  function Button({ children, icon, onClick, ...props }: any) {
-    return React.createElement(
-      "button",
-      { ...props, onClick, "data-testid": "tremor-button" },
-      icon && React.createElement("span", { "data-testid": "tremor-button-icon" }),
-      children,
-    );
-  }
-
-  return {
-    ...actual,
-    TabGroup,
-    TabList,
-    Tab,
-    TabPanels,
-    TabPanel,
-    Card,
-    Grid,
-    Col,
-    Title,
-    Text,
-    Button,
-  };
-});
-
 describe("UsagePage", () => {
   const mockUserDailyActivityAggregatedCall = vi.mocked(networking.userDailyActivityAggregatedCall);
   const mockUserDailyActivityCall = vi.mocked(networking.userDailyActivityCall);
@@ -885,6 +685,26 @@ describe("UsagePage", () => {
   });
 
   describe("admin user selector", () => {
+    // Anchored on the field's own label, so it does not depend on which library draws the control.
+    const userSelectCombobox = (): HTMLElement => {
+      let node: HTMLElement | null = screen.getByText("Filter by user");
+      while (node && !node.querySelector('[role="combobox"]')) {
+        node = node.parentElement;
+      }
+      const combobox = node?.querySelector('[role="combobox"]') ?? null;
+      expect(combobox).not.toBeNull();
+      return combobox as HTMLElement;
+    };
+
+    const openUserSelect = async () => {
+      await userEvent.setup().click(userSelectCombobox());
+    };
+
+    // One library paints the prompt as its own text node and the other leaves it on the input's
+    // placeholder attribute, so either one means the user is being told what to type.
+    const promptsWith = (text: string) =>
+      screen.queryAllByText(text).length + screen.queryAllByPlaceholderText(text).length > 0;
+
     it("should render user selector for admin users in global view", async () => {
       renderWithProviders(<UsagePage {...defaultProps} />);
 
@@ -892,10 +712,8 @@ describe("UsagePage", () => {
         expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
       });
 
-      // Admin should see the user selector select element with the placeholder attribute
-      const userSelects = screen.getAllByRole("combobox");
-      const userSelect = userSelects.find((el) => el.getAttribute("placeholder") === "Select user to filter...");
-      expect(userSelect).toBeDefined();
+      expect(userSelectCombobox()).toBeInTheDocument();
+      expect(promptsWith("Select user to filter...")).toBe(true);
     });
 
     it("should format user options with alias when available", async () => {
@@ -905,6 +723,8 @@ describe("UsagePage", () => {
         expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
       });
 
+      await openUserSelect();
+
       // User with alias should show "alias (id)"
       expect(screen.getByText("Alice (user-001)")).toBeInTheDocument();
       // User without alias but with email should show "email (id)"
@@ -958,6 +778,8 @@ describe("UsagePage", () => {
         expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
       });
 
+      await openUserSelect();
+
       // Duplicate user should appear only once
       const dupElements = screen.getAllByText("DupUser (user-dup)");
       expect(dupElements).toHaveLength(1);
@@ -1003,10 +825,9 @@ describe("UsagePage", () => {
         expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled();
       });
 
-      // Non-admin should not see the user selector
-      const userSelects = screen.getAllByRole("combobox");
-      const userSelect = userSelects.find((el) => el.getAttribute("placeholder") === "Select user to filter...");
-      expect(userSelect).toBeUndefined();
+      // The admin case above proves this label is rendered when the selector exists, so its
+      // absence here is a live assertion rather than a query that can never match.
+      expect(screen.queryByText("Filter by user")).not.toBeInTheDocument();
     });
 
     it("should always pass own userId for non-admin users", async () => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
index dcc0ce06673..80005a20d05 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.test.tsx
@@ -1,86 +1,16 @@
-import { act, fireEvent, render, screen } from "@testing-library/react";
+import { render, screen, within } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { UsageViewSelect } from "./UsageViewSelect";
 
-vi.mock("antd", async () => {
-  const React = await import("react");
+const openMenu = async (user: ReturnType<typeof userEvent.setup>) => {
+  await user.click(screen.getByRole("combobox"));
+};
 
-  function Select(props: any) {
-    const { value, onChange, options, optionRender, labelRender, ...rest } = props;
-
-    const optionElements = options?.map((opt: any) =>
-      React.createElement("option", { key: opt.value, value: opt.value }, opt.label),
-    );
-
-    const optionRenderOutputs = options
-      ?.map((opt: any) => {
-        if (optionRender) {
-          const rendered = optionRender({ value: opt.value, label: opt.label });
-          return React.createElement(
-            "div",
-            {
-              key: `option-render-${opt.value}`,
-              "data-testid": `option-render-${opt.value}`,
-              style: { display: "none" },
-            },
-            rendered,
-          );
-        }
-        return null;
-      })
-      .filter(Boolean);
-
-    return React.createElement(
-      React.Fragment,
-      null,
-      React.createElement(
-        "select",
-        {
-          ...rest,
-          value,
-          onChange: (e: any) => onChange?.(e.target.value),
-          role: "combobox",
-        },
-        optionElements,
-      ),
-      ...(optionRenderOutputs || []),
-    );
-  }
-  (Select as any).displayName = "AntdSelect";
-
-  function Badge(props: any) {
-    const { count, color, children, ...rest } = props;
-    return React.createElement(
-      "span",
-      { ...rest, "data-testid": "antd-badge", "data-color": color },
-      count && React.createElement("span", { "data-testid": "antd-badge-count" }, count),
-      children,
-    );
-  }
-  (Badge as any).displayName = "AntdBadge";
-
-  return { Select, Badge };
-});
-
-vi.mock("@ant-design/icons", async () => {
-  const React = await import("react");
-
-  function Icon(props: any) {
-    return React.createElement("span", { "data-testid": "antd-icon" });
-  }
-
-  return {
-    GlobalOutlined: Icon,
-    BankOutlined: Icon,
-    TeamOutlined: Icon,
-    ShoppingCartOutlined: Icon,
-    TagsOutlined: Icon,
-    RobotOutlined: Icon,
-    UserOutlined: Icon,
-    LineChartOutlined: Icon,
-    BarChartOutlined: Icon,
-  };
-});
+// The listbox is portalled outside the render container in both antd and Base UI, so an
+// option is "offered" when the label appears more times on the page than inside the trigger.
+const offers = (container: HTMLElement, label: string) =>
+  screen.queryAllByText(label).length > within(container).queryAllByText(label).length;
 
 describe("UsageViewSelect", () => {
   const mockOnChange = vi.fn();
@@ -89,53 +19,73 @@ describe("UsageViewSelect", () => {
     mockOnChange.mockClear();
   });
 
-  it("should render", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
+  it("should render", async () => {
+    const user = userEvent.setup();
+    const { container } = render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
 
     expect(screen.getByText("Usage View")).toBeInTheDocument();
     expect(screen.getByText("Select the usage data you want to view")).toBeInTheDocument();
     expect(screen.getByRole("combobox")).toBeInTheDocument();
-    expect(screen.getByRole("option", { name: "Your Usage" })).toBeInTheDocument();
+
+    await openMenu(user);
+    expect(offers(container, "Your Usage")).toBe(true);
   });
 
-  it("should call onChange when value changes", () => {
+  it("should call onChange when value changes", async () => {
+    const user = userEvent.setup();
     render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Admin" />);
 
-    const select = screen.getByRole("combobox");
-    act(() => {
-      fireEvent.change(select, { target: { value: "team" } });
-    });
+    await openMenu(user);
+    const matches = screen.getAllByText("Team Usage");
+    await user.click(matches[matches.length - 1]);
 
-    expect(mockOnChange).toHaveBeenCalledWith("team");
+    expect(mockOnChange).toHaveBeenCalled();
+    expect(mockOnChange.mock.calls[0][0]).toBe("team");
   });
 
-  it("should show Tag Usage for non-admin users with tag usage permission", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
+  it("should show Tag Usage for non-admin users with tag usage permission", async () => {
+    const user = userEvent.setup();
+    const { container } = render(
+      <UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />,
+    );
 
-    expect(screen.getByRole("option", { name: "Tag Usage" })).toBeInTheDocument();
+    await openMenu(user);
+    expect(offers(container, "Tag Usage")).toBe(true);
   });
 
-  it("should hide Tag Usage for non-admin users without tag usage permission", () => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
+  it("should hide Tag Usage for non-admin users without tag usage permission", async () => {
+    const user = userEvent.setup();
+    const { container } = render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" />);
 
-    expect(screen.queryByRole("option", { name: "Tag Usage" })).not.toBeInTheDocument();
+    await openMenu(user);
+    expect(offers(container, "Tag Usage")).toBe(false);
   });
 
-  it.each(["Organization Usage", "Agent Usage (A2A)"])("should show %s to an admin", (optionName) => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Admin" />);
+  it.each(["Organization Usage", "Agent Usage (A2A)"])("should show %s to an admin", async (optionName) => {
+    const user = userEvent.setup();
+    const { container } = render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Admin" />);
 
-    expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument();
+    await openMenu(user);
+    expect(offers(container, optionName)).toBe(true);
   });
 
-  it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", (optionName) => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
+  it.each(["Organization Usage", "Agent Usage (A2A)"])("should hide %s from an internal user", async (optionName) => {
+    const user = userEvent.setup();
+    const { container } = render(
+      <UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />,
+    );
 
-    expect(screen.queryByRole("option", { name: optionName })).not.toBeInTheDocument();
+    await openMenu(user);
+    expect(offers(container, optionName)).toBe(false);
   });
 
-  it.each(["Team Usage", "Tag Usage"])("should keep %s available to an internal user", (optionName) => {
-    render(<UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />);
+  it.each(["Team Usage", "Tag Usage"])("should keep %s available to an internal user", async (optionName) => {
+    const user = userEvent.setup();
+    const { container } = render(
+      <UsageViewSelect value="global" onChange={mockOnChange} userRole="Internal User" canViewTagUsage={true} />,
+    );
 
-    expect(screen.getByRole("option", { name: optionName })).toBeInTheDocument();
+    await openMenu(user);
+    expect(offers(container, optionName)).toBe(true);
   });
 });
diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx
index cd37ea42a90..74d258e2bd0 100644
--- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx
+++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { fireEvent, render, screen } from "@testing-library/react";
 import React from "react";
 import { beforeAll, describe, expect, it, vi } from "vitest";
 import { ActivityMetrics, formatKeyLabel, processActivityData } from "./activity_metrics";
@@ -15,38 +15,12 @@ beforeAll(() => {
   }
 });
 
-vi.mock("@tremor/react", () => ({
-  Card: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
-  Grid: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
-  Text: ({ children }: { children: React.ReactNode }) => <span>{children}</span>,
-  Title: ({ children }: { children: React.ReactNode }) => <h2>{children}</h2>,
-  AreaChart: () => <div>AreaChart</div>,
-  BarChart: () => <div>BarChart</div>,
-}));
-
-vi.mock("antd", () => {
-  const CollapseComponent = ({ children }: { children: React.ReactNode }) => <div>{children}</div>;
-  const PanelComponent = ({ children, header }: { children: React.ReactNode; header: React.ReactNode }) => (
-    <div>
-      <div>{header}</div>
-      <div>{children}</div>
-    </div>
-  );
-  PanelComponent.displayName = "Collapse.Panel";
-  CollapseComponent.Panel = PanelComponent;
-  const TableComponent = ({ dataSource, columns }: { dataSource?: unknown[]; columns?: { title: string }[] }) => (
-    <table>
-      <thead>
-        <tr>{columns?.map((col, i) => <th key={i}>{col.title}</th>)}</tr>
-      </thead>
-      <tbody>{dataSource?.map((_, i) => <tr key={i} />)}</tbody>
-    </table>
-  );
-  return {
-    Collapse: CollapseComponent,
-    Table: TableComponent,
-  };
-});
+// Panel order is a contract; which element the label lands in is not, so compare document order.
+const precedes = (firstLabel: string, secondLabel: string): boolean => {
+  const first = screen.getAllByText(firstLabel)[0];
+  const second = screen.getAllByText(secondLabel)[0];
+  return Boolean(first.compareDocumentPosition(second) & Node.DOCUMENT_POSITION_FOLLOWING);
+};
 
 vi.mock("@/utils/dataUtils", async (importOriginal) => {
   const actual = await importOriginal<typeof import("@/utils/dataUtils")>();
@@ -256,10 +230,7 @@ describe("ActivityMetrics", () => {
     };
 
     render(<ActivityMetrics modelMetrics={multipleModels} />);
-    const headers = screen.getAllByRole("heading", { level: 2 });
-    const gpt4Index = headers.findIndex((h) => h.textContent?.includes("GPT-4"));
-    const gpt35Index = headers.findIndex((h) => h.textContent?.includes("GPT-3.5"));
-    expect(gpt4Index).toBeLessThan(gpt35Index);
+    expect(precedes("GPT-4", "GPT-3.5")).toBe(true);
   });
 
   it("should display model summary cards with correct values", () => {
@@ -387,10 +358,27 @@ describe("ActivityMetrics", () => {
     };
 
     render(<ActivityMetrics modelMetrics={modelsWithEmptyKey} />);
-    const headings = screen.getAllByRole("heading", { level: 2 });
-    const gpt4Index = headings.findIndex((h) => h.textContent?.includes("GPT-4"));
-    const unknownIndex = headings.findIndex((h) => h.textContent?.includes("Unknown"));
-    expect(gpt4Index).toBeLessThan(unknownIndex);
+    expect(precedes("GPT-4", "Unknown")).toBe(true);
+  });
+
+  // A model section owns view-mode state, so collapsing one must not throw its subtree away.
+  it("keeps a model section mounted once it has been expanded", () => {
+    const multipleModels: Record<string, ModelActivityData> = {
+      "gpt-3.5": GPT_35_MODEL_DATA,
+      "gpt-4": { ...mockModelMetrics["gpt-4"], total_spend: 100.5 },
+    };
+
+    render(<ActivityMetrics modelMetrics={multipleModels} />);
+
+    // Only the highest-spend section is expanded initially, so only its body is mounted.
+    const sectionsMounted = () => screen.getAllByText("Spend per day").length;
+    expect(sectionsMounted()).toBe(1);
+
+    fireEvent.click(screen.getAllByText("GPT-3.5")[0]);
+    expect(sectionsMounted()).toBe(2);
+
+    fireEvent.click(screen.getAllByText("GPT-3.5")[0]);
+    expect(sectionsMounted()).toBe(2);
   });
 
   it("should display average tokens per successful request", () => {
diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.test.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.test.tsx
new file mode 100644
index 00000000000..82508472986
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.test.tsx
@@ -0,0 +1,124 @@
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
+import TeamMultiSelect from "./team_multi_select";
+
+vi.mock("@/app/(dashboard)/hooks/teams/useTeams", () => ({
+  useInfiniteTeams: vi.fn(),
+}));
+
+const team = (id: string, alias: string) => ({ team_id: id, team_alias: alias });
+
+const mockTeamsResult = (
+  overrides: Partial<{
+    pages: { teams: ReturnType<typeof team>[] }[];
+    isLoading: boolean;
+    hasNextPage: boolean;
+    isFetchingNextPage: boolean;
+  }> = {},
+) => {
+  const { pages = [{ teams: [team("team-1", "Alpha Team"), team("team-2", "Beta Team")] }], ...rest } = overrides;
+  return {
+    data: { pages },
+    fetchNextPage: vi.fn(),
+    hasNextPage: false,
+    isFetchingNextPage: false,
+    isLoading: false,
+    ...rest,
+  };
+};
+
+describe("TeamMultiSelect", () => {
+  const mockUseInfiniteTeams = vi.mocked(useInfiniteTeams);
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockUseInfiniteTeams.mockReturnValue(mockTeamsResult() as never);
+  });
+
+  const combobox = () => screen.getByRole("combobox");
+
+  // One library paints the prompt as its own text node and the other leaves it on the input's
+  // placeholder attribute, so either one means the user is being told what to type.
+  const promptsWith = (text: string) =>
+    screen.queryAllByText(text).length + screen.queryAllByPlaceholderText(text).length > 0;
+
+  it("renders a search control with the given placeholder", () => {
+    render(<TeamMultiSelect placeholder="Search teams by alias..." />);
+
+    expect(combobox()).toBeInTheDocument();
+    expect(promptsWith("Search teams by alias...")).toBe(true);
+  });
+
+  it("offers every loaded team by alias and id", async () => {
+    const user = userEvent.setup();
+    render(<TeamMultiSelect />);
+
+    await user.click(combobox());
+
+    expect(screen.getByText("Alpha Team")).toBeInTheDocument();
+    expect(screen.getByText("(team-1)")).toBeInTheDocument();
+    expect(screen.getByText("Beta Team")).toBeInTheDocument();
+    expect(screen.getByText("(team-2)")).toBeInTheDocument();
+  });
+
+  it("deduplicates a team that appears on more than one page", async () => {
+    mockUseInfiniteTeams.mockReturnValue(
+      mockTeamsResult({
+        pages: [
+          { teams: [team("team-1", "Alpha Team")] },
+          { teams: [team("team-1", "Alpha Team"), team("team-2", "Beta Team")] },
+        ],
+      }) as never,
+    );
+    const user = userEvent.setup();
+    render(<TeamMultiSelect />);
+
+    await user.click(combobox());
+
+    expect(screen.getAllByText("Alpha Team")).toHaveLength(1);
+    expect(screen.getByText("Beta Team")).toBeInTheDocument();
+  });
+
+  it("reports the picked team id to onChange", async () => {
+    const onChange = vi.fn();
+    const user = userEvent.setup();
+    render(<TeamMultiSelect onChange={onChange} />);
+
+    await user.click(combobox());
+    const matches = screen.getAllByText("Beta Team");
+    await user.click(matches[matches.length - 1]);
+
+    expect(onChange).toHaveBeenCalled();
+    expect(onChange.mock.calls[0][0]).toEqual(["team-2"]);
+  });
+
+  it("does not report a selection while disabled", async () => {
+    const onChange = vi.fn();
+    const user = userEvent.setup();
+    render(<TeamMultiSelect onChange={onChange} disabled />);
+
+    await user.click(combobox());
+
+    expect(screen.queryByText("Alpha Team")).not.toBeInTheDocument();
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
+  it("tells the user when there are no teams to pick", async () => {
+    mockUseInfiniteTeams.mockReturnValue(mockTeamsResult({ pages: [{ teams: [] }] }) as never);
+    const user = userEvent.setup();
+    render(<TeamMultiSelect />);
+
+    await user.click(combobox());
+
+    // Substring match because one library appends an invisible word joiner for its live region.
+    expect(screen.getByText(/No teams found/)).toBeInTheDocument();
+  });
+
+  it("passes the page size and organization filter through to the teams query", () => {
+    render(<TeamMultiSelect pageSize={25} organizationId="org-7" />);
+
+    expect(mockUseInfiniteTeams).toHaveBeenCalledWith(25, undefined, "org-7");
+  });
+});
diff --git a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx
index 66aaf015aca..443ae66af8f 100644
--- a/ui/litellm-dashboard/src/components/per_user_usage.test.tsx
+++ b/ui/litellm-dashboard/src/components/per_user_usage.test.tsx
@@ -29,6 +29,14 @@ const userRow = (userId: string, userAgent: string | null, successfulRequests: n
   spend: 1,
 });
 
+// The distribution panel owns the only chart in this component, so resolving it by slot keeps
+// the assertions independent of how many wrappers the tab library puts around a panel.
+const distributionChart = (): HTMLElement => {
+  const chart = document.querySelector('[data-slot="chart"]');
+  expect(chart).not.toBeNull();
+  return chart as HTMLElement;
+};
+
 describe("PerUserUsage", () => {
   const mockPerUserAnalyticsCall = vi.mocked(networking.perUserAnalyticsCall);
 
@@ -70,6 +78,22 @@ describe("PerUserUsage", () => {
     });
   });
 
+  it("keeps both tab panels mounted so switching tabs does not reset their state", async () => {
+    render(<PerUserUsage {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(screen.getByText("u1")).toBeInTheDocument();
+    });
+
+    // Still on the User Details tab: the distribution panel is mounted alongside it.
+    expect(screen.getByText("User Usage Distribution")).toBeInTheDocument();
+
+    fireEvent.click(screen.getByText("Usage Distribution"));
+
+    // And the details panel survives the switch rather than unmounting.
+    expect(screen.getByText("u1")).toBeInTheDocument();
+  });
+
   it("renders the usage distribution as a stacked bar chart with the explicit palette and users formatter", async () => {
     render(<PerUserUsage {...defaultProps} />);
 
@@ -79,26 +103,22 @@ describe("PerUserUsage", () => {
 
     fireEvent.click(screen.getByText("Usage Distribution"));
 
-    const panel = screen.getByText("User Usage Distribution").closest("div")?.parentElement;
-    expect(panel).not.toBeNull();
-
     await waitFor(() => {
-      expect(panel!.querySelectorAll("path.recharts-rectangle")).toHaveLength(4);
+      expect(distributionChart().querySelectorAll("path.recharts-rectangle")).toHaveLength(4);
     });
 
-    const chart = panel!.querySelector('[data-slot="chart"]');
-    expect(chart).not.toBeNull();
-    expect(chart!.querySelectorAll(".recharts-bar")).toHaveLength(2);
+    const chart = distributionChart();
+    expect(chart.querySelectorAll(".recharts-bar")).toHaveLength(2);
 
-    const rectangles = Array.from(chart!.querySelectorAll("path.recharts-rectangle"));
+    const rectangles = Array.from(chart.querySelectorAll("path.recharts-rectangle"));
     const fills = new Set(rectangles.map((rect) => rect.getAttribute("fill")));
     expect(fills).toEqual(new Set(["var(--color-blue-500, #3b82f6)", "var(--color-green-500, #22c55e)"]));
 
     const xPositions = new Set(rectangles.map((rect) => rect.getAttribute("d")?.match(/^M\s*([\d.]+)/)?.[1]));
     expect(xPositions.size).toBe(3);
 
-    expect(chart!.textContent).toContain("curl/8.0");
-    expect(chart!.textContent).toContain("Unknown");
+    expect(chart.textContent).toContain("curl/8.0");
+    expect(chart.textContent).toContain("Unknown");
     for (const bucket of [
       "1-9 requests",
       "10-99 requests",
@@ -107,10 +127,10 @@ describe("PerUserUsage", () => {
       "10K-99.9K requests",
       "100K+ requests",
     ]) {
-      expect(chart!.textContent).toContain(bucket);
+      expect(chart.textContent).toContain(bucket);
     }
 
-    const tickTexts = Array.from(chart!.querySelectorAll(".recharts-cartesian-axis-tick-value")).map(
+    const tickTexts = Array.from(chart.querySelectorAll(".recharts-cartesian-axis-tick-value")).map(
       (tick) => tick.textContent ?? "",
     );
     expect(tickTexts.some((tick) => / users$/.test(tick))).toBe(true);
diff --git a/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx b/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx
index b2facde3e65..07a03f4addf 100644
--- a/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx
+++ b/ui/litellm-dashboard/src/components/user_agent_activity.test.tsx
@@ -177,40 +177,61 @@ describe("UserAgentActivity", () => {
     // Check that filter label is present
     expect(screen.getByText("Filter by User Agents")).toBeInTheDocument();
 
-    // The Ant Design Select component should be in the document with placeholder
-    const selectElement = screen.getByText("All User Agents");
-    expect(selectElement).toBeInTheDocument();
+    // One library paints the prompt as its own text node and the other leaves it on the input's
+    // placeholder attribute, so either one means the user is being told what the filter does.
+    const prompts =
+      screen.queryAllByText("All User Agents").length + screen.queryAllByPlaceholderText("All User Agents").length;
+    expect(prompts).toBeGreaterThan(0);
   });
 
-  const getPanelForTitle = (title: string): HTMLElement => {
-    // Assumes two wrapper divs between the Tremor <Title> and the panel root; update if Tremor's TabPanel depth changes.
-    const panel = screen.getByText(title).closest("div")?.parentElement;
-    expect(panel).not.toBeNull();
-    return panel!;
+  // Walks up from the panel's heading to the nearest ancestor that owns a chart, so the
+  // assertions do not depend on how many wrappers the tab library puts around a panel.
+  const chartForTitle = (title: string): HTMLElement => {
+    let node: HTMLElement | null = screen.getByText(title);
+    while (node && !node.querySelector('[data-slot="chart"]')) {
+      node = node.parentElement;
+    }
+    const chart = node?.querySelector('[data-slot="chart"]') ?? null;
+    expect(chart).not.toBeNull();
+    return chart as HTMLElement;
   };
 
-  const expectStackedTwoCategoryChart = (panel: HTMLElement, firstBucketLabel: string) => {
-    const chart = panel.querySelector('[data-slot="chart"]');
-    expect(chart).not.toBeNull();
-    expect(chart!.querySelectorAll(".recharts-bar")).toHaveLength(2);
+  const expectStackedTwoCategoryChart = (chart: HTMLElement, firstBucketLabel: string) => {
+    expect(chart.querySelectorAll(".recharts-bar")).toHaveLength(2);
 
-    const rectangles = Array.from(chart!.querySelectorAll("path.recharts-rectangle"));
+    const rectangles = Array.from(chart.querySelectorAll("path.recharts-rectangle"));
     const fills = new Set(rectangles.map((rect) => rect.getAttribute("fill")));
     expect(fills).toEqual(new Set(["var(--color-blue-500, #3b82f6)", "var(--color-cyan-500, #06b6d4)"]));
 
     const xPositions = new Set(rectangles.map((rect) => rect.getAttribute("d")?.match(/^M\s*([\d.]+)/)?.[1]));
     expect(xPositions.size).toBe(1);
 
-    expect(chart!.textContent).toContain("Chrome/1.0");
-    expect(chart!.textContent).toContain("Firefox/2.0");
-    expect(chart!.textContent).toContain(firstBucketLabel);
+    expect(chart.textContent).toContain("Chrome/1.0");
+    expect(chart.textContent).toContain("Firefox/2.0");
+    expect(chart.textContent).toContain(firstBucketLabel);
 
-    const tickTexts = Array.from(chart!.querySelectorAll(".recharts-cartesian-axis-tick-value")).map(
+    const tickTexts = Array.from(chart.querySelectorAll(".recharts-cartesian-axis-tick-value")).map(
       (tick) => tick.textContent ?? "",
     );
     expect(tickTexts.some((tick) => /^\d+K$/.test(tick))).toBe(true);
   };
 
+  it("keeps every tab panel mounted so switching tabs does not reset their state", async () => {
+    render(<UserAgentActivity {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(mockTagDauCall).toHaveBeenCalled();
+    });
+
+    // No tab has been clicked: the inactive DAU/WAU/MAU panels are mounted alongside the active one.
+    expect(screen.getByText("Daily Active Users - Last 7 Days")).toBeInTheDocument();
+    expect(screen.getByText("Weekly Active Users - Last 7 Weeks")).toBeInTheDocument();
+    expect(screen.getByText("Monthly Active Users - Last 7 Months")).toBeInTheDocument();
+
+    // And so is the second panel of the outer tab group.
+    expect(screen.getByText("Per User Usage")).toBeInTheDocument();
+  });
+
   it("renders the DAU chart stacked with default color cycle and abbreviated axis ticks", async () => {
     const firstBucketDate = new Date();
     firstBucketDate.setDate(firstBucketDate.getDate() - 6);
@@ -224,12 +245,16 @@ describe("UserAgentActivity", () => {
 
     render(<UserAgentActivity {...defaultProps} />);
 
-    const panel = getPanelForTitle("Daily Active Users - Last 7 Days");
     await waitFor(() => {
-      expect(panel.querySelectorAll("path.recharts-rectangle")).toHaveLength(2);
+      expect(
+        chartForTitle("Daily Active Users - Last 7 Days").querySelectorAll("path.recharts-rectangle"),
+      ).toHaveLength(2);
     });
 
-    expectStackedTwoCategoryChart(panel, firstBucketDate.toISOString().split("T")[0]);
+    expectStackedTwoCategoryChart(
+      chartForTitle("Daily Active Users - Last 7 Days"),
+      firstBucketDate.toISOString().split("T")[0],
+    );
   });
 
   it("renders the WAU chart stacked with week buckets and abbreviated axis ticks", async () => {
@@ -242,12 +267,13 @@ describe("UserAgentActivity", () => {
 
     render(<UserAgentActivity {...defaultProps} />);
 
-    const panel = getPanelForTitle("Weekly Active Users - Last 7 Weeks");
     await waitFor(() => {
-      expect(panel.querySelectorAll("path.recharts-rectangle")).toHaveLength(2);
+      expect(
+        chartForTitle("Weekly Active Users - Last 7 Weeks").querySelectorAll("path.recharts-rectangle"),
+      ).toHaveLength(2);
     });
 
-    expectStackedTwoCategoryChart(panel, "Week 1");
+    expectStackedTwoCategoryChart(chartForTitle("Weekly Active Users - Last 7 Weeks"), "Week 1");
   });
 
   it("renders the MAU chart stacked with month buckets and abbreviated axis ticks", async () => {
@@ -260,11 +286,12 @@ describe("UserAgentActivity", () => {
 
     render(<UserAgentActivity {...defaultProps} />);
 
-    const panel = getPanelForTitle("Monthly Active Users - Last 7 Months");
     await waitFor(() => {
-      expect(panel.querySelectorAll("path.recharts-rectangle")).toHaveLength(2);
+      expect(
+        chartForTitle("Monthly Active Users - Last 7 Months").querySelectorAll("path.recharts-rectangle"),
+      ).toHaveLength(2);
     });
 
-    expectStackedTwoCategoryChart(panel, "Month 1");
+    expectStackedTwoCategoryChart(chartForTitle("Monthly Active Users - Last 7 Months"), "Month 1");
   });
 });

From c5c98cf706638c2311409e533d0956620b625f3c Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 14:56:58 -0700
Subject: [PATCH 08/31] test(ui): characterise GuardrailsOverview before the
 shadcn migration

Covers the header and export action, all five summary metric cards, the
table toolbar heading, the evaluation settings modal wiring, the busy
state and the request failure message. Every assertion is role, title or
text based so it holds against both the antd markup and its shadcn
replacement, letting the migration commit land without editing this file
---
 .../_components/GuardrailsOverview.test.tsx   | 70 ++++++++++++++++++-
 1 file changed, 68 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
index 9bfa71f77be..b616982d69b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
@@ -1,4 +1,4 @@
-import { render, screen } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
@@ -14,7 +14,7 @@ vi.mock("./ScoreChart", () => ({
 }));
 
 vi.mock("./EvaluationSettingsModal", () => ({
-  EvaluationSettingsModal: () => null,
+  EvaluationSettingsModal: ({ open }: { open: boolean }) => (open ? <div>Evaluation settings modal</div> : null),
 }));
 
 const mockGetGuardrailsUsageOverview = vi.mocked(networking.getGuardrailsUsageOverview);
@@ -28,6 +28,18 @@ function wrapper({ children }: { children: React.ReactNode }) {
   return <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>;
 }
 
+function renderOverview(onSelectGuardrail = vi.fn()) {
+  return render(
+    <GuardrailsOverview
+      accessToken="test-token"
+      startDate="2026-08-01"
+      endDate="2026-08-12"
+      onSelectGuardrail={onSelectGuardrail}
+    />,
+    { wrapper },
+  );
+}
+
 describe("GuardrailsOverview", () => {
   beforeEach(() => {
     vi.clearAllMocks();
@@ -92,4 +104,58 @@ describe("GuardrailsOverview", () => {
 
     expect(onSelectGuardrail).toHaveBeenCalledWith("guardrail-low");
   });
+
+  it("renders the page header and the export action", async () => {
+    renderOverview();
+
+    expect(await screen.findByRole("heading", { name: "Guardrails Monitor", level: 1 })).toBeInTheDocument();
+    expect(screen.getByText("Monitor guardrail performance across all requests")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /Export Data/i })).toBeInTheDocument();
+  });
+
+  it("renders every summary metric card", async () => {
+    renderOverview();
+
+    expect(await screen.findByText("1,500")).toBeInTheDocument();
+    expect(screen.getByText("Total Evaluations")).toBeInTheDocument();
+    expect(screen.getByText("Blocked Requests")).toBeInTheDocument();
+    expect(screen.getByText("84")).toBeInTheDocument();
+    expect(screen.getByText("Pass Rate")).toBeInTheDocument();
+    expect(screen.getByText("94.4%")).toBeInTheDocument();
+    expect(screen.getByText("23ms")).toBeInTheDocument();
+    expect(screen.getByText("Active Guardrails")).toBeInTheDocument();
+    expect(screen.getByText("2")).toBeInTheDocument();
+  });
+
+  it("renders the table toolbar heading and its description", async () => {
+    renderOverview();
+
+    expect(await screen.findByRole("heading", { name: "Guardrail Performance", level: 5 })).toBeInTheDocument();
+    expect(screen.getByText("Click a guardrail to view details, logs, and configuration")).toBeInTheDocument();
+  });
+
+  it("opens the evaluation settings modal from the toolbar action", async () => {
+    const user = userEvent.setup();
+    renderOverview();
+
+    expect(screen.queryByText("Evaluation settings modal")).not.toBeInTheDocument();
+
+    await user.click(await screen.findByTitle("Evaluation settings"));
+
+    expect(await screen.findByText("Evaluation settings modal")).toBeInTheDocument();
+  });
+
+  it("marks the overview busy while the usage request is in flight", async () => {
+    mockGetGuardrailsUsageOverview.mockReturnValue(new Promise(() => {}));
+    renderOverview();
+
+    await waitFor(() => expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument());
+  });
+
+  it("shows a failure message when the usage request rejects", async () => {
+    mockGetGuardrailsUsageOverview.mockRejectedValue(new Error("network down"));
+    renderOverview();
+
+    expect(await screen.findByText("Failed to load data. Try again.")).toBeInTheDocument();
+  });
 });

From b69177712efe0cc7752a5b53b72254fbf60a79bd Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:02:40 -0700
Subject: [PATCH 09/31] refactor(ui): migrate usage to shadcn

Replaces antd and Tremor markup on the usage route with the installed shadcn
base-vega primitives. Behaviour is unchanged: the route's characterisation
tests were rewritten to role and text queries in the previous commit, proven
green against the antd components, and pass through this commit unedited.

Tremor tab panels stayed mounted once rendered, so every migrated TabsContent
and the collapsible model sections in activity_metrics carry keepMounted to
keep view-mode and expansion state alive across tab switches.
---
 ui/litellm-dashboard/eslint-suppressions.json |  54 +-
 .../components/EndpointUsageTable.tsx         |  18 +-
 .../components/EntityUsage/EntityUsage.tsx    | 338 ++++----
 .../EntityUsage/SpendByProvider.tsx           |  44 +-
 .../components/EntityUsage/TopModelView.tsx   |  47 +-
 .../components/UsageAIChatPanel.tsx           |  62 +-
 .../_components/components/UsagePageView.tsx  | 763 +++++++++---------
 .../UsageViewSelect/UsageViewSelect.tsx       |  99 +--
 .../EntityUsageExportModal.tsx                |  81 +-
 .../ExportFormatSelector.tsx                  |  34 +-
 .../EntityUsageExport/ExportTypeSelector.tsx  |  55 +-
 .../EntityUsageExport/UsageExportHeader.tsx   | 116 ++-
 .../src/components/activity_metrics.tsx       | 401 +++++----
 .../common_components/team_multi_select.tsx   | 132 +--
 .../src/components/per_user_usage.tsx         | 225 +++---
 .../src/components/user_agent_activity.tsx    | 320 ++++----
 16 files changed, 1454 insertions(+), 1335 deletions(-)

diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json
index 2f4def9ffa2..c94686b2f1e 100644
--- a/ui/litellm-dashboard/eslint-suppressions.json
+++ b/ui/litellm-dashboard/eslint-suppressions.json
@@ -1658,25 +1658,12 @@
   "src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx": {
     "no-nested-ternary": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx": {
     "local/no-complex-jsx-arrow": {
       "count": 2
     },
-    "no-restricted-imports": {
-      "count": 2
-    }
-  },
-  "src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx": {
-    "no-restricted-imports": {
-      "count": 2
-    }
-  },
-  "src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx": {
     "no-restricted-imports": {
       "count": 1
     }
@@ -1685,9 +1672,6 @@
     "no-nested-ternary": {
       "count": 1
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/immutability": {
       "count": 1
     }
@@ -1703,7 +1687,7 @@
       "count": 1
     },
     "no-restricted-imports": {
-      "count": 2
+      "count": 1
     },
     "react-hooks/purity": {
       "count": 1
@@ -1712,14 +1696,6 @@
       "count": 3
     }
   },
-  "src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx": {
-    "local/no-complex-jsx-arrow": {
-      "count": 2
-    },
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity.ts": {
     "react-hooks/refs": {
       "count": 1
@@ -1982,29 +1958,14 @@
       "count": 1
     }
   },
-  "src/components/EntityUsageExport/EntityUsageExportModal.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/components/EntityUsageExport/ExportFormatSelector.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/components/EntityUsageExport/ExportSummary.tsx": {
     "no-restricted-imports": {
       "count": 1
     }
   },
-  "src/components/EntityUsageExport/ExportTypeSelector.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/components/EntityUsageExport/UsageExportHeader.tsx": {
     "no-restricted-imports": {
-      "count": 3
+      "count": 1
     }
   },
   "src/components/EntityUsageExport/types.ts": {
@@ -2289,9 +2250,6 @@
     },
     "no-nested-ternary": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 2
     }
   },
   "src/components/add_model/AdaptiveRoutingConfig.tsx": {
@@ -2794,9 +2752,6 @@
   "src/components/common_components/team_multi_select.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/components/common_components/user_search_modal.tsx": {
@@ -3171,9 +3126,6 @@
     "local/filename-pascal-case": {
       "count": 1
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/set-state-in-effect": {
       "count": 1
     }
@@ -3714,7 +3666,7 @@
       "count": 1
     },
     "no-restricted-imports": {
-      "count": 3
+      "count": 1
     },
     "react-hooks/set-state-in-effect": {
       "count": 1
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
index 7b6bd57e758..5f72caae437 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
@@ -1,6 +1,6 @@
 import React from "react";
-import { Progress } from "antd";
 import type { ColumnDef } from "@tanstack/react-table";
+import { Meter, MeterIndicator, MeterTrack } from "@/components/ui/meter";
 import { DataTable } from "@/components/shared/DataTable";
 import { MoneyCell } from "@/components/shared/table_cells";
 import { MetricWithMetadata } from "@/components/UsagePage/types";
@@ -53,19 +53,15 @@ const EndpointUsageTable: React.FC<EndpointUsageTableProps> = ({ endpointData })
         const failurePercentage = record.api_requests > 0 ? (record.failed_requests / record.api_requests) * 100 : 0;
         const totalPercentage = successPercentage + failurePercentage;
 
-        const strokeColorConfig: Record<string, string> = {
-          "0%": "#22c55e",
-        };
-        if (successPercentage > 0 && successPercentage < 100) {
-          strokeColorConfig[`${successPercentage}%`] = "#22c55e";
-          strokeColorConfig[`${successPercentage + 0.01}%`] = "#ef4444";
-        }
-        strokeColorConfig["100%"] = failurePercentage > 0 ? "#ef4444" : "#22c55e";
-
         return (
           <div className="flex items-center space-x-3">
             <div className="flex-1 relative">
-              <Progress percent={totalPercentage} size="small" strokeColor={strokeColorConfig} showInfo={false} />
+              {/* The failed share is the track showing through behind the successful share. */}
+              <Meter value={successPercentage} max={totalPercentage || 100} aria-label="Successful requests">
+                <MeterTrack className={failurePercentage > 0 ? "bg-red-500" : undefined}>
+                  <MeterIndicator className="bg-green-500" />
+                </MeterTrack>
+              </Meter>
             </div>
             <div className="flex items-center space-x-2 text-sm min-w-[100px]">
               <span className="text-green-600 font-medium">{record.successful_requests.toLocaleString()}</span>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
index 3a22eb7fd10..2d6dbb823cc 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx
@@ -14,23 +14,13 @@ import { MoneyCell } from "@/components/shared/table_cells";
 import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { hasCapability, type Capability } from "@/utils/capabilities";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
-import {
-  Card,
-  Col,
-  DateRangePickerValue,
-  Grid,
-  Subtitle,
-  Tab,
-  TabGroup,
-  TabList,
-  TabPanel,
-  TabPanels,
-  Text,
-  Title,
-} from "@tremor/react";
-import { DownOutlined, ExportOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons";
+import type { DateRangePickerValue } from "@tremor/react";
+import { ChevronDown, ChevronRight, ExternalLink, Info, Loader2 } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
-import { Alert, Button, Tooltip } from "antd";
+import { Alert, AlertDescription } from "@/components/shared/Alert";
+import { Button } from "@/components/ui/button";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import React, { type ReactNode, useMemo, useState } from "react";
 import TeamMultiSelect from "@/components/common_components/team_multi_select";
 import { ActivityMetrics, processActivityData } from "@/components/activity_metrics";
@@ -341,23 +331,29 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
     [],
   );
 
-  const chev = "text-gray-400 text-xs";
-  const expandIcon = showCostBreakdown ? <DownOutlined className={chev} /> : <RightOutlined className={chev} />;
-  const infoIcon = <InfoCircleOutlined className="text-gray-400 hover:text-gray-600" />;
+  const chev = "size-3 text-gray-400";
+  const expandIcon = showCostBreakdown ? <ChevronDown className={chev} /> : <ChevronRight className={chev} />;
 
   const renderSummaryTile = ({ title, value, className, tooltip, expandable }: SummaryTile) => (
-    <Card
+    <ShadcnCard
       key={title}
       className={expandable ? "cursor-pointer hover:bg-gray-50 transition-colors" : undefined}
       onClick={expandable ? () => setShowCostBreakdown(!showCostBreakdown) : undefined}
     >
-      <div className="flex items-center gap-2">
-        <Title>{title}</Title>
-        {tooltip ? <Tooltip title={tooltip}>{infoIcon}</Tooltip> : null}
-        {expandable ? expandIcon : null}
-      </div>
-      <Text className={`text-2xl font-bold mt-2 ${className ?? ""}`}>{value}</Text>
-    </Card>
+      <CardContent>
+        <div className="flex items-center gap-2">
+          <h3 className="text-lg font-medium text-foreground">{title}</h3>
+          {tooltip ? (
+            <Tooltip>
+              <TooltipTrigger render={<Info className="size-4 text-gray-400 hover:text-gray-600" />} />
+              <TooltipContent>{tooltip}</TooltipContent>
+            </Tooltip>
+          ) : null}
+          {expandable ? expandIcon : null}
+        </div>
+        <p className={`text-2xl font-bold mt-2 ${className ?? ""}`}>{value}</p>
+      </CardContent>
+    </ShadcnCard>
   );
 
   const breakdownTiles = showFlatCost && showCostBreakdown ? buildCostBreakdownTiles(spendData.metadata) : [];
@@ -366,18 +362,18 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
   const modelViewTitle = modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models";
 
   const costPanel = (
-    <Grid numItems={2} className="gap-2 w-full">
-      <Col numColSpan={2}>
-        <Card>
-          <Title>{capitalizedEntityLabel} Spend Overview</Title>
-          <Grid numItems={5} className="gap-4 mt-4">
-            {summaryTiles.map(renderSummaryTile)}
-          </Grid>
-        </Card>
-      </Col>
+    <div className="grid grid-cols-2 gap-2 w-full">
+      <div className="col-span-2">
+        <ShadcnCard>
+          <CardContent>
+            <h3 className="text-lg font-medium text-foreground">{capitalizedEntityLabel} Spend Overview</h3>
+            <div className="grid grid-cols-5 gap-4 mt-4">{summaryTiles.map(renderSummaryTile)}</div>
+          </CardContent>
+        </ShadcnCard>
+      </div>
 
       {/* Daily Spend Chart */}
-      <Col numColSpan={2}>
+      <div className="col-span-2">
         <ShadcnCard>
           <CardHeader>
             <CardTitle className="text-base font-semibold">Daily Spend</CardTitle>
@@ -451,15 +447,15 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
             />
           </CardContent>
         </ShadcnCard>
-      </Col>
+      </div>
 
       {/* Entity Breakdown Section */}
-      <Col numColSpan={2}>
-        <Card>
-          <div className="flex flex-col space-y-4">
+      <div className="col-span-2">
+        <ShadcnCard>
+          <CardContent className="flex flex-col space-y-4">
             <div className="flex flex-col space-y-2">
-              <Title>Spend Per {capitalizedEntityLabel}</Title>
-              <Subtitle className="text-xs">Showing Top 5 by Spend</Subtitle>
+              <h3 className="text-lg font-medium text-foreground">Spend Per {capitalizedEntityLabel}</h3>
+              <p className="text-xs text-muted-foreground">Showing Top 5 by Spend</p>
               <div className="flex items-center text-sm text-gray-500">
                 <span>Get Started by Tracking cost per {capitalizedEntityLabel} </span>
                 <a
@@ -470,8 +466,8 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                 </a>
               </div>
             </div>
-            <Grid numItems={2} className="gap-6">
-              <Col numColSpan={1}>
+            <div className="grid grid-cols-2 gap-6">
+              <div>
                 <BarChart
                   className="mt-4 h-52"
                   data={getProcessedEntityBreakdownForChart()}
@@ -499,8 +495,8 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                     );
                   }}
                 />
-              </Col>
-              <Col numColSpan={1}>
+              </div>
+              <div>
                 <DataTable
                   columns={entityBreakdownColumns}
                   data={getEntityBreakdown().filter((entity) => entity.metrics.spend > 0)}
@@ -509,61 +505,69 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                   noDataMessage={`No ${entityType} spend data`}
                   size="compact"
                 />
-              </Col>
-            </Grid>
-          </div>
-        </Card>
-      </Col>
+              </div>
+            </div>
+          </CardContent>
+        </ShadcnCard>
+      </div>
 
       {/* Top API Keys */}
-      <Col numColSpan={1}>
-        <Card>
-          <Title>Top Virtual Keys</Title>
-          <TopKeyView
-            topKeys={getTopAPIKeys(spendData.results, topKeysLimit)}
-            teams={null}
-            showTags={entityType === "tag"}
-            topKeysLimit={topKeysLimit}
-            setTopKeysLimit={setTopKeysLimit}
-          />
-        </Card>
-      </Col>
+      <div>
+        <ShadcnCard>
+          <CardContent>
+            <h3 className="text-lg font-medium text-foreground">Top Virtual Keys</h3>
+            <TopKeyView
+              topKeys={getTopAPIKeys(spendData.results, topKeysLimit)}
+              teams={null}
+              showTags={entityType === "tag"}
+              topKeysLimit={topKeysLimit}
+              setTopKeysLimit={setTopKeysLimit}
+            />
+          </CardContent>
+        </ShadcnCard>
+      </div>
 
       {/* Top Models */}
-      <Col numColSpan={1}>
-        <Card>
-          <div className="flex justify-between items-center">
-            <Title>{entityType === "agent" ? "Top Agents" : modelViewTitle}</Title>
-            <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
-          </div>
-          <TopModelView
-            topModels={getTopModels(spendData.results, modelBreakdownKey, topModelsLimit)}
-            topModelsLimit={topModelsLimit}
-            setTopModelsLimit={setTopModelsLimit}
-          />
-        </Card>
-      </Col>
+      <div>
+        <ShadcnCard>
+          <CardContent>
+            <div className="flex justify-between items-center">
+              <h3 className="text-lg font-medium text-foreground">
+                {entityType === "agent" ? "Top Agents" : modelViewTitle}
+              </h3>
+              <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
+            </div>
+            <TopModelView
+              topModels={getTopModels(spendData.results, modelBreakdownKey, topModelsLimit)}
+              topModelsLimit={topModelsLimit}
+              setTopModelsLimit={setTopModelsLimit}
+            />
+          </CardContent>
+        </ShadcnCard>
+      </div>
 
       {showAgentBreakdown && (
-        <Col numColSpan={2}>
-          <Card>
-            <Title>Top Agents Driving Spend</Title>
-            <TopModelView
-              topModels={getTopAgents(agentSpendData.results, topAgentsLimit)}
-              topModelsLimit={topAgentsLimit}
-              setTopModelsLimit={setTopAgentsLimit}
-            />
-          </Card>
-        </Col>
+        <div className="col-span-2">
+          <ShadcnCard>
+            <CardContent>
+              <h3 className="text-lg font-medium text-foreground">Top Agents Driving Spend</h3>
+              <TopModelView
+                topModels={getTopAgents(agentSpendData.results, topAgentsLimit)}
+                topModelsLimit={topAgentsLimit}
+                setTopModelsLimit={setTopAgentsLimit}
+              />
+            </CardContent>
+          </ShadcnCard>
+        </div>
       )}
 
       {/* Spend by Provider */}
-      <Col numColSpan={2}>
-        <Card>
-          <div className="flex flex-col space-y-4">
-            <Title>Provider Usage</Title>
-            <Grid numItems={2}>
-              <Col numColSpan={1}>
+      <div className="col-span-2">
+        <ShadcnCard>
+          <CardContent className="flex flex-col space-y-4">
+            <h3 className="text-lg font-medium text-foreground">Provider Usage</h3>
+            <div className="grid grid-cols-2">
+              <div>
                 <DonutChart
                   className="mt-4 h-40"
                   data={providerSpend}
@@ -575,8 +579,8 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                   startAngle={90}
                   endAngle={-270}
                 />
-              </Col>
-              <Col numColSpan={1}>
+              </div>
+              <div>
                 <DataTable
                   columns={providerSpendColumns}
                   data={providerSpend}
@@ -584,12 +588,12 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
                   noDataMessage="No provider usage data"
                   size="compact"
                 />
-              </Col>
-            </Grid>
-          </div>
-        </Card>
-      </Col>
-    </Grid>
+              </div>
+            </div>
+          </CardContent>
+        </ShadcnCard>
+      </div>
+    </div>
   );
 
   const tabs: readonly { key: string; label: string; content: ReactNode }[] = [
@@ -620,80 +624,60 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
   return (
     <div style={{ width: "100%" }} className="relative">
       {isFetchingMore && (
-        <Alert
-          banner
-          type="warning"
-          className="mb-2"
-          message={
-            <div className="flex items-center justify-between">
-              <span>
-                <LoadingOutlined spin className="mr-2" />
-                Currently fetching spend data: fetched {progress.currentPage} / {progress.totalPages} pages. Charts will
-                update periodically as data loads. Moving off of this page will stop and reset this. To continue using
-                the UI in the meantime,{" "}
-                <a href={window.location.href} target="_blank" rel="noopener noreferrer">
-                  open a new tab <ExportOutlined />
-                </a>
-                .
-              </span>
-              <Button type="primary" danger onClick={cancel}>
-                Stop
-              </Button>
-            </div>
-          }
-        />
+        <Alert variant="warning" className="mb-2">
+          <AlertDescription className="flex items-center justify-between text-inherit">
+            <span>
+              <Loader2 className="mr-2 inline size-4 animate-spin align-text-bottom" />
+              Currently fetching spend data: fetched {progress.currentPage} / {progress.totalPages} pages. Charts will
+              update periodically as data loads. Moving off of this page will stop and reset this. To continue using the
+              UI in the meantime,{" "}
+              <a href={window.location.href} target="_blank" rel="noopener noreferrer">
+                open a new tab <ExternalLink className="inline size-3.5 align-text-bottom" />
+              </a>
+              .
+            </span>
+            <Button variant="destructive" onClick={cancel}>
+              Stop
+            </Button>
+          </AlertDescription>
+        </Alert>
       )}
       {cancelled && (
-        <Alert
-          banner
-          type="info"
-          className="mb-2"
-          message={
-            <span>
-              Showing partial data ({progress.currentPage}/{progress.totalPages} pages loaded)
-            </span>
-          }
-        />
+        <Alert variant="info" className="mb-2">
+          <AlertDescription className="text-inherit">
+            Showing partial data ({progress.currentPage}/{progress.totalPages} pages loaded)
+          </AlertDescription>
+        </Alert>
       )}
       {agentIsFetchingMore && showAgentBreakdown && (
-        <Alert
-          banner
-          type="warning"
-          className="mb-2"
-          message={
-            <div className="flex items-center justify-between">
-              <span>
-                <LoadingOutlined spin className="mr-2" />
-                Currently fetching agent data: fetched {agentProgress.currentPage} / {agentProgress.totalPages} pages.
-                Charts will update periodically as data loads. Moving off of this page will stop and reset this. To
-                continue using the UI in the meantime,{" "}
-                <a href={window.location.href} target="_blank" rel="noopener noreferrer">
-                  open a new tab <ExportOutlined />
-                </a>
-                .
-              </span>
-              <Button type="primary" danger onClick={agentCancel}>
-                Stop
-              </Button>
-            </div>
-          }
-        />
+        <Alert variant="warning" className="mb-2">
+          <AlertDescription className="flex items-center justify-between text-inherit">
+            <span>
+              <Loader2 className="mr-2 inline size-4 animate-spin align-text-bottom" />
+              Currently fetching agent data: fetched {agentProgress.currentPage} / {agentProgress.totalPages} pages.
+              Charts will update periodically as data loads. Moving off of this page will stop and reset this. To
+              continue using the UI in the meantime,{" "}
+              <a href={window.location.href} target="_blank" rel="noopener noreferrer">
+                open a new tab <ExternalLink className="inline size-3.5 align-text-bottom" />
+              </a>
+              .
+            </span>
+            <Button variant="destructive" onClick={agentCancel}>
+              Stop
+            </Button>
+          </AlertDescription>
+        </Alert>
       )}
       {agentCancelled && showAgentBreakdown && (
-        <Alert
-          banner
-          type="info"
-          className="mb-2"
-          message={
-            <span>
-              Showing partial agent data ({agentProgress.currentPage}/{agentProgress.totalPages} pages loaded)
-            </span>
-          }
-        />
+        <Alert variant="info" className="mb-2">
+          <AlertDescription className="text-inherit">
+            Showing partial agent data ({agentProgress.currentPage}/{agentProgress.totalPages} pages loaded)
+          </AlertDescription>
+        </Alert>
       )}
       {entityType === "team" && (
         <div className="mb-4">
-          <Text className="mb-2">Filter by team</Text>
+          <p className="mb-2 text-sm text-foreground">Filter by team</p>
           <TeamMultiSelect value={selectedTags} onChange={setSelectedTags} />
         </div>
       )}
@@ -710,18 +694,20 @@ const EntityUsage: React.FC<EntityUsageProps> = ({
         filterMode={entityType === "user" ? "single" : "multiple"}
         teams={teams || []}
       />
-      <TabGroup>
-        <TabList variant="solid" className="mt-1">
+      <Tabs defaultValue={tabs[0].key}>
+        <TabsList className="mt-1">
           {tabs.map(({ key, label }) => (
-            <Tab key={key}>{label}</Tab>
+            <TabsTrigger key={key} value={key} className="flex-none px-3">
+              {label}
+            </TabsTrigger>
           ))}
-        </TabList>
-        <TabPanels>
-          {tabs.map(({ key, content }) => (
-            <TabPanel key={key}>{content}</TabPanel>
-          ))}
-        </TabPanels>
-      </TabGroup>
+        </TabsList>
+        {tabs.map(({ key, content }) => (
+          <TabsContent key={key} value={key} keepMounted>
+            {content}
+          </TabsContent>
+        ))}
+      </Tabs>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx
index ac8ae0e67e1..38f055f65f3 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/SpendByProvider.tsx
@@ -2,10 +2,11 @@ import { DonutChart } from "@/components/shared/charts";
 import { DataTable } from "@/components/shared/DataTable";
 import { MoneyCell } from "@/components/shared/table_cells";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
-import { InfoCircleOutlined } from "@ant-design/icons";
+import { Info } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
-import { Card, Col, Grid, Switch, Title } from "@tremor/react";
-import { Tooltip } from "antd";
+import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Switch } from "@/components/ui/switch";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import React, { useState } from "react";
 import { ProviderLogo } from "@/components/molecules/models/ProviderLogo";
 import { ChartLoader } from "@/components/shared/chart_loader";
@@ -85,29 +86,30 @@ const SpendByProvider: React.FC<SpendByProviderProps> = ({ loading, isDateChangi
 
   return (
     <Card className="h-full">
-      <div className="flex justify-between items-center mb-4">
-        <Title>Spend by Provider</Title>
-        <div className="flex items-center gap-4">
+      <CardHeader>
+        <CardTitle>Spend by Provider</CardTitle>
+        <CardAction className="flex items-center gap-4">
           <div className="flex items-center gap-2">
             <label className="text-sm text-gray-700">Show Zero Spend</label>
-            <Switch checked={includeZeroSpend} onChange={setIncludeZeroSpend} />
+            <Switch checked={includeZeroSpend} onCheckedChange={setIncludeZeroSpend} />
           </div>
           <div className="flex items-center gap-2">
             <div className="flex items-center gap-1">
               <label className="text-sm text-gray-700">Show Unknown</label>
-              <Tooltip title="Requests that failed to route to a provider">
-                <InfoCircleOutlined className="text-gray-400 hover:text-gray-600" />
+              <Tooltip>
+                <TooltipTrigger render={<Info className="size-4 text-gray-400 hover:text-gray-600" />} />
+                <TooltipContent>Requests that failed to route to a provider</TooltipContent>
               </Tooltip>
             </div>
-            <Switch checked={includeUnknown} onChange={setIncludeUnknown} />
+            <Switch checked={includeUnknown} onCheckedChange={setIncludeUnknown} />
           </div>
-        </div>
-      </div>
-      {loading ? (
-        <ChartLoader isDateChanging={isDateChanging} />
-      ) : (
-        <Grid numItems={2}>
-          <Col numColSpan={1}>
+        </CardAction>
+      </CardHeader>
+      <CardContent>
+        {loading ? (
+          <ChartLoader isDateChanging={isDateChanging} />
+        ) : (
+          <div className="grid grid-cols-2">
             <DonutChart
               className="mt-4 h-40"
               data={filteredProviderSpend}
@@ -119,8 +121,6 @@ const SpendByProvider: React.FC<SpendByProviderProps> = ({ loading, isDateChangi
               startAngle={90}
               endAngle={-270}
             />
-          </Col>
-          <Col numColSpan={1}>
             <DataTable
               columns={columns}
               data={filteredProviderSpend}
@@ -128,9 +128,9 @@ const SpendByProvider: React.FC<SpendByProviderProps> = ({ loading, isDateChangi
               noDataMessage="No provider usage data"
               size="compact"
             />
-          </Col>
-        </Grid>
-      )}
+          </div>
+        )}
+      </CardContent>
     </Card>
   );
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx
index 2e2880e8763..2f5b750c118 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/TopModelView.tsx
@@ -1,7 +1,7 @@
 import { BarChart } from "@/components/shared/charts";
 import { DataTable } from "@/components/shared/DataTable";
 import { MoneyCell } from "@/components/shared/table_cells";
-import { Segmented } from "antd";
+import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { useState } from "react";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 
@@ -19,6 +19,8 @@ interface TopModelViewProps {
   setTopModelsLimit: (limit: number) => void;
 }
 
+export const TOP_MODEL_LIMITS = [5, 10, 25, 50];
+
 export default function TopModelView({ topModels, topModelsLimit, setTopModelsLimit }: TopModelViewProps) {
   const [modelViewMode, setModelViewMode] = useState<"chart" | "table">("table");
 
@@ -58,30 +60,25 @@ export default function TopModelView({ topModels, topModelsLimit, setTopModelsLi
   return (
     <>
       <div className="mb-4 flex justify-between items-center">
-        <Segmented
-          options={[
-            { label: "5", value: 5 },
-            { label: "10", value: 10 },
-            { label: "25", value: 25 },
-            { label: "50", value: 50 },
-          ]}
-          value={topModelsLimit}
-          onChange={(value) => setTopModelsLimit(value as number)}
-        />
-        <div className="flex space-x-2">
-          <button
-            onClick={() => setModelViewMode("table")}
-            className={`px-3 py-1 text-sm rounded-md ${modelViewMode === "table" ? "bg-blue-100 text-blue-700" : "bg-gray-100 text-gray-700"}`}
-          >
-            Table View
-          </button>
-          <button
-            onClick={() => setModelViewMode("chart")}
-            className={`px-3 py-1 text-sm rounded-md ${modelViewMode === "chart" ? "bg-blue-100 text-blue-700" : "bg-gray-100 text-gray-700"}`}
-          >
-            Chart View
-          </button>
-        </div>
+        <Tabs value={String(topModelsLimit)} onValueChange={(value: string) => setTopModelsLimit(Number(value))}>
+          <TabsList aria-label="Number of models to show">
+            {TOP_MODEL_LIMITS.map((limit) => (
+              <TabsTrigger key={limit} value={String(limit)} className="flex-none px-3">
+                {limit}
+              </TabsTrigger>
+            ))}
+          </TabsList>
+        </Tabs>
+        <Tabs value={modelViewMode} onValueChange={(value: string) => setModelViewMode(value as "chart" | "table")}>
+          <TabsList aria-label="Top model view mode">
+            <TabsTrigger value="table" className="flex-none px-3">
+              Table View
+            </TabsTrigger>
+            <TabsTrigger value="chart" className="flex-none px-3">
+              Chart View
+            </TabsTrigger>
+          </TabsList>
+        </Tabs>
       </div>
       {modelViewMode === "chart" ? (
         <div className="relative max-h-[600px] overflow-y-auto">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx
index 9d6b09353a2..21314594cb4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageAIChatPanel.tsx
@@ -1,10 +1,18 @@
 import React, { useEffect, useRef, useState } from "react";
-import { Button, Select, Input, Spin } from "antd";
 import ReactMarkdown from "react-markdown";
+import { Button } from "@/components/ui/button";
+import {
+  Combobox,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxInput,
+  ComboboxItem,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { Textarea } from "@/components/ui/textarea";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { modelHubCall, usageAiChatStream, UsageAiToolCallEvent } from "@/components/networking";
 
-const { TextArea } = Input;
-
 interface ToolCallStep {
   tool_name: string;
   tool_label: string;
@@ -41,7 +49,7 @@ const ToolCallDisplay: React.FC<{ step: ToolCallStep }> = ({ step }) => {
     <div className="flex items-start gap-2 px-3 py-2 rounded-lg bg-gray-100 border border-gray-200 text-xs">
       <span className="shrink-0 mt-0.5">
         {step.status === "running" ? (
-          <Spin size="small" />
+          <UiLoadingSpinner className="size-3.5" />
         ) : step.status === "error" ? (
           <span className="text-red-500">✗</span>
         ) : (
@@ -259,18 +267,29 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({ open, onClose, acce
 
       {/* Model selector */}
       <div className="px-5 py-3 border-b border-gray-100 shrink-0">
-        <Select
-          placeholder="Select a model (optional, defaults to gpt-4o-mini)"
-          value={selectedModel}
-          onChange={(value) => setSelectedModel(value)}
-          loading={isLoadingModels}
-          showSearch
-          allowClear
-          size="small"
-          className="w-full"
-          options={availableModels.map((m) => ({ label: m, value: m }))}
-          filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())}
-        />
+        <Combobox
+          items={availableModels}
+          value={selectedModel ?? null}
+          onValueChange={(value: string | null) => setSelectedModel(value ?? undefined)}
+        >
+          <ComboboxInput
+            className="w-full"
+            placeholder="Select a model (optional, defaults to gpt-4o-mini)"
+            aria-label="Select a model (optional, defaults to gpt-4o-mini)"
+            aria-busy={isLoadingModels}
+            showClear={selectedModel !== undefined}
+          />
+          <ComboboxContent>
+            <ComboboxEmpty>{isLoadingModels ? "Loading models…" : "No models found"}</ComboboxEmpty>
+            <ComboboxList>
+              {(model: string) => (
+                <ComboboxItem key={model} value={model}>
+                  {model}
+                </ComboboxItem>
+              )}
+            </ComboboxList>
+          </ComboboxContent>
+        </Combobox>
       </div>
 
       {/* Chat messages */}
@@ -329,7 +348,7 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({ open, onClose, acce
         {/* Status / spinner */}
         {isLoading && !streamingContent && (
           <div className="flex items-center gap-2 px-3 py-2 text-xs text-gray-500">
-            <Spin size="small" />
+            <UiLoadingSpinner className="size-3.5" />
             <span className="italic">{statusMessage || "Thinking..."}</span>
           </div>
         )}
@@ -347,16 +366,17 @@ const UsageAIChatPanel: React.FC<UsageAIChatPanelProps> = ({ open, onClose, acce
       {/* Input area */}
       <div className="px-4 py-3 border-t border-gray-200 bg-white shrink-0">
         <div className="flex gap-2">
-          <TextArea
+          <Textarea
             value={inputText}
             onChange={(e) => setInputText(e.target.value)}
             onKeyDown={handleKeyDown}
             placeholder="Ask about your usage..."
-            autoSize={{ minRows: 1, maxRows: 3 }}
-            className="flex-1"
+            rows={1}
+            className="flex-1 min-h-9 max-h-24"
             disabled={isLoading}
           />
-          <Button type="primary" onClick={handleSend} disabled={!inputText.trim() || isLoading} loading={isLoading}>
+          <Button onClick={handleSend} disabled={!inputText.trim() || isLoading}>
+            {isLoading && <UiLoadingSpinner className="size-4" />}
             Send
           </Button>
         </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
index 494df313ac0..f3e7db83c6a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
@@ -6,27 +6,17 @@
  * Works at 1m+ spend logs, by querying an aggregate table instead.
  */
 
-import { DownOutlined, ExportOutlined, InfoCircleOutlined, LoadingOutlined, RightOutlined } from "@ant-design/icons";
-import { useDebouncedState } from "@tanstack/react-pacer/debouncer";
-import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
-import {
-  Card,
-  Col,
-  DateRangePickerValue,
-  Grid,
-  Tab,
-  TabGroup,
-  TabList,
-  TabPanel,
-  TabPanels,
-  Text,
-  Title,
-} from "@tremor/react";
-import { Alert, Button, Segmented, Select, Tooltip, Typography } from "antd";
-import React, { useCallback, useEffect, useMemo, useRef, useState, type UIEvent } from "react";
+import { ChevronDown, ChevronRight, Download, ExternalLink, Info, Loader2, Sparkles, X } from "lucide-react";
+import type { DateRangePickerValue } from "@tremor/react";
+import React, { useCallback, useEffect, useMemo, useRef, useState } from "react";
 
 import { BarChart } from "@/components/shared/charts";
+import { Alert, AlertAction, AlertDescription, AlertTitle } from "@/components/shared/Alert";
+import { PaginatedSearchSelect } from "@/components/shared/PaginatedSearchSelect";
+import { Button } from "@/components/ui/button";
 import { Card as ShadcnCard, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 
 import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents";
 import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers";
@@ -68,6 +58,7 @@ import EndpointUsage from "./EndpointUsage/EndpointUsage";
 import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage";
 import ModelViewToggle, { ModelViewType } from "./ModelViewToggle";
 import SpendByProvider from "./EntityUsage/SpendByProvider";
+import { TOP_MODEL_LIMITS } from "./EntityUsage/TopModelView";
 import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView";
 import UsageAIChatPanel from "./UsageAIChatPanel";
 import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect";
@@ -113,11 +104,8 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
   const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage");
   const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage");
 
-  // Debounced search for user selector
-  const [userSearchInput, setUserSearchInput] = useState("");
-  const [debouncedUserSearch, setDebouncedUserSearch] = useDebouncedState("", {
-    wait: DEBOUNCE_WAIT_MS,
-  });
+  // The selector debounces its own input, so this holds the already-settled query.
+  const [debouncedUserSearch, setDebouncedUserSearch] = useState("");
 
   const {
     data: usersInfiniteData,
@@ -148,19 +136,6 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
     return result;
   }, [usersInfiniteData]);
 
-  const handleUserSearchChange = (value: string) => {
-    setUserSearchInput(value);
-    setDebouncedUserSearch(value);
-  };
-
-  const handleUserPopupScroll = (e: UIEvent<HTMLDivElement>) => {
-    const target = e.currentTarget;
-    const scrollRatio = (target.scrollTop + target.clientHeight) / target.scrollHeight;
-    if (scrollRatio >= 0.8 && hasNextUsersPage && !isFetchingNextUsersPage) {
-      fetchNextUsersPage();
-    }
-  };
-
   // For admins: null means global view (all users), a string means filter by that user
   // For non-admins: always set to their own user ID
   const [selectedUserId, setSelectedUserId] = useState<string | null>(isAdmin ? null : userID || null);
@@ -522,352 +497,361 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
             <AdvancedDatePicker value={dateValue} onValueChange={handleDateChange} />
           </div>
           {paginatedResult.isFetchingMore && (
-            <Alert
-              banner
-              type="warning"
-              className="mb-2"
-              message={
-                <div className="flex items-center justify-between">
-                  <span>
-                    <LoadingOutlined spin className="mr-2" />
-                    Currently fetching spend data: fetched {paginatedResult.progress.currentPage} /{" "}
-                    {paginatedResult.progress.totalPages} pages. Charts will update periodically as data loads. Moving
-                    off of this page will stop and reset this. To continue using the UI in the meantime,{" "}
-                    <a href={window.location.href} target="_blank" rel="noopener noreferrer">
-                      open a new tab <ExportOutlined />
-                    </a>
-                    .
-                  </span>
-                  <Button type="primary" danger onClick={paginatedResult.cancel}>
-                    Stop
-                  </Button>
-                </div>
-              }
-            />
+            <Alert variant="warning" className="mb-2">
+              <AlertDescription className="flex items-center justify-between text-inherit">
+                <span>
+                  <Loader2 className="mr-2 inline size-4 animate-spin align-text-bottom" />
+                  Currently fetching spend data: fetched {paginatedResult.progress.currentPage} /{" "}
+                  {paginatedResult.progress.totalPages} pages. Charts will update periodically as data loads. Moving off
+                  of this page will stop and reset this. To continue using the UI in the meantime,{" "}
+                  <a href={window.location.href} target="_blank" rel="noopener noreferrer">
+                    open a new tab <ExternalLink className="inline size-3.5 align-text-bottom" />
+                  </a>
+                  .
+                </span>
+                <Button variant="destructive" onClick={paginatedResult.cancel}>
+                  Stop
+                </Button>
+              </AlertDescription>
+            </Alert>
           )}
           {paginatedResult.cancelled && (
-            <Alert
-              banner
-              type="info"
-              className="mb-2"
-              message={
-                <span>
-                  Showing partial data ({paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages}{" "}
-                  pages loaded)
-                </span>
-              }
-            />
+            <Alert variant="info" className="mb-2">
+              <AlertDescription className="text-inherit">
+                Showing partial data ({paginatedResult.progress.currentPage}/{paginatedResult.progress.totalPages} pages
+                loaded)
+              </AlertDescription>
+            </Alert>
           )}
           {/* Your Usage / Global Usage Panel */}
           {(usageView === "global" || usageView === "my-usage") && (
             <>
               {isAdmin && usageView === "global" && (
                 <div className="mb-4">
-                  <Text className="mb-2">Filter by user</Text>
-                  <Select
-                    showSearch
-                    allowClear
-                    style={{ width: "100%" }}
-                    placeholder="Select user to filter..."
-                    value={selectedUserId}
-                    onChange={(value) => setSelectedUserId(value ?? null)}
-                    filterOption={false}
-                    onSearch={handleUserSearchChange}
-                    searchValue={userSearchInput}
-                    onPopupScroll={handleUserPopupScroll}
-                    loading={isLoadingUsers}
-                    notFoundContent={isLoadingUsers ? <LoadingOutlined spin /> : "No users found"}
+                  <p className="mb-2 text-sm text-foreground">Filter by user</p>
+                  <PaginatedSearchSelect
                     options={userOptions}
-                    popupRender={(menu) => (
-                      <>
-                        {menu}
-                        {isFetchingNextUsersPage && (
-                          <div style={{ textAlign: "center", padding: 8 }}>
-                            <LoadingOutlined spin />
-                          </div>
-                        )}
-                      </>
-                    )}
+                    value={selectedUserId ?? undefined}
+                    onValueChange={(value) => setSelectedUserId(value === "" ? null : value)}
+                    onSearchChange={setDebouncedUserSearch}
+                    onLoadMore={fetchNextUsersPage}
+                    hasNextPage={hasNextUsersPage}
+                    isLoading={isLoadingUsers}
+                    isFetchingNextPage={isFetchingNextUsersPage}
+                    placeholder="Select user to filter..."
+                    emptyText="No users found"
                   />
                 </div>
               )}
-              <TabGroup>
+              <Tabs defaultValue="cost">
                 <div className="flex justify-between items-center">
-                  <TabList variant="solid" className="mt-1">
-                    <Tab>Cost</Tab>
-                    <Tab>Model Activity</Tab>
-                    <Tab>Key Activity</Tab>
-                    <Tab>MCP Server Activity</Tab>
-                    <Tab>Endpoint Activity</Tab>
-                  </TabList>
+                  <TabsList className="mt-1">
+                    <TabsTrigger value="cost" className="flex-none px-3">
+                      Cost
+                    </TabsTrigger>
+                    <TabsTrigger value="models" className="flex-none px-3">
+                      Model Activity
+                    </TabsTrigger>
+                    <TabsTrigger value="keys" className="flex-none px-3">
+                      Key Activity
+                    </TabsTrigger>
+                    <TabsTrigger value="mcp" className="flex-none px-3">
+                      MCP Server Activity
+                    </TabsTrigger>
+                    <TabsTrigger value="endpoints" className="flex-none px-3">
+                      Endpoint Activity
+                    </TabsTrigger>
+                  </TabsList>
                   <div className="flex items-center gap-2">
-                    <Button
-                      onClick={() => setIsAiChatOpen(true)}
-                      icon={
-                        <svg className="w-4 h-4" viewBox="0 0 16 16" fill="currentColor">
-                          <path d="M8 1l1.5 3.5L13 6l-3.5 1.5L8 11 6.5 7.5 3 6l3.5-1.5L8 1zm4 7l.75 1.75L14.5 10.5l-1.75.75L12 13l-.75-1.75L9.5 10.5l1.75-.75L12 8zM4 9l.75 1.75L6.5 11.5l-1.75.75L4 14l-.75-1.75L1.5 11.5l1.75-.75L4 9z" />
-                        </svg>
-                      }
-                    >
+                    <Button variant="outline" onClick={() => setIsAiChatOpen(true)}>
+                      <Sparkles />
                       Ask AI
                     </Button>
-                    <Button
-                      onClick={() => setIsGlobalExportModalOpen(true)}
-                      icon={
-                        <svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-                          <path
-                            strokeLinecap="round"
-                            strokeLinejoin="round"
-                            strokeWidth={2}
-                            d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
-                          />
-                        </svg>
-                      }
-                    >
+                    <Button variant="outline" onClick={() => setIsGlobalExportModalOpen(true)}>
+                      <Download />
                       Export Data
                     </Button>
                   </div>
                 </div>
-                <TabPanels>
-                  {/* Cost Panel */}
-                  <TabPanel>
-                    <Grid numItems={2} className="gap-2 w-full">
-                      {/* Total Spend Card */}
-                      <Col numColSpan={2}>
-                        <div className="flex items-center gap-4 mt-2 mb-2">
-                          <Text className="text-tremor-default text-tremor-content dark:text-dark-tremor-content text-lg">
-                            Project Spend{" "}
-                            {dateValue.from && dateValue.to && (
-                              <>
-                                {dateValue.from.toLocaleDateString("en-US", {
-                                  month: "short",
-                                  day: "numeric",
-                                  year:
-                                    dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined,
-                                })}
-                                {" - "}
-                                {dateValue.to.toLocaleDateString("en-US", {
-                                  month: "short",
-                                  day: "numeric",
-                                  year: "numeric",
-                                })}
-                              </>
-                            )}
-                          </Text>
-                        </div>
+                {/* Cost Panel */}
+                <TabsContent value="cost" keepMounted>
+                  <div className="grid grid-cols-2 gap-2 w-full">
+                    {/* Total Spend Card */}
+                    <div className="col-span-2">
+                      <div className="flex items-center gap-4 mt-2 mb-2">
+                        <p className="text-lg text-muted-foreground">
+                          Project Spend{" "}
+                          {dateValue.from && dateValue.to && (
+                            <>
+                              {dateValue.from.toLocaleDateString("en-US", {
+                                month: "short",
+                                day: "numeric",
+                                year:
+                                  dateValue.from.getFullYear() !== dateValue.to.getFullYear() ? "numeric" : undefined,
+                              })}
+                              {" - "}
+                              {dateValue.to.toLocaleDateString("en-US", {
+                                month: "short",
+                                day: "numeric",
+                                year: "numeric",
+                              })}
+                            </>
+                          )}
+                        </p>
+                      </div>
 
-                        <ViewUserSpend
-                          userSpend={totalSpend}
-                          selectedTeam={null}
-                          userMaxBudget={currentUser?.max_budget || null}
-                        />
-                      </Col>
+                      <ViewUserSpend
+                        userSpend={totalSpend}
+                        selectedTeam={null}
+                        userMaxBudget={currentUser?.max_budget || null}
+                      />
+                    </div>
 
-                      <Col numColSpan={2}>
-                        <Card>
-                          <Title>Usage Metrics</Title>
-                          <Grid numItems={5} className="gap-4 mt-4">
-                            <Card>
-                              <Title>Total Requests</Title>
-                              <Text className="text-2xl font-bold mt-2">
-                                {userSpendData.metadata?.total_api_requests?.toLocaleString() || 0}
-                              </Text>
-                            </Card>
-                            <Card>
-                              <div className="flex items-center gap-2">
-                                <Title>Successful Requests</Title>
-                                {gatewayActivity && (
-                                  <Tooltip title="Counted by the gateway when it answers a request, independent of spend logging. Deployment-wide, so it will not match the per-key or per-model breakdowns below.">
-                                    <InfoCircleOutlined className="text-gray-400 hover:text-gray-600" />
+                    <div className="col-span-2">
+                      <ShadcnCard>
+                        <CardContent>
+                          <h3 className="text-lg font-medium text-foreground">Usage Metrics</h3>
+                          <div className="grid grid-cols-5 gap-4 mt-4">
+                            <ShadcnCard>
+                              <CardContent>
+                                <h3 className="text-lg font-medium text-foreground">Total Requests</h3>
+                                <p className="text-2xl font-bold mt-2">
+                                  {userSpendData.metadata?.total_api_requests?.toLocaleString() || 0}
+                                </p>
+                              </CardContent>
+                            </ShadcnCard>
+                            <ShadcnCard>
+                              <CardContent>
+                                <div className="flex items-center gap-2">
+                                  <h3 className="text-lg font-medium text-foreground">Successful Requests</h3>
+                                  {gatewayActivity && (
+                                    <Tooltip>
+                                      <TooltipTrigger
+                                        render={<Info className="size-4 text-gray-400 hover:text-gray-600" />}
+                                      />
+                                      <TooltipContent>
+                                        Counted by the gateway when it answers a request, independent of spend logging.
+                                        Deployment-wide, so it will not match the per-key or per-model breakdowns below.
+                                      </TooltipContent>
+                                    </Tooltip>
+                                  )}
+                                </div>
+                                {/*
+                                  TODO: drop the userSpendData fallback once every deployment
+                                  is writing LiteLLM_DailyGatewayRequests. It covers two cases
+                                  today: a non-admin (who may not read deployment-wide counts)
+                                  and an admin on a proxy whose table is still backfilling.
+                                */}
+                                <p className="text-2xl font-bold mt-2 text-green-600">
+                                  {(
+                                    gatewayActivity?.total_successful_requests ??
+                                    userSpendData.metadata?.total_successful_requests
+                                  )?.toLocaleString() || 0}
+                                </p>
+                              </CardContent>
+                            </ShadcnCard>
+                            <ShadcnCard>
+                              <CardContent>
+                                <div className="flex items-center gap-2">
+                                  <h3 className="text-lg font-medium text-foreground">Failed Requests</h3>
+                                  <Tooltip>
+                                    <TooltipTrigger
+                                      render={<Info className="size-4 text-gray-400 hover:text-gray-600" />}
+                                    />
+                                    <TooltipContent>
+                                      {gatewayActivity
+                                        ? "Counted by the gateway when it answers a request, independent of spend logging. Deployment-wide, so it will not match the per-key or per-model breakdowns below."
+                                        : "Includes requests that failed to route to a provider, tool usage failures, and other request errors where the provider cannot be determined."}
+                                    </TooltipContent>
                                   </Tooltip>
-                                )}
-                              </div>
-                              {/*
-                                TODO: drop the userSpendData fallback once every deployment
-                                is writing LiteLLM_DailyGatewayRequests. It covers two cases
-                                today: a non-admin (who may not read deployment-wide counts)
-                                and an admin on a proxy whose table is still backfilling.
-                              */}
-                              <Text className="text-2xl font-bold mt-2 text-green-600">
-                                {(
-                                  gatewayActivity?.total_successful_requests ??
-                                  userSpendData.metadata?.total_successful_requests
-                                )?.toLocaleString() || 0}
-                              </Text>
-                            </Card>
-                            <Card>
-                              <div className="flex items-center gap-2">
-                                <Title>Failed Requests</Title>
-                                <Tooltip
-                                  title={
-                                    gatewayActivity
-                                      ? "Counted by the gateway when it answers a request, independent of spend logging. Deployment-wide, so it will not match the per-key or per-model breakdowns below."
-                                      : "Includes requests that failed to route to a provider, tool usage failures, and other request errors where the provider cannot be determined."
-                                  }
-                                >
-                                  <InfoCircleOutlined className="text-gray-400 hover:text-gray-600" />
-                                </Tooltip>
-                              </div>
-                              {/* Same source as Successful Requests: the two must agree, or the
-                                  tile disagrees with the endpoint breakdown chart below it. */}
-                              <Text className="text-2xl font-bold mt-2 text-red-600">
-                                {(
-                                  gatewayActivity?.total_failed_requests ??
-                                  userSpendData.metadata?.total_failed_requests
-                                )?.toLocaleString() || 0}
-                              </Text>
-                            </Card>
-                            <Card>
-                              <Title>Average Cost per Request</Title>
-                              <Text className="text-2xl font-bold mt-2">
-                                $
-                                {formatNumberWithCommas(
-                                  (totalSpend || 0) / (userSpendData.metadata?.total_api_requests || 1),
-                                  4,
-                                )}
-                              </Text>
-                            </Card>
-                            <Card
+                                </div>
+                                {/* Same source as Successful Requests: the two must agree, or the
+                                    tile disagrees with the endpoint breakdown chart below it. */}
+                                <p className="text-2xl font-bold mt-2 text-red-600">
+                                  {(
+                                    gatewayActivity?.total_failed_requests ??
+                                    userSpendData.metadata?.total_failed_requests
+                                  )?.toLocaleString() || 0}
+                                </p>
+                              </CardContent>
+                            </ShadcnCard>
+                            <ShadcnCard>
+                              <CardContent>
+                                <h3 className="text-lg font-medium text-foreground">Average Cost per Request</h3>
+                                <p className="text-2xl font-bold mt-2">
+                                  $
+                                  {formatNumberWithCommas(
+                                    (totalSpend || 0) / (userSpendData.metadata?.total_api_requests || 1),
+                                    4,
+                                  )}
+                                </p>
+                              </CardContent>
+                            </ShadcnCard>
+                            <ShadcnCard
                               className="cursor-pointer hover:bg-gray-50 transition-colors"
                               onClick={() => setShowTokenBreakdown(!showTokenBreakdown)}
                             >
-                              <div className="flex items-center gap-2">
-                                <Title>Total Tokens</Title>
-                                {showTokenBreakdown ? (
-                                  <DownOutlined className="text-gray-400 text-xs" />
-                                ) : (
-                                  <RightOutlined className="text-gray-400 text-xs" />
-                                )}
-                              </div>
-                              <Text className="text-2xl font-bold mt-2">
-                                {userSpendData.metadata?.total_tokens?.toLocaleString() || 0}
-                              </Text>
-                            </Card>
-                          </Grid>
+                              <CardContent>
+                                <div className="flex items-center gap-2">
+                                  <h3 className="text-lg font-medium text-foreground">Total Tokens</h3>
+                                  {showTokenBreakdown ? (
+                                    <ChevronDown className="size-3 text-gray-400" />
+                                  ) : (
+                                    <ChevronRight className="size-3 text-gray-400" />
+                                  )}
+                                </div>
+                                <p className="text-2xl font-bold mt-2">
+                                  {userSpendData.metadata?.total_tokens?.toLocaleString() || 0}
+                                </p>
+                              </CardContent>
+                            </ShadcnCard>
+                          </div>
                           {showTokenBreakdown && (
-                            <Grid numItems={4} className="gap-4 mt-4">
-                              <Card>
-                                <Title>Input Tokens</Title>
-                                <Text className="text-2xl font-bold mt-2 text-blue-600">
-                                  {(userSpendData.metadata?.total_prompt_tokens || 0).toLocaleString()}
-                                </Text>
-                              </Card>
-                              <Card>
-                                <Title>Output Tokens</Title>
-                                <Text className="text-2xl font-bold mt-2 text-cyan-600">
-                                  {userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0}
-                                </Text>
-                              </Card>
-                              <Card>
-                                <Title>Cache Read Tokens</Title>
-                                <Text className="text-2xl font-bold mt-2 text-green-600">
-                                  {userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0}
-                                </Text>
-                              </Card>
-                              <Card>
-                                <Title>Cache Write Tokens</Title>
-                                <Text className="text-2xl font-bold mt-2 text-purple-600">
-                                  {userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0}
-                                </Text>
-                              </Card>
-                            </Grid>
+                            <div className="grid grid-cols-4 gap-4 mt-4">
+                              <ShadcnCard>
+                                <CardContent>
+                                  <h3 className="text-lg font-medium text-foreground">Input Tokens</h3>
+                                  <p className="text-2xl font-bold mt-2 text-blue-600">
+                                    {(userSpendData.metadata?.total_prompt_tokens || 0).toLocaleString()}
+                                  </p>
+                                </CardContent>
+                              </ShadcnCard>
+                              <ShadcnCard>
+                                <CardContent>
+                                  <h3 className="text-lg font-medium text-foreground">Output Tokens</h3>
+                                  <p className="text-2xl font-bold mt-2 text-cyan-600">
+                                    {userSpendData.metadata?.total_completion_tokens?.toLocaleString() || 0}
+                                  </p>
+                                </CardContent>
+                              </ShadcnCard>
+                              <ShadcnCard>
+                                <CardContent>
+                                  <h3 className="text-lg font-medium text-foreground">Cache Read Tokens</h3>
+                                  <p className="text-2xl font-bold mt-2 text-green-600">
+                                    {userSpendData.metadata?.total_cache_read_input_tokens?.toLocaleString() || 0}
+                                  </p>
+                                </CardContent>
+                              </ShadcnCard>
+                              <ShadcnCard>
+                                <CardContent>
+                                  <h3 className="text-lg font-medium text-foreground">Cache Write Tokens</h3>
+                                  <p className="text-2xl font-bold mt-2 text-purple-600">
+                                    {userSpendData.metadata?.total_cache_creation_input_tokens?.toLocaleString() || 0}
+                                  </p>
+                                </CardContent>
+                              </ShadcnCard>
+                            </div>
                           )}
-                        </Card>
-                      </Col>
+                        </CardContent>
+                      </ShadcnCard>
+                    </div>
 
-                      {/* Daily Spend Chart */}
-                      <Col numColSpan={2}>
-                        <ShadcnCard>
+                    {/* Daily Spend Chart */}
+                    <div className="col-span-2">
+                      <ShadcnCard>
+                        <CardHeader>
+                          <CardTitle className="text-base font-semibold">Daily Spend</CardTitle>
+                        </CardHeader>
+                        <CardContent>
+                          {loading ? (
+                            <ChartLoader isDateChanging={isDateChanging} />
+                          ) : (
+                            <BarChart
+                              data={sortedDailyResults}
+                              index="date"
+                              categories={["metrics.spend"]}
+                              colors={["cyan"]}
+                              valueFormatter={valueFormatterSpend}
+                              yAxisWidth={100}
+                              showLegend={false}
+                              customTooltip={({ payload, active }) => {
+                                if (!active || !payload?.[0]) return null;
+                                const data = payload[0].payload;
+                                return (
+                                  <div className="bg-white p-4 shadow-lg rounded-lg border">
+                                    <p className="font-bold">{data.date}</p>
+                                    <p className="text-cyan-500">
+                                      Spend: ${formatNumberWithCommas(data.metrics.spend, 2)}
+                                    </p>
+                                    <p className="text-gray-600">Requests: {data.metrics.api_requests}</p>
+                                    <p className="text-gray-600">Successful: {data.metrics.successful_requests}</p>
+                                    <p className="text-gray-600">Failed: {data.metrics.failed_requests}</p>
+                                    <p className="text-gray-600">Tokens: {data.metrics.total_tokens}</p>
+                                  </div>
+                                );
+                              }}
+                            />
+                          )}
+                        </CardContent>
+                      </ShadcnCard>
+                    </div>
+                    {/* Gateway Requests by Endpoint (SGR) */}
+                    {gatewayActivity && gatewayActivity.by_route.length > 0 && (
+                      <div className="col-span-2">
+                        <ShadcnCard data-testid="gateway-requests-by-endpoint">
                           <CardHeader>
-                            <CardTitle className="text-base font-semibold">Daily Spend</CardTitle>
+                            <CardTitle className="text-base font-semibold">
+                              Gateway Requests by Endpoint
+                              <Tooltip>
+                                <TooltipTrigger
+                                  render={<Info className="ml-2 inline size-4 text-gray-400 hover:text-gray-600" />}
+                                />
+                                <TooltipContent>
+                                  Counted by the gateway middleware as each request is answered. Covers LLM, MCP and A2A
+                                  endpoints across the whole deployment.
+                                </TooltipContent>
+                              </Tooltip>
+                            </CardTitle>
                           </CardHeader>
                           <CardContent>
-                            {loading ? (
-                              <ChartLoader isDateChanging={isDateChanging} />
-                            ) : (
-                              <BarChart
-                                data={sortedDailyResults}
-                                index="date"
-                                categories={["metrics.spend"]}
-                                colors={["cyan"]}
-                                valueFormatter={valueFormatterSpend}
-                                yAxisWidth={100}
-                                showLegend={false}
-                                customTooltip={({ payload, active }) => {
-                                  if (!active || !payload?.[0]) return null;
-                                  const data = payload[0].payload;
-                                  return (
-                                    <div className="bg-white p-4 shadow-lg rounded-lg border">
-                                      <p className="font-bold">{data.date}</p>
-                                      <p className="text-cyan-500">
-                                        Spend: ${formatNumberWithCommas(data.metrics.spend, 2)}
-                                      </p>
-                                      <p className="text-gray-600">Requests: {data.metrics.api_requests}</p>
-                                      <p className="text-gray-600">Successful: {data.metrics.successful_requests}</p>
-                                      <p className="text-gray-600">Failed: {data.metrics.failed_requests}</p>
-                                      <p className="text-gray-600">Tokens: {data.metrics.total_tokens}</p>
-                                    </div>
-                                  );
-                                }}
-                              />
-                            )}
+                            <BarChart
+                              data={gatewayRequestsByRoute}
+                              index="route"
+                              categories={["successful_requests", "failed_requests"]}
+                              colors={["green", "red"]}
+                              stack={true}
+                              yAxisWidth={100}
+                              valueFormatter={(value: number) => value.toLocaleString()}
+                            />
                           </CardContent>
                         </ShadcnCard>
-                      </Col>
-                      {/* Gateway Requests by Endpoint (SGR) */}
-                      {gatewayActivity && gatewayActivity.by_route.length > 0 && (
-                        <Col numColSpan={2}>
-                          <ShadcnCard data-testid="gateway-requests-by-endpoint">
-                            <CardHeader>
-                              <CardTitle className="text-base font-semibold">
-                                Gateway Requests by Endpoint
-                                <Tooltip title="Counted by the gateway middleware as each request is answered. Covers LLM, MCP and A2A endpoints across the whole deployment.">
-                                  <InfoCircleOutlined className="ml-2 text-gray-400 hover:text-gray-600" />
-                                </Tooltip>
-                              </CardTitle>
-                            </CardHeader>
-                            <CardContent>
-                              <BarChart
-                                data={gatewayRequestsByRoute}
-                                index="route"
-                                categories={["successful_requests", "failed_requests"]}
-                                colors={["green", "red"]}
-                                stack={true}
-                                yAxisWidth={100}
-                                valueFormatter={(value: number) => value.toLocaleString()}
-                              />
-                            </CardContent>
-                          </ShadcnCard>
-                        </Col>
-                      )}
-                      {/* Top API Keys */}
-                      <Col numColSpan={1}>
-                        <Card className="h-full">
-                          <Title>Top Virtual Keys</Title>
+                      </div>
+                    )}
+                    {/* Top API Keys */}
+                    <div>
+                      <ShadcnCard className="h-full">
+                        <CardContent>
+                          <h3 className="text-lg font-medium text-foreground">Top Virtual Keys</h3>
                           <TopKeyView
                             topKeys={topKeys}
                             teams={null}
                             topKeysLimit={topKeysLimit}
                             setTopKeysLimit={setTopKeysLimit}
                           />
-                        </Card>
-                      </Col>
+                        </CardContent>
+                      </ShadcnCard>
+                    </div>
 
-                      {/* Top Models */}
-                      <Col numColSpan={1}>
-                        <Card className="h-full">
-                          <Title>{modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models"}</Title>
+                    {/* Top Models */}
+                    <div>
+                      <ShadcnCard className="h-full">
+                        <CardContent>
+                          <h3 className="text-lg font-medium text-foreground">
+                            {modelViewType === "groups" ? "Top Public Model Names" : "Top Litellm Models"}
+                          </h3>
                           <div className="flex justify-between items-center mb-4">
-                            <Segmented
-                              options={[
-                                { label: "5", value: 5 },
-                                { label: "10", value: 10 },
-                                { label: "25", value: 25 },
-                                { label: "50", value: 50 },
-                              ]}
-                              value={topModelsLimit}
-                              onChange={(value) => setTopModelsLimit(value as number)}
-                            />
+                            <Tabs
+                              value={String(topModelsLimit)}
+                              onValueChange={(value: string) => setTopModelsLimit(Number(value))}
+                            >
+                              <TabsList>
+                                {TOP_MODEL_LIMITS.map((limit) => (
+                                  <TabsTrigger key={limit} value={String(limit)} className="flex-none px-3">
+                                    {limit}
+                                  </TabsTrigger>
+                                ))}
+                              </TabsList>
+                            </Tabs>
                             <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
                           </div>
                           {loading ? (
@@ -915,40 +899,40 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
                               })()}
                             </div>
                           )}
-                        </Card>
-                      </Col>
-
-                      {/* Spend by Provider */}
-                      <Col numColSpan={2}>
-                        <SpendByProvider
-                          loading={loading}
-                          isDateChanging={isDateChanging}
-                          providerSpend={providerSpend}
-                        />
-                      </Col>
-
-                      {/* Usage Metrics */}
-                    </Grid>
-                  </TabPanel>
-
-                  {/* Activity Panel */}
-                  <TabPanel>
-                    <div className="flex justify-end mt-2 mb-4">
-                      <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
+                        </CardContent>
+                      </ShadcnCard>
                     </div>
-                    <ActivityMetrics modelMetrics={modelMetrics} />
-                  </TabPanel>
-                  <TabPanel>
-                    <ActivityMetrics modelMetrics={keyMetrics} />
-                  </TabPanel>
-                  <TabPanel>
-                    <ActivityMetrics modelMetrics={mcpServerMetrics} />
-                  </TabPanel>
-                  <TabPanel>
-                    <EndpointUsage userSpendData={userSpendData} />
-                  </TabPanel>
-                </TabPanels>
-              </TabGroup>
+
+                    {/* Spend by Provider */}
+                    <div className="col-span-2">
+                      <SpendByProvider
+                        loading={loading}
+                        isDateChanging={isDateChanging}
+                        providerSpend={providerSpend}
+                      />
+                    </div>
+
+                    {/* Usage Metrics */}
+                  </div>
+                </TabsContent>
+
+                {/* Activity Panel */}
+                <TabsContent value="models" keepMounted>
+                  <div className="flex justify-end mt-2 mb-4">
+                    <ModelViewToggle value={modelViewType} onChange={setModelViewType} />
+                  </div>
+                  <ActivityMetrics modelMetrics={modelMetrics} />
+                </TabsContent>
+                <TabsContent value="keys" keepMounted>
+                  <ActivityMetrics modelMetrics={keyMetrics} />
+                </TabsContent>
+                <TabsContent value="mcp" keepMounted>
+                  <ActivityMetrics modelMetrics={mcpServerMetrics} />
+                </TabsContent>
+                <TabsContent value="endpoints" keepMounted>
+                  <EndpointUsage userSpendData={userSpendData} />
+                </TabsContent>
+              </Tabs>
             </>
           )}
           {/* Organization Usage Panel */}
@@ -1009,21 +993,24 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
           {usageView === "tag" && (
             <>
               {showCredentialBanner && (
-                <Alert
-                  banner
-                  type="info"
-                  message="Reusable credentials are automatically tracked as tags"
-                  description={
-                    <Typography.Text>
-                      When a reusable credential is used, it will appear as a tag prefixed with{" "}
-                      <Typography.Text code>Credential: </Typography.Text>
-                      in this view.
-                    </Typography.Text>
-                  }
-                  closable
-                  onClose={() => setShowCredentialBanner(false)}
-                  className="mb-5"
-                />
+                <Alert variant="info" className="mb-5">
+                  <AlertTitle>Reusable credentials are automatically tracked as tags</AlertTitle>
+                  <AlertDescription className="text-inherit">
+                    When a reusable credential is used, it will appear as a tag prefixed with{" "}
+                    <code className="rounded bg-black/5 px-1 py-0.5 font-mono text-xs">Credential: </code>
+                    in this view.
+                  </AlertDescription>
+                  <AlertAction>
+                    <Button
+                      variant="ghost"
+                      size="icon-xs"
+                      aria-label="Close"
+                      onClick={() => setShowCredentialBanner(false)}
+                    >
+                      <X />
+                    </Button>
+                  </AlertAction>
+                </Alert>
               )}
               <EntityUsage
                 accessToken={accessToken}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
index 54c1d5ab7cc..3fa58fc9da8 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsageViewSelect/UsageViewSelect.tsx
@@ -1,16 +1,7 @@
-import {
-  BankOutlined,
-  BarChartOutlined,
-  GlobalOutlined,
-  LineChartOutlined,
-  RobotOutlined,
-  ShoppingCartOutlined,
-  TagsOutlined,
-  TeamOutlined,
-  UserOutlined,
-} from "@ant-design/icons";
-import { Badge, Select } from "antd";
+import { BarChart3, Bot, Building2, Globe, LineChart, ShoppingCart, Tags, User, Users } from "lucide-react";
 import React from "react";
+import { Badge } from "@/components/ui/badge";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { hasCapability, type Capability } from "@/utils/capabilities";
 import { all_admin_roles } from "@/utils/roles";
 export type UsageOption =
@@ -54,61 +45,61 @@ const OPTIONS: OptionConfig[] = [
     description: "View usage across all resources",
     descriptionForAdmin: "View usage across all resources",
     descriptionForNonAdmin: "View your usage",
-    icon: <GlobalOutlined style={{ fontSize: "16px" }} />,
+    icon: <Globe className="size-4" />,
   },
   {
     value: "my-usage",
     label: "Your Usage",
     description: "View your own usage",
-    icon: <UserOutlined style={{ fontSize: "16px" }} />,
+    icon: <User className="size-4" />,
     adminOnly: true,
   },
   {
     value: "organization",
     label: "Organization Usage",
     description: "View usage across all organizations",
-    icon: <BankOutlined style={{ fontSize: "16px" }} />,
+    icon: <Building2 className="size-4" />,
     capability: "viewOrganizationUsage",
   },
   {
     value: "team",
     label: "Team Usage",
     description: "View usage by team",
-    icon: <TeamOutlined style={{ fontSize: "16px" }} />,
+    icon: <Users className="size-4" />,
   },
   {
     value: "customer",
     label: "Customer Usage",
     description: "View usage by customer accounts",
-    icon: <ShoppingCartOutlined style={{ fontSize: "16px" }} />,
+    icon: <ShoppingCart className="size-4" />,
     adminOnly: true,
   },
   {
     value: "tag",
     label: "Tag Usage",
     description: "View usage grouped by tags",
-    icon: <TagsOutlined style={{ fontSize: "16px" }} />,
+    icon: <Tags className="size-4" />,
     adminOnly: true,
   },
   {
     value: "agent",
     label: "Agent Usage (A2A)",
     description: "View usage by AI agents",
-    icon: <RobotOutlined style={{ fontSize: "16px" }} />,
+    icon: <Bot className="size-4" />,
     capability: "viewAgentUsage",
   },
   {
     value: "user",
     label: "User Usage",
     description: "View usage by individual users",
-    icon: <UserOutlined style={{ fontSize: "16px" }} />,
+    icon: <User className="size-4" />,
     adminOnly: true,
   },
   {
     value: "user-agent-activity",
     label: "User Agent Activity",
     description: "View detailed user agent activity logs",
-    icon: <LineChartOutlined style={{ fontSize: "16px" }} />,
+    icon: <LineChart className="size-4" />,
     adminOnly: true,
   },
 ];
@@ -153,12 +144,13 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
     });
   };
   const filteredOptions = getFilteredOptions();
+  const selectedOption = filteredOptions.find((option) => option.value === value);
   return (
     <div className="w-full" data-id={dataId}>
       <div className="flex flex-wrap items-center justify-start gap-4">
         <div className="flex items-stretch gap-2 min-w-0">
           <div className="shrink-0 flex items-center">
-            <BarChartOutlined style={{ fontSize: "32px" }} />
+            <BarChart3 className="size-8" />
           </div>
           <div className="flex-1 min-w-0">
             <h3 className="text-sm font-semibold text-gray-900 mb-0.5 leading-tight">{title}</h3>
@@ -168,42 +160,35 @@ export const UsageViewSelect: React.FC<UsageViewSelectProps> = ({
         <div className="shrink-0">
           <Select
             value={value}
-            onChange={onChange}
-            className="w-54 sm:w-64 md:w-72"
-            size="large"
-            options={filteredOptions.map((opt) => ({
-              value: opt.value,
-              label: opt.label,
-            }))}
-            optionRender={(option) => {
-              const opt = filteredOptions.find((o) => o.value === option.value);
-              if (!opt) return option.label;
-              return (
-                <div className="flex items-center gap-2 py-1">
-                  <div className="shrink-0 mt-0.5">{opt.icon}</div>
-                  <div className="flex-1 min-w-0">
-                    <div className="text-sm font-medium text-gray-900">{opt.label}</div>
-                    <div className="text-xs text-gray-600 mt-0.5">{opt.description}</div>
-                  </div>
-                  {opt.badgeText && (
-                    <div className="items-center">
-                      <Badge color="blue" count={opt.badgeText} />
-                    </div>
-                  )}
-                </div>
-              );
+            onValueChange={(next: UsageOption | null) => {
+              if (next) onChange(next);
             }}
-            labelRender={(props) => {
-              const opt = filteredOptions.find((o) => o.value === props.value);
-              if (!opt) return props.label;
-              return (
-                <div className="flex items-center gap-2">
-                  <div>{opt.icon}</div>
-                  <span className="text-sm">{opt.label}</span>
-                </div>
-              );
-            }}
-          />
+          >
+            <SelectTrigger className="w-54 sm:w-64 md:w-72">
+              <SelectValue>
+                {selectedOption && (
+                  <span className="flex items-center gap-2">
+                    {selectedOption.icon}
+                    <span className="text-sm">{selectedOption.label}</span>
+                  </span>
+                )}
+              </SelectValue>
+            </SelectTrigger>
+            <SelectContent>
+              {filteredOptions.map((option) => (
+                <SelectItem key={option.value} value={option.value}>
+                  <span className="flex items-center gap-2 py-1">
+                    <span className="shrink-0 mt-0.5">{option.icon}</span>
+                    <span className="flex-1 min-w-0">
+                      <span className="block text-sm font-medium text-gray-900">{option.label}</span>
+                      <span className="block text-xs text-gray-600 mt-0.5">{option.description}</span>
+                    </span>
+                    {option.badgeText && <Badge>{option.badgeText}</Badge>}
+                  </span>
+                </SelectItem>
+              ))}
+            </SelectContent>
+          </Select>
         </div>
       </div>
     </div>
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx
index 4c934bbf1d3..6be06914493 100644
--- a/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx
+++ b/ui/litellm-dashboard/src/components/EntityUsageExport/EntityUsageExportModal.tsx
@@ -1,7 +1,10 @@
 import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 import { createTeamAliasMap } from "@/utils/teamUtils";
-import { Button, Modal, Skeleton } from "antd";
+import { Loader2 } from "lucide-react";
 import React, { useMemo, useState } from "react";
+import { Button } from "@/components/ui/button";
+import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Skeleton } from "@/components/ui/skeleton";
 import NotificationsManager from "../molecules/notifications_manager";
 import ExportFormatSelector from "./ExportFormatSelector";
 import ExportSummary from "./ExportSummary";
@@ -49,45 +52,51 @@ const EntityUsageExportModal: React.FC<EntityUsageExportModalProps> = ({
   };
 
   return (
-    <Modal
-      title={<span className="text-base font-semibold">{modalTitle}</span>}
+    <Dialog
       open={isOpen}
-      onCancel={onClose}
-      footer={null}
-      width={480}
+      onOpenChange={(open) => {
+        if (!open) onClose();
+      }}
     >
-      <div className="space-y-5 py-2">
-        {isLoadingTeams ? (
-          <Skeleton active />
-        ) : (
-          <>
-            <ExportSummary dateRange={dateRange} selectedFilters={selectedFilters} />
-            <ExportTypeSelector value={exportScope} onChange={setExportScope} entityType={entityType} />
-            <ExportFormatSelector value={exportFormat} onChange={setExportFormat} />
-          </>
-        )}
-        {isLoadingTeams ? (
+      <DialogContent className="sm:max-w-[480px]">
+        <DialogHeader>
+          <DialogTitle className="text-base font-semibold">{modalTitle}</DialogTitle>
+        </DialogHeader>
+        <div className="space-y-5 py-2">
+          {isLoadingTeams ? (
+            <div className="space-y-3">
+              <Skeleton className="h-4 w-3/4" />
+              <Skeleton className="h-4 w-full" />
+              <Skeleton className="h-4 w-2/3" />
+            </div>
+          ) : (
+            <>
+              <ExportSummary dateRange={dateRange} selectedFilters={selectedFilters} />
+              <ExportTypeSelector value={exportScope} onChange={setExportScope} entityType={entityType} />
+              <ExportFormatSelector value={exportFormat} onChange={setExportFormat} />
+            </>
+          )}
           <div className="flex items-center justify-end gap-2 pt-4 border-t">
-            <Skeleton.Button active />
-            <Skeleton.Button active />
+            {isLoadingTeams ? (
+              <>
+                <Skeleton className="h-9 w-20" />
+                <Skeleton className="h-9 w-28" />
+              </>
+            ) : (
+              <>
+                <Button variant="outline" onClick={onClose} disabled={isExporting}>
+                  Cancel
+                </Button>
+                <Button onClick={() => handleExport()} disabled={isExporting}>
+                  {isExporting && <Loader2 className="animate-spin" />}
+                  {isExporting ? "Exporting..." : `Export ${exportFormat.toUpperCase()}`}
+                </Button>
+              </>
+            )}
           </div>
-        ) : (
-          <div className="flex items-center justify-end gap-2 pt-4 border-t">
-            <Button variant="outlined" onClick={onClose} disabled={isExporting}>
-              Cancel
-            </Button>
-            <Button
-              onClick={() => handleExport()}
-              loading={isExporting || isLoadingTeams}
-              disabled={isExporting || isLoadingTeams}
-              type="primary"
-            >
-              {isExporting ? "Exporting..." : `Export ${exportFormat.toUpperCase()}`}
-            </Button>
-          </div>
-        )}
-      </div>
-    </Modal>
+        </div>
+      </DialogContent>
+    </Dialog>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.tsx
index 2adc4a0eca3..8048585b06c 100644
--- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.tsx
+++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportFormatSelector.tsx
@@ -1,5 +1,5 @@
 import React from "react";
-import { Select } from "antd";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import type { ExportFormat } from "./types";
 
 interface ExportFormatSelectorProps {
@@ -7,25 +7,27 @@ interface ExportFormatSelectorProps {
   onChange: (value: ExportFormat) => void;
 }
 
+const FORMAT_LABELS: Record<ExportFormat, string> = {
+  csv: "CSV (Excel, Google Sheets)",
+  json: "JSON (includes metadata)",
+};
+
 const ExportFormatSelector: React.FC<ExportFormatSelectorProps> = ({ value, onChange }) => {
   return (
     <div>
       <label className="text-sm font-medium text-gray-700 block mb-2">Format</label>
-      <Select
-        value={value}
-        onChange={onChange}
-        className="w-full"
-        options={[
-          {
-            value: "csv",
-            label: "CSV (Excel, Google Sheets)",
-          },
-          {
-            value: "json",
-            label: "JSON (includes metadata)",
-          },
-        ]}
-      />
+      <Select value={value} onValueChange={(next: ExportFormat | null) => next && onChange(next)}>
+        <SelectTrigger className="w-full">
+          <SelectValue>{FORMAT_LABELS[value]}</SelectValue>
+        </SelectTrigger>
+        <SelectContent>
+          {(Object.keys(FORMAT_LABELS) as ExportFormat[]).map((format) => (
+            <SelectItem key={format} value={format}>
+              {FORMAT_LABELS[format]}
+            </SelectItem>
+          ))}
+        </SelectContent>
+      </Select>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx
index 6bc85948608..498041fd77e 100644
--- a/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx
+++ b/ui/litellm-dashboard/src/components/EntityUsageExport/ExportTypeSelector.tsx
@@ -1,5 +1,5 @@
 import React from "react";
-import { Radio } from "antd";
+import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
 import type { ExportScope, EntityType } from "./types";
 
 interface ExportTypeSelectorProps {
@@ -9,36 +9,41 @@ interface ExportTypeSelectorProps {
 }
 
 const ExportTypeSelector: React.FC<ExportTypeSelectorProps> = ({ value, onChange, entityType }) => {
+  const scopes: { value: ExportScope; title: string; description: string }[] = [
+    {
+      value: "daily",
+      title: `Day-by-day breakdown by ${entityType}`,
+      description: `Daily metrics for each ${entityType}`,
+    },
+    {
+      value: "daily_with_keys",
+      title: `Day-by-day breakdown by ${entityType} and key`,
+      description: `Daily metrics for each ${entityType}, split by API key`,
+    },
+    {
+      value: "daily_with_models",
+      title: `Day-by-day by ${entityType} and model`,
+      description: "Daily metrics split by model",
+    },
+  ];
+
   return (
     <div>
       <label className="text-sm font-medium text-gray-700 block mb-2">Export type</label>
-      <Radio.Group value={value} onChange={(e) => onChange(e.target.value)} className="w-full">
-        <div className="space-y-2">
-          <label className="flex items-start p-3 border border-gray-200 rounded-lg hover:bg-gray-50 cursor-pointer transition-colors">
-            <Radio value="daily" className="mt-0.5" />
+      <RadioGroup value={value} onValueChange={(next) => onChange(next as ExportScope)} className="gap-2">
+        {scopes.map((scope) => (
+          <label
+            key={scope.value}
+            className="flex items-start p-3 border border-gray-200 rounded-lg hover:bg-gray-50 cursor-pointer transition-colors"
+          >
+            <RadioGroupItem value={scope.value} className="mt-0.5" />
             <div className="ml-3 flex-1">
-              <div className="font-medium text-sm">Day-by-day breakdown by {entityType}</div>
-              <div className="text-xs text-gray-500 mt-0.5">Daily metrics for each {entityType}</div>
+              <div className="font-medium text-sm">{scope.title}</div>
+              <div className="text-xs text-gray-500 mt-0.5">{scope.description}</div>
             </div>
           </label>
-
-          <label className="flex items-start p-3 border border-gray-200 rounded-lg hover:bg-gray-50 cursor-pointer transition-colors">
-            <Radio value="daily_with_keys" className="mt-0.5" />
-            <div className="ml-3 flex-1">
-              <div className="font-medium text-sm">Day-by-day breakdown by {entityType} and key</div>
-              <div className="text-xs text-gray-500 mt-0.5">Daily metrics for each {entityType}, split by API key</div>
-            </div>
-          </label>
-
-          <label className="flex items-start p-3 border border-gray-200 rounded-lg hover:bg-gray-50 cursor-pointer transition-colors">
-            <Radio value="daily_with_models" className="mt-0.5" />
-            <div className="ml-3 flex-1">
-              <div className="font-medium text-sm">Day-by-day by {entityType} and model</div>
-              <div className="text-xs text-gray-500 mt-0.5">Daily metrics split by model</div>
-            </div>
-          </label>
-        </div>
-      </Radio.Group>
+        ))}
+      </RadioGroup>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx
index ee7857a1544..78781df5948 100644
--- a/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx
+++ b/ui/litellm-dashboard/src/components/EntityUsageExport/UsageExportHeader.tsx
@@ -1,7 +1,20 @@
 import type { DateRangePickerValue } from "@tremor/react";
-import { Button, Text } from "@tremor/react";
-import { Select } from "antd";
+import { Download } from "lucide-react";
 import React, { useState } from "react";
+import { Button } from "@/components/ui/button";
+import {
+  Combobox,
+  ComboboxChip,
+  ComboboxChips,
+  ComboboxChipsInput,
+  ComboboxClear,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxInput,
+  ComboboxItem,
+  ComboboxList,
+  ComboboxValue,
+} from "@/components/ui/combobox";
 import EntityUsageExportModal from "./EntityUsageExportModal";
 import type { EntitySpendData, EntityType } from "./types";
 import type { Team } from "@/components/key_team_helpers/key_list";
@@ -40,13 +53,22 @@ const UsageExportHeader: React.FC<UsageExportHeaderProps> = ({
 }) => {
   const [isExportModalOpen, setIsExportModalOpen] = useState(false);
 
-  // Determine grid layout based on what's visible
-  const getGridCols = () => {
-    const hasFilters = showFilters && filterOptions.length > 0;
+  const hasFilters = showFilters && filterOptions.length > 0;
+  const optionValues = filterOptions.map((option) => option.value);
+  const labelOf = (value: string) => filterOptions.find((option) => option.value === value)?.label ?? value;
 
-    if (hasFilters) return "grid-cols-[1fr_auto]";
-    return "grid-cols-[auto]";
-  };
+  const filterList = (
+    <ComboboxContent>
+      <ComboboxEmpty>No options found</ComboboxEmpty>
+      <ComboboxList>
+        {(value: string) => (
+          <ComboboxItem key={value} value={value}>
+            {labelOf(value)}
+          </ComboboxItem>
+        )}
+      </ComboboxList>
+    </ComboboxContent>
+  );
 
   return (
     <>
@@ -56,42 +78,58 @@ const UsageExportHeader: React.FC<UsageExportHeaderProps> = ({
          * align to the same baseline regardless of label heights. This removes
          * vertical drift when the right column has a label above the input.
          */}
-        <div className={`grid ${getGridCols()} items-end gap-4`}>
-          {showFilters && filterOptions.length > 0 && (
+        <div className={`grid ${hasFilters ? "grid-cols-[1fr_auto]" : "grid-cols-[auto]"} items-end gap-4`}>
+          {hasFilters && (
             <div>
-              {filterLabel && <Text className="mb-2">{filterLabel}</Text>}
-              <Select
-                mode={filterMode === "single" ? undefined : "multiple"}
-                style={{ width: "100%" }}
-                placeholder={filterPlaceholder}
-                value={filterMode === "single" ? selectedFilters[0] ?? undefined : selectedFilters}
-                onChange={(value: any) => {
-                  if (filterMode === "single") {
-                    onFiltersChange?.(value ? [value] : []);
-                  } else {
-                    onFiltersChange?.(value);
-                  }
-                }}
-                options={filterOptions}
-                allowClear
-              />
+              {filterLabel && <label className="text-sm font-medium text-gray-700 block mb-2">{filterLabel}</label>}
+              {filterMode === "single" ? (
+                <Combobox
+                  items={optionValues}
+                  value={selectedFilters[0] ?? null}
+                  onValueChange={(next: string | null) => onFiltersChange?.(next ? [next] : [])}
+                  itemToStringLabel={labelOf}
+                >
+                  <ComboboxInput
+                    className="w-full"
+                    placeholder={filterPlaceholder}
+                    aria-label={filterPlaceholder}
+                    showClear={selectedFilters.length > 0}
+                  />
+                  {filterList}
+                </Combobox>
+              ) : (
+                <Combobox
+                  multiple
+                  items={optionValues}
+                  value={selectedFilters}
+                  onValueChange={(next: string[]) => onFiltersChange?.(next)}
+                >
+                  <ComboboxChips className="w-full">
+                    <ComboboxValue>
+                      {(selected: string[]) =>
+                        selected.map((value) => (
+                          <ComboboxChip key={value} aria-label={labelOf(value)}>
+                            {labelOf(value)}
+                          </ComboboxChip>
+                        ))
+                      }
+                    </ComboboxValue>
+                    <ComboboxChipsInput
+                      className="border-0 bg-transparent"
+                      placeholder={filterPlaceholder}
+                      aria-label={filterPlaceholder}
+                    />
+                    {selectedFilters.length > 0 && <ComboboxClear aria-label={`Clear ${filterLabel ?? "filters"}`} />}
+                  </ComboboxChips>
+                  {filterList}
+                </Combobox>
+              )}
             </div>
           )}
 
           <div className="justify-self-end">
-            <Button
-              onClick={() => setIsExportModalOpen(true)}
-              icon={() => (
-                <svg className="w-4 h-4" fill="none" stroke="currentColor" viewBox="0 0 24 24">
-                  <path
-                    strokeLinecap="round"
-                    strokeLinejoin="round"
-                    strokeWidth={2}
-                    d="M4 16v1a3 3 0 003 3h10a3 3 0 003-3v-1m-4-4l-4 4m0 0l-4-4m4 4V4"
-                  />
-                </svg>
-              )}
-            >
+            <Button onClick={() => setIsExportModalOpen(true)}>
+              <Download />
               Export Data
             </Button>
           </div>
diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx
index b8199ee29b6..288d07378b7 100644
--- a/ui/litellm-dashboard/src/components/activity_metrics.tsx
+++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx
@@ -1,9 +1,10 @@
 import { AreaChart, BarChart, CustomLegend, CustomTooltip } from "@/components/shared/charts";
 import { formatNumberWithCommas } from "@/utils/dataUtils";
 import { resolveTeamAliasFromTeamID } from "@/utils/teamUtils";
-import { Card, Grid, Text, Title } from "@tremor/react";
-import { Collapse } from "antd";
-import React from "react";
+import { Card, CardContent } from "@/components/ui/card";
+import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
+import { ChevronDown } from "lucide-react";
+import React, { useState } from "react";
 import { Team } from "./key_team_helpers/key_list";
 import KeyModelUsageView from "./UsagePage/components/KeyModelUsageView";
 import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData, TopModelData } from "./UsagePage/types";
@@ -26,50 +27,65 @@ const ModelSection = ({
   return (
     <div className="space-y-2">
       {/* Summary Cards */}
-      <Grid numItems={4} className="gap-4">
+      <div className="grid grid-cols-4 gap-4">
         <Card>
-          <Text>Total Requests</Text>
-          <Title>{metrics.total_requests.toLocaleString()}</Title>
+          <CardContent>
+            <p className="text-sm text-muted-foreground">Total Requests</p>
+            <h3 className="text-lg font-medium text-foreground">{metrics.total_requests.toLocaleString()}</h3>
+          </CardContent>
         </Card>
         <Card>
-          <Text>Total Successful Requests</Text>
-          <Title>{metrics.total_successful_requests.toLocaleString()}</Title>
+          <CardContent>
+            <p className="text-sm text-muted-foreground">Total Successful Requests</p>
+            <h3 className="text-lg font-medium text-foreground">
+              {metrics.total_successful_requests.toLocaleString()}
+            </h3>
+          </CardContent>
         </Card>
         <Card>
-          <Text>Total Tokens</Text>
-          <Title>{metrics.total_tokens.toLocaleString()}</Title>
-          <Text>{Math.round(metrics.total_tokens / metrics.total_successful_requests)} avg per successful request</Text>
+          <CardContent>
+            <p className="text-sm text-muted-foreground">Total Tokens</p>
+            <h3 className="text-lg font-medium text-foreground">{metrics.total_tokens.toLocaleString()}</h3>
+            <p className="text-sm text-muted-foreground">
+              {Math.round(metrics.total_tokens / metrics.total_successful_requests)} avg per successful request
+            </p>
+          </CardContent>
         </Card>
         <Card>
-          <Text>Total Spend</Text>
-          <Title>${formatNumberWithCommas(metrics.total_spend, 2)}</Title>
-          <Text>
-            ${formatNumberWithCommas(metrics.total_spend / metrics.total_successful_requests, 3)} per successful request
-          </Text>
+          <CardContent>
+            <p className="text-sm text-muted-foreground">Total Spend</p>
+            <h3 className="text-lg font-medium text-foreground">${formatNumberWithCommas(metrics.total_spend, 2)}</h3>
+            <p className="text-sm text-muted-foreground">
+              ${formatNumberWithCommas(metrics.total_spend / metrics.total_successful_requests, 3)} per successful
+              request
+            </p>
+          </CardContent>
         </Card>
-      </Grid>
+      </div>
 
       {metrics.top_api_keys && metrics.top_api_keys.length > 0 && (
         <Card className="mt-4">
-          <Title>Top Virtual Keys by Spend</Title>
-          <div className="mt-3">
-            <div className="grid grid-cols-1 gap-2">
-              {metrics.top_api_keys.map((keyData, index) => (
-                <div key={keyData.api_key} className="flex justify-between items-center p-3 bg-gray-50 rounded-lg">
-                  <div>
-                    <Text className="font-medium">{keyData.key_alias || `${keyData.api_key.substring(0, 10)}...`}</Text>
-                    {keyData.team_id && <Text className="text-xs text-gray-500">Team: {keyData.team_id}</Text>}
+          <CardContent>
+            <h3 className="text-lg font-medium text-foreground">Top Virtual Keys by Spend</h3>
+            <div className="mt-3">
+              <div className="grid grid-cols-1 gap-2">
+                {metrics.top_api_keys.map((keyData) => (
+                  <div key={keyData.api_key} className="flex justify-between items-center p-3 bg-gray-50 rounded-lg">
+                    <div>
+                      <p className="font-medium">{keyData.key_alias || `${keyData.api_key.substring(0, 10)}...`}</p>
+                      {keyData.team_id && <p className="text-xs text-gray-500">Team: {keyData.team_id}</p>}
+                    </div>
+                    <div className="text-right">
+                      <p className="font-medium">${formatNumberWithCommas(keyData.spend, 2)}</p>
+                      <p className="text-xs text-gray-500">
+                        {keyData.requests.toLocaleString()} requests | {keyData.tokens.toLocaleString()} tokens
+                      </p>
+                    </div>
                   </div>
-                  <div className="text-right">
-                    <Text className="font-medium">${formatNumberWithCommas(keyData.spend, 2)}</Text>
-                    <Text className="text-xs text-gray-500">
-                      {keyData.requests.toLocaleString()} requests | {keyData.tokens.toLocaleString()} tokens
-                    </Text>
-                  </div>
-                </div>
-              ))}
+                ))}
+              </div>
             </div>
-          </div>
+          </CardContent>
         </Card>
       )}
 
@@ -77,110 +93,158 @@ const ModelSection = ({
 
       {/* Spend per day - Full width card */}
       <Card className="mt-4">
-        <div className="flex justify-between items-center">
-          <Title>Spend per day</Title>
-          <CustomLegend categories={["metrics.spend"]} colors={["green"]} />
-        </div>
-        <BarChart
-          className="mt-4"
-          data={metrics.daily_data}
-          index="date"
-          categories={["metrics.spend"]}
-          colors={["green"]}
-          valueFormatter={(value: number) => `$${formatNumberWithCommas(value, 2, true)}`}
-          yAxisWidth={72}
-        />
-      </Card>
-
-      {/* Charts */}
-      <Grid numItems={2} className="gap-4 mt-4">
-        <Card>
+        <CardContent>
           <div className="flex justify-between items-center">
-            <Title>Total Tokens</Title>
-            <CustomLegend
-              categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
-              colors={["blue", "cyan", "indigo"]}
-            />
-          </div>
-          <AreaChart
-            className="mt-4"
-            data={metrics.daily_data}
-            index="date"
-            categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
-            colors={["blue", "cyan", "indigo"]}
-            valueFormatter={valueFormatter}
-            customTooltip={CustomTooltip}
-            showLegend={false}
-          />
-        </Card>
-
-        <Card>
-          <div className="flex justify-between items-center">
-            <Title>Requests per day</Title>
-            <CustomLegend categories={["metrics.api_requests"]} colors={["blue"]} />
+            <h3 className="text-lg font-medium text-foreground">Spend per day</h3>
+            <CustomLegend categories={["metrics.spend"]} colors={["green"]} />
           </div>
           <BarChart
             className="mt-4"
             data={metrics.daily_data}
             index="date"
-            categories={["metrics.api_requests"]}
-            colors={["blue"]}
-            valueFormatter={valueFormatter}
-            customTooltip={CustomTooltip}
-            showLegend={false}
+            categories={["metrics.spend"]}
+            colors={["green"]}
+            valueFormatter={(value: number) => `$${formatNumberWithCommas(value, 2, true)}`}
+            yAxisWidth={72}
           />
-        </Card>
+        </CardContent>
+      </Card>
 
+      {/* Charts */}
+      <div className="grid grid-cols-2 gap-4 mt-4">
         <Card>
-          <div className="flex justify-between items-center">
-            <Title>Success vs Failed Requests</Title>
-            <CustomLegend
-              categories={["metrics.successful_requests", "metrics.failed_requests"]}
-              colors={["green", "red"]}
-            />
-          </div>
-          <AreaChart
-            className="mt-4"
-            data={metrics.daily_data}
-            index="date"
-            categories={["metrics.successful_requests", "metrics.failed_requests"]}
-            colors={["green", "red"]}
-            valueFormatter={valueFormatter}
-            customTooltip={CustomTooltip}
-            showLegend={false}
-          />
-        </Card>
-
-        {!hidePromptCachingMetrics && (
-          <Card>
+          <CardContent>
             <div className="flex justify-between items-center">
-              <Title>Prompt Caching Metrics</Title>
+              <h3 className="text-lg font-medium text-foreground">Total Tokens</h3>
               <CustomLegend
-                categories={["metrics.cache_read_input_tokens", "metrics.cache_creation_input_tokens"]}
-                colors={["cyan", "purple"]}
+                categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
+                colors={["blue", "cyan", "indigo"]}
               />
             </div>
-            <div className="mb-2">
-              <Text>Cache Read: {metrics.total_cache_read_input_tokens?.toLocaleString() || 0} tokens</Text>
-              <Text>Cache Creation: {metrics.total_cache_creation_input_tokens?.toLocaleString() || 0} tokens</Text>
-            </div>
             <AreaChart
               className="mt-4"
               data={metrics.daily_data}
               index="date"
-              categories={["metrics.cache_read_input_tokens", "metrics.cache_creation_input_tokens"]}
-              colors={["cyan", "purple"]}
+              categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
+              colors={["blue", "cyan", "indigo"]}
               valueFormatter={valueFormatter}
               customTooltip={CustomTooltip}
               showLegend={false}
             />
+          </CardContent>
+        </Card>
+
+        <Card>
+          <CardContent>
+            <div className="flex justify-between items-center">
+              <h3 className="text-lg font-medium text-foreground">Requests per day</h3>
+              <CustomLegend categories={["metrics.api_requests"]} colors={["blue"]} />
+            </div>
+            <BarChart
+              className="mt-4"
+              data={metrics.daily_data}
+              index="date"
+              categories={["metrics.api_requests"]}
+              colors={["blue"]}
+              valueFormatter={valueFormatter}
+              customTooltip={CustomTooltip}
+              showLegend={false}
+            />
+          </CardContent>
+        </Card>
+
+        <Card>
+          <CardContent>
+            <div className="flex justify-between items-center">
+              <h3 className="text-lg font-medium text-foreground">Success vs Failed Requests</h3>
+              <CustomLegend
+                categories={["metrics.successful_requests", "metrics.failed_requests"]}
+                colors={["green", "red"]}
+              />
+            </div>
+            <AreaChart
+              className="mt-4"
+              data={metrics.daily_data}
+              index="date"
+              categories={["metrics.successful_requests", "metrics.failed_requests"]}
+              colors={["green", "red"]}
+              valueFormatter={valueFormatter}
+              customTooltip={CustomTooltip}
+              showLegend={false}
+            />
+          </CardContent>
+        </Card>
+
+        {!hidePromptCachingMetrics && (
+          <Card>
+            <CardContent>
+              <div className="flex justify-between items-center">
+                <h3 className="text-lg font-medium text-foreground">Prompt Caching Metrics</h3>
+                <CustomLegend
+                  categories={["metrics.cache_read_input_tokens", "metrics.cache_creation_input_tokens"]}
+                  colors={["cyan", "purple"]}
+                />
+              </div>
+              <div className="mb-2">
+                <p className="text-sm">
+                  Cache Read: {metrics.total_cache_read_input_tokens?.toLocaleString() || 0} tokens
+                </p>
+                <p className="text-sm">
+                  Cache Creation: {metrics.total_cache_creation_input_tokens?.toLocaleString() || 0} tokens
+                </p>
+              </div>
+              <AreaChart
+                className="mt-4"
+                data={metrics.daily_data}
+                index="date"
+                categories={["metrics.cache_read_input_tokens", "metrics.cache_creation_input_tokens"]}
+                colors={["cyan", "purple"]}
+                valueFormatter={valueFormatter}
+                customTooltip={CustomTooltip}
+                showLegend={false}
+              />
+            </CardContent>
           </Card>
         )}
-      </Grid>
+      </div>
     </div>
   );
 };
 
+// Mirrors the antd Collapse this replaced: a panel's subtree is not built until it is first
+// opened, and it stays mounted afterwards so its view-mode state survives being collapsed.
+const ModelCollapsible = ({
+  defaultOpen,
+  header,
+  children,
+}: {
+  defaultOpen: boolean;
+  header: React.ReactNode;
+  children: React.ReactNode;
+}) => {
+  const [open, setOpen] = useState(defaultOpen);
+  const [everOpened, setEverOpened] = useState(defaultOpen);
+
+  return (
+    <Collapsible
+      open={open}
+      onOpenChange={(next: boolean) => {
+        setOpen(next);
+        if (next) setEverOpened(true);
+      }}
+      className="border-b last:border-b-0"
+    >
+      <CollapsibleTrigger className="flex w-full items-center gap-2 px-4 py-3 text-left">
+        <ChevronDown className={`size-4 shrink-0 text-gray-400 transition-transform ${open ? "" : "-rotate-90"}`} />
+        {header}
+      </CollapsibleTrigger>
+      <CollapsibleContent keepMounted={everOpened} className="px-4 pb-4">
+        {children}
+      </CollapsibleContent>
+    </Collapsible>
+  );
+};
+
 export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics, hidePromptCachingMetrics = false }) => {
   const modelNames = Object.keys(modelMetrics).sort((a, b) => {
     if (a === "") return 1;
@@ -257,78 +321,97 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics,
     <div className="space-y-8">
       {/* Global Summary */}
       <div className="border rounded-lg p-4">
-        <Title>Overall Usage</Title>
-        <Grid numItems={4} className="gap-4 mb-4">
+        <h3 className="text-lg font-medium text-foreground">Overall Usage</h3>
+        <div className="grid grid-cols-4 gap-4 mb-4">
           <Card>
-            <Text>Total Requests</Text>
-            <Title>{totalMetrics.total_requests.toLocaleString()}</Title>
+            <CardContent>
+              <p className="text-sm text-muted-foreground">Total Requests</p>
+              <h3 className="text-lg font-medium text-foreground">{totalMetrics.total_requests.toLocaleString()}</h3>
+            </CardContent>
           </Card>
           <Card>
-            <Text>Total Successful Requests</Text>
-            <Title>{totalMetrics.total_successful_requests.toLocaleString()}</Title>
+            <CardContent>
+              <p className="text-sm text-muted-foreground">Total Successful Requests</p>
+              <h3 className="text-lg font-medium text-foreground">
+                {totalMetrics.total_successful_requests.toLocaleString()}
+              </h3>
+            </CardContent>
           </Card>
           <Card>
-            <Text>Total Tokens</Text>
-            <Title>{totalMetrics.total_tokens.toLocaleString()}</Title>
+            <CardContent>
+              <p className="text-sm text-muted-foreground">Total Tokens</p>
+              <h3 className="text-lg font-medium text-foreground">{totalMetrics.total_tokens.toLocaleString()}</h3>
+            </CardContent>
           </Card>
           <Card>
-            <Text>Total Spend</Text>
-            <Title>${formatNumberWithCommas(totalMetrics.total_spend, 2)}</Title>
+            <CardContent>
+              <p className="text-sm text-muted-foreground">Total Spend</p>
+              <h3 className="text-lg font-medium text-foreground">
+                ${formatNumberWithCommas(totalMetrics.total_spend, 2)}
+              </h3>
+            </CardContent>
           </Card>
-        </Grid>
+        </div>
 
-        <Grid numItems={2} className="gap-4">
+        <div className="grid grid-cols-2 gap-4">
           <Card>
-            <div className="flex justify-between items-center">
-              <Title>Total Tokens Over Time</Title>
-              <CustomLegend
+            <CardContent>
+              <div className="flex justify-between items-center">
+                <h3 className="text-lg font-medium text-foreground">Total Tokens Over Time</h3>
+                <CustomLegend
+                  categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
+                  colors={["blue", "cyan", "indigo"]}
+                />
+              </div>
+              <AreaChart
+                className="mt-4"
+                data={sortedDailyData}
+                index="date"
                 categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
                 colors={["blue", "cyan", "indigo"]}
+                valueFormatter={valueFormatter}
+                customTooltip={CustomTooltip}
+                showLegend={false}
+                yAxisWidth={80}
               />
-            </div>
-            <AreaChart
-              className="mt-4"
-              data={sortedDailyData}
-              index="date"
-              categories={["metrics.prompt_tokens", "metrics.completion_tokens", "metrics.total_tokens"]}
-              colors={["blue", "cyan", "indigo"]}
-              valueFormatter={valueFormatter}
-              customTooltip={CustomTooltip}
-              showLegend={false}
-              yAxisWidth={80}
-            />
+            </CardContent>
           </Card>
           <Card>
-            <div className="flex justify-between items-center">
-              <Title>Total Requests Over Time</Title>
-              <CustomLegend
+            <CardContent>
+              <div className="flex justify-between items-center">
+                <h3 className="text-lg font-medium text-foreground">Total Requests Over Time</h3>
+                <CustomLegend
+                  categories={["metrics.successful_requests", "metrics.failed_requests"]}
+                  colors={["emerald", "red"]}
+                />
+              </div>
+              <AreaChart
+                className="mt-4"
+                data={sortedDailyData}
+                index="date"
                 categories={["metrics.successful_requests", "metrics.failed_requests"]}
                 colors={["emerald", "red"]}
+                valueFormatter={valueFormatter}
+                customTooltip={CustomTooltip}
+                showLegend={false}
+                yAxisWidth={80}
               />
-            </div>
-            <AreaChart
-              className="mt-4"
-              data={sortedDailyData}
-              index="date"
-              categories={["metrics.successful_requests", "metrics.failed_requests"]}
-              colors={["emerald", "red"]}
-              valueFormatter={valueFormatter}
-              customTooltip={CustomTooltip}
-              showLegend={false}
-              yAxisWidth={80}
-            />
+            </CardContent>
           </Card>
-        </Grid>
+        </div>
       </div>
 
       {/* Individual Model Sections */}
-      <Collapse defaultActiveKey={modelNames[0]}>
+      <div className="rounded-lg border">
         {modelNames.map((modelName) => (
-          <Collapse.Panel
+          <ModelCollapsible
             key={modelName}
+            defaultOpen={modelName === modelNames[0]}
             header={
               <div className="flex justify-between items-center w-full">
-                <Title>{modelMetrics[modelName].label || "Unknown Item"}</Title>
+                <h3 className="text-lg font-medium text-foreground">
+                  {modelMetrics[modelName].label || "Unknown Item"}
+                </h3>
                 <div className="flex space-x-4 text-sm text-gray-500">
                   <span>${formatNumberWithCommas(modelMetrics[modelName].total_spend, 2)}</span>
                   <span>{modelMetrics[modelName].total_requests.toLocaleString()} requests</span>
@@ -341,9 +424,9 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics,
               metrics={modelMetrics[modelName]}
               hidePromptCachingMetrics={hidePromptCachingMetrics}
             />
-          </Collapse.Panel>
+          </ModelCollapsible>
         ))}
-      </Collapse>
+      </div>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx
index 3a48b5f7b50..12a2031798b 100644
--- a/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx
+++ b/ui/litellm-dashboard/src/components/common_components/team_multi_select.tsx
@@ -1,13 +1,22 @@
 import React, { useMemo, useState, type UIEvent } from "react";
-import { Select, Typography } from "antd";
-import { LoadingOutlined } from "@ant-design/icons";
-import { useDebouncedState } from "@tanstack/react-pacer/debouncer";
+import { Loader2 } from "lucide-react";
+import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer";
+import {
+  Combobox,
+  ComboboxChip,
+  ComboboxChips,
+  ComboboxChipsInput,
+  ComboboxClear,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxItem,
+  ComboboxList,
+  ComboboxValue,
+} from "@/components/ui/combobox";
 import { useInfiniteTeams } from "@/app/(dashboard)/hooks/teams/useTeams";
 import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
 import { Team } from "../key_team_helpers/key_list";
 
-const { Text } = Typography;
-
 interface TeamMultiSelectProps {
   value?: string[];
   onChange?: (value: string[]) => void;
@@ -27,77 +36,82 @@ const TeamMultiSelect: React.FC<TeamMultiSelectProps> = ({
   pageSize = 20,
   placeholder = "Search teams by alias...",
 }) => {
-  const [searchInput, setSearchInput] = useState("");
-  const [debouncedSearch, setDebouncedSearch] = useDebouncedState("", {
-    wait: DEBOUNCE_WAIT_MS,
-  });
+  const [search, setSearch] = useState("");
+  const debouncedSetSearch = useDebouncedCallback(setSearch, { wait: DEBOUNCE_WAIT_MS });
 
   const { data, fetchNextPage, hasNextPage, isFetchingNextPage, isLoading } = useInfiniteTeams(
     pageSize,
-    debouncedSearch || undefined,
+    search || undefined,
     organizationId,
   );
 
-  const teams = useMemo(() => {
-    if (!data?.pages) return [];
-    const seen = new Set<string>();
-    const result: Team[] = [];
-    for (const page of data.pages) {
-      for (const team of page.teams) {
-        if (seen.has(team.team_id)) continue;
-        seen.add(team.team_id);
-        result.push(team);
-      }
-    }
-    return result;
-  }, [data]);
+  const teamById = useMemo(
+    () =>
+      new Map<string, Team>(
+        (data?.pages ?? []).flatMap((page) => page.teams).map((team) => [team.team_id, team] as const),
+      ),
+    [data],
+  );
+  const teamIds = useMemo(() => Array.from(teamById.keys()), [teamById]);
 
-  const handlePopupScroll = (e: UIEvent<HTMLDivElement>) => {
-    const target = e.currentTarget;
+  const aliasOf = (teamId: string) => teamById.get(teamId)?.team_alias ?? teamId;
+
+  const handleScroll = (event: UIEvent<HTMLDivElement>) => {
+    const target = event.currentTarget;
+    if (target.scrollHeight === 0) return;
     const scrollRatio = (target.scrollTop + target.clientHeight) / target.scrollHeight;
     if (scrollRatio >= SCROLL_THRESHOLD && hasNextPage && !isFetchingNextPage) {
       fetchNextPage();
     }
   };
 
-  const handleSearch = (val: string) => {
-    setSearchInput(val);
-    setDebouncedSearch(val);
-  };
-
   return (
-    <Select
-      mode="multiple"
-      showSearch
-      placeholder={placeholder}
+    <Combobox
+      multiple
+      items={teamIds}
       value={value}
-      onChange={(val: string[]) => onChange?.(val)}
+      onValueChange={(next: string[]) => onChange?.(next)}
+      filter={null}
+      onInputValueChange={debouncedSetSearch}
       disabled={disabled}
-      allowClear
-      filterOption={false}
-      onSearch={handleSearch}
-      searchValue={searchInput}
-      onPopupScroll={handlePopupScroll}
-      loading={isLoading}
-      notFoundContent={isLoading ? <LoadingOutlined spin /> : "No teams found"}
-      style={{ width: "100%" }}
-      popupRender={(menu) => (
-        <>
-          {menu}
-          {isFetchingNextPage && (
-            <div style={{ textAlign: "center", padding: 8 }}>
-              <LoadingOutlined spin />
-            </div>
-          )}
-        </>
-      )}
     >
-      {teams.map((team) => (
-        <Select.Option key={team.team_id} value={team.team_id}>
-          <span className="font-medium">{team.team_alias}</span> <Text type="secondary">({team.team_id})</Text>
-        </Select.Option>
-      ))}
-    </Select>
+      <ComboboxChips className="w-full" aria-busy={isLoading}>
+        <ComboboxValue>
+          {(selected: string[]) =>
+            selected.map((teamId) => (
+              <ComboboxChip key={teamId} aria-label={aliasOf(teamId)}>
+                {aliasOf(teamId)}
+              </ComboboxChip>
+            ))
+          }
+        </ComboboxValue>
+        <ComboboxChipsInput
+          className="border-0 bg-transparent"
+          placeholder={placeholder}
+          aria-label={placeholder}
+          disabled={disabled}
+        />
+        {value.length > 0 && <ComboboxClear aria-label="Clear all teams" disabled={disabled} />}
+      </ComboboxChips>
+      <ComboboxContent>
+        <ComboboxEmpty>
+          {isLoading ? <Loader2 className="size-4 animate-spin text-muted-foreground" /> : "No teams found"}
+        </ComboboxEmpty>
+        <ComboboxList onScroll={handleScroll}>
+          {(teamId: string) => (
+            <ComboboxItem key={teamId} value={teamId}>
+              <span className="font-medium">{aliasOf(teamId)}</span>{" "}
+              <span className="text-muted-foreground">({teamId})</span>
+            </ComboboxItem>
+          )}
+        </ComboboxList>
+        {isFetchingNextPage && (
+          <div className="flex justify-center py-2">
+            <Loader2 className="size-4 animate-spin text-muted-foreground" />
+          </div>
+        )}
+      </ComboboxContent>
+    </Combobox>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/components/per_user_usage.tsx b/ui/litellm-dashboard/src/components/per_user_usage.tsx
index cef3fa5130b..b239da83766 100644
--- a/ui/litellm-dashboard/src/components/per_user_usage.tsx
+++ b/ui/litellm-dashboard/src/components/per_user_usage.tsx
@@ -1,8 +1,9 @@
 import React, { useState, useEffect } from "react";
-import { Title, Subtitle, Text, Button, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
 import type { ColumnDef } from "@tanstack/react-table";
 import { BarChart } from "@/components/shared/charts";
 import { DataTable } from "@/components/shared/DataTable";
+import { Button } from "@/components/ui/button";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { perUserAnalyticsCall } from "./networking";
 
 interface PerUserMetrics {
@@ -119,131 +120,133 @@ const PerUserUsage: React.FC<PerUserUsageProps> = ({ accessToken, selectedTags,
 
   return (
     <div className="mb-6">
-      <Title>Per User Usage</Title>
-      <Subtitle>Individual developer usage metrics</Subtitle>
+      <h3 className="text-lg font-medium text-foreground">Per User Usage</h3>
+      <p className="text-sm text-muted-foreground">Individual developer usage metrics</p>
 
-      <TabGroup>
-        <TabList className="mb-6">
-          <Tab>User Details</Tab>
-          <Tab>Usage Distribution</Tab>
-        </TabList>
+      <Tabs defaultValue="details">
+        <TabsList className="mb-6">
+          <TabsTrigger value="details" className="flex-none px-3">
+            User Details
+          </TabsTrigger>
+          <TabsTrigger value="distribution" className="flex-none px-3">
+            Usage Distribution
+          </TabsTrigger>
+        </TabsList>
 
-        <TabPanels>
-          {/* Tab 1: Existing User Details Table */}
-          <TabPanel>
-            <DataTable
-              columns={columns}
-              data={perUserData.results.slice(0, 10)}
-              getRowId={(row) => row.user_id}
-              noDataMessage="No per-user usage data"
-              size="compact"
-            />
+        {/* Tab 1: Existing User Details Table */}
+        <TabsContent value="details" keepMounted>
+          <DataTable
+            columns={columns}
+            data={perUserData.results.slice(0, 10)}
+            getRowId={(row) => row.user_id}
+            noDataMessage="No per-user usage data"
+            size="compact"
+          />
 
-            {perUserData.results.length > 10 && (
-              <div className="mt-4 flex justify-between items-center">
-                <Text className="text-sm text-gray-500">Showing 10 of {perUserData.total_count} results</Text>
-                <div className="flex gap-2">
-                  <Button size="sm" variant="secondary" onClick={handlePrevPage} disabled={currentPage === 1}>
-                    Previous
-                  </Button>
-                  <Button
-                    size="sm"
-                    variant="secondary"
-                    onClick={handleNextPage}
-                    disabled={currentPage >= perUserData.total_pages}
-                  >
-                    Next
-                  </Button>
-                </div>
+          {perUserData.results.length > 10 && (
+            <div className="mt-4 flex justify-between items-center">
+              <p className="text-sm text-gray-500">Showing 10 of {perUserData.total_count} results</p>
+              <div className="flex gap-2">
+                <Button size="sm" variant="secondary" onClick={handlePrevPage} disabled={currentPage === 1}>
+                  Previous
+                </Button>
+                <Button
+                  size="sm"
+                  variant="secondary"
+                  onClick={handleNextPage}
+                  disabled={currentPage >= perUserData.total_pages}
+                >
+                  Next
+                </Button>
               </div>
-            )}
-          </TabPanel>
-
-          {/* Tab 2: Usage Distribution Histogram */}
-          <TabPanel>
-            <div className="mb-4">
-              <Title className="text-lg">User Usage Distribution</Title>
-              <Subtitle>Number of users by successful request frequency</Subtitle>
             </div>
+          )}
+        </TabsContent>
 
-            <BarChart
-              data={(() => {
-                // Get top user agents by frequency first
-                const userAgentCounts = new Map<string, number>();
-                perUserData.results.forEach((item: PerUserMetrics) => {
-                  const agent = item.user_agent || "Unknown";
-                  userAgentCounts.set(agent, (userAgentCounts.get(agent) || 0) + 1);
-                });
+        {/* Tab 2: Usage Distribution Histogram */}
+        <TabsContent value="distribution" keepMounted>
+          <div className="mb-4">
+            <h4 className="text-lg font-medium text-foreground">User Usage Distribution</h4>
+            <p className="text-sm text-muted-foreground">Number of users by successful request frequency</p>
+          </div>
 
-                const topUserAgents = Array.from(userAgentCounts.entries())
-                  .sort(([, a], [, b]) => b - a)
-                  .slice(0, MAX_USER_AGENTS)
-                  .map(([agent]) => agent);
+          <BarChart
+            data={(() => {
+              // Get top user agents by frequency first
+              const userAgentCounts = new Map<string, number>();
+              perUserData.results.forEach((item: PerUserMetrics) => {
+                const agent = item.user_agent || "Unknown";
+                userAgentCounts.set(agent, (userAgentCounts.get(agent) || 0) + 1);
+              });
 
-                // Categorize users by successful request count and user agent
-                const categories = {
-                  "1-9 requests": { range: [1, 9], agents: {} as Record<string, number> },
-                  "10-99 requests": { range: [10, 99], agents: {} as Record<string, number> },
-                  "100-999 requests": { range: [100, 999], agents: {} as Record<string, number> },
-                  "1K-9.9K requests": { range: [1000, 9999], agents: {} as Record<string, number> },
-                  "10K-99.9K requests": { range: [10000, 99999], agents: {} as Record<string, number> },
-                  "100K+ requests": { range: [100000, Infinity], agents: {} as Record<string, number> },
-                };
+              const topUserAgents = Array.from(userAgentCounts.entries())
+                .sort(([, a], [, b]) => b - a)
+                .slice(0, MAX_USER_AGENTS)
+                .map(([agent]) => agent);
 
-                // Count users in each category by user agent (only for top user agents)
-                perUserData.results.forEach((item: PerUserMetrics) => {
-                  const successCount = item.successful_requests;
-                  const userAgent = item.user_agent || "Unknown";
+              // Categorize users by successful request count and user agent
+              const categories = {
+                "1-9 requests": { range: [1, 9], agents: {} as Record<string, number> },
+                "10-99 requests": { range: [10, 99], agents: {} as Record<string, number> },
+                "100-999 requests": { range: [100, 999], agents: {} as Record<string, number> },
+                "1K-9.9K requests": { range: [1000, 9999], agents: {} as Record<string, number> },
+                "10K-99.9K requests": { range: [10000, 99999], agents: {} as Record<string, number> },
+                "100K+ requests": { range: [100000, Infinity], agents: {} as Record<string, number> },
+              };
 
-                  // Only process if this is one of the top user agents
-                  if (topUserAgents.includes(userAgent)) {
-                    Object.entries(categories).forEach(([categoryName, category]) => {
-                      if (successCount >= category.range[0] && successCount <= category.range[1]) {
-                        if (!category.agents[userAgent]) {
-                          category.agents[userAgent] = 0;
-                        }
-                        category.agents[userAgent]++;
+              // Count users in each category by user agent (only for top user agents)
+              perUserData.results.forEach((item: PerUserMetrics) => {
+                const successCount = item.successful_requests;
+                const userAgent = item.user_agent || "Unknown";
+
+                // Only process if this is one of the top user agents
+                if (topUserAgents.includes(userAgent)) {
+                  Object.entries(categories).forEach(([categoryName, category]) => {
+                    if (successCount >= category.range[0] && successCount <= category.range[1]) {
+                      if (!category.agents[userAgent]) {
+                        category.agents[userAgent] = 0;
                       }
-                    });
-                  }
-                });
-
-                // Convert to chart data format for stacked bar chart
-                return Object.entries(categories).map(([categoryName, category]) => {
-                  const dataPoint: Record<string, any> = { category: categoryName };
-
-                  // Add count for each top user agent
-                  topUserAgents.forEach((agent) => {
-                    dataPoint[agent] = category.agents[agent] || 0;
+                      category.agents[userAgent]++;
+                    }
                   });
+                }
+              });
 
-                  return dataPoint;
-                });
-              })()}
-              index="category"
-              categories={(() => {
-                // Count user agents by frequency and get top ones
-                const userAgentCounts = new Map<string, number>();
-                perUserData.results.forEach((item: PerUserMetrics) => {
-                  const agent = item.user_agent || "Unknown";
-                  userAgentCounts.set(agent, (userAgentCounts.get(agent) || 0) + 1);
+              // Convert to chart data format for stacked bar chart
+              return Object.entries(categories).map(([categoryName, category]) => {
+                const dataPoint: Record<string, any> = { category: categoryName };
+
+                // Add count for each top user agent
+                topUserAgents.forEach((agent) => {
+                  dataPoint[agent] = category.agents[agent] || 0;
                 });
 
-                // Sort by frequency (most common first) and limit to top MAX_USER_AGENTS
-                return Array.from(userAgentCounts.entries())
-                  .sort(([, a], [, b]) => b - a)
-                  .slice(0, MAX_USER_AGENTS)
-                  .map(([agent]) => agent);
-              })()}
-              colors={["blue", "green", "orange", "red", "purple", "yellow", "pink", "indigo"]}
-              valueFormatter={(value: number) => `${value} users`}
-              yAxisWidth={80}
-              showLegend={true}
-              stack={true}
-            />
-          </TabPanel>
-        </TabPanels>
-      </TabGroup>
+                return dataPoint;
+              });
+            })()}
+            index="category"
+            categories={(() => {
+              // Count user agents by frequency and get top ones
+              const userAgentCounts = new Map<string, number>();
+              perUserData.results.forEach((item: PerUserMetrics) => {
+                const agent = item.user_agent || "Unknown";
+                userAgentCounts.set(agent, (userAgentCounts.get(agent) || 0) + 1);
+              });
+
+              // Sort by frequency (most common first) and limit to top MAX_USER_AGENTS
+              return Array.from(userAgentCounts.entries())
+                .sort(([, a], [, b]) => b - a)
+                .slice(0, MAX_USER_AGENTS)
+                .map(([agent]) => agent);
+            })()}
+            colors={["blue", "green", "orange", "red", "purple", "yellow", "pink", "indigo"]}
+            valueFormatter={(value: number) => `${value} users`}
+            yAxisWidth={80}
+            showLegend={true}
+            stack={true}
+          />
+        </TabsContent>
+      </Tabs>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/components/user_agent_activity.tsx b/ui/litellm-dashboard/src/components/user_agent_activity.tsx
index 29a5336e26f..a05db3313ab 100644
--- a/ui/litellm-dashboard/src/components/user_agent_activity.tsx
+++ b/ui/litellm-dashboard/src/components/user_agent_activity.tsx
@@ -1,10 +1,23 @@
 import React, { useState, useEffect } from "react";
-import { Card, Title, Text, Grid, Metric, Subtitle, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react";
-import { Select, Tooltip } from "antd";
+import { Card, CardContent } from "@/components/ui/card";
+import {
+  Combobox,
+  ComboboxChip,
+  ComboboxChips,
+  ComboboxChipsInput,
+  ComboboxClear,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxItem,
+  ComboboxList,
+  ComboboxValue,
+} from "@/components/ui/combobox";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { BarChart } from "@/components/shared/charts";
 import { userAgentSummaryCall, tagDauCall, tagWauCall, tagMauCall, tagDistinctCall } from "./networking";
 import PerUserUsage from "./per_user_usage";
-import { DateRangePickerValue } from "@tremor/react";
+import type { DateRangePickerValue } from "@tremor/react";
 import { ChartLoader } from "./shared/chart_loader";
 
 // New interfaces for the updated API response
@@ -356,39 +369,53 @@ const UserAgentActivity: React.FC<UserAgentActivityProps> = ({ accessToken, user
     <div className="space-y-6 mt-6">
       {/* Summary Section Card */}
       <Card>
-        <div className="space-y-6">
+        <CardContent className="space-y-6">
           <div className="flex justify-between items-start">
             <div>
-              <Title>Summary by User Agent</Title>
-              <Subtitle>Performance metrics for different user agents</Subtitle>
+              <h3 className="text-lg font-medium text-foreground">Summary by User Agent</h3>
+              <p className="text-sm text-muted-foreground">Performance metrics for different user agents</p>
             </div>
 
             {/* User Agent Filter */}
             <div className="w-96">
-              <Text className="text-sm font-medium block mb-2">Filter by User Agents</Text>
-              <Select
-                mode="multiple"
-                placeholder="All User Agents"
+              <label className="text-sm font-medium block mb-2">Filter by User Agents</label>
+              <Combobox
+                multiple
+                items={availableTags}
                 value={selectedTags}
-                onChange={setSelectedTags}
-                style={{ width: "100%" }}
-                showSearch={true}
-                allowClear={true}
-                loading={tagsLoading}
-                optionFilterProp="label"
-                className="rounded-md"
-                maxTagCount="responsive"
+                onValueChange={(next: string[]) => setSelectedTags(next)}
               >
-                {availableTags.map((tag) => {
-                  const userAgent = extractUserAgent(tag);
-                  const displayName = userAgent.length > 50 ? `${userAgent.substring(0, 50)}...` : userAgent;
-                  return (
-                    <Select.Option key={tag} value={tag} label={displayName} title={userAgent}>
-                      {displayName}
-                    </Select.Option>
-                  );
-                })}
-              </Select>
+                <ComboboxChips className="w-full" aria-busy={tagsLoading}>
+                  <ComboboxValue>
+                    {(selected: string[]) =>
+                      selected.map((tag) => (
+                        <ComboboxChip key={tag} aria-label={extractUserAgent(tag)}>
+                          {truncateUserAgent(extractUserAgent(tag))}
+                        </ComboboxChip>
+                      ))
+                    }
+                  </ComboboxValue>
+                  <ComboboxChipsInput
+                    className="border-0 bg-transparent"
+                    placeholder="All User Agents"
+                    aria-label="All User Agents"
+                  />
+                  {selectedTags.length > 0 && <ComboboxClear aria-label="Clear user agent filter" />}
+                </ComboboxChips>
+                <ComboboxContent>
+                  <ComboboxEmpty>No user agents found</ComboboxEmpty>
+                  <ComboboxList>
+                    {(tag: string) => {
+                      const userAgent = extractUserAgent(tag);
+                      return (
+                        <ComboboxItem key={tag} value={tag} title={userAgent}>
+                          {userAgent.length > 50 ? `${userAgent.substring(0, 50)}...` : userAgent}
+                        </ComboboxItem>
+                      );
+                    }}
+                  </ComboboxList>
+                </ComboboxContent>
+              </Combobox>
             </div>
           </div>
 
@@ -398,151 +425,166 @@ const UserAgentActivity: React.FC<UserAgentActivityProps> = ({ accessToken, user
           {summaryLoading ? (
             <ChartLoader isDateChanging={false} />
           ) : (
-            <Grid numItems={4} className="gap-4">
+            <div className="grid grid-cols-4 gap-4">
               {(summaryData.results || []).slice(0, 4).map((tag, index) => {
                 const userAgent = extractUserAgent(tag.tag);
                 const displayName = truncateUserAgent(userAgent);
                 return (
                   <Card key={index}>
-                    <Tooltip title={userAgent} placement="top">
-                      <Title className="truncate">{displayName}</Title>
-                    </Tooltip>
-                    <div className="mt-4 space-y-3">
-                      <div>
-                        <Text className="text-sm text-gray-600">Success Requests</Text>
-                        <Metric className="text-lg">{formatAbbreviatedNumber(tag.successful_requests)}</Metric>
+                    <CardContent>
+                      <Tooltip>
+                        <TooltipTrigger
+                          render={<h4 className="truncate text-lg font-medium text-foreground">{displayName}</h4>}
+                        />
+                        <TooltipContent side="top">{userAgent}</TooltipContent>
+                      </Tooltip>
+                      <div className="mt-4 space-y-3">
+                        <div>
+                          <p className="text-sm text-gray-600">Success Requests</p>
+                          <p className="text-lg font-semibold">{formatAbbreviatedNumber(tag.successful_requests)}</p>
+                        </div>
+                        <div>
+                          <p className="text-sm text-gray-600">Total Tokens</p>
+                          <p className="text-lg font-semibold">{formatAbbreviatedNumber(tag.total_tokens)}</p>
+                        </div>
+                        <div>
+                          <p className="text-sm text-gray-600">Total Cost</p>
+                          <p className="text-lg font-semibold">${formatAbbreviatedNumber(tag.total_spend, 4)}</p>
+                        </div>
                       </div>
-                      <div>
-                        <Text className="text-sm text-gray-600">Total Tokens</Text>
-                        <Metric className="text-lg">{formatAbbreviatedNumber(tag.total_tokens)}</Metric>
-                      </div>
-                      <div>
-                        <Text className="text-sm text-gray-600">Total Cost</Text>
-                        <Metric className="text-lg">${formatAbbreviatedNumber(tag.total_spend, 4)}</Metric>
-                      </div>
-                    </div>
+                    </CardContent>
                   </Card>
                 );
               })}
               {/* Fill remaining slots if less than 4 agents */}
               {Array.from({ length: Math.max(0, 4 - (summaryData.results || []).length) }).map((_, index) => (
                 <Card key={`empty-${index}`}>
-                  <Title>No Data</Title>
-                  <div className="mt-4 space-y-3">
-                    <div>
-                      <Text className="text-sm text-gray-600">Success Requests</Text>
-                      <Metric className="text-lg">-</Metric>
+                  <CardContent>
+                    <h4 className="text-lg font-medium text-foreground">No Data</h4>
+                    <div className="mt-4 space-y-3">
+                      <div>
+                        <p className="text-sm text-gray-600">Success Requests</p>
+                        <p className="text-lg font-semibold">-</p>
+                      </div>
+                      <div>
+                        <p className="text-sm text-gray-600">Total Tokens</p>
+                        <p className="text-lg font-semibold">-</p>
+                      </div>
+                      <div>
+                        <p className="text-sm text-gray-600">Total Cost</p>
+                        <p className="text-lg font-semibold">-</p>
+                      </div>
                     </div>
-                    <div>
-                      <Text className="text-sm text-gray-600">Total Tokens</Text>
-                      <Metric className="text-lg">-</Metric>
-                    </div>
-                    <div>
-                      <Text className="text-sm text-gray-600">Total Cost</Text>
-                      <Metric className="text-lg">-</Metric>
-                    </div>
-                  </div>
+                  </CardContent>
                 </Card>
               ))}
-            </Grid>
+            </div>
           )}
-        </div>
+        </CardContent>
       </Card>
 
-      {/* Main TabGroup for DAU/WAU/MAU vs Per User Usage */}
+      {/* Main tabs for DAU/WAU/MAU vs Per User Usage */}
       <Card>
-        <TabGroup>
-          <TabList className="mb-6">
-            <Tab>DAU/WAU/MAU</Tab>
-            <Tab>Per User Usage (Last 30 Days)</Tab>
-          </TabList>
+        <CardContent>
+          <Tabs defaultValue="active-users">
+            <TabsList className="mb-6">
+              <TabsTrigger value="active-users" className="flex-none px-3">
+                DAU/WAU/MAU
+              </TabsTrigger>
+              <TabsTrigger value="per-user" className="flex-none px-3">
+                Per User Usage (Last 30 Days)
+              </TabsTrigger>
+            </TabsList>
 
-          <TabPanels>
             {/* DAU/WAU/MAU Tab Panel */}
-            <TabPanel>
+            <TabsContent value="active-users" keepMounted>
               <div className="mb-6">
-                <Title>DAU, WAU & MAU per Agent</Title>
-                <Subtitle>Active users across different time periods</Subtitle>
+                <h3 className="text-lg font-medium text-foreground">DAU, WAU &amp; MAU per Agent</h3>
+                <p className="text-sm text-muted-foreground">Active users across different time periods</p>
               </div>
 
-              <TabGroup>
-                <TabList className="mb-6">
-                  <Tab>DAU</Tab>
-                  <Tab>WAU</Tab>
-                  <Tab>MAU</Tab>
-                </TabList>
+              <Tabs defaultValue="dau">
+                <TabsList className="mb-6">
+                  <TabsTrigger value="dau" className="flex-none px-3">
+                    DAU
+                  </TabsTrigger>
+                  <TabsTrigger value="wau" className="flex-none px-3">
+                    WAU
+                  </TabsTrigger>
+                  <TabsTrigger value="mau" className="flex-none px-3">
+                    MAU
+                  </TabsTrigger>
+                </TabsList>
 
-                <TabPanels>
-                  <TabPanel>
-                    <div className="mb-4">
-                      <Title className="text-lg">Daily Active Users - Last 7 Days</Title>
-                    </div>
-                    {dauLoading ? (
-                      <ChartLoader isDateChanging={false} />
-                    ) : (
-                      <BarChart
-                        data={dailyChartData}
-                        index="date"
-                        categories={allDauTags.map(extractUserAgent)}
-                        valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
-                        yAxisWidth={60}
-                        showLegend={true}
-                        stack={true}
-                      />
-                    )}
-                  </TabPanel>
+                <TabsContent value="dau" keepMounted>
+                  <div className="mb-4">
+                    <h4 className="text-lg font-medium text-foreground">Daily Active Users - Last 7 Days</h4>
+                  </div>
+                  {dauLoading ? (
+                    <ChartLoader isDateChanging={false} />
+                  ) : (
+                    <BarChart
+                      data={dailyChartData}
+                      index="date"
+                      categories={allDauTags.map(extractUserAgent)}
+                      valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
+                      yAxisWidth={60}
+                      showLegend={true}
+                      stack={true}
+                    />
+                  )}
+                </TabsContent>
 
-                  <TabPanel>
-                    <div className="mb-4">
-                      <Title className="text-lg">Weekly Active Users - Last 7 Weeks</Title>
-                    </div>
-                    {wauLoading ? (
-                      <ChartLoader isDateChanging={false} />
-                    ) : (
-                      <BarChart
-                        data={weeklyChartData}
-                        index="week"
-                        categories={allWauTags.map(extractUserAgent)}
-                        valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
-                        yAxisWidth={60}
-                        showLegend={true}
-                        stack={true}
-                      />
-                    )}
-                  </TabPanel>
+                <TabsContent value="wau" keepMounted>
+                  <div className="mb-4">
+                    <h4 className="text-lg font-medium text-foreground">Weekly Active Users - Last 7 Weeks</h4>
+                  </div>
+                  {wauLoading ? (
+                    <ChartLoader isDateChanging={false} />
+                  ) : (
+                    <BarChart
+                      data={weeklyChartData}
+                      index="week"
+                      categories={allWauTags.map(extractUserAgent)}
+                      valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
+                      yAxisWidth={60}
+                      showLegend={true}
+                      stack={true}
+                    />
+                  )}
+                </TabsContent>
 
-                  <TabPanel>
-                    <div className="mb-4">
-                      <Title className="text-lg">Monthly Active Users - Last 7 Months</Title>
-                    </div>
-                    {mauLoading ? (
-                      <ChartLoader isDateChanging={false} />
-                    ) : (
-                      <BarChart
-                        data={monthlyChartData}
-                        index="month"
-                        categories={allMauTags.map(extractUserAgent)}
-                        valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
-                        yAxisWidth={60}
-                        showLegend={true}
-                        stack={true}
-                      />
-                    )}
-                  </TabPanel>
-                </TabPanels>
-              </TabGroup>
-            </TabPanel>
+                <TabsContent value="mau" keepMounted>
+                  <div className="mb-4">
+                    <h4 className="text-lg font-medium text-foreground">Monthly Active Users - Last 7 Months</h4>
+                  </div>
+                  {mauLoading ? (
+                    <ChartLoader isDateChanging={false} />
+                  ) : (
+                    <BarChart
+                      data={monthlyChartData}
+                      index="month"
+                      categories={allMauTags.map(extractUserAgent)}
+                      valueFormatter={(value: number) => formatAbbreviatedNumber(value)}
+                      yAxisWidth={60}
+                      showLegend={true}
+                      stack={true}
+                    />
+                  )}
+                </TabsContent>
+              </Tabs>
+            </TabsContent>
 
             {/* Per User Usage Tab Panel */}
-            <TabPanel>
+            <TabsContent value="per-user" keepMounted>
               <PerUserUsage
                 accessToken={accessToken}
                 selectedTags={selectedTags}
                 formatAbbreviatedNumber={formatAbbreviatedNumber}
               />
-            </TabPanel>
-          </TabPanels>
-        </TabGroup>
+            </TabsContent>
+          </Tabs>
+        </CardContent>
       </Card>
     </div>
   );

From 771c36c1eec0d67fc9ee98150e14244e0a71ecce Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:05:09 -0700
Subject: [PATCH 10/31] refactor(ui): migrate guardrails route to shadcn

Ports the 21 route-owned guardrails components off antd and Tremor onto the
installed shadcn base-vega primitives. Forms, tables and shared components stay
where they are: the route analyzer buckets add_guardrail_form, guardrail_info
and six others as DEFERRED, and eleven components under src/components as
SHARED, so none of them are touched here.

Three behaviours needed explicit handling because Base UI does not match antd by
default. Base UI Tabs mounts only the active panel, so the playground draft and
every panel's local state would reset on a tab switch; TabsContent now carries
keepMounted. ComboboxList takes a function child to render the filtered items,
and mapping over the array instead renders everything and silently disables the
type to filter behaviour antd gave for free. SelectContent gets
alignItemWithTrigger={false} so popups anchor below the trigger the way antd's
did, which also removes a focus race that made the template picker
intermittently unclickable.

The three dialogs reached from inside the still-antd wizard need a layer above
its z-index of 1000 to be visible at all.

Adds three tests rather than editing the characterisation tests from the
previous commit: a guard that every tab panel stays mounted, and one filtering
regression each for the pattern and mode dropdowns.

Prunes the guardrails entries from eslint-suppressions.json.
---
 ui/litellm-dashboard/eslint-suppressions.json |  86 +---
 .../_components/GuardrailTestPanel.tsx        |  98 +++--
 .../_components/GuardrailTestPlayground.tsx   | 207 +++++----
 .../_components/GuardrailTestResults.tsx      | 107 ++---
 .../_components/GuardrailsPanel.test.tsx      |   7 +
 .../_components/GuardrailsPanel.tsx           | 219 +++++-----
 .../content_filter/CategoryTable.tsx          |  71 +--
 .../ContentCategoryConfiguration.tsx          | 320 +++++++-------
 .../ContentFilterConfiguration.tsx            | 113 +++--
 .../content_filter/ContentFilterDisplay.tsx   |  67 +--
 .../content_filter/ContentFilterManager.tsx   |  28 +-
 .../content_filter/CustomPatternModal.tsx     | 110 +++--
 .../content_filter/KeywordModal.tsx           | 109 +++--
 .../content_filter/KeywordTable.tsx           |  30 +-
 .../content_filter/PatternModal.test.tsx      |  14 +
 .../content_filter/PatternModal.tsx           | 152 ++++---
 .../content_filter/PatternTable.tsx           |  42 +-
 .../content_filter/action_options.ts          |  10 +
 .../content_filter/dialog_layering.ts         |   1 +
 .../custom_code/CustomCodeModal.test.tsx      |  11 +
 .../custom_code/CustomCodeModal.tsx           | 403 +++++++++---------
 .../_components/guardrail_garden.tsx          |  68 ++-
 .../_components/guardrail_garden_card.tsx     |  38 +-
 .../_components/guardrail_garden_detail.tsx   |  31 +-
 .../guardrails/_components/pii_components.tsx | 229 +++++-----
 .../_components/pii_configuration.tsx         |   9 +-
 .../ToolPermissionRulesEditor.tsx             | 348 ++++++++-------
 27 files changed, 1490 insertions(+), 1438 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/action_options.ts
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts

diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json
index 2f4def9ffa2..8ba5d7aab3f 100644
--- a/ui/litellm-dashboard/eslint-suppressions.json
+++ b/ui/litellm-dashboard/eslint-suppressions.json
@@ -326,28 +326,12 @@
       "count": 1
     }
   },
-  "src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx": {
-    "no-restricted-imports": {
-      "count": 2
-    }
-  },
   "src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx": {
     "no-nested-ternary": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx": {
-    "no-restricted-imports": {
-      "count": 2
     }
   },
   "src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/set-state-in-effect": {
       "count": 1
     }
@@ -386,11 +370,6 @@
       "count": 2
     }
   },
-  "src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/guardrails/_components/content_filter/CompetitorIntentConfiguration.tsx": {
     "no-nested-ternary": {
       "count": 1
@@ -403,14 +382,8 @@
     }
   },
   "src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx": {
-    "local/no-complex-jsx-arrow": {
-      "count": 1
-    },
     "no-nested-ternary": {
-      "count": 3
-    },
-    "no-restricted-imports": {
-      "count": 1
+      "count": 2
     },
     "react-hooks/set-state-in-effect": {
       "count": 1
@@ -419,62 +392,20 @@
   "src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx": {
     "local/no-complex-jsx-arrow": {
       "count": 3
-    },
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx": {
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx": {
     "max-params": {
       "count": 2
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/set-state-in-effect": {
       "count": 1
     }
   },
-  "src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx": {
-    "local/no-complex-jsx-arrow": {
-      "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
-  "src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx": {
     "no-nested-ternary": {
       "count": 6
     },
-    "no-restricted-imports": {
-      "count": 2
-    },
     "react-hooks/set-state-in-effect": {
       "count": 1
     }
@@ -487,9 +418,6 @@
   "src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx": {
@@ -500,9 +428,6 @@
   "src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/guardrail_info.tsx": {
@@ -568,23 +493,14 @@
   "src/app/(dashboard)/guardrails/_components/pii_components.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/pii_configuration.tsx": {
     "local/filename-pascal-case": {
       "count": 1
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx": {
-    "no-restricted-imports": {
-      "count": 2
-    },
     "react-hooks/purity": {
       "count": 1
     }
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx
index e310da55063..8297c1e1e3b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx
@@ -1,13 +1,12 @@
 import React, { useState } from "react";
-import { Button } from "@tremor/react";
-import { Input, Typography, Tooltip } from "antd";
-import { CopyOutlined, InfoCircleOutlined } from "@ant-design/icons";
+import { Copy, Info } from "lucide-react";
 import NotificationsManager from "@/components/molecules/notifications_manager";
+import { Button } from "@/components/ui/button";
+import { Textarea } from "@/components/ui/textarea";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import GuardrailTestResults from "./GuardrailTestResults";
 
-const { TextArea } = Input;
-const { Text } = Typography;
-
 interface GuardrailTestPanelProps {
   guardrailNames: string[];
   onSubmit: (text: string, metadata?: Record<string, unknown> | null) => void;
@@ -108,23 +107,23 @@ export function GuardrailTestPanel({
   return (
     <div className="space-y-4 h-full flex flex-col">
       {/* Header */}
-      <div className="flex items-center justify-between pb-3 border-b border-gray-200">
+      <div className="flex items-center justify-between border-b border-border pb-3">
         <div className="flex items-center space-x-3">
           <div className="flex-1 min-w-0">
-            <div className="flex items-center space-x-2 mb-1">
-              <h2 className="text-lg font-semibold text-gray-900">Test Guardrails:</h2>
+            <div className="mb-1 flex items-center space-x-2">
+              <h2 className="text-lg font-semibold">Test Guardrails:</h2>
               <div className="flex flex-wrap gap-2">
                 {guardrailNames.map((name) => (
                   <div
                     key={name}
-                    className="inline-flex items-center space-x-1 bg-blue-50 px-3 py-1 rounded-md border border-blue-200"
+                    className="inline-flex items-center space-x-1 rounded-md border border-blue-200 bg-blue-50 px-3 py-1"
                   >
-                    <span className="font-mono text-blue-700 font-medium text-sm">{name}</span>
+                    <span className="font-mono text-sm font-medium text-blue-700">{name}</span>
                   </div>
                 ))}
               </div>
             </div>
-            <p className="text-sm text-gray-500">
+            <p className="text-sm text-muted-foreground">
               Test {guardrailNames.length > 1 ? "guardrails" : "guardrail"} and compare results
             </p>
           </div>
@@ -135,46 +134,63 @@ export function GuardrailTestPanel({
       <div className="flex-1 overflow-auto space-y-4">
         <div className="space-y-3">
           <div>
-            <div className="flex justify-between items-center mb-2">
+            <div className="mb-2 flex items-center justify-between">
               <div className="flex items-center gap-2">
-                <label className="text-sm font-medium text-gray-700">Input Text</label>
-                <Tooltip title="Press Enter to submit. Use Shift+Enter for new line.">
-                  <InfoCircleOutlined className="text-gray-400 cursor-help" />
+                <label className="text-sm font-medium">Input Text</label>
+                <Tooltip>
+                  <TooltipTrigger
+                    render={
+                      <span className="cursor-help text-muted-foreground">
+                        <Info className="size-3.5" />
+                      </span>
+                    }
+                  />
+                  <TooltipContent>Press Enter to submit. Use Shift+Enter for new line.</TooltipContent>
                 </Tooltip>
               </div>
               {inputText && (
-                <Button size="xs" variant="secondary" icon={CopyOutlined} onClick={handleCopyInput}>
+                <Button size="sm" variant="secondary" onClick={handleCopyInput}>
+                  <Copy />
                   Copy Input
                 </Button>
               )}
             </div>
-            <TextArea
+            <Textarea
               value={inputText}
               onChange={(e) => setInputText(e.target.value)}
               onKeyDown={handleKeyDown}
               placeholder="Enter text to test with guardrails..."
               rows={8}
-              className="font-mono text-sm"
+              className="font-mono text-sm field-sizing-fixed"
             />
-            <div className="flex justify-between items-center mt-1">
-              <Text className="text-xs text-gray-500">
-                Press <kbd className="px-1 py-0.5 bg-gray-100 border border-gray-300 rounded-sm text-xs">Enter</kbd> to
-                submit •{" "}
-                <kbd className="px-1 py-0.5 bg-gray-100 border border-gray-300 rounded-sm text-xs">Shift+Enter</kbd> for
-                new line
-              </Text>
-              <Text className="text-xs text-gray-500">Characters: {inputText.length}</Text>
+            <div className="mt-1 flex items-center justify-between">
+              <span className="text-xs text-muted-foreground">
+                Press <kbd className="rounded-sm border border-border bg-muted px-1 py-0.5 text-xs">Enter</kbd> to
+                submit • <kbd className="rounded-sm border border-border bg-muted px-1 py-0.5 text-xs">Shift+Enter</kbd>{" "}
+                for new line
+              </span>
+              <span className="text-xs text-muted-foreground">Characters: {inputText.length}</span>
             </div>
           </div>
 
           <div>
-            <div className="flex items-center gap-2 mb-2">
-              <label className="text-sm font-medium text-gray-700">Metadata (optional)</label>
-              <Tooltip title="JSON object forwarded to the guardrail as request_data['metadata']. Custom guardrails can read per-request configuration from it.">
-                <InfoCircleOutlined className="text-gray-400 cursor-help" />
+            <div className="mb-2 flex items-center gap-2">
+              <label className="text-sm font-medium">Metadata (optional)</label>
+              <Tooltip>
+                <TooltipTrigger
+                  render={
+                    <span className="cursor-help text-muted-foreground">
+                      <Info className="size-3.5" />
+                    </span>
+                  }
+                />
+                <TooltipContent>
+                  JSON object forwarded to the guardrail as request_data[&apos;metadata&apos;]. Custom guardrails can
+                  read per-request configuration from it.
+                </TooltipContent>
               </Tooltip>
             </div>
-            <TextArea
+            <Textarea
               value={metadataText}
               onChange={(e) => {
                 setMetadataText(e.target.value);
@@ -184,18 +200,20 @@ export function GuardrailTestPanel({
               }}
               placeholder='{"forbidden_topics": ["tax", "finance"]}'
               rows={3}
-              className="font-mono text-sm"
-              status={metadataError ? "error" : undefined}
+              className="font-mono text-sm field-sizing-fixed"
+              aria-invalid={metadataError ? true : undefined}
             />
-            {metadataError && (
-              <Text type="danger" className="text-xs">
-                {metadataError}
-              </Text>
-            )}
+            {metadataError && <span className="text-xs text-destructive">{metadataError}</span>}
           </div>
 
           <div className="pt-2">
-            <Button onClick={handleSubmit} loading={isLoading} disabled={!inputText.trim()} className="w-full">
+            <Button
+              onClick={handleSubmit}
+              disabled={!inputText.trim() || isLoading}
+              aria-busy={isLoading}
+              className="w-full"
+            >
+              {isLoading && <UiLoadingSpinner className="size-4" />}
               {isLoading
                 ? `Testing ${guardrailNames.length} guardrail${guardrailNames.length > 1 ? "s" : ""}...`
                 : `Test ${guardrailNames.length} guardrail${guardrailNames.length > 1 ? "s" : ""}`}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx
index d3d8901101a..08916b03312 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestPlayground.tsx
@@ -1,9 +1,11 @@
 import React, { useState } from "react";
-import { Card, List, Empty, Spin, Input, Typography } from "antd";
-import { ExperimentOutlined, SearchOutlined } from "@ant-design/icons";
+import { FlaskConical, Search } from "lucide-react";
 import GuardrailTestPanel from "./GuardrailTestPanel";
 import { applyGuardrail } from "@/components/networking";
 import NotificationsManager from "@/components/molecules/notifications_manager";
+import { Card, CardContent } from "@/components/ui/card";
+import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 
 interface GuardrailItem {
   guardrail_id?: string;
@@ -112,115 +114,110 @@ const GuardrailTestPlayground: React.FC<GuardrailTestPlaygroundProps> = ({
 
   return (
     <div className="w-full h-[calc(100vh-200px)]">
-      <Card className="h-full" styles={{ body: { padding: 0, height: "100%" } }}>
-        <div className="flex h-full">
-          {/* Left Sidebar - Guardrails List */}
-          <div className="w-1/4 border-r border-gray-200 flex flex-col overflow-hidden">
-            <div className="p-4 border-b border-gray-200">
-              <div className="mb-3">
-                <h3 className="text-lg font-semibold mb-3">Guardrails</h3>
-                <Input
-                  prefix={<SearchOutlined />}
-                  placeholder="Search guardrails..."
-                  value={searchQuery}
-                  onChange={(e) => setSearchQuery(e.target.value)}
-                />
+      <Card className="h-full overflow-hidden py-0">
+        <CardContent className="h-full p-0">
+          <div className="flex h-full">
+            {/* Left Sidebar - Guardrails List */}
+            <div className="flex w-1/4 flex-col overflow-hidden border-r border-border">
+              <div className="border-b border-border p-4">
+                <div className="mb-3">
+                  <h3 className="mb-3 text-lg font-semibold">Guardrails</h3>
+                  <InputGroup>
+                    <InputGroupAddon>
+                      <Search className="size-4 text-muted-foreground" />
+                    </InputGroupAddon>
+                    <InputGroupInput
+                      placeholder="Search guardrails..."
+                      value={searchQuery}
+                      onChange={(e) => setSearchQuery(e.target.value)}
+                    />
+                  </InputGroup>
+                </div>
+              </div>
+
+              <div className="flex-1 overflow-auto">
+                {isLoading ? (
+                  <div className="flex h-32 items-center justify-center" aria-busy="true">
+                    <UiLoadingSpinner className="size-6 text-muted-foreground" />
+                  </div>
+                ) : filteredGuardrails.length === 0 ? (
+                  <div className="p-4 text-center text-muted-foreground">
+                    {searchQuery ? "No guardrails match your search" : "No guardrails available"}
+                  </div>
+                ) : (
+                  <ul className="m-0 list-none p-0">
+                    {filteredGuardrails.map((guardrail) => (
+                      <li
+                        key={guardrail.guardrail_id ?? guardrail.guardrail_name}
+                        onClick={() => {
+                          if (guardrail.guardrail_name) {
+                            toggleGuardrailSelection(guardrail.guardrail_name);
+                          }
+                        }}
+                        className={`cursor-pointer border-b border-border py-3 pr-4 pl-6 transition-colors hover:bg-muted/40 ${
+                          selectedGuardrails.has(guardrail.guardrail_name || "")
+                            ? "border-l-4 border-l-primary bg-accent"
+                            : "border-l-4 border-l-transparent"
+                        }`}
+                      >
+                        <div className="flex items-center space-x-2">
+                          <FlaskConical className="size-4 text-muted-foreground" />
+                          <span className="font-medium">{guardrail.guardrail_name}</span>
+                        </div>
+                        <div className="mt-1 space-y-1 text-xs">
+                          <div>
+                            <span className="font-medium">Type: </span>
+                            <span className="text-muted-foreground">{guardrail.litellm_params.guardrail}</span>
+                          </div>
+                          <div>
+                            <span className="font-medium">Mode: </span>
+                            <span className="text-muted-foreground">{guardrail.litellm_params.mode}</span>
+                          </div>
+                        </div>
+                      </li>
+                    ))}
+                  </ul>
+                )}
+              </div>
+
+              <div className="border-t border-border bg-muted/40 p-3">
+                <span className="text-xs text-muted-foreground">
+                  {selectedGuardrails.size} of {filteredGuardrails.length} selected
+                </span>
               </div>
             </div>
 
-            <div className="flex-1 overflow-auto">
-              {isLoading ? (
-                <div className="flex items-center justify-center h-32">
-                  <Spin />
-                </div>
-              ) : filteredGuardrails.length === 0 ? (
-                <div className="p-4">
-                  <Empty description={searchQuery ? "No guardrails match your search" : "No guardrails available"} />
-                </div>
-              ) : (
-                <List
-                  dataSource={filteredGuardrails}
-                  renderItem={(guardrail) => (
-                    <List.Item
-                      onClick={() => {
-                        if (guardrail.guardrail_name) {
-                          toggleGuardrailSelection(guardrail.guardrail_name);
-                        }
-                      }}
-                      style={{ paddingLeft: 24, paddingRight: 16 }}
-                      className={`cursor-pointer hover:bg-gray-50 transition-colors ${
-                        selectedGuardrails.has(guardrail.guardrail_name || "")
-                          ? "bg-blue-50 border-l-4 border-l-blue-500"
-                          : "border-l-4 border-l-transparent"
-                      }`}
-                    >
-                      <List.Item.Meta
-                        title={
-                          <div className="flex items-center space-x-2">
-                            <ExperimentOutlined className="text-gray-400" />
-                            <span className="font-medium text-gray-900">{guardrail.guardrail_name}</span>
-                          </div>
-                        }
-                        description={
-                          <div className="text-xs space-y-1 mt-1">
-                            <div>
-                              <span className="font-medium">Type: </span>
-                              <span className="text-gray-600">{guardrail.litellm_params.guardrail}</span>
-                            </div>
-                            <div>
-                              <span className="font-medium">Mode: </span>
-                              <span className="text-gray-600">{guardrail.litellm_params.mode}</span>
-                            </div>
-                          </div>
-                        }
-                      />
-                    </List.Item>
-                  )}
-                />
-              )}
-            </div>
+            {/* Right Panel - Test Area */}
+            <div className="flex w-3/4 flex-col">
+              <div className="flex items-center justify-between border-b border-border p-4">
+                <h2 className="mb-0 text-xl font-semibold">Guardrail Testing Playground</h2>
+              </div>
 
-            <div className="p-3 border-t border-gray-200 bg-gray-50">
-              <Typography.Text className="text-xs text-gray-600">
-                {selectedGuardrails.size} of {filteredGuardrails.length} selected
-              </Typography.Text>
+              <div className="flex-1 overflow-auto p-4">
+                {selectedGuardrails.size === 0 ? (
+                  <div className="flex h-full flex-col items-center justify-center text-muted-foreground">
+                    <FlaskConical className="mb-4 size-12" />
+                    <p className="mb-2 text-lg font-medium">Select Guardrails to Test</p>
+                    <p className="max-w-md text-center">
+                      Choose one or more guardrails from the left sidebar to start testing and comparing results.
+                    </p>
+                  </div>
+                ) : (
+                  <div className="h-full">
+                    <GuardrailTestPanel
+                      guardrailNames={Array.from(selectedGuardrails)}
+                      onSubmit={handleTestGuardrails}
+                      results={testResults.length > 0 ? testResults : null}
+                      errors={testErrors.length > 0 ? testErrors : null}
+                      isLoading={isTesting}
+                      onClose={() => setSelectedGuardrails(new Set())}
+                    />
+                  </div>
+                )}
+              </div>
             </div>
           </div>
-
-          {/* Right Panel - Test Area */}
-          <div className="w-3/4 flex flex-col bg-white">
-            <div className="p-4 border-b border-gray-200 flex justify-between items-center">
-              <Typography.Title level={2} className="text-xl font-semibold mb-0">
-                Guardrail Testing Playground
-              </Typography.Title>
-            </div>
-
-            <div className="flex-1 overflow-auto p-4">
-              {selectedGuardrails.size === 0 ? (
-                <div className="h-full flex flex-col items-center justify-center text-gray-400">
-                  <ExperimentOutlined style={{ fontSize: "48px", marginBottom: "16px" }} />
-                  <Typography.Paragraph className="text-lg font-medium text-gray-600 mb-2">
-                    Select Guardrails to Test
-                  </Typography.Paragraph>
-                  <Typography.Paragraph className="text-center text-gray-500 max-w-md">
-                    Choose one or more guardrails from the left sidebar to start testing and comparing results.
-                  </Typography.Paragraph>
-                </div>
-              ) : (
-                <div className="h-full">
-                  <GuardrailTestPanel
-                    guardrailNames={Array.from(selectedGuardrails)}
-                    onSubmit={handleTestGuardrails}
-                    results={testResults.length > 0 ? testResults : null}
-                    errors={testErrors.length > 0 ? testErrors : null}
-                    isLoading={isTesting}
-                    onClose={() => setSelectedGuardrails(new Set())}
-                  />
-                </div>
-              )}
-            </div>
-          </div>
-        </div>
+        </CardContent>
       </Card>
     </div>
   );
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx
index 10ce709244b..ba410a8a800 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailTestResults.tsx
@@ -1,7 +1,8 @@
 import React, { useState } from "react";
-import { Button, Card } from "@tremor/react";
-import { CopyOutlined, CheckCircleOutlined, ClockCircleOutlined, DownOutlined, RightOutlined } from "@ant-design/icons";
+import { Check, ChevronDown, ChevronRight, Clock, Copy } from "lucide-react";
 import NotificationsManager from "@/components/molecules/notifications_manager";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent } from "@/components/ui/card";
 
 interface TestResult {
   guardrailName: string;
@@ -66,39 +67,38 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
   }
 
   return (
-    <div className="space-y-3 pt-4 border-t border-gray-200">
-      <h3 className="text-sm font-semibold text-gray-900">Results</h3>
+    <div className="space-y-3 border-t border-border pt-4">
+      <h3 className="text-sm font-semibold">Results</h3>
 
       {/* Success Results */}
       {results &&
         results.map((result) => {
           const isCollapsed = collapsedResults.has(result.guardrailName);
           return (
-            <Card key={result.guardrailName} className="bg-green-50 border-green-200">
-              <div className="space-y-3">
+            <Card key={result.guardrailName} className="border-emerald-200 bg-emerald-50">
+              <CardContent className="space-y-3">
                 <div className="flex items-center justify-between">
                   <div
-                    className="flex items-center space-x-2 cursor-pointer flex-1"
+                    className="flex flex-1 cursor-pointer items-center space-x-2"
                     onClick={() => toggleResultCollapse(result.guardrailName)}
                   >
                     {isCollapsed ? (
-                      <RightOutlined className="text-gray-500 text-xs" />
+                      <ChevronRight className="size-3 text-muted-foreground" />
                     ) : (
-                      <DownOutlined className="text-gray-500 text-xs" />
+                      <ChevronDown className="size-3 text-muted-foreground" />
                     )}
-                    <CheckCircleOutlined className="text-green-600 text-lg" />
-                    <span className="text-sm font-medium text-green-800">{result.guardrailName}</span>
+                    <Check className="size-4 text-emerald-600" />
+                    <span className="text-sm font-medium text-emerald-800">{result.guardrailName}</span>
                   </div>
                   <div className="flex items-center gap-3">
-                    <div className="flex items-center space-x-1 text-xs text-gray-600">
-                      <ClockCircleOutlined />
+                    <div className="flex items-center space-x-1 text-xs text-muted-foreground">
+                      <Clock className="size-3" />
                       <span className="font-medium">{result.latency}ms</span>
                     </div>
                     {!isCollapsed && (
                       <Button
-                        size="xs"
+                        size="sm"
                         variant="secondary"
-                        icon={CopyOutlined}
                         onClick={async () => {
                           const success = await copyToClipboard(result.response_text);
                           if (success) {
@@ -108,6 +108,7 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
                           }
                         }}
                       >
+                        <Copy />
                         Copy
                       </Button>
                     )}
@@ -115,18 +116,18 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
                 </div>
                 {!isCollapsed && (
                   <>
-                    <div className="bg-white border border-green-200 rounded-sm p-3">
-                      <label className="text-xs font-medium text-gray-600 mb-2 block">Output Text</label>
-                      <div className="font-mono text-sm text-gray-900 whitespace-pre-wrap wrap-break-word">
+                    <div className="rounded-sm border border-emerald-200 bg-background p-3">
+                      <label className="mb-2 block text-xs font-medium text-muted-foreground">Output Text</label>
+                      <div className="font-mono text-sm whitespace-pre-wrap wrap-break-word">
                         {result.response_text}
                       </div>
                     </div>
-                    <div className="text-xs text-gray-600">
+                    <div className="text-xs text-muted-foreground">
                       <span className="font-medium">Characters:</span> {result.response_text.length}
                     </div>
                   </>
                 )}
-              </div>
+              </CardContent>
             </Card>
           );
         })}
@@ -136,40 +137,42 @@ export function GuardrailTestResults({ results, errors }: GuardrailTestResultsPr
         errors.map((errorItem) => {
           const isCollapsed = collapsedResults.has(errorItem.guardrailName);
           return (
-            <Card key={errorItem.guardrailName} className="bg-red-50 border-red-200">
-              <div className="flex items-start space-x-2">
-                <div className="cursor-pointer mt-0.5" onClick={() => toggleResultCollapse(errorItem.guardrailName)}>
-                  {isCollapsed ? (
-                    <RightOutlined className="text-gray-500 text-xs" />
-                  ) : (
-                    <DownOutlined className="text-gray-500 text-xs" />
-                  )}
-                </div>
-                <div className="text-red-600 mt-0.5">
-                  <svg className="w-5 h-5" fill="currentColor" viewBox="0 0 20 20">
-                    <path
-                      fillRule="evenodd"
-                      d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
-                      clipRule="evenodd"
-                    />
-                  </svg>
-                </div>
-                <div className="flex-1">
-                  <div className="flex items-center justify-between mb-1">
-                    <p
-                      className="text-sm font-medium text-red-800 cursor-pointer"
-                      onClick={() => toggleResultCollapse(errorItem.guardrailName)}
-                    >
-                      {errorItem.guardrailName} - Error
-                    </p>
-                    <div className="flex items-center space-x-1 text-xs text-gray-600">
-                      <ClockCircleOutlined />
-                      <span className="font-medium">{errorItem.latency}ms</span>
-                    </div>
+            <Card key={errorItem.guardrailName} className="border-red-200 bg-red-50">
+              <CardContent>
+                <div className="flex items-start space-x-2">
+                  <div className="mt-0.5 cursor-pointer" onClick={() => toggleResultCollapse(errorItem.guardrailName)}>
+                    {isCollapsed ? (
+                      <ChevronRight className="size-3 text-muted-foreground" />
+                    ) : (
+                      <ChevronDown className="size-3 text-muted-foreground" />
+                    )}
+                  </div>
+                  <div className="mt-0.5 text-destructive">
+                    <svg className="w-5 h-5" fill="currentColor" viewBox="0 0 20 20">
+                      <path
+                        fillRule="evenodd"
+                        d="M10 18a8 8 0 100-16 8 8 0 000 16zM8.707 7.293a1 1 0 00-1.414 1.414L8.586 10l-1.293 1.293a1 1 0 101.414 1.414L10 11.414l1.293 1.293a1 1 0 001.414-1.414L11.414 10l1.293-1.293a1 1 0 00-1.414-1.414L10 8.586 8.707 7.293z"
+                        clipRule="evenodd"
+                      />
+                    </svg>
+                  </div>
+                  <div className="flex-1">
+                    <div className="mb-1 flex items-center justify-between">
+                      <p
+                        className="cursor-pointer text-sm font-medium text-red-800"
+                        onClick={() => toggleResultCollapse(errorItem.guardrailName)}
+                      >
+                        {errorItem.guardrailName} - Error
+                      </p>
+                      <div className="flex items-center space-x-1 text-xs text-muted-foreground">
+                        <Clock className="size-3" />
+                        <span className="font-medium">{errorItem.latency}ms</span>
+                      </div>
+                    </div>
+                    {!isCollapsed && <p className="mt-1 text-sm text-red-700">{errorItem.error.message}</p>}
                   </div>
-                  {!isCollapsed && <p className="text-sm text-red-700 mt-1">{errorItem.error.message}</p>}
                 </div>
-              </div>
+              </CardContent>
             </Card>
           );
         })}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
index 88e52e0cdd1..d165d766ad0 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.test.tsx
@@ -138,6 +138,13 @@ describe("GuardrailsPanel", () => {
     expect(mockGetGuardrailsList).toHaveBeenCalledTimes(2);
   });
 
+  it("should mount every tab panel up front so panel state survives tab switches", async () => {
+    render(<GuardrailsPanel {...defaultProps} />);
+
+    expect(await screen.findByLabelText("playground draft")).toBeInTheDocument();
+    expect(screen.getByText("Mock Team Guardrails Tab")).toBeInTheDocument();
+  });
+
   it("should keep test playground state when switching tabs away and back", async () => {
     render(<GuardrailsPanel {...defaultProps} />);
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
index 33c6933634c..9df2fa04c04 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/GuardrailsPanel.tsx
@@ -1,5 +1,5 @@
 import React, { useState, useEffect } from "react";
-import { Tabs } from "antd";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
 import { ChevronDown, Code, Plus } from "lucide-react";
 import { getGuardrailsList, deleteGuardrailCall } from "@/components/networking";
 import { buttonVariants } from "@/components/ui/button";
@@ -125,118 +125,119 @@ const GuardrailsPanel: React.FC<GuardrailsPanelProps> = ({ accessToken, userRole
 
   return (
     <div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
-      <Tabs
-        defaultActiveKey="guardrails"
-        items={[
-          ...(isAdmin
-            ? [
-                {
-                  key: "garden",
-                  label: "Guardrail Garden",
-                  children: <GuardrailGarden accessToken={accessToken} onGuardrailCreated={handleSuccess} />,
-                },
-                {
-                  key: "guardrails",
-                  label: "Guardrails",
-                  children: (
-                    <>
-                      <div className="flex justify-between items-center mb-4">
-                        <DropdownMenu>
-                          <DropdownMenuTrigger
-                            disabled={!accessToken}
-                            className={cn(buttonVariants({ variant: "default" }))}
-                          >
-                            <Plus />
-                            Add New Guardrail
-                            <ChevronDown />
-                          </DropdownMenuTrigger>
-                          <DropdownMenuContent align="start" className="w-56">
-                            <DropdownMenuItem onClick={handleAddGuardrail}>
-                              <Plus />
-                              Add Provider Guardrail
-                            </DropdownMenuItem>
-                            <DropdownMenuItem onClick={handleAddCustomCodeGuardrail}>
-                              <Code />
-                              Create Custom Code Guardrail
-                            </DropdownMenuItem>
-                          </DropdownMenuContent>
-                        </DropdownMenu>
-                      </div>
+      <Tabs defaultValue="guardrails">
+        <TabsList>
+          {isAdmin && (
+            <>
+              <TabsTrigger value="garden" className="flex-none">
+                Guardrail Garden
+              </TabsTrigger>
+              <TabsTrigger value="guardrails" className="flex-none">
+                Guardrails
+              </TabsTrigger>
+              <TabsTrigger value="playground" className="flex-none" disabled={!accessToken}>
+                Test Playground
+              </TabsTrigger>
+            </>
+          )}
+          <TabsTrigger value="submitted" className="flex-none">
+            Submitted Guardrails
+          </TabsTrigger>
+        </TabsList>
 
-                      {selectedGuardrailId ? (
-                        <GuardrailInfoView
-                          guardrailId={selectedGuardrailId}
-                          onClose={() => setSelectedGuardrailId(null)}
-                          accessToken={accessToken}
-                          isAdmin={isAdmin}
-                        />
-                      ) : (
-                        <GuardrailTable
-                          guardrailsList={guardrailsList}
-                          isLoading={isLoading}
-                          onDeleteClick={handleDeleteClick}
-                          onGuardrailClick={(id) => setSelectedGuardrailId(id)}
-                        />
-                      )}
+        {isAdmin && (
+          <>
+            <TabsContent value="garden" keepMounted>
+              <GuardrailGarden accessToken={accessToken} onGuardrailCreated={handleSuccess} />
+            </TabsContent>
 
-                      <AddGuardrailForm
-                        visible={isAddModalVisible}
-                        onClose={handleCloseModal}
-                        accessToken={accessToken}
-                        onSuccess={handleSuccess}
-                      />
+            <TabsContent value="guardrails" keepMounted>
+              <div className="flex justify-between items-center mb-4">
+                <DropdownMenu>
+                  <DropdownMenuTrigger disabled={!accessToken} className={cn(buttonVariants({ variant: "default" }))}>
+                    <Plus />
+                    Add New Guardrail
+                    <ChevronDown />
+                  </DropdownMenuTrigger>
+                  <DropdownMenuContent align="start" className="w-56">
+                    <DropdownMenuItem onClick={handleAddGuardrail}>
+                      <Plus />
+                      Add Provider Guardrail
+                    </DropdownMenuItem>
+                    <DropdownMenuItem onClick={handleAddCustomCodeGuardrail}>
+                      <Code />
+                      Create Custom Code Guardrail
+                    </DropdownMenuItem>
+                  </DropdownMenuContent>
+                </DropdownMenu>
+              </div>
 
-                      <CustomCodeModal
-                        visible={isCustomCodeModalVisible}
-                        onClose={handleCloseCustomCodeModal}
-                        accessToken={accessToken}
-                        onSuccess={handleSuccess}
-                      />
+              {selectedGuardrailId ? (
+                <GuardrailInfoView
+                  guardrailId={selectedGuardrailId}
+                  onClose={() => setSelectedGuardrailId(null)}
+                  accessToken={accessToken}
+                  isAdmin={isAdmin}
+                />
+              ) : (
+                <GuardrailTable
+                  guardrailsList={guardrailsList}
+                  isLoading={isLoading}
+                  onDeleteClick={handleDeleteClick}
+                  onGuardrailClick={(id) => setSelectedGuardrailId(id)}
+                />
+              )}
 
-                      <DeleteResourceModal
-                        isOpen={isDeleteModalOpen}
-                        title="Delete Guardrail"
-                        message={`Are you sure you want to delete guardrail: ${guardrailToDelete?.guardrail_name}? This action cannot be undone.`}
-                        resourceInformationTitle="Guardrail Information"
-                        resourceInformation={[
-                          { label: "Name", value: guardrailToDelete?.guardrail_name },
-                          { label: "ID", value: guardrailToDelete?.guardrail_id, code: true },
-                          { label: "Provider", value: providerDisplayName },
-                          { label: "Mode", value: guardrailToDelete?.litellm_params.mode },
-                          {
-                            label: "Default On",
-                            value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No",
-                          },
-                        ]}
-                        onCancel={handleDeleteCancel}
-                        onOk={handleDeleteConfirm}
-                        confirmLoading={isDeleting}
-                      />
-                    </>
-                  ),
-                },
-                {
-                  key: "playground",
-                  label: "Test Playground",
-                  disabled: !accessToken,
-                  children: (
-                    <GuardrailTestPlayground
-                      guardrailsList={guardrailsList}
-                      isLoading={isLoading}
-                      accessToken={accessToken}
-                      onClose={() => {}}
-                    />
-                  ),
-                },
-              ]
-            : []),
-          {
-            key: "submitted",
-            label: "Submitted Guardrails",
-            children: <TeamGuardrailsTab accessToken={accessToken} />,
-          },
-        ]}
-      />
+              <AddGuardrailForm
+                visible={isAddModalVisible}
+                onClose={handleCloseModal}
+                accessToken={accessToken}
+                onSuccess={handleSuccess}
+              />
+
+              <CustomCodeModal
+                visible={isCustomCodeModalVisible}
+                onClose={handleCloseCustomCodeModal}
+                accessToken={accessToken}
+                onSuccess={handleSuccess}
+              />
+
+              <DeleteResourceModal
+                isOpen={isDeleteModalOpen}
+                title="Delete Guardrail"
+                message={`Are you sure you want to delete guardrail: ${guardrailToDelete?.guardrail_name}? This action cannot be undone.`}
+                resourceInformationTitle="Guardrail Information"
+                resourceInformation={[
+                  { label: "Name", value: guardrailToDelete?.guardrail_name },
+                  { label: "ID", value: guardrailToDelete?.guardrail_id, code: true },
+                  { label: "Provider", value: providerDisplayName },
+                  { label: "Mode", value: guardrailToDelete?.litellm_params.mode },
+                  {
+                    label: "Default On",
+                    value: guardrailToDelete?.litellm_params.default_on ? "Yes" : "No",
+                  },
+                ]}
+                onCancel={handleDeleteCancel}
+                onOk={handleDeleteConfirm}
+                confirmLoading={isDeleting}
+              />
+            </TabsContent>
+
+            <TabsContent value="playground" keepMounted>
+              <GuardrailTestPlayground
+                guardrailsList={guardrailsList}
+                isLoading={isLoading}
+                accessToken={accessToken}
+                onClose={() => {}}
+              />
+            </TabsContent>
+          </>
+        )}
+
+        <TabsContent value="submitted" keepMounted>
+          <TeamGuardrailsTab accessToken={accessToken} />
+        </TabsContent>
+      </Tabs>
     </div>
   );
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx
index cdc85bf0c0f..fa40ce6ab54 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CategoryTable.tsx
@@ -1,11 +1,11 @@
 import React from "react";
-import { Typography, Select, Tag, Button } from "antd";
-import { DeleteOutlined } from "@ant-design/icons";
+import { Trash2 } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
 import { DataTable } from "@/components/shared/DataTable";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS, SEVERITY_ITEMS } from "./action_options";
 
 interface ContentCategory {
   id: string;
@@ -38,14 +38,8 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
         const { category, display_name: displayName } = row.original;
         return (
           <div>
-            <Text strong>{displayName}</Text>
-            {displayName !== category && (
-              <div>
-                <Text type="secondary" style={{ fontSize: 12 }}>
-                  {category}
-                </Text>
-              </div>
-            )}
+            <span className="font-semibold">{displayName}</span>
+            {displayName !== category && <div className="text-xs text-muted-foreground">{category}</div>}
           </div>
         );
       },
@@ -57,23 +51,26 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
       cell: ({ row }) => {
         const { id, severity_threshold: severity } = row.original;
         if (readOnly) {
-          const colorMap = {
-            high: "red",
-            medium: "orange",
-            low: "yellow",
-          } as const;
-          return <Tag color={colorMap[severity as keyof typeof colorMap]}>{severity.toUpperCase()}</Tag>;
+          return <Badge variant={severity === "high" ? "destructive" : "secondary"}>{severity.toUpperCase()}</Badge>;
         }
         return (
           <Select
+            items={SEVERITY_ITEMS}
             value={severity}
-            onChange={(value) => onSeverityChange?.(id, value as "high" | "medium" | "low")}
-            style={{ width: 150 }}
-            size="small"
+            onValueChange={(value: string | null) =>
+              value && onSeverityChange?.(id, value as "high" | "medium" | "low")
+            }
           >
-            <Option value="high">High</Option>
-            <Option value="medium">Medium</Option>
-            <Option value="low">Low</Option>
+            <SelectTrigger size="sm" className="w-[150px]" aria-label="Severity Threshold">
+              <SelectValue />
+            </SelectTrigger>
+            <SelectContent alignItemWithTrigger={false}>
+              {SEVERITY_ITEMS.map((item) => (
+                <SelectItem key={item.value} value={item.value}>
+                  {item.label}
+                </SelectItem>
+              ))}
+            </SelectContent>
           </Select>
         );
       },
@@ -85,17 +82,24 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
       cell: ({ row }) => {
         const { action, id } = row.original;
         if (readOnly) {
-          return <Tag color={action === "BLOCK" ? "red" : "blue"}>{action}</Tag>;
+          return <Badge variant={action === "BLOCK" ? "destructive" : "secondary"}>{action}</Badge>;
         }
         return (
           <Select
+            items={ACTION_ITEMS}
             value={action}
-            onChange={(value) => onActionChange?.(id, value as "BLOCK" | "MASK")}
-            style={{ width: 120 }}
-            size="small"
+            onValueChange={(value: string | null) => value && onActionChange?.(id, value as "BLOCK" | "MASK")}
           >
-            <Option value="BLOCK">Block</Option>
-            <Option value="MASK">Mask</Option>
+            <SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
+              <SelectValue />
+            </SelectTrigger>
+            <SelectContent alignItemWithTrigger={false}>
+              {ACTION_ITEMS.map((item) => (
+                <SelectItem key={item.value} value={item.value}>
+                  {item.label}
+                </SelectItem>
+              ))}
+            </SelectContent>
           </Select>
         );
       },
@@ -108,7 +112,8 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
       id: "actions",
       size: 100,
       cell: ({ row }) => (
-        <Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove?.(row.original.id)}>
+        <Button variant="ghost" size="sm" onClick={() => onRemove?.(row.original.id)}>
+          <Trash2 />
           Delete
         </Button>
       ),
@@ -116,7 +121,7 @@ const CategoryTable: React.FC<CategoryTableProps> = ({
   }
 
   if (categories.length === 0) {
-    return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No categories configured.</div>;
+    return <div className="py-10 text-center text-muted-foreground">No categories configured.</div>;
   }
 
   return <DataTable data={categories} columns={columns} getRowId={(row) => row.id} size="compact" />;
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx
index d0e0a4c6c8b..1924133a6cc 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentCategoryConfiguration.tsx
@@ -1,12 +1,22 @@
 import React from "react";
-import { Card, Typography, Select, Tag, Collapse, Button } from "antd";
-import { DeleteOutlined, PlusOutlined, FileTextOutlined } from "@ant-design/icons";
+import { ChevronRight, FileText, Plus, Trash2 } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
 import { getCategoryYaml } from "@/components/networking";
 import { DataTable } from "@/components/shared/DataTable";
-
-const { Title, Text } = Typography;
-const { Option } = Select;
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
+import {
+  Combobox,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxInput,
+  ComboboxItem,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS, SEVERITY_ITEMS } from "./action_options";
 
 interface ContentCategory {
   name: string;
@@ -170,10 +180,8 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
         const category = availableCategories.find((c) => c.name === row.original.category);
         return (
           <div>
-            <div style={{ fontWeight: 500 }}>{row.original.display_name}</div>
-            {category?.description && (
-              <div style={{ fontSize: "12px", color: "#888", marginTop: "4px" }}>{category.description}</div>
-            )}
+            <div className="font-medium">{row.original.display_name}</div>
+            {category?.description && <div className="mt-1 text-xs text-muted-foreground">{category.description}</div>}
           </div>
         );
       },
@@ -184,16 +192,20 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
       size: 150,
       cell: ({ row }) => (
         <Select
+          items={ACTION_ITEMS}
           value={row.original.action}
-          onChange={(value) => onCategoryUpdate(row.original.id, "action", value)}
-          style={{ width: "100%" }}
+          onValueChange={(value: string | null) => value && onCategoryUpdate(row.original.id, "action", value)}
         >
-          <Option value="BLOCK">
-            <Tag color="red">BLOCK</Tag>
-          </Option>
-          <Option value="MASK">
-            <Tag color="orange">MASK</Tag>
-          </Option>
+          <SelectTrigger size="sm" className="w-full" aria-label="Action">
+            <SelectValue />
+          </SelectTrigger>
+          <SelectContent alignItemWithTrigger={false}>
+            {ACTION_ITEMS.map((item) => (
+              <SelectItem key={item.value} value={item.value}>
+                <Badge variant={item.value === "BLOCK" ? "destructive" : "secondary"}>{item.value}</Badge>
+              </SelectItem>
+            ))}
+          </SelectContent>
         </Select>
       ),
     },
@@ -203,13 +215,22 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
       size: 180,
       cell: ({ row }) => (
         <Select
+          items={SEVERITY_ITEMS}
           value={row.original.severity_threshold}
-          onChange={(value) => onCategoryUpdate(row.original.id, "severity_threshold", value)}
-          style={{ width: "100%" }}
+          onValueChange={(value: string | null) =>
+            value && onCategoryUpdate(row.original.id, "severity_threshold", value)
+          }
         >
-          <Option value="low">Low</Option>
-          <Option value="medium">Medium</Option>
-          <Option value="high">High</Option>
+          <SelectTrigger size="sm" className="w-full" aria-label="Severity Threshold">
+            <SelectValue />
+          </SelectTrigger>
+          <SelectContent alignItemWithTrigger={false}>
+            {SEVERITY_ITEMS.map((item) => (
+              <SelectItem key={item.value} value={item.value}>
+                {item.label}
+              </SelectItem>
+            ))}
+          </SelectContent>
         </Select>
       ),
     },
@@ -218,7 +239,8 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
       id: "actions",
       size: 80,
       cell: ({ row }) => (
-        <Button icon={<DeleteOutlined />} onClick={() => onCategoryRemove(row.original.id)} size="small">
+        <Button variant="outline" size="sm" onClick={() => onCategoryRemove(row.original.id)}>
+          <Trash2 />
           Remove
         </Button>
       ),
@@ -228,172 +250,120 @@ const ContentCategoryConfiguration: React.FC<ContentCategoryConfigurationProps>
   const unselectedCategories = availableCategories.filter(
     (cat) => !selectedCategories.some((sel) => sel.category === cat.name),
   );
+  const pendingCategory = availableCategories.find((c) => c.name === selectedCategoryName) ?? null;
 
   return (
-    <Card
-      title={
-        <div
-          style={{ display: "flex", justifyContent: "space-between", alignItems: "center", flexWrap: "wrap", gap: 8 }}
-        >
-          <Title level={5} style={{ margin: 0 }}>
-            Blocked topics
-          </Title>
-          <Text type="secondary" style={{ fontSize: 12, fontWeight: 400 }}>
+    <Card>
+      <CardHeader>
+        <div className="flex flex-wrap items-center justify-between gap-2">
+          <CardTitle>Blocked topics</CardTitle>
+          <p className="text-xs font-normal text-muted-foreground">
             Select topics to block using keyword and semantic analysis
-          </Text>
+          </p>
+        </div>
+      </CardHeader>
+      <CardContent>
+        <div className="mb-4 flex gap-2">
+          <Combobox
+            items={unselectedCategories}
+            value={pendingCategory}
+            onValueChange={(category: ContentCategory | null) => setSelectedCategoryName(category?.name ?? "")}
+            itemToStringLabel={(category: ContentCategory) => category.display_name}
+          >
+            <ComboboxInput className="w-full" placeholder="Select a content category" />
+            <ComboboxContent>
+              <ComboboxEmpty>No matching categories</ComboboxEmpty>
+              <ComboboxList>
+                {(cat: ContentCategory) => (
+                  <ComboboxItem key={cat.name} value={cat}>
+                    <div>
+                      <div className="font-medium">{cat.display_name}</div>
+                      <div className="mt-0.5 text-xs text-muted-foreground">{cat.description}</div>
+                    </div>
+                  </ComboboxItem>
+                )}
+              </ComboboxList>
+            </ComboboxContent>
+          </Combobox>
+          <Button onClick={handleAddCategory} disabled={!selectedCategoryName}>
+            <Plus />
+            Add
+          </Button>
         </div>
-      }
-      size="small"
-    >
-      <div style={{ marginBottom: 16, display: "flex", gap: 8 }}>
-        <Select
-          placeholder="Select a content category"
-          value={selectedCategoryName || undefined}
-          onChange={setSelectedCategoryName}
-          style={{ flex: 1 }}
-          showSearch
-          optionLabelProp="label"
-          filterOption={(input, option) =>
-            (option?.label?.toString().toLowerCase() ?? "").includes(input.toLowerCase())
-          }
-        >
-          {unselectedCategories.map((cat) => (
-            <Option key={cat.name} value={cat.name} label={cat.display_name}>
-              <div>
-                <div style={{ fontWeight: 500 }}>{cat.display_name}</div>
-                <div style={{ fontSize: "12px", color: "#666", marginTop: "2px" }}>{cat.description}</div>
-              </div>
-            </Option>
-          ))}
-        </Select>
-        <Button type="primary" onClick={handleAddCategory} disabled={!selectedCategoryName} icon={<PlusOutlined />}>
-          Add
-        </Button>
-      </div>
 
-      {/* Preview box - shown when category is selected but not yet added */}
-      {selectedCategoryName && (
-        <div
-          style={{
-            marginBottom: 16,
-            padding: "12px",
-            background: "#f9f9f9",
-            border: "1px solid #e0e0e0",
-            borderRadius: "4px",
-          }}
-        >
-          <div style={{ marginBottom: 8, fontWeight: 500, fontSize: "14px" }}>
-            Preview: {availableCategories.find((c) => c.name === selectedCategoryName)?.display_name}
-            {categoryFileTypes[selectedCategoryName] && (
-              <span style={{ marginLeft: 8, fontSize: "12px", color: "#888", fontWeight: 400 }}>
-                ({categoryFileTypes[selectedCategoryName]?.toUpperCase()})
-              </span>
+        {/* Preview box - shown when category is selected but not yet added */}
+        {selectedCategoryName && (
+          <div className="mb-4 rounded-md border border-border bg-muted/40 p-3">
+            <div className="mb-2 text-sm font-medium">
+              Preview: {availableCategories.find((c) => c.name === selectedCategoryName)?.display_name}
+              {categoryFileTypes[selectedCategoryName] && (
+                <span className="ml-2 text-xs font-normal text-muted-foreground">
+                  ({categoryFileTypes[selectedCategoryName]?.toUpperCase()})
+                </span>
+              )}
+            </div>
+            {loadingPreviewYaml ? (
+              <div className="p-4 text-center text-muted-foreground">Loading content...</div>
+            ) : previewYaml ? (
+              <pre className="m-0 max-h-[300px] max-w-full overflow-auto rounded-md border border-border bg-background p-3 text-xs leading-relaxed break-words whitespace-pre-wrap">
+                <code>{previewYaml}</code>
+              </pre>
+            ) : (
+              <div className="p-2 text-center text-xs text-muted-foreground">Unable to load category content</div>
             )}
           </div>
-          {loadingPreviewYaml ? (
-            <div style={{ padding: "16px", textAlign: "center", color: "#888" }}>Loading content...</div>
-          ) : previewYaml ? (
-            <pre
-              style={{
-                background: "#fff",
-                padding: "12px",
-                borderRadius: "4px",
-                overflow: "auto",
-                maxHeight: "300px",
-                maxWidth: "100%",
-                fontSize: "12px",
-                lineHeight: "1.5",
-                margin: 0,
-                border: "1px solid #e0e0e0",
-                whiteSpace: "pre-wrap",
-                wordBreak: "break-word",
-              }}
-            >
-              <code>{previewYaml}</code>
-            </pre>
-          ) : (
-            <div style={{ padding: "8px", textAlign: "center", color: "#888", fontSize: "12px" }}>
-              Unable to load category content
-            </div>
-          )}
-        </div>
-      )}
+        )}
 
-      {selectedCategories.length > 0 ? (
-        <>
-          <DataTable data={selectedCategories} columns={columns} getRowId={(row) => row.id} size="compact" />
-          <div style={{ marginTop: 16 }}>
-            <Collapse
-              activeKey={expandedYamlCategories}
-              onChange={(keys) => {
-                const keyArray = Array.isArray(keys) ? keys : keys ? [keys] : [];
-                const oldExpanded = new Set(expandedYamlCategories);
-
-                // Find newly expanded categories and fetch their YAML
-                keyArray.forEach((key) => {
-                  const categoryName = key as string;
-                  if (!oldExpanded.has(categoryName) && !categoryYaml[categoryName]) {
-                    fetchCategoryYaml(categoryName);
-                  }
-                });
-
-                setExpandedYamlCategories(keyArray as string[]);
-              }}
-              ghost
-              items={selectedCategories.map((category) => {
+        {selectedCategories.length > 0 ? (
+          <>
+            <DataTable data={selectedCategories} columns={columns} getRowId={(row) => row.id} size="compact" />
+            <div className="mt-4 space-y-2">
+              {selectedCategories.map((category) => {
                 const fileType = categoryFileTypes[category.category] || "yaml";
-                const fileTypeLabel = fileType.toUpperCase();
+                const isExpanded = expandedYamlCategories.includes(category.category);
 
-                return {
-                  key: category.category,
-                  label: (
-                    <div style={{ display: "flex", alignItems: "center", gap: 8 }}>
-                      <FileTextOutlined />
+                return (
+                  <Collapsible
+                    key={category.category}
+                    open={isExpanded}
+                    onOpenChange={(open) => {
+                      if (open && !categoryYaml[category.category]) {
+                        fetchCategoryYaml(category.category);
+                      }
+                      setExpandedYamlCategories((prev) =>
+                        open ? [...prev, category.category] : prev.filter((name) => name !== category.category),
+                      );
+                    }}
+                  >
+                    <CollapsibleTrigger className="flex items-center gap-2 text-sm">
+                      <ChevronRight className={`size-4 transition-transform ${isExpanded ? "rotate-90" : ""}`} />
+                      <FileText className="size-4" />
                       <span>
-                        View {fileTypeLabel} for {category.display_name}
+                        View {fileType.toUpperCase()} for {category.display_name}
                       </span>
-                    </div>
-                  ),
-                  children: loadingYaml[category.category] ? (
-                    <div style={{ padding: "16px", textAlign: "center", color: "#888" }}>Loading content...</div>
-                  ) : categoryYaml[category.category] ? (
-                    <pre
-                      style={{
-                        background: "#f5f5f5",
-                        padding: "16px",
-                        borderRadius: "4px",
-                        overflow: "auto",
-                        maxHeight: "400px",
-                        fontSize: "12px",
-                        lineHeight: "1.5",
-                        margin: 0,
-                      }}
-                    >
-                      <code>{categoryYaml[category.category]}</code>
-                    </pre>
-                  ) : (
-                    <div style={{ padding: "16px", textAlign: "center", color: "#888" }}>
-                      Content will load when expanded
-                    </div>
-                  ),
-                };
+                    </CollapsibleTrigger>
+                    <CollapsibleContent>
+                      {loadingYaml[category.category] ? (
+                        <div className="p-4 text-center text-muted-foreground">Loading content...</div>
+                      ) : categoryYaml[category.category] ? (
+                        <pre className="m-0 max-h-[400px] overflow-auto rounded-md bg-muted p-4 text-xs leading-relaxed">
+                          <code>{categoryYaml[category.category]}</code>
+                        </pre>
+                      ) : (
+                        <div className="p-4 text-center text-muted-foreground">Content will load when expanded</div>
+                      )}
+                    </CollapsibleContent>
+                  </Collapsible>
+                );
               })}
-            />
+            </div>
+          </>
+        ) : (
+          <div className="rounded-md border border-dashed border-border p-6 text-center text-muted-foreground">
+            No blocked topics selected. Add topics to detect and block harmful content.
           </div>
-        </>
-      ) : (
-        <div
-          style={{
-            textAlign: "center",
-            padding: "24px",
-            color: "#888",
-            border: "1px dashed #d9d9d9",
-            borderRadius: "4px",
-          }}
-        >
-          No blocked topics selected. Add topics to detect and block harmful content.
-        </div>
-      )}
+        )}
+      </CardContent>
     </Card>
   );
 };
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx
index 337520fac7d..a8c4ca3b2e7 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterConfiguration.tsx
@@ -1,8 +1,10 @@
-import React, { useState } from "react";
-import { Typography, Space, Upload, Card, Button } from "antd";
-import { PlusOutlined, UploadOutlined } from "@ant-design/icons";
+import React, { useRef, useState } from "react";
+import { Plus, Upload } from "lucide-react";
 import { validateBlockedWordsFile } from "@/components/networking";
 import NotificationsManager from "@/components/molecules/notifications_manager";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import PatternModal from "./PatternModal";
 import CustomPatternModal from "./CustomPatternModal";
 import KeywordModal from "./KeywordModal";
@@ -11,8 +13,6 @@ import KeywordTable from "./KeywordTable";
 import ContentCategoryConfiguration from "./ContentCategoryConfiguration";
 import CompetitorIntentConfiguration, { CompetitorIntentConfig } from "./CompetitorIntentConfiguration";
 
-const { Title, Text } = Typography;
-
 interface PrebuiltPattern {
   name: string;
   display_name: string;
@@ -115,6 +115,7 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
   const [newKeywordAction, setNewKeywordAction] = useState<"BLOCK" | "MASK">("BLOCK");
   const [newKeywordDescription, setNewKeywordDescription] = useState<string>("");
   const [uploadValidating, setUploadValidating] = useState(false);
+  const fileInputRef = useRef<HTMLInputElement>(null);
 
   const handleAddPrebuiltPattern = () => {
     if (!selectedPatternName) {
@@ -201,6 +202,14 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
     return false;
   };
 
+  const handleFileInputChange = (event: React.ChangeEvent<HTMLInputElement>) => {
+    const file = event.target.files?.[0];
+    event.target.value = "";
+    if (file) {
+      handleFileUpload(file);
+    }
+  };
+
   const showPatterns = !showStep || showStep === "patterns";
   const showKeywords = !showStep || showStep === "keywords";
   const showCategories = !showStep || showStep === "categories";
@@ -210,68 +219,78 @@ const ContentFilterConfiguration: React.FC<ContentFilterConfigurationProps> = ({
     <div className="space-y-6">
       {!showStep && (
         <div>
-          <Text type="secondary">
+          <p className="text-muted-foreground">
             Configure patterns, keywords, and content categories to detect and filter sensitive information in requests
             and responses.
-          </Text>
+          </p>
         </div>
       )}
 
       {showPatterns && (
-        <Card
-          title={
-            <div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
-              <Title level={5} style={{ margin: 0 }}>
-                Pattern Detection
-              </Title>
-              <Text type="secondary" style={{ fontSize: 14, fontWeight: 400 }}>
+        <Card>
+          <CardHeader>
+            <div className="flex flex-wrap items-center justify-between gap-2">
+              <CardTitle>Pattern Detection</CardTitle>
+              <p className="text-sm font-normal text-muted-foreground">
                 Detect sensitive information using regex patterns (SSN, credit cards, API keys, etc.)
-              </Text>
+              </p>
             </div>
-          }
-          size="small"
-        >
-          <div style={{ marginBottom: 16 }}>
-            <Space>
-              <Button type="primary" onClick={() => setPatternModalVisible(true)} icon={<PlusOutlined />}>
+          </CardHeader>
+          <CardContent>
+            <div className="mb-4 flex flex-wrap gap-2">
+              <Button onClick={() => setPatternModalVisible(true)}>
+                <Plus />
                 Add prebuilt pattern
               </Button>
-              <Button onClick={() => setCustomPatternModalVisible(true)} icon={<PlusOutlined />}>
+              <Button variant="outline" onClick={() => setCustomPatternModalVisible(true)}>
+                <Plus />
                 Add custom regex
               </Button>
-            </Space>
-          </div>
-          <PatternTable patterns={selectedPatterns} onActionChange={onPatternActionChange} onRemove={onPatternRemove} />
+            </div>
+            <PatternTable
+              patterns={selectedPatterns}
+              onActionChange={onPatternActionChange}
+              onRemove={onPatternRemove}
+            />
+          </CardContent>
         </Card>
       )}
 
       {showKeywords && (
-        <Card
-          title={
-            <div style={{ display: "flex", justifyContent: "space-between", alignItems: "center" }}>
-              <Title level={5} style={{ margin: 0 }}>
-                Blocked Keywords
-              </Title>
-              <Text type="secondary" style={{ fontSize: 14, fontWeight: 400 }}>
+        <Card>
+          <CardHeader>
+            <div className="flex flex-wrap items-center justify-between gap-2">
+              <CardTitle>Blocked Keywords</CardTitle>
+              <p className="text-sm font-normal text-muted-foreground">
                 Block or mask specific sensitive terms and phrases
-              </Text>
+              </p>
             </div>
-          }
-          size="small"
-        >
-          <div style={{ marginBottom: 16 }}>
-            <Space>
-              <Button type="primary" onClick={() => setKeywordModalVisible(true)} icon={<PlusOutlined />}>
+          </CardHeader>
+          <CardContent>
+            <div className="mb-4 flex flex-wrap gap-2">
+              <Button onClick={() => setKeywordModalVisible(true)}>
+                <Plus />
                 Add keyword
               </Button>
-              <Upload beforeUpload={handleFileUpload} accept=".yaml,.yml" showUploadList={false}>
-                <Button icon={<UploadOutlined />} loading={uploadValidating}>
-                  Upload YAML file
-                </Button>
-              </Upload>
-            </Space>
-          </div>
-          <KeywordTable keywords={blockedWords} onActionChange={onBlockedWordUpdate} onRemove={onBlockedWordRemove} />
+              <input
+                ref={fileInputRef}
+                type="file"
+                accept=".yaml,.yml"
+                className="hidden"
+                onChange={handleFileInputChange}
+              />
+              <Button
+                variant="outline"
+                disabled={uploadValidating}
+                aria-busy={uploadValidating}
+                onClick={() => fileInputRef.current?.click()}
+              >
+                {uploadValidating ? <UiLoadingSpinner className="size-4" /> : <Upload />}
+                Upload YAML file
+              </Button>
+            </div>
+            <KeywordTable keywords={blockedWords} onActionChange={onBlockedWordUpdate} onRemove={onBlockedWordRemove} />
+          </CardContent>
         </Card>
       )}
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx
index 3e7ee600943..d34f2920bb3 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterDisplay.tsx
@@ -1,5 +1,6 @@
 import React from "react";
-import { Card, Text, Badge } from "@tremor/react";
+import { Badge } from "@/components/ui/badge";
+import { Card, CardContent } from "@/components/ui/card";
 import PatternTable from "./PatternTable";
 import KeywordTable from "./KeywordTable";
 import CategoryTable from "./CategoryTable";
@@ -66,45 +67,51 @@ const ContentFilterDisplay: React.FC<ContentFilterDisplayProps> = ({
     <>
       {categories.length > 0 && (
         <Card className="mt-6">
-          <div className="flex justify-between items-center mb-4">
-            <Text className="text-lg font-semibold">Content Categories</Text>
-            <Badge color="blue">{categories.length} categories configured</Badge>
-          </div>
-          <CategoryTable
-            categories={categories}
-            onActionChange={readOnly ? undefined : onCategoryActionChange}
-            onSeverityChange={readOnly ? undefined : onCategorySeverityChange}
-            onRemove={readOnly ? undefined : onCategoryRemove}
-            readOnly={readOnly}
-          />
+          <CardContent>
+            <div className="mb-4 flex items-center justify-between">
+              <p className="text-lg font-semibold">Content Categories</p>
+              <Badge variant="secondary">{categories.length} categories configured</Badge>
+            </div>
+            <CategoryTable
+              categories={categories}
+              onActionChange={readOnly ? undefined : onCategoryActionChange}
+              onSeverityChange={readOnly ? undefined : onCategorySeverityChange}
+              onRemove={readOnly ? undefined : onCategoryRemove}
+              readOnly={readOnly}
+            />
+          </CardContent>
         </Card>
       )}
 
       {patterns.length > 0 && (
         <Card className="mt-6">
-          <div className="flex justify-between items-center mb-4">
-            <Text className="text-lg font-semibold">Pattern Detection</Text>
-            <Badge color="blue">{patterns.length} patterns configured</Badge>
-          </div>
-          <PatternTable
-            patterns={patterns}
-            onActionChange={readOnly ? noOp : onPatternActionChange || noOp}
-            onRemove={readOnly ? noOp : onPatternRemove || noOp}
-          />
+          <CardContent>
+            <div className="mb-4 flex items-center justify-between">
+              <p className="text-lg font-semibold">Pattern Detection</p>
+              <Badge variant="secondary">{patterns.length} patterns configured</Badge>
+            </div>
+            <PatternTable
+              patterns={patterns}
+              onActionChange={readOnly ? noOp : onPatternActionChange || noOp}
+              onRemove={readOnly ? noOp : onPatternRemove || noOp}
+            />
+          </CardContent>
         </Card>
       )}
 
       {blockedWords.length > 0 && (
         <Card className="mt-6">
-          <div className="flex justify-between items-center mb-4">
-            <Text className="text-lg font-semibold">Blocked Keywords</Text>
-            <Badge color="blue">{blockedWords.length} keywords configured</Badge>
-          </div>
-          <KeywordTable
-            keywords={blockedWords}
-            onActionChange={readOnly ? noOp : onBlockedWordUpdate || noOp}
-            onRemove={readOnly ? noOp : onBlockedWordRemove || noOp}
-          />
+          <CardContent>
+            <div className="mb-4 flex items-center justify-between">
+              <p className="text-lg font-semibold">Blocked Keywords</p>
+              <Badge variant="secondary">{blockedWords.length} keywords configured</Badge>
+            </div>
+            <KeywordTable
+              keywords={blockedWords}
+              onActionChange={readOnly ? noOp : onBlockedWordUpdate || noOp}
+              onRemove={readOnly ? noOp : onBlockedWordRemove || noOp}
+            />
+          </CardContent>
         </Card>
       )}
     </>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx
index 5f63cb3a2d4..b6d4f4b1970 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/ContentFilterManager.tsx
@@ -1,11 +1,11 @@
-import { Alert, Divider, Typography } from "antd";
+import { TriangleAlert } from "lucide-react";
 import React, { useEffect, useState } from "react";
+import { Alert, AlertDescription } from "@/components/shared/Alert";
+import { Separator } from "@/components/ui/separator";
 import ContentFilterConfiguration from "./ContentFilterConfiguration";
 import ContentFilterDisplay from "./ContentFilterDisplay";
 import type { CompetitorIntentConfig } from "./CompetitorIntentConfiguration";
 
-const { Text } = Typography;
-
 interface Pattern {
   id: string;
   type: "prebuilt" | "custom";
@@ -233,19 +233,17 @@ const ContentFilterManager: React.FC<ContentFilterManagerProps> = ({
   // Edit mode
   return (
     <>
-      <Divider orientation="left">Content Filter Configuration</Divider>
+      <div className="my-6 flex items-center gap-4">
+        <span className="shrink-0 font-medium">Content Filter Configuration</span>
+        <Separator className="flex-1" />
+      </div>
       {hasUnsavedChanges && (
-        <Alert
-          type="warning"
-          showIcon
-          className="mb-4"
-          message={
-            <Text>
-              You have unsaved changes to patterns or keywords. Remember to click &quot;Save Changes&quot; at the
-              bottom.
-            </Text>
-          }
-        />
+        <Alert variant="warning" className="mb-4">
+          <TriangleAlert />
+          <AlertDescription>
+            You have unsaved changes to patterns or keywords. Remember to click &quot;Save Changes&quot; at the bottom.
+          </AlertDescription>
+        </Alert>
       )}
       <div className="mb-6">
         {guardrailSettings && guardrailSettings.content_filter_settings && (
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx
index 6e2f8b07847..20bdddd3b74 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/CustomPatternModal.tsx
@@ -1,8 +1,10 @@
 import React from "react";
-import { Typography, Select, Modal, Space, Button, Input } from "antd";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Button } from "@/components/ui/button";
+import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS } from "./action_options";
+import { ABOVE_ANTD_MODAL } from "./dialog_layering";
 
 interface CustomPatternModalProps {
   visible: boolean;
@@ -28,50 +30,66 @@ const CustomPatternModal: React.FC<CustomPatternModalProps> = ({
   onCancel,
 }) => {
   return (
-    <Modal title="Add custom regex pattern" open={visible} onCancel={onCancel} footer={null} width={800}>
-      <Space direction="vertical" style={{ width: "100%" }} size="large">
-        <div>
-          <Text strong>Pattern name</Text>
-          <Input
-            placeholder="e.g., internal_id, employee_code"
-            value={patternName}
-            onChange={(e) => onNameChange(e.target.value)}
-            style={{ marginTop: 8 }}
-          />
+    <Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
+      <DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
+        <DialogHeader>
+          <DialogTitle>Add custom regex pattern</DialogTitle>
+        </DialogHeader>
+
+        <div className="space-y-6">
+          <div>
+            <p className="font-semibold">Pattern name</p>
+            <Input
+              className="mt-2"
+              placeholder="e.g., internal_id, employee_code"
+              value={patternName}
+              onChange={(e) => onNameChange(e.target.value)}
+            />
+          </div>
+
+          <div>
+            <p className="font-semibold">Regex pattern</p>
+            <Input
+              className="mt-2"
+              placeholder="e.g., ID-[0-9]{6}"
+              value={patternRegex}
+              onChange={(e) => onRegexChange(e.target.value)}
+            />
+            <p className="text-xs text-muted-foreground">Enter a valid regular expression to match sensitive data</p>
+          </div>
+
+          <div>
+            <p className="font-semibold">Action</p>
+            <p className="mt-1 mb-2 text-muted-foreground">
+              Choose what action the guardrail should take when this pattern is detected
+            </p>
+            <Select
+              items={ACTION_ITEMS}
+              value={patternAction}
+              onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
+            >
+              <SelectTrigger className="w-full" aria-label="Action">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                {ACTION_ITEMS.map((item) => (
+                  <SelectItem key={item.value} value={item.value}>
+                    {item.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
         </div>
 
-        <div>
-          <Text strong>Regex pattern</Text>
-          <Input
-            placeholder="e.g., ID-[0-9]{6}"
-            value={patternRegex}
-            onChange={(e) => onRegexChange(e.target.value)}
-            style={{ marginTop: 8 }}
-          />
-          <Text type="secondary" style={{ fontSize: 12 }}>
-            Enter a valid regular expression to match sensitive data
-          </Text>
-        </div>
-
-        <div>
-          <Text strong>Action</Text>
-          <Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
-            Choose what action the guardrail should take when this pattern is detected
-          </Text>
-          <Select value={patternAction} onChange={onActionChange} style={{ width: "100%" }}>
-            <Option value="BLOCK">Block</Option>
-            <Option value="MASK">Mask</Option>
-          </Select>
-        </div>
-      </Space>
-
-      <div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
-        <Button onClick={onCancel}>Cancel</Button>
-        <Button type="primary" onClick={onAdd}>
-          Add
-        </Button>
-      </div>
-    </Modal>
+        <DialogFooter>
+          <Button variant="outline" onClick={onCancel}>
+            Cancel
+          </Button>
+          <Button onClick={onAdd}>Add</Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx
index 56b721adda5..45627a2f101 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordModal.tsx
@@ -1,8 +1,11 @@
 import React from "react";
-import { Typography, Select, Modal, Space, Button, Input } from "antd";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Button } from "@/components/ui/button";
+import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Textarea } from "@/components/ui/textarea";
+import { ACTION_ITEMS } from "./action_options";
+import { ABOVE_ANTD_MODAL } from "./dialog_layering";
 
 interface KeywordModalProps {
   visible: boolean;
@@ -28,48 +31,66 @@ const KeywordModal: React.FC<KeywordModalProps> = ({
   onCancel,
 }) => {
   return (
-    <Modal title="Add blocked keyword" open={visible} onCancel={onCancel} footer={null} width={800}>
-      <Space direction="vertical" style={{ width: "100%" }} size="large">
-        <div>
-          <Text strong>Keyword</Text>
-          <Input
-            placeholder="Enter sensitive keyword or phrase"
-            value={keyword}
-            onChange={(e) => onKeywordChange(e.target.value)}
-            style={{ marginTop: 8 }}
-          />
+    <Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
+      <DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
+        <DialogHeader>
+          <DialogTitle>Add blocked keyword</DialogTitle>
+        </DialogHeader>
+
+        <div className="space-y-6">
+          <div>
+            <p className="font-semibold">Keyword</p>
+            <Input
+              className="mt-2"
+              placeholder="Enter sensitive keyword or phrase"
+              value={keyword}
+              onChange={(e) => onKeywordChange(e.target.value)}
+            />
+          </div>
+
+          <div>
+            <p className="font-semibold">Action</p>
+            <p className="mt-1 mb-2 text-muted-foreground">
+              Choose what action the guardrail should take when this keyword is detected
+            </p>
+            <Select
+              items={ACTION_ITEMS}
+              value={action}
+              onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
+            >
+              <SelectTrigger className="w-full" aria-label="Action">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                {ACTION_ITEMS.map((item) => (
+                  <SelectItem key={item.value} value={item.value}>
+                    {item.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
+
+          <div>
+            <p className="font-semibold">Description (optional)</p>
+            <Textarea
+              className="mt-2 field-sizing-fixed"
+              placeholder="Explain why this keyword is sensitive"
+              value={description}
+              onChange={(e) => onDescriptionChange(e.target.value)}
+              rows={3}
+            />
+          </div>
         </div>
 
-        <div>
-          <Text strong>Action</Text>
-          <Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
-            Choose what action the guardrail should take when this keyword is detected
-          </Text>
-          <Select value={action} onChange={onActionChange} style={{ width: "100%" }}>
-            <Option value="BLOCK">Block</Option>
-            <Option value="MASK">Mask</Option>
-          </Select>
-        </div>
-
-        <div>
-          <Text strong>Description (optional)</Text>
-          <Input.TextArea
-            placeholder="Explain why this keyword is sensitive"
-            value={description}
-            onChange={(e) => onDescriptionChange(e.target.value)}
-            rows={3}
-            style={{ marginTop: 8 }}
-          />
-        </div>
-      </Space>
-
-      <div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
-        <Button onClick={onCancel}>Cancel</Button>
-        <Button type="primary" onClick={onAdd}>
-          Add
-        </Button>
-      </div>
-    </Modal>
+        <DialogFooter>
+          <Button variant="outline" onClick={onCancel}>
+            Cancel
+          </Button>
+          <Button onClick={onAdd}>Add</Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx
index d16e157f819..5c7e3ef3ab8 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/KeywordTable.tsx
@@ -1,10 +1,10 @@
-import { DeleteOutlined } from "@ant-design/icons";
+import { Trash2 } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
-import { Button, Select } from "antd";
 import React from "react";
 import { DataTable } from "@/components/shared/DataTable";
-
-const { Option } = Select;
+import { Button } from "@/components/ui/button";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS } from "./action_options";
 
 interface BlockedWord {
   id: string;
@@ -31,13 +31,20 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
       size: 150,
       cell: ({ row }) => (
         <Select
+          items={ACTION_ITEMS}
           value={row.original.action}
-          onChange={(value) => onActionChange(row.original.id, "action", value)}
-          style={{ width: 120 }}
-          size="small"
+          onValueChange={(value: string | null) => value && onActionChange(row.original.id, "action", value)}
         >
-          <Option value="BLOCK">Block</Option>
-          <Option value="MASK">Mask</Option>
+          <SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
+            <SelectValue />
+          </SelectTrigger>
+          <SelectContent alignItemWithTrigger={false}>
+            {ACTION_ITEMS.map((item) => (
+              <SelectItem key={item.value} value={item.value}>
+                {item.label}
+              </SelectItem>
+            ))}
+          </SelectContent>
         </Select>
       ),
     },
@@ -51,7 +58,8 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
       id: "actions",
       size: 100,
       cell: ({ row }) => (
-        <Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove(row.original.id)}>
+        <Button variant="ghost" size="sm" onClick={() => onRemove(row.original.id)}>
+          <Trash2 />
           Delete
         </Button>
       ),
@@ -59,7 +67,7 @@ const KeywordTable: React.FC<KeywordTableProps> = ({ keywords, onActionChange, o
   ];
 
   if (keywords.length === 0) {
-    return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No keywords added.</div>;
+    return <div className="py-10 text-center text-muted-foreground">No keywords added.</div>;
   }
 
   return <DataTable data={keywords} columns={columns} getRowId={(row) => row.id} size="compact" />;
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
index 4311f4e13aa..77cbcf6211d 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
@@ -70,6 +70,20 @@ describe("PatternModal", () => {
     expect(mockOnPatternNameChange.mock.calls[0][0]).toBe("us_ssn");
   });
 
+  it("should narrow the pattern options to the typed search text", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
+
+    await user.click(screen.getAllByRole("combobox")[0]);
+    await user.keyboard("visa");
+
+    expect(await screen.findByText("Visa card")).toBeInTheDocument();
+    expect(screen.queryByText("Email address")).not.toBeInTheDocument();
+    expect(screen.queryByText("AWS access key")).not.toBeInTheDocument();
+  });
+
   it("should report the chosen action", async () => {
     const user = userEvent.setup();
     renderModal();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
index 32265fb3888..df87c511cf4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
@@ -1,8 +1,20 @@
 import React from "react";
-import { Typography, Select, Modal, Space, Button } from "antd";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Button } from "@/components/ui/button";
+import {
+  Combobox,
+  ComboboxCollection,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxGroup,
+  ComboboxInput,
+  ComboboxItem,
+  ComboboxLabel,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS } from "./action_options";
+import { ABOVE_ANTD_MODAL } from "./dialog_layering";
 
 interface PrebuiltPattern {
   name: string;
@@ -11,6 +23,11 @@ interface PrebuiltPattern {
   description: string;
 }
 
+interface PatternGroup {
+  category: string;
+  items: PrebuiltPattern[];
+}
+
 interface PatternModalProps {
   visible: boolean;
   prebuiltPatterns: PrebuiltPattern[];
@@ -34,64 +51,83 @@ const PatternModal: React.FC<PatternModalProps> = ({
   onAdd,
   onCancel,
 }) => {
+  const selectedPattern = prebuiltPatterns.find((pattern) => pattern.name === selectedPatternName) ?? null;
+  const patternGroups = categories
+    .map((category) => ({
+      category,
+      items: prebuiltPatterns.filter((pattern) => pattern.category === category),
+    }))
+    .filter((group) => group.items.length > 0);
+
   return (
-    <Modal title="Add prebuilt pattern" open={visible} onCancel={onCancel} footer={null} width={800}>
-      <Space direction="vertical" style={{ width: "100%" }} size="large">
-        <div>
-          <Text strong>Pattern type</Text>
-          <Select
-            placeholder="Choose pattern type"
-            value={selectedPatternName}
-            onChange={onPatternNameChange}
-            style={{ width: "100%", marginTop: 8 }}
-            showSearch
-            filterOption={(input, option) => {
-              const pattern = prebuiltPatterns.find((p) => p.name === option?.value);
-              if (pattern) {
-                return (
-                  pattern.display_name.toLowerCase().includes(input.toLowerCase()) ||
-                  pattern.name.toLowerCase().includes(input.toLowerCase())
-                );
-              }
-              return false;
-            }}
-          >
-            {categories.map((category) => {
-              const categoryPatterns = prebuiltPatterns.filter((p) => p.category === category);
-              if (categoryPatterns.length === 0) return null;
+    <Dialog open={visible} onOpenChange={(open) => !open && onCancel()}>
+      <DialogContent className={`max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[800px] ${ABOVE_ANTD_MODAL}`}>
+        <DialogHeader>
+          <DialogTitle>Add prebuilt pattern</DialogTitle>
+        </DialogHeader>
 
-              return (
-                <Select.OptGroup key={category} label={category}>
-                  {categoryPatterns.map((pattern) => (
-                    <Option key={pattern.name} value={pattern.name}>
-                      {pattern.display_name}
-                    </Option>
-                  ))}
-                </Select.OptGroup>
-              );
-            })}
-          </Select>
+        <div className="space-y-6">
+          <div>
+            <p className="font-semibold">Pattern type</p>
+            <Combobox
+              items={patternGroups}
+              value={selectedPattern}
+              onValueChange={(pattern: PrebuiltPattern | null) => pattern && onPatternNameChange(pattern.name)}
+              itemToStringLabel={(pattern: PrebuiltPattern) => pattern.display_name}
+            >
+              <ComboboxInput className="mt-2 w-full" placeholder="Choose pattern type" />
+              <ComboboxContent>
+                <ComboboxEmpty>No matching patterns</ComboboxEmpty>
+                <ComboboxList>
+                  {(group: PatternGroup) => (
+                    <ComboboxGroup key={group.category} items={group.items}>
+                      <ComboboxLabel>{group.category}</ComboboxLabel>
+                      <ComboboxCollection>
+                        {(pattern: PrebuiltPattern) => (
+                          <ComboboxItem key={pattern.name} value={pattern}>
+                            {pattern.display_name}
+                          </ComboboxItem>
+                        )}
+                      </ComboboxCollection>
+                    </ComboboxGroup>
+                  )}
+                </ComboboxList>
+              </ComboboxContent>
+            </Combobox>
+          </div>
+
+          <div>
+            <p className="font-semibold">Action</p>
+            <p className="mt-1 mb-2 text-muted-foreground">
+              Choose what action the guardrail should take when this pattern is detected
+            </p>
+            <Select
+              items={ACTION_ITEMS}
+              value={patternAction}
+              onValueChange={(value: string | null) => value && onActionChange(value as "BLOCK" | "MASK")}
+            >
+              <SelectTrigger className="w-full" aria-label="Action">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                {ACTION_ITEMS.map((item) => (
+                  <SelectItem key={item.value} value={item.value}>
+                    {item.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
         </div>
 
-        <div>
-          <Text strong>Action</Text>
-          <Text type="secondary" style={{ display: "block", marginTop: 4, marginBottom: 8 }}>
-            Choose what action the guardrail should take when this pattern is detected
-          </Text>
-          <Select value={patternAction} onChange={onActionChange} style={{ width: "100%" }}>
-            <Option value="BLOCK">Block</Option>
-            <Option value="MASK">Mask</Option>
-          </Select>
-        </div>
-      </Space>
-
-      <div style={{ display: "flex", justifyContent: "flex-end", gap: "8px", marginTop: "24px" }}>
-        <Button onClick={onCancel}>Cancel</Button>
-        <Button type="primary" onClick={onAdd}>
-          Add
-        </Button>
-      </div>
-    </Modal>
+        <DialogFooter>
+          <Button variant="outline" onClick={onCancel}>
+            Cancel
+          </Button>
+          <Button onClick={onAdd}>Add</Button>
+        </DialogFooter>
+      </DialogContent>
+    </Dialog>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx
index 403dbec25cc..6dd266f07a0 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternTable.tsx
@@ -1,11 +1,11 @@
 import React from "react";
-import { Typography, Select, Tag, Button } from "antd";
-import { DeleteOutlined } from "@ant-design/icons";
+import { Trash2 } from "lucide-react";
 import type { ColumnDef } from "@tanstack/react-table";
 import { DataTable } from "@/components/shared/DataTable";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { ACTION_ITEMS } from "./action_options";
 
 interface Pattern {
   id: string;
@@ -28,11 +28,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
       header: "Type",
       accessorKey: "type",
       size: 100,
-      cell: ({ row }) => (
-        <Tag color={row.original.type === "prebuilt" ? "blue" : "green"}>
-          {row.original.type === "prebuilt" ? "Prebuilt" : "Custom"}
-        </Tag>
-      ),
+      cell: ({ row }) => <Badge variant="secondary">{row.original.type === "prebuilt" ? "Prebuilt" : "Custom"}</Badge>,
     },
     {
       header: "Pattern name",
@@ -44,9 +40,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
       accessorKey: "pattern",
       cell: ({ row }) =>
         row.original.pattern ? (
-          <Text code style={{ fontSize: 12 }}>
-            {row.original.pattern.substring(0, 40)}...
-          </Text>
+          <code className="rounded-sm bg-muted px-1 py-0.5 text-xs">{row.original.pattern.substring(0, 40)}...</code>
         ) : (
           "-"
         ),
@@ -57,13 +51,20 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
       size: 150,
       cell: ({ row }) => (
         <Select
+          items={ACTION_ITEMS}
           value={row.original.action}
-          onChange={(value) => onActionChange(row.original.id, value as "BLOCK" | "MASK")}
-          style={{ width: 120 }}
-          size="small"
+          onValueChange={(value: string | null) => value && onActionChange(row.original.id, value as "BLOCK" | "MASK")}
         >
-          <Option value="BLOCK">Block</Option>
-          <Option value="MASK">Mask</Option>
+          <SelectTrigger size="sm" className="w-[120px]" aria-label="Action">
+            <SelectValue />
+          </SelectTrigger>
+          <SelectContent alignItemWithTrigger={false}>
+            {ACTION_ITEMS.map((item) => (
+              <SelectItem key={item.value} value={item.value}>
+                {item.label}
+              </SelectItem>
+            ))}
+          </SelectContent>
         </Select>
       ),
     },
@@ -72,7 +73,8 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
       id: "actions",
       size: 100,
       cell: ({ row }) => (
-        <Button type="text" danger size="small" icon={<DeleteOutlined />} onClick={() => onRemove(row.original.id)}>
+        <Button variant="ghost" size="sm" onClick={() => onRemove(row.original.id)}>
+          <Trash2 />
           Delete
         </Button>
       ),
@@ -80,7 +82,7 @@ const PatternTable: React.FC<PatternTableProps> = ({ patterns, onActionChange, o
   ];
 
   if (patterns.length === 0) {
-    return <div style={{ textAlign: "center", padding: "40px 0", color: "#999" }}>No patterns added.</div>;
+    return <div className="py-10 text-center text-muted-foreground">No patterns added.</div>;
   }
 
   return <DataTable data={patterns} columns={columns} getRowId={(row) => row.id} size="compact" />;
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/action_options.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/action_options.ts
new file mode 100644
index 00000000000..d959da32756
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/action_options.ts
@@ -0,0 +1,10 @@
+export const ACTION_ITEMS = [
+  { value: "BLOCK", label: "Block" },
+  { value: "MASK", label: "Mask" },
+] as const;
+
+export const SEVERITY_ITEMS = [
+  { value: "high", label: "High" },
+  { value: "medium", label: "Medium" },
+  { value: "low", label: "Low" },
+] as const;
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts
new file mode 100644
index 00000000000..f31ab788beb
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/dialog_layering.ts
@@ -0,0 +1 @@
+export const ABOVE_ANTD_MODAL = "z-[1100]";
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
index ff3432a0d82..865e72ffa86 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.test.tsx
@@ -123,6 +123,17 @@ describe("CustomCodeModal", () => {
     expect(await screen.findByDisplayValue(/SSN detected/)).toBeInTheDocument();
   });
 
+  it("should narrow the mode options to the typed search text", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    await user.click(screen.getAllByRole("combobox")[0]);
+    await user.keyboard("mcp");
+
+    expect(await screen.findByText("pre_mcp_call (Before MCP Tool Call)")).toBeInTheDocument();
+    expect(screen.queryByText("logging_only")).not.toBeInTheDocument();
+  });
+
   it("should expand the test section and run a test against the backend", async () => {
     const user = userEvent.setup();
     mockTest.mockResolvedValue({ success: true, result: { action: "allow" } } as never);
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx
index 79cd9e19e6f..f2d2ab724f4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx
@@ -1,21 +1,34 @@
 import React, { useState, useRef, useEffect } from "react";
-import { Modal, Select, Switch, Collapse, Input, Divider } from "antd";
-import { Button, TextInput } from "@tremor/react";
-import {
-  CodeOutlined,
-  PlayCircleOutlined,
-  CheckCircleOutlined,
-  CloseCircleOutlined,
-  CaretRightOutlined,
-  SaveOutlined,
-  UsergroupAddOutlined,
-  ExportOutlined,
-} from "@ant-design/icons";
+import { CheckCircle2, ChevronRight, Code, ExternalLink, PlayCircle, Save, Users, XCircle } from "lucide-react";
 import { createGuardrailCall, updateGuardrailCall, testCustomCodeGuardrail } from "@/components/networking";
 import NotificationsManager from "@/components/molecules/notifications_manager";
-
-const { Panel } = Collapse;
-const { TextArea } = Input;
+import { Button } from "@/components/ui/button";
+import { Collapsible, CollapsibleContent, CollapsibleTrigger } from "@/components/ui/collapsible";
+import {
+  Combobox,
+  ComboboxChip,
+  ComboboxChips,
+  ComboboxChipsInput,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxItem,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { Dialog, DialogContent, DialogDescription, DialogHeader, DialogTitle } from "@/components/ui/dialog";
+import { Input } from "@/components/ui/input";
+import {
+  Select,
+  SelectContent,
+  SelectGroup,
+  SelectItem,
+  SelectLabel,
+  SelectSeparator,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { Switch } from "@/components/ui/switch";
+import { Textarea } from "@/components/ui/textarea";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 
 // Code templates
 const CODE_TEMPLATES = {
@@ -144,6 +157,17 @@ const MODE_OPTIONS = [
   { value: "during_mcp_call", label: "during_mcp_call (During MCP Tool Call)" },
 ];
 
+const TEMPLATE_ITEMS = Object.entries(CODE_TEMPLATES).map(([key, template]) => ({
+  value: key,
+  label: template.name,
+}));
+
+type ModeOption = (typeof MODE_OPTIONS)[number];
+
+const MODE_OPTION_BY_VALUE: Record<string, ModeOption> = Object.fromEntries(
+  MODE_OPTIONS.map((option) => [option.value, option]),
+);
+
 // Data for editing an existing guardrail
 export interface EditGuardrailData {
   guardrail_id: string;
@@ -470,105 +494,104 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
   };
 
   const lineCount = code.split("\n").length;
+  const selectedModeOptions = mode.map((value) => MODE_OPTION_BY_VALUE[value]).filter(Boolean);
 
   return (
-    <Modal
-      open={visible}
-      onCancel={onClose}
-      footer={null}
-      width={1400}
-      className="custom-code-modal"
-      closable={true}
-      destroyOnClose
-    >
-      <div className="flex flex-col h-[80vh]">
-        {/* Header */}
-        <div className="pb-4 border-b border-gray-200">
-          <h2 className="text-xl font-semibold text-gray-900">
+    <Dialog open={visible} onOpenChange={(open) => !open && onClose()}>
+      <DialogContent className="max-h-[calc(100dvh-2rem)] overflow-y-auto sm:max-w-[1400px]">
+        <DialogHeader>
+          <DialogTitle className="text-xl font-semibold">
             {isEditMode ? "Edit Custom Guardrail" : "Create Custom Guardrail"}
-          </h2>
-          <p className="text-sm text-gray-500 mt-1">Define custom logic using Python-like syntax</p>
-        </div>
+          </DialogTitle>
+          <DialogDescription>Define custom logic using Python-like syntax</DialogDescription>
+        </DialogHeader>
 
         {/* Top Controls */}
-        <div className="flex items-center gap-4 py-4 border-b border-gray-100">
-          <div className="flex-1 max-w-[200px]">
-            <label className="block text-xs font-medium text-gray-600 mb-1">Guardrail Name</label>
-            <TextInput value={guardrailName} onValueChange={setGuardrailName} placeholder="e.g., block-pii-custom" />
-          </div>
-          <div className="w-[280px]">
-            <label className="block text-xs font-medium text-gray-600 mb-1">Mode (can select multiple)</label>
-            <Select
-              mode="multiple"
-              value={mode}
-              onChange={setMode}
-              options={MODE_OPTIONS}
-              className="w-full"
-              size="middle"
-              placeholder="Select modes"
+        <div className="flex items-center gap-4 border-b border-border py-4">
+          <div className="max-w-[200px] flex-1">
+            <label className="mb-1 block text-xs font-medium text-muted-foreground">Guardrail Name</label>
+            <Input
+              value={guardrailName}
+              onChange={(e) => setGuardrailName(e.target.value)}
+              placeholder="e.g., block-pii-custom"
             />
           </div>
-          <div className="w-[180px]">
-            <label className="block text-xs font-medium text-gray-600 mb-1">Template</label>
-            <Select
-              value={selectedTemplate}
-              onChange={handleTemplateChange}
-              className="w-full"
-              size="middle"
-              dropdownRender={(menu) => (
-                <>
-                  {menu}
-                  <Divider style={{ margin: "8px 0" }} />
-                  <div
-                    style={{
-                      padding: "8px 12px",
-                      cursor: "pointer",
-                      color: "#1890ff",
-                      fontSize: "12px",
-                      display: "flex",
-                      alignItems: "center",
-                      gap: "4px",
-                    }}
-                    onClick={(e) => {
-                      e.preventDefault();
-                      window.open("https://models.litellm.ai/guardrails", "_blank");
-                    }}
-                    onMouseEnter={(e) => {
-                      e.currentTarget.style.backgroundColor = "#f0f0f0";
-                    }}
-                    onMouseLeave={(e) => {
-                      e.currentTarget.style.backgroundColor = "transparent";
-                    }}
-                  >
-                    <UsergroupAddOutlined />
-                    <span>Browse Community templates</span>
-                    <ExportOutlined style={{ fontSize: "10px" }} />
-                  </div>
-                </>
-              )}
+          <div className="w-[280px]">
+            <label className="mb-1 block text-xs font-medium text-muted-foreground">Mode (can select multiple)</label>
+            <Combobox
+              items={MODE_OPTIONS}
+              value={selectedModeOptions}
+              onValueChange={(options: ModeOption[]) => setMode(options.map((option) => option.value))}
+              multiple
             >
-              <Select.OptGroup label="STANDARD">
-                {Object.entries(CODE_TEMPLATES).map(([key, template]) => (
-                  <Select.Option key={key} value={key}>
-                    {template.name}
-                  </Select.Option>
+              <ComboboxChips className="w-full">
+                {selectedModeOptions.map((option) => (
+                  <ComboboxChip key={option.value} aria-label={option.label}>
+                    {option.label}
+                  </ComboboxChip>
                 ))}
-              </Select.OptGroup>
+                <ComboboxChipsInput
+                  className="border-0 bg-transparent"
+                  placeholder={mode.length === 0 ? "Select modes" : undefined}
+                />
+              </ComboboxChips>
+              <ComboboxContent>
+                <ComboboxEmpty>No matching modes</ComboboxEmpty>
+                <ComboboxList>
+                  {(option: ModeOption) => (
+                    <ComboboxItem key={option.value} value={option}>
+                      {option.label}
+                    </ComboboxItem>
+                  )}
+                </ComboboxList>
+              </ComboboxContent>
+            </Combobox>
+          </div>
+          <div className="w-[180px]">
+            <label className="mb-1 block text-xs font-medium text-muted-foreground">Template</label>
+            <Select
+              items={TEMPLATE_ITEMS}
+              value={selectedTemplate}
+              onValueChange={(value: string | null) => value && handleTemplateChange(value)}
+            >
+              <SelectTrigger className="w-full" aria-label="Template">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                <SelectGroup>
+                  <SelectLabel>STANDARD</SelectLabel>
+                  {TEMPLATE_ITEMS.map((template) => (
+                    <SelectItem key={template.value} value={template.value}>
+                      {template.label}
+                    </SelectItem>
+                  ))}
+                </SelectGroup>
+                <SelectSeparator />
+                <button
+                  type="button"
+                  onClick={() => window.open("https://models.litellm.ai/guardrails", "_blank")}
+                  className="flex w-full items-center gap-1 rounded-sm px-2 py-1.5 text-xs text-primary hover:bg-accent"
+                >
+                  <Users className="size-3.5" />
+                  <span>Browse Community templates</span>
+                  <ExternalLink className="size-2.5" />
+                </button>
+              </SelectContent>
             </Select>
           </div>
           <div className="flex items-center gap-2 pt-5">
-            <span className="text-sm text-gray-600">Default On</span>
-            <Switch checked={defaultOn} onChange={setDefaultOn} />
+            <span className="text-sm text-muted-foreground">Default On</span>
+            <Switch checked={defaultOn} onCheckedChange={setDefaultOn} aria-label="Default On" />
           </div>
         </div>
 
         {/* Main Content */}
-        <div className="flex flex-1 overflow-hidden mt-4 gap-6">
+        <div className="mt-4 flex gap-6">
           {/* Code Editor */}
-          <div className="flex-2 flex flex-col min-w-0 overflow-y-auto">
-            <div className="flex items-center justify-between mb-2 shrink-0">
-              <span className="text-xs font-semibold text-gray-500 uppercase tracking-wide">Python Logic</span>
-              <span className="text-xs text-gray-400">Restricted environment (no imports)</span>
+          <div className="flex min-w-0 flex-1 flex-col">
+            <div className="mb-2 flex shrink-0 items-center justify-between">
+              <span className="text-xs font-semibold tracking-wide text-muted-foreground uppercase">Python Logic</span>
+              <span className="text-xs text-muted-foreground">Restricted environment (no imports)</span>
             </div>
             <div
               className="relative rounded-lg overflow-hidden border border-gray-700 bg-[#1e1e1e] shrink-0"
@@ -607,27 +630,23 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
             </div>
 
             {/* Test Section */}
-            <Collapse
-              activeKey={testExpanded ? ["test"] : []}
-              onChange={(keys) => setTestExpanded(keys.includes("test"))}
-              className="mt-3 bg-white border border-gray-200 rounded-lg shrink-0"
-              expandIcon={({ isActive }) => <CaretRightOutlined rotate={isActive ? 90 : 0} />}
+            <Collapsible
+              open={testExpanded}
+              onOpenChange={setTestExpanded}
+              className="mt-3 shrink-0 rounded-lg border border-border"
             >
-              <Panel
-                header={
-                  <span className="flex items-center gap-2 text-sm font-medium">
-                    <PlayCircleOutlined className="text-blue-500" />
-                    Test Your Guardrail
-                  </span>
-                }
-                key="test"
-              >
+              <CollapsibleTrigger className="flex w-full items-center gap-2 p-3 text-sm font-medium">
+                <ChevronRight className={`size-4 transition-transform ${testExpanded ? "rotate-90" : ""}`} />
+                <PlayCircle className="size-4 text-muted-foreground" />
+                Test Your Guardrail
+              </CollapsibleTrigger>
+              <CollapsibleContent className="p-3 pt-0">
                 <div className="space-y-3">
                   <div>
                     <div className="flex items-center justify-between mb-2">
-                      <label className="block text-xs font-medium text-gray-600">Test Input (JSON)</label>
+                      <label className="block text-xs font-medium text-muted-foreground">Test Input (JSON)</label>
                       <div className="flex items-center gap-2">
-                        <span className="text-xs text-gray-500">Load example:</span>
+                        <span className="text-xs text-muted-foreground">Load example:</span>
                         <button
                           type="button"
                           onClick={() => setTestInput(JSON.stringify(TEST_INPUT_EXAMPLES.pre_call.data, null, 2))}
@@ -651,7 +670,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
                         </button>
                       </div>
                     </div>
-                    <div className="mb-2 p-2 bg-gray-50 rounded-sm text-xs text-gray-600 border border-gray-200">
+                    <div className="mb-2 rounded-sm border border-border bg-muted/40 p-2 text-xs text-muted-foreground">
                       <div className="grid grid-cols-2 gap-x-4 gap-y-1">
                         <div>
                           <strong>texts</strong>: Message content (always)
@@ -676,16 +695,17 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
                         </div>
                       </div>
                     </div>
-                    <TextArea
+                    <Textarea
                       value={testInput}
                       onChange={(e) => setTestInput(e.target.value)}
                       rows={8}
-                      className="font-mono text-xs"
+                      className="font-mono text-xs field-sizing-fixed"
                       placeholder='{"texts": ["test message"], ...}'
                     />
                   </div>
                   <div className="flex items-center gap-3">
-                    <Button size="xs" onClick={handleTest} disabled={isTesting} icon={PlayCircleOutlined}>
+                    <Button size="sm" onClick={handleTest} disabled={isTesting} aria-busy={isTesting}>
+                      {isTesting ? <UiLoadingSpinner className="size-4" /> : <PlayCircle />}
                       {isTesting ? "Running..." : "Run Test"}
                     </Button>
                     {testResult && (
@@ -702,7 +722,7 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
                       >
                         {testResult.error ? (
                           <>
-                            <CloseCircleOutlined />
+                            <XCircle className="size-4" />
                             <span>
                               {testResult.error_type && <span className="font-medium">[{testResult.error_type}] </span>}
                               {testResult.error}
@@ -710,140 +730,117 @@ const CustomCodeModal: React.FC<CustomCodeModalProps> = ({ visible, onClose, onS
                           </>
                         ) : testResult.action === "allow" ? (
                           <>
-                            <CheckCircleOutlined /> Allowed
+                            <CheckCircle2 className="size-4" /> Allowed
                           </>
                         ) : testResult.action === "block" ? (
                           <>
-                            <CloseCircleOutlined /> Blocked: {testResult.reason}
+                            <XCircle className="size-4" /> Blocked: {testResult.reason}
                           </>
                         ) : testResult.action === "modify" ? (
                           <>
-                            <CheckCircleOutlined /> Modified
+                            <CheckCircle2 className="size-4" /> Modified
                             {testResult.texts && testResult.texts.length > 0 && (
-                              <span className="text-xs text-gray-500 ml-1">
-                                → {testResult.texts[0].substring(0, 50)}
+                              <span className="ml-1 text-xs text-muted-foreground">
+                                -&gt; {testResult.texts[0].substring(0, 50)}
                                 {testResult.texts[0].length > 50 ? "..." : ""}
                               </span>
                             )}
                           </>
                         ) : (
                           <>
-                            <CheckCircleOutlined /> {testResult.action || "Unknown"}
+                            <CheckCircle2 className="size-4" /> {testResult.action || "Unknown"}
                           </>
                         )}
                       </div>
                     )}
                   </div>
                 </div>
-              </Panel>
-            </Collapse>
+              </CollapsibleContent>
+            </Collapsible>
             {/* Contribution CTA Banner */}
-            <div className="mt-3 p-4 bg-linear-to-r from-blue-50 to-indigo-50 border border-blue-200 rounded-lg flex items-center justify-between shrink-0">
+            <div className="mt-3 flex shrink-0 items-center justify-between rounded-lg border border-blue-200 bg-linear-to-r from-blue-50 to-indigo-50 p-4">
               <div className="flex items-center gap-3">
-                <div className="bg-blue-100 rounded-full p-2">
-                  <UsergroupAddOutlined className="text-blue-600 text-lg" />
+                <div className="rounded-full bg-blue-100 p-2">
+                  <Users className="size-5 text-blue-600" />
                 </div>
                 <div>
-                  <div className="text-sm font-medium text-gray-900">Built a useful guardrail?</div>
-                  <div className="text-xs text-gray-600">Share it with the community and help others build faster</div>
+                  <div className="text-sm font-medium">Built a useful guardrail?</div>
+                  <div className="text-xs text-muted-foreground">
+                    Share it with the community and help others build faster
+                  </div>
                 </div>
               </div>
-              <Button
-                size="xs"
-                onClick={() => window.open("https://github.com/BerriAI/litellm-guardrails", "_blank")}
-                icon={ExportOutlined}
-                className="bg-blue-600 hover:bg-blue-700 text-white border-0"
-              >
+              <Button size="sm" onClick={() => window.open("https://github.com/BerriAI/litellm-guardrails", "_blank")}>
+                <ExternalLink />
                 Contribute Template
               </Button>
             </div>
           </div>
 
           {/* Primitives Panel */}
-          <div className="w-[300px] shrink-0 overflow-auto border-l border-gray-200 pl-6">
-            <div className="flex items-center gap-2 mb-3">
-              <CodeOutlined className="text-blue-500" />
-              <span className="font-semibold text-gray-700">Available Primitives</span>
+          <div className="w-[300px] shrink-0 overflow-auto border-l border-border pl-6">
+            <div className="mb-3 flex items-center gap-2">
+              <Code className="size-4 text-muted-foreground" />
+              <span className="font-semibold">Available Primitives</span>
             </div>
-            <p className="text-xs text-gray-500 mb-3">Click to copy functions to clipboard</p>
+            <p className="mb-3 text-xs text-muted-foreground">Click to copy functions to clipboard</p>
 
-            <Collapse
-              defaultActiveKey={["Return Values"]}
-              className="primitives-collapse bg-transparent border-0"
-              expandIconPosition="end"
-            >
+            <div className="space-y-2">
               {Object.entries(PRIMITIVES).map(([category, primitives]) => (
-                <Panel
-                  header={<span className="text-sm font-medium text-gray-700">{category}</span>}
+                <Collapsible
                   key={category}
-                  className="bg-white mb-2 rounded-lg border border-gray-200"
+                  defaultOpen={category === "Return Values"}
+                  className="rounded-lg border border-border"
                 >
-                  <div className="space-y-2">
-                    {primitives.map((p) => (
-                      <button
-                        key={p.name}
-                        onClick={() => copyPrimitive(p.name)}
-                        className={`w-full text-left px-2 py-2 rounded transition-colors ${
-                          copiedPrimitive === p.name ? "bg-green-100" : "bg-gray-50 hover:bg-blue-50"
-                        }`}
-                      >
-                        {copiedPrimitive === p.name ? (
-                          <span className="flex items-center gap-1 text-xs font-mono text-green-700">
-                            <CheckCircleOutlined /> Copied!
-                          </span>
-                        ) : (
-                          <>
-                            <div className="text-xs font-mono text-gray-800">{p.name}</div>
-                            <div className="text-[10px] text-gray-500 mt-0.5">{p.desc}</div>
-                          </>
-                        )}
-                      </button>
-                    ))}
-                  </div>
-                </Panel>
+                  <CollapsibleTrigger className="group flex w-full items-center justify-between px-3 py-2 text-sm font-medium">
+                    {category}
+                    <ChevronRight className="size-4 transition-transform group-data-panel-open:rotate-90" />
+                  </CollapsibleTrigger>
+                  <CollapsibleContent className="px-3 pb-3">
+                    <div className="space-y-2">
+                      {primitives.map((p) => (
+                        <button
+                          key={p.name}
+                          onClick={() => copyPrimitive(p.name)}
+                          className={`w-full rounded-sm px-2 py-2 text-left transition-colors ${
+                            copiedPrimitive === p.name ? "bg-accent" : "bg-muted/40 hover:bg-accent"
+                          }`}
+                        >
+                          {copiedPrimitive === p.name ? (
+                            <span className="flex items-center gap-1 font-mono text-xs">
+                              <CheckCircle2 className="size-3.5" /> Copied!
+                            </span>
+                          ) : (
+                            <>
+                              <div className="font-mono text-xs">{p.name}</div>
+                              <div className="mt-0.5 text-[10px] text-muted-foreground">{p.desc}</div>
+                            </>
+                          )}
+                        </button>
+                      ))}
+                    </div>
+                  </CollapsibleContent>
+                </Collapsible>
               ))}
-            </Collapse>
+            </div>
           </div>
         </div>
 
         {/* Footer */}
-        <div className="flex items-center justify-between pt-4 mt-4 border-t border-gray-200">
-          <span className="text-xs text-gray-400">Changes are auto-saved to local draft</span>
+        <div className="mt-4 flex items-center justify-between border-t border-border pt-4">
+          <span className="text-xs text-muted-foreground">Changes are auto-saved to local draft</span>
           <div className="flex items-center gap-3">
             <Button variant="secondary" onClick={onClose}>
               Cancel
             </Button>
-            <Button
-              onClick={handleSave}
-              loading={isSaving}
-              disabled={isSaving || !guardrailName.trim()}
-              icon={SaveOutlined}
-            >
+            <Button onClick={handleSave} disabled={isSaving || !guardrailName.trim()} aria-busy={isSaving}>
+              {isSaving ? <UiLoadingSpinner className="size-4" /> : <Save />}
               {isEditMode ? "Update Guardrail" : "Save Guardrail"}
             </Button>
           </div>
         </div>
-      </div>
-
-      <style>{`
-        .custom-code-modal .ant-modal-content {
-          padding: 24px;
-        }
-        .custom-code-modal .ant-modal-close {
-          top: 20px;
-          right: 20px;
-        }
-        .primitives-collapse .ant-collapse-item {
-          border: none !important;
-        }
-        .primitives-collapse .ant-collapse-header {
-          padding: 8px 12px !important;
-        }
-        .primitives-collapse .ant-collapse-content-box {
-          padding: 8px 12px !important;
-        }
-      `}</style>
-    </Modal>
+      </DialogContent>
+    </Dialog>
   );
 };
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx
index de7c29dc31c..a09ce9435bb 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden.tsx
@@ -1,6 +1,6 @@
 import React, { useState } from "react";
-import { Input } from "antd";
-import { SearchOutlined, ArrowRightOutlined } from "@ant-design/icons";
+import { ArrowRight, Search } from "lucide-react";
+import { InputGroup, InputGroupAddon, InputGroupInput } from "@/components/ui/input-group";
 import { GuardrailCardInfo, ALL_CARDS } from "./guardrail_garden_data";
 import GuardrailCard from "./guardrail_garden_card";
 import GuardrailDetailView from "./guardrail_garden_detail";
@@ -43,72 +43,52 @@ const GuardrailGarden: React.FC<GuardrailGardenProps> = ({ accessToken, onGuardr
 
   return (
     <div>
-      {/* Search Bar */}
-      <div style={{ marginBottom: 24 }}>
-        <Input
-          size="large"
-          placeholder="Search guardrails"
-          prefix={<SearchOutlined style={{ color: "#9ca3af" }} />}
-          value={searchQuery}
-          onChange={(e) => setSearchQuery(e.target.value)}
-          style={{ borderRadius: 8 }}
-        />
+      <div className="mb-6">
+        <InputGroup>
+          <InputGroupAddon>
+            <Search className="size-4 text-muted-foreground" />
+          </InputGroupAddon>
+          <InputGroupInput
+            placeholder="Search guardrails"
+            value={searchQuery}
+            onChange={(e) => setSearchQuery(e.target.value)}
+          />
+        </InputGroup>
       </div>
 
-      {/* LiteLLM Content Filter Section */}
-      <div style={{ marginBottom: 40 }}>
-        <div style={{ display: "flex", alignItems: "center", justifyContent: "space-between", marginBottom: 4 }}>
-          <h2 style={{ fontSize: 20, fontWeight: 600, color: "#111827", margin: 0 }}>LiteLLM Content Filter</h2>
+      <div className="mb-10">
+        <div className="mb-1 flex items-center justify-between">
+          <h2 className="m-0 text-xl font-semibold text-foreground">LiteLLM Content Filter</h2>
           <span
-            style={{
-              display: "inline-flex",
-              alignItems: "center",
-              gap: 6,
-              fontSize: 14,
-              color: "#1a73e8",
-              cursor: "pointer",
-            }}
+            className="inline-flex cursor-pointer items-center gap-1.5 text-sm text-primary"
             onClick={() => setShowAllLitellm(!showAllLitellm)}
           >
             {showAllLitellm ? (
               <>Show less</>
             ) : (
               <>
-                <ArrowRightOutlined style={{ fontSize: 12 }} />
+                <ArrowRight className="size-3" />
                 {`Show all (${litellmCards.length})`}
               </>
             )}
           </span>
         </div>
-        <p style={{ fontSize: 13, color: "#6b7280", margin: "4px 0 20px 0" }}>
+        <p className="mt-1 mb-5 text-[13px] text-muted-foreground">
           Built-in guardrails powered by LiteLLM. Zero latency, no external dependencies, no additional cost.
         </p>
-        <div
-          style={{
-            display: "grid",
-            gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))",
-            gap: 16,
-          }}
-        >
+        <div className="grid grid-cols-[repeat(auto-fill,minmax(220px,1fr))] gap-4">
           {(showAllLitellm ? litellmCards : litellmCards.slice(0, CARDS_PER_ROW * VISIBLE_ROWS)).map((card) => (
             <GuardrailCard key={card.id} card={card} onClick={() => setSelectedCard(card)} />
           ))}
         </div>
       </div>
 
-      {/* Partner Guardrails Section */}
-      <div style={{ marginBottom: 40 }}>
-        <h2 style={{ fontSize: 20, fontWeight: 600, color: "#111827", margin: "0 0 4px 0" }}>Partner Guardrails</h2>
-        <p style={{ fontSize: 13, color: "#6b7280", margin: "4px 0 20px 0" }}>
+      <div className="mb-10">
+        <h2 className="mt-0 mb-1 text-xl font-semibold text-foreground">Partner Guardrails</h2>
+        <p className="mt-1 mb-5 text-[13px] text-muted-foreground">
           Third-party guardrail integrations from leading AI security providers.
         </p>
-        <div
-          style={{
-            display: "grid",
-            gridTemplateColumns: "repeat(auto-fill, minmax(220px, 1fr))",
-            gap: 16,
-          }}
-        >
+        <div className="grid grid-cols-[repeat(auto-fill,minmax(220px,1fr))] gap-4">
           {partnerCards.map((card) => (
             <GuardrailCard key={card.id} card={card} onClick={() => setSelectedCard(card)} />
           ))}
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx
index 53abf3eb81c..357c714fc76 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_card.tsx
@@ -1,45 +1,25 @@
-import React, { useState } from "react";
-import { CheckCircleFilled } from "@ant-design/icons";
+import React from "react";
+import { CircleCheck } from "lucide-react";
 import { GuardrailCardInfo } from "./guardrail_garden_data";
 import { Logo } from "@/components/molecules/logo/Logo";
 
 const GuardrailCard: React.FC<{ card: GuardrailCardInfo; onClick: () => void }> = ({ card, onClick }) => {
-  const [hovered, setHovered] = useState(false);
-
   return (
     <div
       onClick={onClick}
-      onMouseEnter={() => setHovered(true)}
-      onMouseLeave={() => setHovered(false)}
-      style={{
-        borderRadius: 12,
-        border: hovered ? "1px solid #93c5fd" : "1px solid #e5e7eb",
-        backgroundColor: "#ffffff",
-        padding: "20px 20px 16px 20px",
-        cursor: "pointer",
-        transition: "border-color 0.15s, box-shadow 0.15s",
-        display: "flex",
-        flexDirection: "column",
-        minHeight: 170,
-        boxShadow: hovered ? "0 1px 6px rgba(59,130,246,0.08)" : "none",
-      }}
+      className="flex min-h-[170px] cursor-pointer flex-col rounded-xl border border-border bg-card px-5 pt-5 pb-4 transition-[border-color,box-shadow] hover:border-primary/40 hover:shadow-sm"
     >
-      {/* Icon + Name row */}
-      <div style={{ display: "flex", alignItems: "center", gap: 10, marginBottom: 10 }}>
+      <div className="mb-2.5 flex items-center gap-2.5">
         <Logo src={card.logo} label={card.name} className="w-7 h-7 rounded-md object-contain shrink-0" />
-        <span style={{ fontSize: 14, fontWeight: 600, color: "#111827", lineHeight: 1.3 }}>{card.name}</span>
+        <span className="text-sm leading-tight font-semibold text-foreground">{card.name}</span>
       </div>
 
-      {/* Description */}
-      <p className="line-clamp-3" style={{ fontSize: 12, color: "#6b7280", lineHeight: 1.6, margin: 0, flex: 1 }}>
-        {card.description}
-      </p>
+      <p className="line-clamp-3 m-0 flex-1 text-xs leading-relaxed text-muted-foreground">{card.description}</p>
 
-      {/* Eval badge */}
       {card.eval && (
-        <div style={{ marginTop: 10, display: "flex", alignItems: "center", gap: 4 }}>
-          <CheckCircleFilled style={{ color: "#16a34a", fontSize: 12 }} />
-          <span style={{ fontSize: 11, color: "#16a34a", fontWeight: 500 }}>
+        <div className="mt-2.5 flex items-center gap-1 text-emerald-600">
+          <CircleCheck className="size-3" />
+          <span className="text-[11px] font-medium">
             F1: {card.eval.f1}% &middot; {card.eval.testCases} test cases
           </span>
         </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx
index 71c7a527614..0a0f70b6bc2 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/guardrail_garden_detail.tsx
@@ -1,6 +1,6 @@
 import React, { useState } from "react";
-import { Button } from "antd";
-import { ArrowLeftOutlined } from "@ant-design/icons";
+import { ArrowLeft } from "lucide-react";
+import { Button } from "@/components/ui/button";
 import AddGuardrailForm from "./add_guardrail_form";
 import { Logo } from "@/components/molecules/logo/Logo";
 import { GUARDRAIL_PRESETS } from "./guardrail_garden_configs";
@@ -44,17 +44,9 @@ const GuardrailDetailView: React.FC<GuardrailDetailViewProps> = ({ card, onBack,
       {/* Back link */}
       <div
         onClick={onBack}
-        style={{
-          display: "inline-flex",
-          alignItems: "center",
-          gap: 6,
-          color: "#5f6368",
-          cursor: "pointer",
-          fontSize: 14,
-          marginBottom: 24,
-        }}
+        className="mb-6 inline-flex cursor-pointer items-center gap-1.5 text-sm text-muted-foreground"
       >
-        <ArrowLeftOutlined style={{ fontSize: 11 }} />
+        <ArrowLeft className="size-3" />
         <span>{card.name}</span>
       </div>
 
@@ -67,19 +59,8 @@ const GuardrailDetailView: React.FC<GuardrailDetailViewProps> = ({ card, onBack,
       <p style={{ fontSize: 14, color: "#5f6368", margin: "0 0 20px 0", lineHeight: 1.6 }}>{card.description}</p>
 
       {/* Action buttons — outlined style like Vertex */}
-      <div style={{ display: "flex", gap: 10, marginBottom: 32 }}>
-        <Button
-          onClick={() => setIsAddFormVisible(true)}
-          style={{
-            borderRadius: 20,
-            padding: "4px 20px",
-            height: 36,
-            borderColor: "#dadce0",
-            color: "#1a73e8",
-            fontWeight: 500,
-            fontSize: 14,
-          }}
-        >
+      <div className="mb-8 flex gap-2.5">
+        <Button variant="outline" className="rounded-full" onClick={() => setIsAddFormVisible(true)}>
           Create Guardrail
         </Button>
       </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx
index 3043832ddc4..da5e61b2f3d 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_components.tsx
@@ -1,10 +1,21 @@
 import React from "react";
-import { Typography, Select, Button, Checkbox, Tooltip, Tag } from "antd";
-import { CloseOutlined, EyeInvisibleOutlined, StopOutlined, FilterOutlined } from "@ant-design/icons";
+import { EyeOff, Filter, Info, Ban, X } from "lucide-react";
 import { PiiEntityCategory } from "@/components/guardrails/types";
-
-const { Text } = Typography;
-const { Option } = Select;
+import { Badge } from "@/components/ui/badge";
+import { Button } from "@/components/ui/button";
+import { Checkbox } from "@/components/ui/checkbox";
+import {
+  Combobox,
+  ComboboxChip,
+  ComboboxChips,
+  ComboboxChipsInput,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxItem,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 
 // Helper functions
 export const formatEntityName = (name: string) => {
@@ -14,9 +25,9 @@ export const formatEntityName = (name: string) => {
 export const getActionIcon = (action: string) => {
   switch (action) {
     case "MASK":
-      return <EyeInvisibleOutlined style={{ marginRight: 4 }} />;
+      return <EyeOff className="mr-1 size-3.5" />;
     case "BLOCK":
-      return <StopOutlined style={{ marginRight: 4 }} />;
+      return <Ban className="mr-1 size-3.5" />;
     default:
       return null;
   }
@@ -30,34 +41,37 @@ export interface CategoryFilterProps {
 }
 
 export const CategoryFilter: React.FC<CategoryFilterProps> = ({ categories, selectedCategories, onChange }) => {
+  const categoryNames = categories.map((cat) => cat.category);
+
   return (
     <div>
-      <div className="flex items-center mb-2">
-        <FilterOutlined className="text-gray-500 mr-1" />
-        <Text className="text-gray-500 font-medium">Filter by category</Text>
+      <div className="mb-2 flex items-center">
+        <Filter className="mr-1 size-4 text-muted-foreground" />
+        <span className="font-medium text-muted-foreground">Filter by category</span>
       </div>
-      <Select
-        mode="multiple"
-        placeholder="Select categories to filter by"
-        style={{ width: "100%" }}
-        onChange={onChange}
-        value={selectedCategories}
-        allowClear
-        showSearch
-        optionFilterProp="children"
-        className="mb-4"
-        tagRender={(props) => (
-          <Tag color="blue" closable={props.closable} onClose={props.onClose} className="mr-2 mb-2">
-            {props.label}
-          </Tag>
-        )}
-      >
-        {categories.map((cat) => (
-          <Option key={cat.category} value={cat.category}>
-            {cat.category}
-          </Option>
-        ))}
-      </Select>
+      <Combobox items={categoryNames} value={selectedCategories} onValueChange={onChange} multiple>
+        <ComboboxChips className="mb-4 w-full">
+          {selectedCategories.map((category) => (
+            <ComboboxChip key={category} aria-label={category}>
+              {category}
+            </ComboboxChip>
+          ))}
+          <ComboboxChipsInput
+            className="border-0 bg-transparent"
+            placeholder={selectedCategories.length === 0 ? "Select categories to filter by" : undefined}
+          />
+        </ComboboxChips>
+        <ComboboxContent>
+          <ComboboxEmpty>No matching categories</ComboboxEmpty>
+          <ComboboxList>
+            {(category: string) => (
+              <ComboboxItem key={category} value={category}>
+                {category}
+              </ComboboxItem>
+            )}
+          </ComboboxList>
+        </ComboboxContent>
+      </Combobox>
     </div>
   );
 };
@@ -71,46 +85,34 @@ export interface QuickActionsProps {
 
 export const QuickActions: React.FC<QuickActionsProps> = ({ onSelectAll, onUnselectAll, hasSelectedEntities }) => {
   return (
-    <div className="bg-gray-50 p-5 rounded-lg mb-6 border border-gray-200 shadow-xs">
-      <div className="flex items-center justify-between mb-4">
+    <div className="mb-6 rounded-lg border border-border bg-muted/40 p-5 shadow-xs">
+      <div className="mb-4 flex items-center justify-between">
         <div className="flex items-center">
-          <Text strong className="text-gray-700 text-base">
-            Quick Actions
-          </Text>
-          <Tooltip title="Apply action to all PII types at once">
-            <div className="ml-2 text-gray-400 cursor-help text-xs">ⓘ</div>
+          <span className="text-base font-semibold">Quick Actions</span>
+          <Tooltip>
+            <TooltipTrigger
+              render={
+                <span className="ml-2 cursor-help text-muted-foreground">
+                  <Info className="size-3.5" />
+                </span>
+              }
+            />
+            <TooltipContent>Apply action to all PII types at once</TooltipContent>
           </Tooltip>
         </div>
-        <Button
-          color="danger"
-          variant="outlined"
-          onClick={onUnselectAll}
-          disabled={!hasSelectedEntities}
-          icon={<CloseOutlined />}
-        >
+        <Button variant="outline" onClick={onUnselectAll} disabled={!hasSelectedEntities}>
+          <X />
           Unselect All
         </Button>
       </div>
       <div className="grid grid-cols-2 gap-4">
-        <Button
-          color="primary"
-          variant="outlined"
-          onClick={() => onSelectAll("MASK")}
-          className="h-10"
-          block
-          icon={<EyeInvisibleOutlined />}
-        >
-          Select All & Mask
+        <Button variant="outline" className="h-10 w-full" onClick={() => onSelectAll("MASK")}>
+          <EyeOff />
+          Select All &amp; Mask
         </Button>
-        <Button
-          color="danger"
-          variant="outlined"
-          onClick={() => onSelectAll("BLOCK")}
-          className="h-10 hover:bg-red-100"
-          block
-          icon={<StopOutlined />}
-        >
-          Select All & Block
+        <Button variant="outline" className="h-10 w-full" onClick={() => onSelectAll("BLOCK")}>
+          <Ban />
+          Select All &amp; Block
         </Button>
       </div>
     </div>
@@ -138,62 +140,59 @@ export const PiiEntityList: React.FC<PiiEntityListProps> = ({
   entityToCategoryMap,
 }) => {
   return (
-    <div className="border rounded-lg overflow-hidden shadow-xs">
-      <div className="bg-gray-50 px-5 py-3 border-b flex">
-        <Text strong className="flex-1 text-gray-700">
-          PII Type
-        </Text>
-        <Text strong className="w-32 text-right text-gray-700">
-          Action
-        </Text>
+    <div className="overflow-hidden rounded-lg border border-border shadow-xs">
+      <div className="flex border-b border-border bg-muted/40 px-5 py-3">
+        <span className="flex-1 font-semibold">PII Type</span>
+        <span className="w-32 text-right font-semibold">Action</span>
       </div>
       <div className="max-h-[400px] overflow-y-auto">
         {entities.length === 0 ? (
-          <div className="py-10 text-center text-gray-500">No PII types match your filter criteria</div>
+          <div className="py-10 text-center text-muted-foreground">No PII types match your filter criteria</div>
         ) : (
-          entities.map((entity) => (
-            <div
-              key={entity}
-              className={`px-5 py-3 flex items-center justify-between hover:bg-gray-50 border-b ${
-                selectedEntities.includes(entity) ? "bg-blue-50" : ""
-              }`}
-            >
-              <div className="flex items-center flex-1">
-                <Checkbox
-                  checked={selectedEntities.includes(entity)}
-                  onChange={() => onEntitySelect(entity)}
-                  className="mr-3"
-                />
-                <Text className={selectedEntities.includes(entity) ? "font-medium text-gray-900" : "text-gray-700"}>
-                  {formatEntityName(entity)}
-                </Text>
-                {entityToCategoryMap.get(entity) && (
-                  <Tag className="ml-2 text-xs" color="blue">
-                    {entityToCategoryMap.get(entity)}
-                  </Tag>
-                )}
+          entities.map((entity) => {
+            const isSelected = selectedEntities.includes(entity);
+            return (
+              <div
+                key={entity}
+                className={`flex items-center justify-between border-b border-border px-5 py-3 hover:bg-muted/40 ${
+                  isSelected ? "bg-accent" : ""
+                }`}
+              >
+                <div className="flex flex-1 items-center">
+                  <Checkbox className="mr-3" checked={isSelected} onCheckedChange={() => onEntitySelect(entity)} />
+                  <span className={isSelected ? "font-medium text-foreground" : "text-muted-foreground"}>
+                    {formatEntityName(entity)}
+                  </span>
+                  {entityToCategoryMap.get(entity) && (
+                    <Badge variant="secondary" className="ml-2">
+                      {entityToCategoryMap.get(entity)}
+                    </Badge>
+                  )}
+                </div>
+                <div className="w-32">
+                  <Select
+                    value={isSelected ? selectedActions[entity] || "MASK" : "MASK"}
+                    onValueChange={(value: string | null) => value && onActionSelect(entity, value)}
+                    disabled={!isSelected}
+                  >
+                    <SelectTrigger className={`w-[120px] ${isSelected ? "" : "opacity-50"}`} aria-label="Action">
+                      <SelectValue />
+                    </SelectTrigger>
+                    <SelectContent alignItemWithTrigger={false}>
+                      {actions.map((action) => (
+                        <SelectItem key={action} value={action}>
+                          <span className="flex items-center">
+                            {getActionIcon(action)}
+                            {action}
+                          </span>
+                        </SelectItem>
+                      ))}
+                    </SelectContent>
+                  </Select>
+                </div>
               </div>
-              <div className="w-32">
-                <Select
-                  value={selectedEntities.includes(entity) ? selectedActions[entity] || "MASK" : "MASK"}
-                  onChange={(value) => onActionSelect(entity, value)}
-                  style={{ width: 120 }}
-                  disabled={!selectedEntities.includes(entity)}
-                  className={`${!selectedEntities.includes(entity) ? "opacity-50" : ""}`}
-                  dropdownMatchSelectWidth={false}
-                >
-                  {actions.map((action) => (
-                    <Option key={action} value={action}>
-                      <div className="flex items-center">
-                        {getActionIcon(action)}
-                        {action}
-                      </div>
-                    </Option>
-                  ))}
-                </Select>
-              </div>
-            </div>
-          ))
+            );
+          })
         )}
       </div>
     </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
index 14408b121ee..cba29ed615d 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/pii_configuration.tsx
@@ -1,10 +1,7 @@
-import { Typography } from "antd";
 import React, { useState } from "react";
 import { CategoryFilter, PiiEntityList, QuickActions } from "./pii_components";
 import { PiiConfigurationProps } from "@/components/guardrails/types";
 
-const { Title, Text } = Typography;
-
 /**
  * A reusable component for rendering PII entity selection and action configuration
  * Used in both add and edit guardrail forms
@@ -57,11 +54,9 @@ const PiiConfiguration: React.FC<PiiConfigurationProps> = ({
     <div className="pii-configuration">
       <div className="flex justify-between items-center mb-5">
         <div className="flex items-center">
-          <Title level={4} className="m-0! font-semibold text-gray-800">
-            Configure PII Protection
-          </Title>
+          <h4 className="m-0 text-lg font-semibold text-foreground">Configure PII Protection</h4>
         </div>
-        <Text className="text-gray-500">{selectedEntities.length} items selected</Text>
+        <span className="text-muted-foreground">{selectedEntities.length} items selected</span>
       </div>
 
       <div className="mb-6">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx
index 7083e32f73e..8fe2bf5bf21 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/tool_permission/ToolPermissionRulesEditor.tsx
@@ -1,7 +1,12 @@
 import React from "react";
-import { Card, Text } from "@tremor/react";
-import { Button, Divider, Empty, Input, Select, Space, Tooltip } from "antd";
-import { InfoCircleOutlined, PlusOutlined, DeleteOutlined } from "@ant-design/icons";
+import { Info, Plus, Trash2 } from "lucide-react";
+import { Button } from "@/components/ui/button";
+import { Card, CardContent } from "@/components/ui/card";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Separator } from "@/components/ui/separator";
+import { Textarea } from "@/components/ui/textarea";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 
 export type ToolPermissionDecision = "allow" | "deny";
 export type ToolPermissionDefaultAction = "allow" | "deny";
@@ -28,6 +33,16 @@ interface ToolPermissionRulesEditorProps {
   disabled?: boolean;
 }
 
+const DECISION_ITEMS = [
+  { value: "allow", label: "Allow" },
+  { value: "deny", label: "Deny" },
+] as const;
+
+const ON_DISALLOWED_ITEMS = [
+  { value: "block", label: "Block" },
+  { value: "rewrite", label: "Rewrite" },
+] as const;
+
 const DEFAULT_CONFIG: ToolPermissionConfig = {
   rules: [],
   default_action: "deny",
@@ -115,8 +130,9 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
     if (entries.length === 0) {
       return (
         <Button
+          variant="outline"
           disabled={disabled}
-          size="small"
+          size="sm"
           onClick={() => updateRule(index, { allowed_param_patterns: { "": "" } })}
         >
           + Restrict tool arguments (optional)
@@ -126,9 +142,9 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
 
     return (
       <div className="space-y-2">
-        <Text className="text-sm text-gray-600">Argument constraints (dot or array paths)</Text>
+        <p className="text-sm text-muted-foreground">Argument constraints (dot or array paths)</p>
         {entries.map(([path, pattern], patternIndex) => (
-          <Space key={`${rule.id || index}-${patternIndex}`} align="start">
+          <div key={`${rule.id || index}-${patternIndex}`} className="flex items-start gap-2">
             <Input
               disabled={disabled}
               placeholder="messages[0].content"
@@ -142,20 +158,24 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
               onChange={(e) => updateAllowedParamPattern(index, patternIndex, e.target.value)}
             />
             <Button
+              variant="outline"
+              size="icon"
+              aria-label="Remove constraint"
               disabled={disabled}
-              icon={<DeleteOutlined />}
-              danger
               onClick={() =>
                 updateAllowedParamEntries(index, (entries) => {
                   entries.splice(patternIndex, 1);
                 })
               }
-            />
-          </Space>
+            >
+              <Trash2 />
+            </Button>
+          </div>
         ))}
         <Button
+          variant="outline"
           disabled={disabled}
-          size="small"
+          size="sm"
           onClick={() =>
             updateRule(index, {
               allowed_param_patterns: {
@@ -173,148 +193,186 @@ const ToolPermissionRulesEditor: React.FC<ToolPermissionRulesEditorProps> = ({ v
 
   return (
     <Card>
-      <div className="flex items-center justify-between">
-        <div>
-          <Text className="text-lg font-semibold">LiteLLM Tool Permission Guardrail</Text>
-          <Text className="text-sm text-gray-500">
-            Provide regex patterns (e.g., ^mcp__github_.*$) for tool names or types and optionally constrain payload
-            fields.
-          </Text>
+      <CardContent>
+        <div className="flex items-center justify-between">
+          <div>
+            <p className="text-lg font-semibold">LiteLLM Tool Permission Guardrail</p>
+            <p className="text-sm text-muted-foreground">
+              Provide regex patterns (e.g., ^mcp__github_.*$) for tool names or types and optionally constrain payload
+              fields.
+            </p>
+          </div>
+          {!disabled && (
+            <Button onClick={addRule}>
+              <Plus />
+              Add Rule
+            </Button>
+          )}
         </div>
-        {!disabled && (
-          <Button
-            icon={<PlusOutlined />}
-            type="primary"
-            onClick={addRule}
-            className="bg-blue-600! text-white! hover:bg-blue-500!"
-          >
-            Add Rule
-          </Button>
+
+        <Separator className="my-4" />
+
+        {config.rules.length === 0 ? (
+          <div className="py-10 text-center text-muted-foreground">No tool rules added yet</div>
+        ) : (
+          <div className="space-y-4">
+            {config.rules.map((rule, index) => (
+              <Card key={rule.id || index} className="bg-muted/40">
+                <CardContent>
+                  <div className="mb-3 flex items-center justify-between">
+                    <p className="font-semibold">Rule {index + 1}</p>
+                    <Button variant="ghost" disabled={disabled} onClick={() => removeRule(index)}>
+                      <Trash2 />
+                      Remove
+                    </Button>
+                  </div>
+                  <div className="grid grid-cols-1 gap-4 md:grid-cols-2">
+                    <div>
+                      <p className="text-sm font-medium">Rule ID</p>
+                      <Input
+                        disabled={disabled}
+                        placeholder="unique_rule_id"
+                        value={rule.id}
+                        onChange={(e) => updateRule(index, { id: e.target.value })}
+                      />
+                    </div>
+                    <div>
+                      <p className="text-sm font-medium">Tool Name (optional)</p>
+                      <Input
+                        disabled={disabled}
+                        placeholder="^mcp__github_.*$"
+                        value={rule.tool_name ?? ""}
+                        onChange={(e) =>
+                          updateRule(index, {
+                            tool_name: e.target.value.trim() === "" ? undefined : e.target.value,
+                          })
+                        }
+                      />
+                    </div>
+                  </div>
+
+                  <div className="mt-4 grid grid-cols-1 gap-4 md:grid-cols-2">
+                    <div>
+                      <p className="text-sm font-medium">Tool Type (optional)</p>
+                      <Input
+                        disabled={disabled}
+                        placeholder="^function$"
+                        value={rule.tool_type ?? ""}
+                        onChange={(e) =>
+                          updateRule(index, {
+                            tool_type: e.target.value.trim() === "" ? undefined : e.target.value,
+                          })
+                        }
+                      />
+                    </div>
+                  </div>
+
+                  <div className="mt-4 flex flex-col gap-2">
+                    <p className="text-sm font-medium">Decision</p>
+                    <Select
+                      items={DECISION_ITEMS}
+                      disabled={disabled}
+                      value={rule.decision}
+                      onValueChange={(value: string | null) =>
+                        value && updateRule(index, { decision: value as ToolPermissionDecision })
+                      }
+                    >
+                      <SelectTrigger className="w-[200px]" aria-label="Decision">
+                        <SelectValue />
+                      </SelectTrigger>
+                      <SelectContent alignItemWithTrigger={false}>
+                        {DECISION_ITEMS.map((item) => (
+                          <SelectItem key={item.value} value={item.value}>
+                            {item.label}
+                          </SelectItem>
+                        ))}
+                      </SelectContent>
+                    </Select>
+                  </div>
+
+                  <div className="mt-4">{renderAllowedParamPatterns(rule, index)}</div>
+                </CardContent>
+              </Card>
+            ))}
+          </div>
         )}
-      </div>
 
-      <Divider />
+        <Separator className="my-4" />
 
-      {config.rules.length === 0 ? (
-        <Empty description="No tool rules added yet" />
-      ) : (
-        <div className="space-y-4">
-          {config.rules.map((rule, index) => (
-            <Card key={rule.id || index} className="bg-gray-50">
-              <div className="flex items-center justify-between mb-3">
-                <Text className="font-semibold">Rule {index + 1}</Text>
-                <Button
-                  icon={<DeleteOutlined />}
-                  danger
-                  type="text"
-                  disabled={disabled}
-                  onClick={() => removeRule(index)}
-                >
-                  Remove
-                </Button>
-              </div>
-              <div className="grid grid-cols-1 gap-4 md:grid-cols-2">
-                <div>
-                  <Text className="text-sm font-medium">Rule ID</Text>
-                  <Input
-                    disabled={disabled}
-                    placeholder="unique_rule_id"
-                    value={rule.id}
-                    onChange={(e) => updateRule(index, { id: e.target.value })}
-                  />
-                </div>
-                <div>
-                  <Text className="text-sm font-medium">Tool Name (optional)</Text>
-                  <Input
-                    disabled={disabled}
-                    placeholder="^mcp__github_.*$"
-                    value={rule.tool_name ?? ""}
-                    onChange={(e) =>
-                      updateRule(index, {
-                        tool_name: e.target.value.trim() === "" ? undefined : e.target.value,
-                      })
-                    }
-                  />
-                </div>
-              </div>
-
-              <div className="grid grid-cols-1 gap-4 md:grid-cols-2 mt-4">
-                <div>
-                  <Text className="text-sm font-medium">Tool Type (optional)</Text>
-                  <Input
-                    disabled={disabled}
-                    placeholder="^function$"
-                    value={rule.tool_type ?? ""}
-                    onChange={(e) =>
-                      updateRule(index, {
-                        tool_type: e.target.value.trim() === "" ? undefined : e.target.value,
-                      })
-                    }
-                  />
-                </div>
-              </div>
-
-              <div className="mt-4 flex flex-col gap-2">
-                <Text className="text-sm font-medium">Decision</Text>
-                <Select
-                  disabled={disabled}
-                  value={rule.decision}
-                  style={{ width: 200 }}
-                  onChange={(value) => updateRule(index, { decision: value as ToolPermissionDecision })}
-                >
-                  <Select.Option value="allow">Allow</Select.Option>
-                  <Select.Option value="deny">Deny</Select.Option>
-                </Select>
-              </div>
-
-              <div className="mt-4">{renderAllowedParamPatterns(rule, index)}</div>
-            </Card>
-          ))}
+        <div className="grid gap-4 md:grid-cols-2">
+          <div>
+            <p className="text-sm font-medium">Default action</p>
+            <Select
+              items={DECISION_ITEMS}
+              disabled={disabled}
+              value={config.default_action}
+              onValueChange={(value: string | null) =>
+                value && updateConfig({ default_action: value as ToolPermissionDefaultAction })
+              }
+            >
+              <SelectTrigger className="w-full" aria-label="Default action">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                {DECISION_ITEMS.map((item) => (
+                  <SelectItem key={item.value} value={item.value}>
+                    {item.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
+          <div>
+            <p className="flex items-center gap-1 text-sm font-medium">
+              On disallowed action
+              <Tooltip>
+                <TooltipTrigger
+                  render={
+                    <span className="cursor-help text-muted-foreground">
+                      <Info className="size-3.5" />
+                    </span>
+                  }
+                />
+                <TooltipContent>
+                  Block returns an error when a forbidden tool is invoked. Rewrite strips the tool call but lets the
+                  rest of the response continue.
+                </TooltipContent>
+              </Tooltip>
+            </p>
+            <Select
+              items={ON_DISALLOWED_ITEMS}
+              disabled={disabled}
+              value={config.on_disallowed_action}
+              onValueChange={(value: string | null) =>
+                value && updateConfig({ on_disallowed_action: value as ToolPermissionOnDisallowedAction })
+              }
+            >
+              <SelectTrigger className="w-full" aria-label="On disallowed action">
+                <SelectValue />
+              </SelectTrigger>
+              <SelectContent alignItemWithTrigger={false}>
+                {ON_DISALLOWED_ITEMS.map((item) => (
+                  <SelectItem key={item.value} value={item.value}>
+                    {item.label}
+                  </SelectItem>
+                ))}
+              </SelectContent>
+            </Select>
+          </div>
         </div>
-      )}
 
-      <Divider />
-
-      <div className="grid gap-4 md:grid-cols-2">
-        <div>
-          <Text className="text-sm font-medium">Default action</Text>
-          <Select
+        <div className="mt-4">
+          <p className="text-sm font-medium">Violation message (optional)</p>
+          <Textarea
+            className="field-sizing-fixed"
             disabled={disabled}
-            value={config.default_action}
-            onChange={(value) => updateConfig({ default_action: value as ToolPermissionDefaultAction })}
-          >
-            <Select.Option value="allow">Allow</Select.Option>
-            <Select.Option value="deny">Deny</Select.Option>
-          </Select>
+            rows={3}
+            placeholder="This violates our org policy..."
+            value={config.violation_message_template}
+            onChange={(e) => updateConfig({ violation_message_template: e.target.value })}
+          />
         </div>
-        <div>
-          <Text className="text-sm font-medium flex items-center gap-1">
-            On disallowed action
-            <Tooltip title="Block returns an error when a forbidden tool is invoked. Rewrite strips the tool call but lets the rest of the response continue.">
-              <InfoCircleOutlined />
-            </Tooltip>
-          </Text>
-          <Select
-            disabled={disabled}
-            value={config.on_disallowed_action}
-            onChange={(value) => updateConfig({ on_disallowed_action: value as ToolPermissionOnDisallowedAction })}
-          >
-            <Select.Option value="block">Block</Select.Option>
-            <Select.Option value="rewrite">Rewrite</Select.Option>
-          </Select>
-        </div>
-      </div>
-
-      <div className="mt-4">
-        <Text className="text-sm font-medium">Violation message (optional)</Text>
-        <Input.TextArea
-          disabled={disabled}
-          rows={3}
-          placeholder="This violates our org policy..."
-          value={config.violation_message_template}
-          onChange={(e) => updateConfig({ violation_message_template: e.target.value })}
-        />
-      </div>
+      </CardContent>
     </Card>
   );
 };

From 7c9cd2b3dfdaacc0fb9e4e7daed75bdd327a61c1 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:21:54 -0700
Subject: [PATCH 11/31] refactor(ui): migrate guardrails-monitor to shadcn

Replaces the antd Button, Row, Col, Spin and Typography usage in
GuardrailsOverview with the shadcn button, a CSS grid and the house
loading wrapper, and swaps the ant-design icons for lucide equivalents
that the already-migrated GuardrailDetail sibling uses. The busy state
keeps the aria-busy attribute the antd Spin exposed, so the loading
contract is unchanged for assistive tech.

The route's other analyzer-listed file, GuardrailsMonitorView, keeps its
DateRangePickerValue type import because the shared AdvancedDatePicker
prop contract requires it, so there is nothing to migrate there.

Retires the file's now unused no-restricted-imports suppression
---
 ui/litellm-dashboard/eslint-suppressions.json |  3 -
 .../_components/GuardrailsOverview.tsx        | 84 +++++++++----------
 2 files changed, 40 insertions(+), 47 deletions(-)

diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json
index d7f71a5840d..5df4a5af0ae 100644
--- a/ui/litellm-dashboard/eslint-suppressions.json
+++ b/ui/litellm-dashboard/eslint-suppressions.json
@@ -321,9 +321,6 @@
   "src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx": {
     "no-nested-ternary": {
       "count": 5
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/guardrails/_components/GuardrailTestPanel.tsx": {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
index 8d45b4a4ee8..3d91fc1b9b6 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
@@ -1,11 +1,12 @@
-import { DownloadOutlined, RiseOutlined, SafetyOutlined, SettingOutlined, WarningOutlined } from "@ant-design/icons";
 import { useQuery } from "@tanstack/react-query";
 import type { ColumnDef, OnChangeFn, SortingState } from "@tanstack/react-table";
-import { Button, Col, Row, Spin, Typography } from "antd";
+import { Download, Settings, Shield, TrendingUp, TriangleAlert } from "lucide-react";
 import React, { useMemo, useState } from "react";
 import { DataTable, DataTableSortHeader } from "@/components/shared/DataTable";
 import { getGuardrailsUsageOverview } from "@/components/networking";
 import { type PerformanceRow } from "@/components/GuardrailsMonitor/mockData";
+import { Button } from "@/components/ui/button";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { EvaluationSettingsModal } from "./EvaluationSettingsModal";
 import { MetricCard } from "@/components/GuardrailsMonitor/MetricCard";
 import { ScoreChart } from "./ScoreChart";
@@ -197,51 +198,42 @@ export function GuardrailsOverview({
       <div className="flex items-start justify-between mb-5">
         <div>
           <div className="flex items-center gap-2 mb-1">
-            <SafetyOutlined className="text-lg text-indigo-500" />
+            <Shield className="size-5 text-indigo-500" />
             <h1 className="text-xl font-semibold text-gray-900">Guardrails Monitor</h1>
           </div>
           <p className="text-sm text-gray-500">Monitor guardrail performance across all requests</p>
         </div>
         <div className="flex items-center gap-3">
-          <Button type="default" icon={<DownloadOutlined />} title="Coming soon">
+          <Button variant="outline" title="Coming soon">
+            <Download className="size-4" />
             Export Data
           </Button>
         </div>
       </div>
 
-      <Row gutter={[16, 16]} className="mb-6">
-        <Col xs={12} sm={12} md={8} flex="1 0 20%">
-          <MetricCard label="Total Evaluations" value={metrics.totalRequests.toLocaleString()} />
-        </Col>
-        <Col xs={12} sm={12} md={8} flex="1 0 20%">
-          <MetricCard
-            label="Blocked Requests"
-            value={metrics.totalBlocked.toLocaleString()}
-            valueColor="text-red-600"
-            icon={<WarningOutlined className="text-red-400" />}
-          />
-        </Col>
-        <Col xs={12} sm={12} md={8} flex="1 0 20%">
-          <MetricCard
-            label="Pass Rate"
-            value={`${metrics.passRate}%`}
-            valueColor="text-green-600"
-            icon={<RiseOutlined className="text-green-400" />}
-          />
-        </Col>
-        <Col xs={12} sm={12} md={8} flex="1 0 20%">
-          <MetricCard
-            label="Avg. latency added"
-            value={`${metrics.avgLatency}ms`}
-            valueColor={
-              metrics.avgLatency > 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600"
-            }
-          />
-        </Col>
-        <Col xs={12} sm={12} md={8} flex="1 0 20%">
-          <MetricCard label="Active Guardrails" value={metrics.count} />
-        </Col>
-      </Row>
+      <div className="mb-6 grid grid-cols-5 gap-4">
+        <MetricCard label="Total Evaluations" value={metrics.totalRequests.toLocaleString()} />
+        <MetricCard
+          label="Blocked Requests"
+          value={metrics.totalBlocked.toLocaleString()}
+          valueColor="text-red-600"
+          icon={<TriangleAlert className="size-4 text-red-400" />}
+        />
+        <MetricCard
+          label="Pass Rate"
+          value={`${metrics.passRate}%`}
+          valueColor="text-green-600"
+          icon={<TrendingUp className="size-4 text-green-400" />}
+        />
+        <MetricCard
+          label="Avg. latency added"
+          value={`${metrics.avgLatency}ms`}
+          valueColor={
+            metrics.avgLatency > 150 ? "text-red-600" : metrics.avgLatency > 50 ? "text-amber-600" : "text-green-600"
+          }
+        />
+        <MetricCard label="Active Guardrails" value={metrics.count} />
+      </div>
 
       <div className="mb-6">
         <ScoreChart data={chartData} />
@@ -250,7 +242,11 @@ export function GuardrailsOverview({
       <div>
         {(isLoading || error) && (
           <div className="mb-2 flex items-center gap-2">
-            {isLoading && <Spin size="small" />}
+            {isLoading && (
+              <span role="status" aria-busy="true" aria-label="Loading" className="inline-flex">
+                <UiLoadingSpinner className="size-4 text-primary" />
+              </span>
+            )}
             {error && <span className="text-sm text-red-600">Failed to load data. Try again.</span>}
           </div>
         )}
@@ -270,20 +266,20 @@ export function GuardrailsOverview({
           toolbar={() => (
             <div className="flex items-start justify-between gap-4">
               <div>
-                <Typography.Title level={5} className="mb-0! text-gray-900">
-                  Guardrail Performance
-                </Typography.Title>
+                <h5 className="mb-0 text-base font-semibold text-gray-900">Guardrail Performance</h5>
                 <p className="text-xs text-gray-500 mt-0.5">
                   Click a guardrail to view details, logs, and configuration
                 </p>
               </div>
               <div className="flex items-center gap-2">
                 <Button
-                  type="default"
-                  icon={<SettingOutlined />}
+                  variant="outline"
+                  size="icon"
                   onClick={() => setEvaluationModalOpen(true)}
                   title="Evaluation settings"
-                />
+                >
+                  <Settings className="size-4" />
+                </Button>
               </div>
             </div>
           )}

From 355ef22257af73e8589bf6ab707e1c0e658f97a3 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:22:03 -0700
Subject: [PATCH 12/31] refactor(ui): drop narration comments from the migrated
 usage components

---
 .../EndpointUsage/components/EndpointUsageTable.tsx        | 1 -
 .../usage/_components/components/UsagePageView.tsx         | 7 +++----
 ui/litellm-dashboard/src/components/activity_metrics.tsx   | 2 --
 3 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
index 5f72caae437..01a0ac7460a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EndpointUsage/components/EndpointUsageTable.tsx
@@ -56,7 +56,6 @@ const EndpointUsageTable: React.FC<EndpointUsageTableProps> = ({ endpointData })
         return (
           <div className="flex items-center space-x-3">
             <div className="flex-1 relative">
-              {/* The failed share is the track showing through behind the successful share. */}
               <Meter value={successPercentage} max={totalPercentage || 100} aria-label="Successful requests">
                 <MeterTrack className={failurePercentage > 0 ? "bg-red-500" : undefined}>
                   <MeterIndicator className="bg-green-500" />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
index f3e7db83c6a..a716d59bdc4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx
@@ -104,8 +104,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
   const canViewOrganizationUsage = hasCapability(userRole, "viewOrganizationUsage");
   const canViewAgentUsage = hasCapability(userRole, "viewAgentUsage");
 
-  // The selector debounces its own input, so this holds the already-settled query.
-  const [debouncedUserSearch, setDebouncedUserSearch] = useState("");
+  const [settledUserSearch, setSettledUserSearch] = useState("");
 
   const {
     data: usersInfiniteData,
@@ -113,7 +112,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
     hasNextPage: hasNextUsersPage,
     isFetchingNextPage: isFetchingNextUsersPage,
     isLoading: isLoadingUsers,
-  } = useInfiniteUsers(50, debouncedUserSearch || undefined);
+  } = useInfiniteUsers(50, settledUserSearch || undefined);
 
   const userOptions = useMemo(() => {
     if (!usersInfiniteData?.pages) return [];
@@ -533,7 +532,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
                     options={userOptions}
                     value={selectedUserId ?? undefined}
                     onValueChange={(value) => setSelectedUserId(value === "" ? null : value)}
-                    onSearchChange={setDebouncedUserSearch}
+                    onSearchChange={setSettledUserSearch}
                     onLoadMore={fetchNextUsersPage}
                     hasNextPage={hasNextUsersPage}
                     isLoading={isLoadingUsers}
diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx
index 288d07378b7..b53c6d2c3bf 100644
--- a/ui/litellm-dashboard/src/components/activity_metrics.tsx
+++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx
@@ -211,8 +211,6 @@ const ModelSection = ({
   );
 };
 
-// Mirrors the antd Collapse this replaced: a panel's subtree is not built until it is first
-// opened, and it stays mounted afterwards so its view-mode state survives being collapsed.
 const ModelCollapsible = ({
   defaultOpen,
   header,

From 923a49a6f8cf85da15588475341dbba379a0b35f Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:27:05 -0700
Subject: [PATCH 13/31] test(ui): pin playground behaviour with
 library-agnostic queries

Rewrite the playground tests that reached for antd class names so they
locate controls by role, text, placeholder or lucide icon instead. Add
characterisation suites for AgentBuilderView and RealtimePlayground,
which had none, including tab state that must survive a round trip
through another tab.

Every assertion here passes against the current antd components.
---
 .../chat_ui/AgentBuilderView.test.tsx         | 326 ++++++++++++++++++
 .../chat_ui/RealtimePlayground.test.tsx       | 267 ++++++++++++++
 .../components/compareUI/CompareUI.test.tsx   |   7 +-
 .../components/ComparisonPanel.test.tsx       | 139 +++++++-
 .../components/MessageInput.test.tsx          | 101 ++++--
 .../components/UnifiedSelector.test.tsx       | 204 +++++------
 6 files changed, 885 insertions(+), 159 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
 create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
new file mode 100644
index 00000000000..179b7b8557a
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
@@ -0,0 +1,326 @@
+import { render, screen, waitFor, within } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { useState } from "react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import AgentBuilderView from "./AgentBuilderView";
+import type { AgentModel } from "../../llm_calls/fetch_agents";
+
+const modelCreateCall = vi.fn().mockResolvedValue({ model_id: "id-new" });
+const modelPatchUpdateCall = vi.fn().mockResolvedValue({});
+const modelDeleteCall = vi.fn().mockResolvedValue({});
+const keyCreateCall = vi.fn().mockResolvedValue({ key: "sk-agent-key" });
+const fetchMCPServers = vi.fn().mockResolvedValue([]);
+const fetchAvailableAgentModels = vi.fn();
+const fetchAvailableModels = vi.fn().mockResolvedValue([{ model_group: "gpt-4o" }, { model_group: "claude-sonnet-4" }]);
+
+vi.mock("@/components/networking", () => ({
+  proxyBaseUrl: "https://proxy.example.com",
+  modelCreateCall: (...args: unknown[]) => modelCreateCall(...args),
+  modelPatchUpdateCall: (...args: unknown[]) => modelPatchUpdateCall(...args),
+  modelDeleteCall: (...args: unknown[]) => modelDeleteCall(...args),
+  keyCreateCall: (...args: unknown[]) => keyCreateCall(...args),
+  fetchMCPServers: (...args: unknown[]) => fetchMCPServers(...args),
+}));
+
+vi.mock("../../llm_calls/fetch_agents", () => ({
+  fetchAvailableAgentModels: (...args: unknown[]) => fetchAvailableAgentModels(...args),
+}));
+
+vi.mock("@/components/llm_calls/fetch_models", () => ({
+  fetchAvailableModels: (...args: unknown[]) => fetchAvailableModels(...args),
+}));
+
+vi.mock("@/components/CodeBlock", () => ({
+  default: ({ code }: { code: string }) => <pre data-testid="code-block">{code}</pre>,
+}));
+
+vi.mock("@/components/molecules/notifications_manager", () => ({
+  default: { success: vi.fn(), fromBackend: vi.fn() },
+}));
+
+// Both tab bodies keep their own state, which is what the tab container has to
+// preserve when the user moves away and comes back.
+const StatefulPanel = ({ label }: { label: string }) => {
+  const [draft, setDraft] = useState("");
+  return <input aria-label={label} value={draft} onChange={(event) => setDraft(event.target.value)} />;
+};
+
+vi.mock("./ChatUI", () => ({
+  default: () => <StatefulPanel label="chat scratch" />,
+}));
+
+vi.mock("../complianceUI/ComplianceUI", () => ({
+  default: () => <StatefulPanel label="batch scratch" />,
+}));
+
+const AGENTS: AgentModel[] = [
+  {
+    model_name: "support-agent",
+    litellm_params: { model: "litellm_agent/gpt-4o", litellm_system_prompt: "Be helpful.", temperature: 0.3 },
+    model_info: { id: "agent-1" },
+  },
+  {
+    model_name: "research-agent",
+    litellm_params: { model: "litellm_agent/claude-sonnet-4" },
+    model_info: { id: "agent-2" },
+  },
+];
+
+const props = {
+  accessToken: "sk-access",
+  token: "tok",
+  userID: "u1",
+  userRole: "Admin",
+};
+
+// The labels in the configure panel are not wired to their controls, so reach
+// the control through the field that the label heads.
+const controlUnder = (label: string): HTMLElement =>
+  within(screen.getByText(label).parentElement!).getByRole("combobox");
+
+const renderView = () => render(<AgentBuilderView {...props} />);
+
+const waitForRoster = async () => {
+  await screen.findByRole("button", { name: "support-agent litellm_agent" });
+};
+
+beforeEach(() => {
+  vi.clearAllMocks();
+  fetchAvailableAgentModels.mockResolvedValue(AGENTS);
+  fetchAvailableModels.mockResolvedValue([{ model_group: "gpt-4o" }, { model_group: "claude-sonnet-4" }]);
+  fetchMCPServers.mockResolvedValue([]);
+  modelCreateCall.mockResolvedValue({ model_id: "id-new" });
+  Object.defineProperty(window, "matchMedia", {
+    writable: true,
+    value: (query: string) => ({
+      matches: false,
+      media: query,
+      onchange: null,
+      addListener: () => {},
+      removeListener: () => {},
+      addEventListener: () => {},
+      removeEventListener: () => {},
+      dispatchEvent: () => false,
+    }),
+  });
+});
+
+describe("AgentBuilderView", () => {
+  it("asks the visitor to sign in when there is no session", () => {
+    render(<AgentBuilderView accessToken={null} token={null} userID={null} userRole={null} />);
+
+    expect(screen.getByText("Sign in to use Agent Builder.")).toBeInTheDocument();
+    expect(fetchAvailableAgentModels).not.toHaveBeenCalled();
+  });
+
+  it("marks itself busy while the roster loads", async () => {
+    let release: (agents: AgentModel[]) => void = () => {};
+    fetchAvailableAgentModels.mockReturnValue(
+      new Promise<AgentModel[]>((resolve) => {
+        release = resolve;
+      }),
+    );
+
+    renderView();
+
+    expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument();
+
+    release(AGENTS);
+    await waitForRoster();
+    expect(document.querySelector('[aria-busy="true"]')).not.toBeInTheDocument();
+  });
+
+  it("lists every agent and opens the first one's configuration", async () => {
+    renderView();
+    await waitForRoster();
+
+    expect(screen.getByRole("button", { name: "research-agent litellm_agent" })).toBeInTheDocument();
+    expect(screen.getByText("Agent Builder")).toBeInTheDocument();
+    await waitFor(() => expect(screen.getByDisplayValue("support-agent")).toBeInTheDocument());
+    expect(screen.getByDisplayValue("Be helpful.")).toBeInTheDocument();
+    expect(screen.getByDisplayValue("0.3")).toBeInTheDocument();
+  });
+
+  it("loads the configuration of whichever agent is picked", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("button", { name: "research-agent litellm_agent" }));
+
+    await waitFor(() => expect(screen.getByDisplayValue("research-agent")).toBeInTheDocument());
+  });
+
+  it("offers a blank draft and a save control for a new agent", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("button", { name: /New agent/i }));
+
+    expect(screen.getByRole("button", { name: /Save Agent/i })).toBeInTheDocument();
+    expect(screen.getByDisplayValue("You are a helpful assistant.")).toBeInTheDocument();
+    expect(screen.queryByRole("button", { name: /Update Agent/i })).not.toBeInTheDocument();
+  });
+
+  it("creates the agent under the litellm_agent prefix", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("button", { name: /New agent/i }));
+    await user.type(screen.getByPlaceholderText("My Agent"), "billing-agent");
+    await user.click(screen.getByRole("button", { name: /Save Agent/i }));
+
+    await waitFor(() => expect(modelCreateCall).toHaveBeenCalled());
+    const payload = modelCreateCall.mock.calls[0][1];
+    expect(payload.model_name).toBe("billing-agent");
+    expect(payload.litellm_params.model).toBe("litellm_agent/gpt-4o");
+  });
+
+  it("will not save a draft without a name", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("button", { name: /New agent/i }));
+    await user.click(screen.getByRole("button", { name: /Save Agent/i }));
+
+    expect(modelCreateCall).not.toHaveBeenCalled();
+  });
+
+  it("updates the selected agent through its model id", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+    await screen.findByDisplayValue("support-agent");
+
+    await user.click(screen.getByRole("button", { name: /Update Agent/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(modelPatchUpdateCall.mock.calls[0][2]).toBe("agent-1");
+  });
+
+  it("deletes only after the warning is confirmed", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+    await screen.findByDisplayValue("support-agent");
+
+    await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]);
+
+    expect(await screen.findByText(/Are you sure you want to delete "support-agent"/)).toBeInTheDocument();
+    expect(modelDeleteCall).not.toHaveBeenCalled();
+
+    // The trigger and the confirmation share a label; the confirmation is the
+    // one rendered last, in the overlay.
+    const confirmations = screen.getAllByRole("button", { name: /Delete$/ });
+    await user.click(confirmations[confirmations.length - 1]);
+
+    await waitFor(() => expect(modelDeleteCall).toHaveBeenCalledWith("sk-access", "agent-1"));
+  });
+
+  it("abandons the delete when the warning is dismissed", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+    await screen.findByDisplayValue("support-agent");
+
+    await user.click(screen.getAllByRole("button", { name: /Delete$/ })[0]);
+    await screen.findByText(/Are you sure you want to delete "support-agent"/);
+    await user.click(screen.getAllByRole("button", { name: /Cancel$/ }).at(-1)!);
+
+    await waitFor(() =>
+      expect(screen.queryByText(/Are you sure you want to delete "support-agent"/)).not.toBeInTheDocument(),
+    );
+    expect(modelDeleteCall).not.toHaveBeenCalled();
+  });
+
+  it("shows a ready-to-run curl example on the Connect tab", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("tab", { name: /Connect/i }));
+
+    const snippet = await screen.findByTestId("code-block");
+    expect(snippet).toHaveTextContent("https://proxy.example.com/v1/chat/completions");
+    expect(snippet).toHaveTextContent('"model": "support-agent"');
+  });
+
+  it("mints a key scoped to the selected agent", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("tab", { name: /Connect/i }));
+    await user.click(await screen.findByRole("button", { name: /Create key for this agent/i }));
+
+    await waitFor(() => expect(keyCreateCall).toHaveBeenCalled());
+    expect(keyCreateCall.mock.calls[0][2].models).toEqual(["support-agent"]);
+    expect(await screen.findByTestId("code-block")).toHaveTextContent("Bearer sk-agent-key");
+  });
+
+  it("keeps a tab's own state alive while the user works in another tab", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("tab", { name: /Chat/i }));
+    const scratch = await screen.findByLabelText("chat scratch");
+    await user.type(scratch, "half a thought");
+    expect(scratch).toHaveValue("half a thought");
+
+    await user.click(screen.getByRole("tab", { name: /Configure/i }));
+    await screen.findByDisplayValue("support-agent");
+
+    await user.click(screen.getByRole("tab", { name: /Chat/i }));
+    expect(await screen.findByLabelText("chat scratch")).toHaveValue("half a thought");
+  });
+
+  it("keeps the batch tab's state alive across a round trip too", async () => {
+    const user = userEvent.setup();
+    renderView();
+    await waitForRoster();
+
+    await user.click(screen.getByRole("tab", { name: /Batch Test/i }));
+    await user.type(await screen.findByLabelText("batch scratch"), "seven cases");
+
+    await user.click(screen.getByRole("tab", { name: /Connect/i }));
+    await screen.findByTestId("code-block");
+
+    await user.click(screen.getByRole("tab", { name: /Batch Test/i }));
+    expect(await screen.findByLabelText("batch scratch")).toHaveValue("seven cases");
+  });
+
+  it("attaches the MCP servers the agent should reach", async () => {
+    const user = userEvent.setup();
+    fetchMCPServers.mockResolvedValue([{ server_id: "srv-1", alias: "github", server_name: "github-mcp" }]);
+    renderView();
+    await waitForRoster();
+    await screen.findByDisplayValue("support-agent");
+
+    await user.click(controlUnder("MCP servers"));
+    const options = await screen.findAllByText("github");
+    await user.click(options[options.length - 1]);
+    await user.keyboard("{Escape}");
+
+    await user.click(screen.getByRole("button", { name: /Update Agent/i }));
+
+    await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalled());
+    expect(modelPatchUpdateCall.mock.calls[0][1].litellm_params.tools).toEqual([
+      { type: "mcp", server_label: "litellm", server_url: "litellm_proxy/mcp/github", require_approval: "never" },
+    ]);
+  });
+
+  it("warns that the builder is experimental", async () => {
+    renderView();
+    await waitForRoster();
+
+    expect(screen.getByText(/Agent Builder is experimental/)).toBeInTheDocument();
+    expect(within(screen.getByText(/Agent Builder is experimental/)).getByRole("link")).toHaveAttribute(
+      "href",
+      "mailto:product@berri.ai",
+    );
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx
new file mode 100644
index 00000000000..a72a02a59e9
--- /dev/null
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.test.tsx
@@ -0,0 +1,267 @@
+import { act, render, screen, waitFor } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import RealtimePlayground from "./RealtimePlayground";
+
+vi.mock("@/components/networking", () => ({
+  getProxyBaseUrl: () => "https://proxy.example.com",
+}));
+
+class FakeSocket {
+  static instances: FakeSocket[] = [];
+  static OPEN = 1;
+
+  readyState = 0;
+  sent: string[] = [];
+  onopen: (() => void) | null = null;
+  onmessage: ((event: { data: string }) => void) | null = null;
+  onerror: (() => void) | null = null;
+  onclose: (() => void) | null = null;
+  close = vi.fn(() => {
+    this.readyState = 3;
+    this.onclose?.();
+  });
+
+  constructor(
+    public url: string,
+    public protocols?: string[],
+  ) {
+    FakeSocket.instances.push(this);
+  }
+
+  send(payload: string) {
+    this.sent.push(payload);
+  }
+
+  open() {
+    this.readyState = 1;
+    this.onopen?.();
+  }
+
+  emit(message: Record<string, unknown>) {
+    this.onmessage?.({ data: JSON.stringify(message) });
+  }
+}
+
+const latestSocket = () => FakeSocket.instances[FakeSocket.instances.length - 1];
+
+const connect = async (user: ReturnType<typeof userEvent.setup>) => {
+  await user.click(screen.getByRole("button", { name: /Connect/i }));
+  await act(async () => {
+    latestSocket().open();
+  });
+};
+
+const props = {
+  accessToken: "sk-realtime",
+  selectedModel: "gpt-realtime",
+};
+
+beforeEach(() => {
+  FakeSocket.instances = [];
+  vi.stubGlobal("WebSocket", FakeSocket);
+  vi.stubGlobal(
+    "AudioContext",
+    class {
+      currentTime = 0;
+      destination = {};
+      close = vi.fn();
+      createBuffer = vi.fn(() => ({ getChannelData: () => new Float32Array(1), duration: 0 }));
+      createBufferSource = vi.fn(() => ({ connect: vi.fn(), start: vi.fn(), buffer: null }));
+    },
+  );
+  Object.defineProperty(window, "matchMedia", {
+    writable: true,
+    value: (query: string) => ({
+      matches: false,
+      media: query,
+      onchange: null,
+      addListener: () => {},
+      removeListener: () => {},
+      addEventListener: () => {},
+      removeEventListener: () => {},
+      dispatchEvent: () => false,
+    }),
+  });
+  Element.prototype.scrollIntoView = vi.fn();
+});
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+describe("RealtimePlayground", () => {
+  it("opens disconnected, with the invitation to connect", () => {
+    render(<RealtimePlayground {...props} />);
+
+    expect(screen.getByText("Realtime Voice Chat")).toBeInTheDocument();
+    expect(screen.getByText("Disconnected")).toBeInTheDocument();
+    expect(screen.getByText("Realtime Voice Playground")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /Connect/i })).toBeInTheDocument();
+  });
+
+  it("hides the composer until a session exists", () => {
+    render(<RealtimePlayground {...props} />);
+
+    expect(screen.queryByPlaceholderText("Type a message or use the mic...")).not.toBeInTheDocument();
+  });
+
+  it("dials the realtime endpoint for the selected model, carrying the key as a protocol", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await user.click(screen.getByRole("button", { name: /Connect/i }));
+
+    expect(latestSocket().url).toBe("wss://proxy.example.com/v1/realtime?model=gpt-realtime");
+    expect(latestSocket().protocols).toEqual(["realtime", "openai-insecure-api-key.sk-realtime"]);
+  });
+
+  it("passes a custom proxy base url through instead of the default", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} customProxyBaseUrl="https://tenant.example.com" />);
+
+    await user.click(screen.getByRole("button", { name: /Connect/i }));
+
+    expect(latestSocket().url).toContain("wss://tenant.example.com/v1/realtime");
+  });
+
+  it("appends the selected guardrails to the session url", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} selectedGuardrails={["pii", "toxicity"]} />);
+
+    await user.click(screen.getByRole("button", { name: /Connect/i }));
+
+    expect(latestSocket().url).toContain("guardrails=pii%2Ctoxicity");
+  });
+
+  it("refuses to dial without a model and says why", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} selectedModel="" />);
+
+    await user.click(screen.getByRole("button", { name: /Connect/i }));
+
+    expect(FakeSocket.instances).toHaveLength(0);
+    expect(screen.getByText("Please select a model first")).toBeInTheDocument();
+  });
+
+  it("reveals the composer and the disconnect control once the session opens", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+
+    expect(screen.getByText("Connected")).toBeInTheDocument();
+    expect(screen.getByText("Connected to realtime API")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("Type a message or use the mic...")).toBeInTheDocument();
+    expect(screen.getByRole("button", { name: /Disconnect/i })).toBeInTheDocument();
+    expect(screen.getByTitle("Start recording")).toBeInTheDocument();
+  });
+
+  it("configures the session against the chosen voice when it is created", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await act(async () => {
+      latestSocket().emit({ type: "session.created" });
+    });
+
+    const update = JSON.parse(latestSocket().sent[0]);
+    expect(update.type).toBe("session.update");
+    expect(update.session.voice).toBe("alloy");
+    expect(update.session.type).toBe("realtime");
+  });
+
+  it("sends what was typed and then asks for a response", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await user.type(screen.getByPlaceholderText("Type a message or use the mic..."), "hello there");
+    await user.click(screen.getByRole("button", { name: /send/i }));
+
+    const payloads = latestSocket().sent.map((raw) => JSON.parse(raw));
+    expect(payloads[0].item.content[0].text).toBe("hello there");
+    expect(payloads[1]).toEqual({ type: "response.create" });
+    expect(screen.getByText("hello there")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("Type a message or use the mic...")).toHaveValue("");
+  });
+
+  it("will not send an empty message", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await user.click(screen.getByRole("button", { name: /send/i }));
+
+    expect(latestSocket().sent).toHaveLength(0);
+  });
+
+  it("streams assistant text deltas into a single reply", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await act(async () => {
+      latestSocket().emit({ type: "response.output_text.delta", delta: "Hel" });
+      latestSocket().emit({ type: "response.output_text.delta", delta: "lo!" });
+    });
+
+    expect(screen.getByText("Hello!")).toBeInTheDocument();
+  });
+
+  it("falls back to the completed response when no delta arrived", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await act(async () => {
+      latestSocket().emit({
+        type: "response.done",
+        response: { output: [{ content: [{ type: "output_audio", transcript: "spoken reply" }] }] },
+      });
+    });
+
+    expect(screen.getByText("spoken reply")).toBeInTheDocument();
+  });
+
+  it("shows what the microphone heard", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await act(async () => {
+      latestSocket().emit({
+        type: "conversation.item.input_audio_transcription.completed",
+        transcript: "what is the weather",
+      });
+    });
+
+    expect(screen.getByText("what is the weather")).toBeInTheDocument();
+  });
+
+  it("surfaces an error frame in the transcript", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    await act(async () => {
+      latestSocket().emit({ type: "error", error: { message: "rate limited" } });
+    });
+
+    expect(screen.getByText("Error: rate limited")).toBeInTheDocument();
+  });
+
+  it("closes the socket and returns to the disconnected state", async () => {
+    const user = userEvent.setup();
+    render(<RealtimePlayground {...props} />);
+
+    await connect(user);
+    const socket = latestSocket();
+    await user.click(screen.getByRole("button", { name: /Disconnect/i }));
+
+    expect(socket.close).toHaveBeenCalled();
+    await waitFor(() => expect(screen.getByRole("button", { name: /Connect/i })).toBeInTheDocument());
+    expect(screen.queryByPlaceholderText("Type a message or use the mic...")).not.toBeInTheDocument();
+  });
+});
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx
index 4278cb6a0e4..a68513b5204 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.test.tsx
@@ -1,4 +1,4 @@
-import { render, waitFor } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import CompareUI from "./CompareUI";
@@ -108,10 +108,7 @@ describe("CompareUI", () => {
     let comparisonPanels = container.querySelectorAll('[data-testid^="comparison-panel-"]');
     expect(comparisonPanels).toHaveLength(2);
 
-    const addButtons = Array.from(container.querySelectorAll('button[class*="ant-btn"]'));
-    const addComparisonButton = addButtons.find((btn) => btn.textContent?.includes("Add Comparison"));
-    expect(addComparisonButton).toBeInTheDocument();
-    await user.click(addComparisonButton!);
+    await user.click(screen.getByRole("button", { name: /Add Comparison/i }));
 
     // Wait for the new comparison panel to be added (should have 3 total now)
     await waitFor(() => {
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
index c07ad367606..f18ea86b8d3 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
@@ -1,4 +1,4 @@
-import { render } from "@testing-library/react";
+import { render, screen, waitFor } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { ComparisonInstance } from "../CompareUI";
@@ -76,20 +76,139 @@ const mockProps = {
   apiKey: "test-api-key",
 };
 
+// Both the settings and the remove control are icon-only buttons whose lucide
+// icon survives the migration, so locate them by that icon.
+const buttonWithIcon = (icon: string): HTMLButtonElement => {
+  const match = Array.from(document.querySelectorAll("button")).find((button) =>
+    button.querySelector(`svg.lucide-${icon}`),
+  );
+  if (!match) throw new Error(`no button carrying the ${icon} icon`);
+  return match;
+};
+
+const openSettings = async (user: ReturnType<typeof userEvent.setup>) => {
+  await user.click(buttonWithIcon("settings"));
+  await screen.findByText("General Settings");
+};
+
 describe("ComparisonPanel", () => {
-  it("should render", () => {
-    const { getByTestId } = render(<ComparisonPanel {...mockProps} />);
-    expect(getByTestId("unified-selector")).toBeInTheDocument();
-    expect(getByTestId("message-display")).toBeInTheDocument();
+  it("renders the selector and the transcript", () => {
+    render(<ComparisonPanel {...mockProps} />);
+
+    expect(screen.getByTestId("unified-selector")).toBeInTheDocument();
+    expect(screen.getByTestId("message-display")).toBeInTheDocument();
   });
 
-  it("should call onRemove when remove button is clicked", async () => {
+  it("removes the panel when the remove control is used", async () => {
     const user = userEvent.setup();
     const onRemove = vi.fn();
-    const { container } = render(<ComparisonPanel {...mockProps} onRemove={onRemove} />);
-    const removeButton = container.querySelector('button[class*="text-red-600"]');
-    expect(removeButton).toBeInTheDocument();
-    await user.click(removeButton!);
+    render(<ComparisonPanel {...mockProps} onRemove={onRemove} />);
+
+    await user.click(buttonWithIcon("x"));
+
     expect(onRemove).toHaveBeenCalledTimes(1);
   });
+
+  it("hides the remove control on the last remaining panel", () => {
+    render(<ComparisonPanel {...mockProps} canRemove={false} />);
+
+    expect(() => buttonWithIcon("x")).toThrow();
+  });
+
+  it("keeps the settings out of sight until the gear is used", async () => {
+    const user = userEvent.setup();
+    render(<ComparisonPanel {...mockProps} />);
+
+    expect(screen.queryByText("General Settings")).not.toBeInTheDocument();
+
+    await openSettings(user);
+
+    expect(screen.getByText("General Settings")).toBeInTheDocument();
+    expect(screen.getByText("Advanced Settings")).toBeInTheDocument();
+    expect(screen.getByTestId("tag-selector")).toBeInTheDocument();
+    expect(screen.getByTestId("vector-store-selector")).toBeInTheDocument();
+    expect(screen.getByTestId("guardrail-selector")).toBeInTheDocument();
+  });
+
+  it("shows the current temperature and token ceiling", async () => {
+    const user = userEvent.setup();
+    render(<ComparisonPanel {...mockProps} />);
+
+    await openSettings(user);
+
+    expect(screen.getByText("Temperature")).toBeInTheDocument();
+    expect(screen.getByText("1.00")).toBeInTheDocument();
+    expect(screen.getByText("Max Tokens")).toBeInTheDocument();
+    expect(screen.getByText("2048")).toBeInTheDocument();
+
+    const ranges = Array.from(document.querySelectorAll("[aria-valuenow]"));
+    expect(ranges.map((range) => range.getAttribute("aria-valuenow"))).toEqual(["1", "2048"]);
+  });
+
+  it("pushes the whole parameter set to every panel when sync is switched on", async () => {
+    const user = userEvent.setup();
+    const onUpdate = vi.fn();
+    render(<ComparisonPanel {...mockProps} onUpdate={onUpdate} />);
+
+    await openSettings(user);
+    await user.click(screen.getByRole("checkbox", { name: /Sync Settings Across Models/i }));
+
+    await waitFor(() => expect(onUpdate).toHaveBeenCalled());
+    const [updates, options] = onUpdate.mock.calls[0];
+    expect(updates.applyAcrossModels).toBe(true);
+    expect(updates.temperature).toBe(1);
+    expect(updates.maxTokens).toBe(2048);
+    expect(options.applyToAll).toBe(true);
+    expect(options.keysToApply).toContain("temperature");
+    expect(options.keysToApply).toContain("maxTokens");
+  });
+
+  it("turns sync off without resetting the values it was sharing", async () => {
+    const user = userEvent.setup();
+    const onUpdate = vi.fn();
+    render(
+      <ComparisonPanel
+        {...mockProps}
+        comparison={{ ...mockComparison, applyAcrossModels: true }}
+        onUpdate={onUpdate}
+      />,
+    );
+
+    await openSettings(user);
+    await user.click(screen.getByRole("checkbox", { name: /Sync Settings Across Models/i }));
+
+    await waitFor(() => expect(onUpdate).toHaveBeenCalled());
+    expect(onUpdate.mock.calls[0][0]).toEqual({ applyAcrossModels: false });
+  });
+
+  it("keeps an advanced-parameter toggle local while sync is off", async () => {
+    const user = userEvent.setup();
+    const onUpdate = vi.fn();
+    render(<ComparisonPanel {...mockProps} onUpdate={onUpdate} />);
+
+    await openSettings(user);
+    await user.click(screen.getByRole("checkbox", { name: /Use Advanced Parameters/i }));
+
+    await waitFor(() => expect(onUpdate).toHaveBeenCalled());
+    expect(onUpdate.mock.calls[0][0]).toEqual({ useAdvancedParams: true });
+    expect(onUpdate.mock.calls[0][1]).toBeUndefined();
+  });
+
+  it("fans an advanced-parameter toggle out to every panel while sync is on", async () => {
+    const user = userEvent.setup();
+    const onUpdate = vi.fn();
+    render(
+      <ComparisonPanel
+        {...mockProps}
+        comparison={{ ...mockComparison, applyAcrossModels: true }}
+        onUpdate={onUpdate}
+      />,
+    );
+
+    await openSettings(user);
+    await user.click(screen.getByRole("checkbox", { name: /Use Advanced Parameters/i }));
+
+    await waitFor(() => expect(onUpdate).toHaveBeenCalled());
+    expect(onUpdate.mock.calls[0][1]).toEqual({ applyToAll: true, keysToApply: ["useAdvancedParams"] });
+  });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx
index eb157904165..049b70b982c 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.test.tsx
@@ -1,42 +1,95 @@
-import { render } from "@testing-library/react";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
 import { describe, expect, it, vi } from "vitest";
 import { MessageInput } from "./MessageInput";
 
+const PLACEHOLDER = "Type your message... (Shift+Enter for new line)";
+
 describe("MessageInput", () => {
-  it("should render", () => {
-    const onChange = vi.fn();
-    const onSend = vi.fn();
-    const { container } = render(<MessageInput value="" onChange={onChange} onSend={onSend} />);
-    const textarea = container.querySelector("textarea");
-    const button = container.querySelector("button");
-    expect(textarea).toBeInTheDocument();
-    expect(button).toBeInTheDocument();
+  it("renders a message box and a send control", () => {
+    render(<MessageInput value="" onChange={vi.fn()} onSend={vi.fn()} />);
+
+    expect(screen.getByPlaceholderText(PLACEHOLDER)).toBeInTheDocument();
+    expect(screen.getByRole("button")).toBeInTheDocument();
   });
 
-  it("should disable send button initially", () => {
+  it("reports every keystroke through onChange", async () => {
+    const user = userEvent.setup();
     const onChange = vi.fn();
-    const onSend = vi.fn();
-    const { container } = render(<MessageInput value="" onChange={onChange} onSend={onSend} />);
-    const button = container.querySelector("button") as HTMLButtonElement;
+    render(<MessageInput value="" onChange={onChange} onSend={vi.fn()} />);
 
-    expect(button).toBeDisabled();
+    await user.type(screen.getByPlaceholderText(PLACEHOLDER), "hi");
+
+    expect(onChange).toHaveBeenCalledWith("h");
   });
 
-  it("should enable send button when hasAttachment is true even with empty value", () => {
-    const onChange = vi.fn();
+  it("refuses to send while the box is empty and nothing is attached", async () => {
+    const user = userEvent.setup();
     const onSend = vi.fn();
-    const uploadComponent = <div data-testid="upload-component">Upload</div>;
-    const { container, getByTestId } = render(
+    render(<MessageInput value="" onChange={vi.fn()} onSend={onSend} />);
+
+    const send = screen.getByRole("button");
+    expect(send).toBeDisabled();
+
+    await user.click(send);
+    expect(onSend).not.toHaveBeenCalled();
+  });
+
+  it("sends once the box holds text", async () => {
+    const user = userEvent.setup();
+    const onSend = vi.fn();
+    render(<MessageInput value="hello" onChange={vi.fn()} onSend={onSend} />);
+
+    const send = screen.getByRole("button");
+    expect(send).not.toBeDisabled();
+
+    await user.click(send);
+    expect(onSend).toHaveBeenCalledTimes(1);
+  });
+
+  it("sends on an attachment alone and surfaces the upload control", async () => {
+    const user = userEvent.setup();
+    const onSend = vi.fn();
+    render(
       <MessageInput
         value=""
-        onChange={onChange}
+        onChange={vi.fn()}
         onSend={onSend}
-        hasAttachment={true}
-        uploadComponent={uploadComponent}
+        hasAttachment
+        uploadComponent={<span data-testid="upload-component">Upload</span>}
       />,
     );
-    const button = container.querySelector("button") as HTMLButtonElement;
-    expect(getByTestId("upload-component")).toBeInTheDocument();
-    expect(button).not.toBeDisabled();
+
+    expect(screen.getByTestId("upload-component")).toBeInTheDocument();
+
+    const send = screen.getByRole("button");
+    expect(send).not.toBeDisabled();
+
+    await user.click(send);
+    expect(onSend).toHaveBeenCalledTimes(1);
+  });
+
+  it("stays inert while disabled even with text present", async () => {
+    const user = userEvent.setup();
+    const onSend = vi.fn();
+    render(<MessageInput value="hello" onChange={vi.fn()} onSend={onSend} disabled />);
+
+    await user.click(screen.getByRole("button"));
+    expect(onSend).not.toHaveBeenCalled();
+  });
+
+  it("sends on Enter but adds a newline on Shift+Enter", async () => {
+    const user = userEvent.setup();
+    const onSend = vi.fn();
+    render(<MessageInput value="hello" onChange={vi.fn()} onSend={onSend} />);
+
+    const box = screen.getByPlaceholderText(PLACEHOLDER);
+    await user.click(box);
+
+    await user.keyboard("{Shift>}{Enter}{/Shift}");
+    expect(onSend).not.toHaveBeenCalled();
+
+    await user.keyboard("{Enter}");
+    expect(onSend).toHaveBeenCalledTimes(1);
   });
 });
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
index 1c6951fc421..4686804afec 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
@@ -4,158 +4,122 @@ import { describe, expect, it, vi } from "vitest";
 import { UnifiedSelector } from "./UnifiedSelector";
 import { EndpointId, ENDPOINT_CONFIGS } from "../endpoint_config";
 
+const CHAT = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
+const AGENTS = ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS];
+
+// A select's prompt text is a rendered node in one library and an input
+// placeholder attribute in the other, so accept either.
+const promptIsVisible = (text: string) =>
+  screen.queryByText(text) !== null || screen.queryByPlaceholderText(text) !== null;
+
+const openList = async (user: ReturnType<typeof userEvent.setup>) => {
+  await user.click(screen.getByRole("combobox"));
+};
+
 describe("UnifiedSelector", () => {
-  it("should render", () => {
-    const onChange = vi.fn();
-    const options = [
-      { value: "option1", label: "Option 1" },
-      { value: "option2", label: "Option 2" },
-    ];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
-
-    render(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
-
-    const select = screen.getByRole("combobox");
-    expect(select).toBeInTheDocument();
-  });
-
-  it("should display placeholder when not loading", () => {
-    const onChange = vi.fn();
-    const options = [{ value: "option1", label: "Option 1" }];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
-
-    const { container } = render(
-      <UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />,
+  it("renders a combobox", () => {
+    render(
+      <UnifiedSelector
+        value=""
+        options={[{ value: "option1", label: "Option One" }]}
+        loading={false}
+        config={CHAT}
+        onChange={vi.fn()}
+      />,
     );
 
-    const placeholder = container.querySelector(".ant-select-selection-placeholder");
-    expect(placeholder).toHaveTextContent(config.selectorPlaceholder);
+    expect(screen.getByRole("combobox")).toBeInTheDocument();
   });
 
-  it("should display loading placeholder when loading", () => {
-    const onChange = vi.fn();
-    const options = [{ value: "option1", label: "Option 1" }];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
+  it("prompts with the endpoint's own selector copy", () => {
+    render(<UnifiedSelector value="" options={[]} loading={false} config={CHAT} onChange={vi.fn()} />);
 
-    const { container } = render(
-      <UnifiedSelector value="" options={options} loading={true} config={config} onChange={onChange} />,
-    );
-
-    const placeholder = container.querySelector(".ant-select-selection-placeholder");
-    expect(placeholder).toHaveTextContent(`Loading ${config.selectorLabel.toLowerCase()}s...`);
+    expect(promptIsVisible(CHAT.selectorPlaceholder)).toBe(true);
   });
 
-  it("should call onChange when option is selected", async () => {
+  it("prompts with the agent endpoint's copy when configured for agents", () => {
+    render(<UnifiedSelector value="" options={[]} loading={false} config={AGENTS} onChange={vi.fn()} />);
+
+    expect(promptIsVisible(AGENTS.selectorPlaceholder)).toBe(true);
+  });
+
+  it("swaps the prompt for a loading message while options are in flight", () => {
+    render(<UnifiedSelector value="" options={[]} loading config={CHAT} onChange={vi.fn()} />);
+
+    expect(promptIsVisible(`Loading ${CHAT.selectorLabel.toLowerCase()}s...`)).toBe(true);
+  });
+
+  it("reports the chosen option's value, not its label", async () => {
     const user = userEvent.setup();
     const onChange = vi.fn();
-    const options = [
-      { value: "option1", label: "Option 1" },
-      { value: "option2", label: "Option 2" },
-    ];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
-
-    render(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
-
-    const select = screen.getByRole("combobox");
-    await user.click(select);
-
-    await waitFor(() => {
-      const option = screen.getByText("Option 1");
-      expect(option).toBeInTheDocument();
-    });
-
-    const option = screen.getByText("Option 1");
-    await user.click(option);
-
-    await waitFor(() => {
-      expect(onChange).toHaveBeenCalled();
-    });
-    const callArgs = onChange.mock.calls[0];
-    expect(callArgs[0]).toBe("option1");
-  });
-
-  it("should display selected value", () => {
-    const onChange = vi.fn();
-    const options = [
-      { value: "option1", label: "Option 1" },
-      { value: "option2", label: "Option 2" },
-    ];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
-
-    const { container } = render(
-      <UnifiedSelector value="option1" options={options} loading={false} config={config} onChange={onChange} />,
+    render(
+      <UnifiedSelector
+        value=""
+        options={[
+          { value: "option1", label: "Option One" },
+          { value: "option2", label: "Option Two" },
+        ]}
+        loading={false}
+        config={CHAT}
+        onChange={onChange}
+      />,
     );
 
-    const selectedValue = container.querySelector(".ant-select-selection-item");
-    expect(selectedValue).toHaveTextContent("Option 1");
+    await openList(user);
+
+    // antd renders a hidden measurement copy of each option, so take the last
+    // match: it is the live one, and it is the only one when there is no copy.
+    const matches = await screen.findAllByText("Option Two");
+    await user.click(matches[matches.length - 1]);
+
+    await waitFor(() => expect(onChange).toHaveBeenCalled());
+    expect(onChange.mock.calls[0][0]).toBe("option2");
   });
 
-  it("should filter options by search input", async () => {
+  it("narrows the list as the user searches", async () => {
     const user = userEvent.setup();
-    const onChange = vi.fn();
-    const options = [
-      { value: "option1", label: "Option One" },
-      { value: "option2", label: "Option Two" },
-      { value: "option3", label: "Different" },
-    ];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
+    render(
+      <UnifiedSelector
+        value=""
+        options={[
+          { value: "option1", label: "Option One" },
+          { value: "option2", label: "Option Two" },
+          { value: "option3", label: "Different" },
+        ]}
+        loading={false}
+        config={CHAT}
+        onChange={vi.fn()}
+      />,
+    );
 
-    render(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
-
-    const select = screen.getByRole("combobox");
-    await user.click(select);
-    await user.type(select, "One");
+    const combobox = screen.getByRole("combobox");
+    await user.click(combobox);
+    await user.type(combobox, "One");
 
     await waitFor(() => {
-      expect(screen.getByText("Option One")).toBeInTheDocument();
+      expect(screen.getAllByText("Option One").length).toBeGreaterThan(0);
       expect(screen.queryByText("Option Two")).not.toBeInTheDocument();
       expect(screen.queryByText("Different")).not.toBeInTheDocument();
     });
   });
 
-  it("should show loading spinner in notFoundContent when loading", async () => {
+  it("shows a busy indicator in the empty list while loading", async () => {
     const user = userEvent.setup();
-    const onChange = vi.fn();
-    const options: { value: string; label: string }[] = [];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
+    render(<UnifiedSelector value="" options={[]} loading config={CHAT} onChange={vi.fn()} />);
 
-    render(<UnifiedSelector value="" options={options} loading={true} config={config} onChange={onChange} />);
-
-    const select = screen.getByRole("combobox");
-    await user.click(select);
+    await openList(user);
 
     await waitFor(() => {
-      const spin = document.querySelector(".ant-spin");
-      expect(spin).toBeInTheDocument();
+      expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument();
     });
   });
 
-  it("should show no options message when not loading and no options", async () => {
+  it("says so when there is nothing to pick and nothing is loading", async () => {
     const user = userEvent.setup();
-    const onChange = vi.fn();
-    const options: { value: string; label: string }[] = [];
-    const config = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
+    render(<UnifiedSelector value="" options={[]} loading={false} config={CHAT} onChange={vi.fn()} />);
 
-    render(<UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />);
+    await openList(user);
 
-    const select = screen.getByRole("combobox");
-    await user.click(select);
-
-    await waitFor(() => {
-      expect(screen.getByText(`No ${config.selectorLabel.toLowerCase()}s available`)).toBeInTheDocument();
-    });
-  });
-
-  it("should work with agent endpoint config", () => {
-    const onChange = vi.fn();
-    const options = [{ value: "agent1", label: "Agent One" }];
-    const config = ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS];
-
-    const { container } = render(
-      <UnifiedSelector value="" options={options} loading={false} config={config} onChange={onChange} />,
-    );
-
-    const placeholder = container.querySelector(".ant-select-selection-placeholder");
-    expect(placeholder).toHaveTextContent(config.selectorPlaceholder);
+    expect(await screen.findByText(`No ${CHAT.selectorLabel.toLowerCase()}s available`)).toBeInTheDocument();
   });
 });

From b341a223399874314f151c0667b50b4b4255ee02 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:57:03 -0700
Subject: [PATCH 14/31] fix(ui): keep the guardrails metric cards wrapping at
 narrow widths

The antd Row and Col the migration replaced never squeezed the summary
cards below their content width; they wrapped onto a second line instead.
Measured on a live dashboard, antd laid out five per row at 1280 and 1024,
then four plus one at 900 and three plus two at 820, never narrower than
about 126px. A fixed grid-cols-5 kept all five on one line and compressed
them to 86px at 820, so the metric values overflowed their cards.

An auto-fit track with a 7rem minimum reproduces antd's wrap points and
card widths exactly at all four measured viewports, and is identical at
1280, so the route's visual baseline is unchanged.

The added test fails against grid-cols-5 and passes against the track
---
 .../_components/GuardrailsOverview.test.tsx              | 9 +++++++++
 .../_components/GuardrailsOverview.tsx                   | 2 +-
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
index b616982d69b..2aff2083870 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
@@ -152,6 +152,15 @@ describe("GuardrailsOverview", () => {
     await waitFor(() => expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument());
   });
 
+  it("wraps the metric cards on a flexible track instead of a fixed column count", async () => {
+    renderOverview();
+
+    const grid = (await screen.findByText("Total Evaluations")).closest("div.grid");
+
+    expect(grid).not.toBeNull();
+    expect(grid!.className).toMatch(/grid-cols-\[repeat\(auto-fit,minmax\([^)]+,1fr\)\)\]/);
+  });
+
   it("shows a failure message when the usage request rejects", async () => {
     mockGetGuardrailsUsageOverview.mockRejectedValue(new Error("network down"));
     renderOverview();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
index 3d91fc1b9b6..e1048c5322f 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.tsx
@@ -211,7 +211,7 @@ export function GuardrailsOverview({
         </div>
       </div>
 
-      <div className="mb-6 grid grid-cols-5 gap-4">
+      <div className="mb-6 grid grid-cols-[repeat(auto-fit,minmax(7rem,1fr))] gap-4">
         <MetricCard label="Total Evaluations" value={metrics.totalRequests.toLocaleString()} />
         <MetricCard
           label="Blocked Requests"

From 9a26da90a47459d249030b6873ec99619161fdc9 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 15:59:53 -0700
Subject: [PATCH 15/31] refactor(ui): migrate playground to shadcn

Replaces antd and Tremor with shadcn primitives across the six route-owned,
form-free playground components: the compare view and its panel, message input
and unified selector, plus the realtime playground and the agent builder.

Markup only, no behaviour change. The characterisation tests added in the
previous commit are untouched here and stay green through the swap.

Adds ui/slider.tsx via the shadcn CLI and retires the six antd
no-restricted-imports suppressions the migration made obsolete.
---
 ui/litellm-dashboard/eslint-suppressions.json |  30 +-
 .../components/chat_ui/AgentBuilderView.tsx   | 530 +++++++++---------
 .../components/chat_ui/RealtimePlayground.tsx |  77 +--
 .../components/compareUI/CompareUI.tsx        |  76 ++-
 .../compareUI/components/ComparisonPanel.tsx  |  75 ++-
 .../compareUI/components/MessageInput.tsx     |  33 +-
 .../compareUI/components/UnifiedSelector.tsx  |  66 ++-
 .../src/components/ui/slider.tsx              |  41 ++
 8 files changed, 514 insertions(+), 414 deletions(-)
 create mode 100644 ui/litellm-dashboard/src/components/ui/slider.tsx

diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json
index d7f71a5840d..0c5bcc65702 100644
--- a/ui/litellm-dashboard/eslint-suppressions.json
+++ b/ui/litellm-dashboard/eslint-suppressions.json
@@ -1055,9 +1055,6 @@
     "no-nested-ternary": {
       "count": 2
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/set-state-in-effect": {
       "count": 5
     }
@@ -1098,9 +1095,6 @@
     "no-nested-ternary": {
       "count": 2
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/immutability": {
       "count": 2
     },
@@ -1115,9 +1109,6 @@
     "no-nested-ternary": {
       "count": 4
     },
-    "no-restricted-imports": {
-      "count": 1
-    },
     "react-hooks/set-state-in-effect": {
       "count": 1
     }
@@ -1125,9 +1116,6 @@
   "src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx": {
     "local/no-complex-jsx-arrow": {
       "count": 2
-    },
-    "no-restricted-imports": {
-      "count": 1
     }
   },
   "src/app/(dashboard)/playground/components/compareUI/components/MessageDisplay.tsx": {
@@ -1135,21 +1123,11 @@
       "count": 1
     }
   },
-  "src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/playground/components/compareUI/components/ModelSelector.tsx": {
     "no-restricted-imports": {
       "count": 2
     }
   },
-  "src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx": {
-    "no-restricted-imports": {
-      "count": 1
-    }
-  },
   "src/app/(dashboard)/playground/components/complianceUI/ComplianceUI.tsx": {
     "local/no-complex-jsx-arrow": {
       "count": 2
@@ -3582,6 +3560,14 @@
       "count": 1
     }
   },
+  "src/components/ui/slider.tsx": {
+    "local/filename-pascal-case": {
+      "count": 1
+    },
+    "no-nested-ternary": {
+      "count": 1
+    }
+  },
   "src/components/ui/switch.tsx": {
     "local/filename-pascal-case": {
       "count": 1
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
index d4333b95c62..44995707b10 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
@@ -1,16 +1,24 @@
 "use client";
 
-import {
-  CommentOutlined,
-  DeleteOutlined,
-  ExperimentOutlined,
-  LinkOutlined,
-  PlusOutlined,
-  RobotOutlined,
-  SaveOutlined,
-} from "@ant-design/icons";
-import { Button, Input, Modal, Select, Spin, Tabs } from "antd";
+import { Bot, FlaskConical, Link as LinkIcon, MessageSquare, Plus, Save, Trash2 } from "lucide-react";
 import React, { useCallback, useEffect, useState } from "react";
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
+import { Textarea } from "@/components/ui/textarea";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
+import { MultiSelect } from "@/components/shared/MultiSelect";
+import { useVisitedTabs } from "@/hooks/useVisitedTabs";
 import CodeBlock from "@/components/CodeBlock";
 import NotificationsManager from "@/components/molecules/notifications_manager";
 import {
@@ -27,8 +35,6 @@ import { fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_m
 import ComplianceUI from "../complianceUI/ComplianceUI";
 import ChatUI from "./ChatUI";
 
-const { TextArea } = Input;
-
 export interface AgentBuilderViewProps {
   accessToken: string | null;
   token: string | null;
@@ -45,6 +51,8 @@ export interface AgentBuilderViewProps {
 
 const NEW_AGENT_ID = "__new__";
 
+type AgentTab = "configure" | "chat" | "test" | "connect";
+
 function getConnectTabBaseUrl(
   proxySettings: AgentBuilderViewProps["proxySettings"],
   customProxyBaseUrl?: string,
@@ -116,7 +124,7 @@ function ConnectTabContent({
           Create a virtual key that can only call this agent. The key will be scoped to you (user_id) and restricted to
           the model <span className="font-mono text-gray-800">{agentName}</span>.
         </p>
-        <Button type="primary" onClick={onCreateKey} loading={creatingKey} disabled={disabledPersonalKeyCreation}>
+        <Button onClick={onCreateKey} disabled={creatingKey || disabledPersonalKeyCreation}>
           Create key for this agent
         </Button>
         {disabledPersonalKeyCreation && (
@@ -190,7 +198,12 @@ export default function AgentBuilderView({
   const [modelGroups, setModelGroups] = useState<ModelGroup[]>([]);
   const [loadingAgents, setLoadingAgents] = useState(true);
   const [selectedId, setSelectedId] = useState<string | null>(null);
-  const [activeTab, setActiveTab] = useState<"configure" | "chat" | "test" | "connect">("configure");
+  const [activeTab, setActiveTab] = useState<AgentTab>("configure");
+  const { onTabChange, hasVisited } = useVisitedTabs("configure");
+  const goToTab = (tab: AgentTab) => {
+    setActiveTab(tab);
+    onTabChange(tab);
+  };
   const [creatingKey, setCreatingKey] = useState(false);
   const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
 
@@ -207,6 +220,7 @@ export default function AgentBuilderView({
 
   const [saving, setSaving] = useState(false);
   const [deleting, setDeleting] = useState(false);
+  const [confirmingDelete, setConfirmingDelete] = useState(false);
 
   const effectiveApiKey = apiKey || accessToken || "";
   const selectedAgent =
@@ -314,7 +328,7 @@ export default function AgentBuilderView({
     setDraftTemperature(0.7);
     setDraftMaxTokens(4096);
     setDraftTools([]);
-    setActiveTab("configure");
+    goToTab("configure");
   };
 
   const handleSaveAgent = async () => {
@@ -344,7 +358,7 @@ export default function AgentBuilderView({
         ? list.find((a) => getAgentModelId(a) === createdId) ?? list.find((a) => a.model_name === draftName.trim())
         : list.find((a) => a.model_name === draftName.trim());
       setSelectedId(created ? getAgentSelectionKey(created) : list[0] ? getAgentSelectionKey(list[0]) : null);
-      setActiveTab("chat");
+      goToTab("chat");
     } catch (e) {
       NotificationsManager.fromBackend("Failed to save agent");
     } finally {
@@ -411,27 +425,24 @@ export default function AgentBuilderView({
 
   const handleDeleteAgent = () => {
     if (!selectedAgent || !selectedAgentModelId || !accessToken) return;
-    Modal.confirm({
-      title: "Delete agent",
-      content: `Are you sure you want to delete "${selectedAgent.model_name}"? This cannot be undone.`,
-      okText: "Delete",
-      okType: "danger",
-      cancelText: "Cancel",
-      onOk: async () => {
-        setDeleting(true);
-        try {
-          await modelDeleteCall(accessToken, selectedAgentModelId);
-          NotificationsManager.success("Agent deleted");
-          const list = await loadAgents();
-          const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId);
-          setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null);
-        } catch (e) {
-          NotificationsManager.fromBackend("Failed to delete agent");
-        } finally {
-          setDeleting(false);
-        }
-      },
-    });
+    setConfirmingDelete(true);
+  };
+
+  const handleConfirmDelete = async () => {
+    if (!selectedAgent || !selectedAgentModelId || !accessToken) return;
+    setDeleting(true);
+    try {
+      await modelDeleteCall(accessToken, selectedAgentModelId);
+      NotificationsManager.success("Agent deleted");
+      const list = await loadAgents();
+      const remaining = list.filter((a) => getAgentModelId(a) !== selectedAgentModelId);
+      setSelectedId(remaining.length > 0 ? getAgentSelectionKey(remaining[0]) : null);
+    } catch (e) {
+      NotificationsManager.fromBackend("Failed to delete agent");
+    } finally {
+      setDeleting(false);
+      setConfirmingDelete(false);
+    }
   };
 
   if (!accessToken || !userID || !userRole) {
@@ -446,13 +457,8 @@ export default function AgentBuilderView({
         <div className="flex h-12 items-center justify-between px-4">
           <span className="text-sm font-medium text-gray-900">Agent Builder</span>
           {isNewAgent ? (
-            <Button
-              type="primary"
-              icon={<SaveOutlined />}
-              onClick={handleSaveAgent}
-              loading={saving}
-              disabled={!draftName?.trim() || !draftUnderlyingModel}
-            >
+            <Button onClick={handleSaveAgent} disabled={saving || !draftName?.trim() || !draftUnderlyingModel}>
+              <Save />
               Save Agent
             </Button>
           ) : (
@@ -460,7 +466,7 @@ export default function AgentBuilderView({
           )}
         </div>
         <div className="flex items-center gap-2 border-t border-amber-200 bg-amber-50 px-4 py-2 text-xs text-amber-800">
-          <ExperimentOutlined className="shrink-0 text-amber-600" />
+          <FlaskConical className="size-4 shrink-0 text-amber-600" />
           <span>
             Agent Builder is experimental and may change or be removed without notice. We’d love your feedback—email us
             at{" "}
@@ -477,12 +483,14 @@ export default function AgentBuilderView({
         <div className="w-60 shrink-0 border-r border-gray-200 bg-white flex flex-col">
           <div className="flex items-center justify-between border-b border-gray-200 p-3">
             <span className="text-xs font-semibold uppercase tracking-wide text-gray-500">Agents</span>
-            <Button type="text" size="small" icon={<PlusOutlined />} onClick={handleAddAgent} aria-label="Add agent" />
+            <Button variant="ghost" size="icon-sm" onClick={handleAddAgent} aria-label="Add agent">
+              <Plus />
+            </Button>
           </div>
           <div className="flex-1 overflow-y-auto p-2">
             {loadingAgents ? (
-              <div className="flex justify-center py-4">
-                <Spin size="small" />
+              <div className="flex justify-center py-4" aria-busy="true">
+                <UiLoadingSpinner className="size-4 text-gray-400" />
               </div>
             ) : (
               <>
@@ -509,7 +517,7 @@ export default function AgentBuilderView({
                   onClick={handleAddAgent}
                   className="mb-1 w-full rounded-md border border-dashed border-gray-300 px-3 py-2 text-left text-sm text-gray-500 hover:border-blue-400 hover:bg-blue-50/50 hover:text-gray-700"
                 >
-                  <PlusOutlined className="mr-1" /> New agent
+                  <Plus className="mr-1 inline size-4" /> New agent
                 </button>
               </>
             )}
@@ -526,227 +534,231 @@ export default function AgentBuilderView({
           {(selectedId !== null || isNewAgent) && (
             <>
               <Tabs
-                activeKey={activeTab}
-                onChange={(k) => setActiveTab(k as "configure" | "chat" | "test" | "connect")}
-                className="flex-1 overflow-hidden [&_.ant-tabs-content]:h-full [&_.ant-tabs-tabpane]:h-full [&_.ant-tabs-nav]:pl-4"
-                items={[
-                  {
-                    key: "configure",
-                    label: (
-                      <span>
-                        <RobotOutlined className="mr-1" /> Configure
-                      </span>
-                    ),
-                    children: (
-                      <div className="h-full overflow-y-auto p-6">
-                        {isNewAgent || selectedAgent ? (
-                          <div className="mx-auto max-w-xl space-y-4">
-                            {!selectedAgentModelId && selectedAgent && (
-                              <div className="rounded-sm border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-800">
-                                This agent cannot be updated or deleted here (missing model id). Manage it from Models
-                                &amp; Endpoints.
-                              </div>
-                            )}
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">Agent name</label>
-                              <Input
-                                value={draftName}
-                                onChange={(e) => setDraftName(e.target.value)}
-                                placeholder="My Agent"
-                              />
-                            </div>
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">System prompt</label>
-                              <TextArea
-                                value={draftSystemPrompt}
-                                onChange={(e) => setDraftSystemPrompt(e.target.value)}
-                                placeholder="You are a helpful assistant..."
-                                rows={6}
-                              />
-                            </div>
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">Underlying LLM</label>
-                              <Select
-                                value={draftUnderlyingModel}
-                                onChange={setDraftUnderlyingModel}
-                                className="w-full"
-                                options={modelGroups.map((m) => ({ value: m.model_group, label: m.model_group }))}
-                                placeholder="Select model"
-                              />
-                            </div>
-                            <div className="grid grid-cols-2 gap-4">
-                              <div>
-                                <label className="mb-1 block text-sm font-medium text-gray-700">Temperature</label>
-                                <Input
-                                  type="number"
-                                  min={0}
-                                  max={2}
-                                  step={0.1}
-                                  value={draftTemperature}
-                                  onChange={(e) => setDraftTemperature(Number(e.target.value))}
-                                />
-                              </div>
-                              <div>
-                                <label className="mb-1 block text-sm font-medium text-gray-700">Max tokens</label>
-                                <Input
-                                  type="number"
-                                  min={1}
-                                  value={draftMaxTokens}
-                                  onChange={(e) => setDraftMaxTokens(Number(e.target.value))}
-                                />
-                              </div>
-                            </div>
-                            <div>
-                              <label className="mb-1 block text-sm font-medium text-gray-700">MCP servers</label>
-                              <Select
-                                mode="multiple"
-                                placeholder="Select MCP servers to attach (same format as chat completions API)"
-                                value={selectedMCPServerIds}
-                                onChange={handleMCPServerChange}
-                                loading={loadingMCPServers}
-                                className="w-full"
-                                allowClear
-                                showSearch
-                                optionFilterProp="label"
-                                options={mcpServers.map((s) => ({
-                                  value: s.server_id,
-                                  label: s.alias || s.server_name || s.server_id,
-                                }))}
-                              />
-                              {selectedAgent && draftTools.length > 0 && (
-                                <p className="mt-1 text-xs text-gray-500">
-                                  {draftTools.length} MCP server{draftTools.length !== 1 ? "s" : ""} saved. Use the same{" "}
-                                  <code className="rounded-sm bg-gray-100 px-1">tools</code> array in chat completions
-                                  when calling this agent.
-                                </p>
-                              )}
-                            </div>
-                            {selectedAgent && (
-                              <div className="flex flex-wrap items-center gap-2 pt-2">
-                                {selectedAgentModelId && (
-                                  <>
-                                    <Button
-                                      type="primary"
-                                      icon={<SaveOutlined />}
-                                      onClick={handleUpdateAgent}
-                                      loading={saving}
-                                      disabled={!draftName?.trim() || !draftUnderlyingModel}
-                                    >
-                                      Update Agent
-                                    </Button>
-                                    <Button
-                                      type="default"
-                                      danger
-                                      icon={<DeleteOutlined />}
-                                      onClick={handleDeleteAgent}
-                                      loading={deleting}
-                                    >
-                                      Delete
-                                    </Button>
-                                  </>
-                                )}
-                                <Button type="primary" icon={<CommentOutlined />} onClick={() => setActiveTab("chat")}>
-                                  Test in Chat
+                value={activeTab}
+                onValueChange={(value) => goToTab(value as AgentTab)}
+                className="flex flex-1 flex-col overflow-hidden"
+              >
+                <TabsList variant="line" className="h-auto w-full justify-start rounded-none border-b p-0 pl-4">
+                  <TabsTrigger value="configure" className="flex-none rounded-none px-4 py-2">
+                    <Bot />
+                    Configure
+                  </TabsTrigger>
+                  <TabsTrigger value="chat" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
+                    <MessageSquare />
+                    Chat
+                  </TabsTrigger>
+                  <TabsTrigger value="test" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
+                    <FlaskConical />
+                    Batch Test
+                  </TabsTrigger>
+                  <TabsTrigger value="connect" disabled={isNewAgent} className="flex-none rounded-none px-4 py-2">
+                    <LinkIcon />
+                    Connect
+                  </TabsTrigger>
+                </TabsList>
+
+                <TabsContent
+                  value="configure"
+                  keepMounted={hasVisited("configure")}
+                  className="min-h-0 overflow-hidden"
+                >
+                  <div className="h-full overflow-y-auto p-6">
+                    {isNewAgent || selectedAgent ? (
+                      <div className="mx-auto max-w-xl space-y-4">
+                        {!selectedAgentModelId && selectedAgent && (
+                          <div className="rounded-sm border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-800">
+                            This agent cannot be updated or deleted here (missing model id). Manage it from Models &amp;
+                            Endpoints.
+                          </div>
+                        )}
+                        <div>
+                          <label className="mb-1 block text-sm font-medium text-gray-700">Agent name</label>
+                          <Input
+                            value={draftName}
+                            onChange={(e) => setDraftName(e.target.value)}
+                            placeholder="My Agent"
+                          />
+                        </div>
+                        <div>
+                          <label className="mb-1 block text-sm font-medium text-gray-700">System prompt</label>
+                          <Textarea
+                            value={draftSystemPrompt}
+                            onChange={(e) => setDraftSystemPrompt(e.target.value)}
+                            placeholder="You are a helpful assistant..."
+                            rows={6}
+                            className="field-sizing-fixed"
+                          />
+                        </div>
+                        <div>
+                          <label className="mb-1 block text-sm font-medium text-gray-700">Underlying LLM</label>
+                          <Select
+                            value={draftUnderlyingModel ?? null}
+                            onValueChange={(model: string | null) => setDraftUnderlyingModel(model ?? undefined)}
+                          >
+                            <SelectTrigger className="w-full" aria-label="Underlying LLM">
+                              <SelectValue placeholder="Select model" />
+                            </SelectTrigger>
+                            <SelectContent>
+                              {modelGroups.map((m) => (
+                                <SelectItem key={m.model_group} value={m.model_group}>
+                                  {m.model_group}
+                                </SelectItem>
+                              ))}
+                            </SelectContent>
+                          </Select>
+                        </div>
+                        <div className="grid grid-cols-2 gap-4">
+                          <div>
+                            <label className="mb-1 block text-sm font-medium text-gray-700">Temperature</label>
+                            <Input
+                              type="number"
+                              min={0}
+                              max={2}
+                              step={0.1}
+                              value={draftTemperature}
+                              onChange={(e) => setDraftTemperature(Number(e.target.value))}
+                            />
+                          </div>
+                          <div>
+                            <label className="mb-1 block text-sm font-medium text-gray-700">Max tokens</label>
+                            <Input
+                              type="number"
+                              min={1}
+                              value={draftMaxTokens}
+                              onChange={(e) => setDraftMaxTokens(Number(e.target.value))}
+                            />
+                          </div>
+                        </div>
+                        <div>
+                          <label className="mb-1 block text-sm font-medium text-gray-700">MCP servers</label>
+                          <MultiSelect
+                            placeholder="Select MCP servers to attach (same format as chat completions API)"
+                            value={selectedMCPServerIds}
+                            onValueChange={handleMCPServerChange}
+                            loading={loadingMCPServers}
+                            className="w-full"
+                            options={mcpServers.map((s) => ({
+                              value: s.server_id,
+                              label: s.alias || s.server_name || s.server_id,
+                            }))}
+                          />
+                          {selectedAgent && draftTools.length > 0 && (
+                            <p className="mt-1 text-xs text-gray-500">
+                              {draftTools.length} MCP server{draftTools.length !== 1 ? "s" : ""} saved. Use the same{" "}
+                              <code className="rounded-sm bg-gray-100 px-1">tools</code> array in chat completions when
+                              calling this agent.
+                            </p>
+                          )}
+                        </div>
+                        {selectedAgent && (
+                          <div className="flex flex-wrap items-center gap-2 pt-2">
+                            {selectedAgentModelId && (
+                              <>
+                                <Button
+                                  onClick={handleUpdateAgent}
+                                  disabled={saving || !draftName?.trim() || !draftUnderlyingModel}
+                                >
+                                  <Save />
+                                  Update Agent
                                 </Button>
-                              </div>
+                                <Button variant="destructive" onClick={handleDeleteAgent} disabled={deleting}>
+                                  <Trash2 />
+                                  Delete
+                                </Button>
+                              </>
                             )}
-                          </div>
-                        ) : null}
-                      </div>
-                    ),
-                  },
-                  {
-                    key: "chat",
-                    label: (
-                      <span>
-                        <CommentOutlined className="mr-1" /> Chat
-                      </span>
-                    ),
-                    disabled: isNewAgent,
-                    children: (
-                      <div className="flex h-full flex-col min-h-0">
-                        {selectedAgent ? (
-                          <ChatUI
-                            key={selectedAgent.model_name}
-                            simplified
-                            fixedModel={selectedAgent.model_name}
-                            accessToken={accessToken}
-                            token={token}
-                            userRole={userRole}
-                            userID={userID}
-                            disabledPersonalKeyCreation={disabledPersonalKeyCreation}
-                            proxySettings={proxySettings}
-                          />
-                        ) : (
-                          <div className="flex flex-1 items-center justify-center text-gray-500">
-                            Save an agent first to test in Chat.
+                            <Button onClick={() => goToTab("chat")}>
+                              <MessageSquare />
+                              Test in Chat
+                            </Button>
                           </div>
                         )}
                       </div>
-                    ),
-                  },
-                  {
-                    key: "test",
-                    label: (
-                      <span>
-                        <ExperimentOutlined className="mr-1" /> Batch Test
-                      </span>
-                    ),
-                    disabled: isNewAgent,
-                    children: (
-                      <div className="flex h-full flex-col min-h-0">
-                        {selectedAgent ? (
-                          <ComplianceUI
-                            accessToken={accessToken}
-                            disabledPersonalKeyCreation={disabledPersonalKeyCreation}
-                            backendMode="chat_completions"
-                            fixedModel={selectedAgent.model_name}
-                            proxySettings={proxySettings}
-                          />
-                        ) : (
-                          <div className="flex flex-1 items-center justify-center text-gray-500">
-                            Select an agent to run batch tests.
-                          </div>
-                        )}
+                    ) : null}
+                  </div>
+                </TabsContent>
+                <TabsContent value="chat" keepMounted={hasVisited("chat")} className="min-h-0 overflow-hidden">
+                  <div className="flex h-full flex-col min-h-0">
+                    {selectedAgent ? (
+                      <ChatUI
+                        key={selectedAgent.model_name}
+                        simplified
+                        fixedModel={selectedAgent.model_name}
+                        accessToken={accessToken}
+                        token={token}
+                        userRole={userRole}
+                        userID={userID}
+                        disabledPersonalKeyCreation={disabledPersonalKeyCreation}
+                        proxySettings={proxySettings}
+                      />
+                    ) : (
+                      <div className="flex flex-1 items-center justify-center text-gray-500">
+                        Save an agent first to test in Chat.
                       </div>
-                    ),
-                  },
-                  {
-                    key: "connect",
-                    label: (
-                      <span>
-                        <LinkOutlined className="mr-1" /> Connect
-                      </span>
-                    ),
-                    disabled: isNewAgent,
-                    children: (
-                      <div className="h-full overflow-y-auto p-6">
-                        {selectedAgent ? (
-                          <ConnectTabContent
-                            agentName={selectedAgent.model_name}
-                            proxySettings={proxySettings}
-                            customProxyBaseUrl={customProxyBaseUrl}
-                            accessToken={accessToken}
-                            userID={userID}
-                            disabledPersonalKeyCreation={disabledPersonalKeyCreation}
-                            creatingKey={creatingKey}
-                            createdKeyValue={createdKeyValue}
-                            onCreateKey={handleCreateKeyForAgent}
-                          />
-                        ) : (
-                          <div className="flex flex-1 items-center justify-center text-gray-500">
-                            Select an agent to see how to connect.
-                          </div>
-                        )}
+                    )}
+                  </div>
+                </TabsContent>
+                <TabsContent value="test" keepMounted={hasVisited("test")} className="min-h-0 overflow-hidden">
+                  <div className="flex h-full flex-col min-h-0">
+                    {selectedAgent ? (
+                      <ComplianceUI
+                        accessToken={accessToken}
+                        disabledPersonalKeyCreation={disabledPersonalKeyCreation}
+                        backendMode="chat_completions"
+                        fixedModel={selectedAgent.model_name}
+                        proxySettings={proxySettings}
+                      />
+                    ) : (
+                      <div className="flex flex-1 items-center justify-center text-gray-500">
+                        Select an agent to run batch tests.
                       </div>
-                    ),
-                  },
-                ]}
-              />
+                    )}
+                  </div>
+                </TabsContent>
+                <TabsContent value="connect" keepMounted={hasVisited("connect")} className="min-h-0 overflow-hidden">
+                  <div className="h-full overflow-y-auto p-6">
+                    {selectedAgent ? (
+                      <ConnectTabContent
+                        agentName={selectedAgent.model_name}
+                        proxySettings={proxySettings}
+                        customProxyBaseUrl={customProxyBaseUrl}
+                        accessToken={accessToken}
+                        userID={userID}
+                        disabledPersonalKeyCreation={disabledPersonalKeyCreation}
+                        creatingKey={creatingKey}
+                        createdKeyValue={createdKeyValue}
+                        onCreateKey={handleCreateKeyForAgent}
+                      />
+                    ) : (
+                      <div className="flex flex-1 items-center justify-center text-gray-500">
+                        Select an agent to see how to connect.
+                      </div>
+                    )}
+                  </div>
+                </TabsContent>
+              </Tabs>
             </>
           )}
         </div>
       </div>
+
+      <AlertDialog open={confirmingDelete} onOpenChange={setConfirmingDelete}>
+        <AlertDialogContent>
+          <AlertDialogHeader>
+            <AlertDialogTitle>Delete agent</AlertDialogTitle>
+            <AlertDialogDescription>
+              Are you sure you want to delete &quot;{selectedAgent?.model_name}&quot;? This cannot be undone.
+            </AlertDialogDescription>
+          </AlertDialogHeader>
+          <AlertDialogFooter>
+            <AlertDialogAction variant="outline">Cancel</AlertDialogAction>
+            {/* Not AlertDialogAction: that closes on click, and the dialog must stay up while the delete is in flight */}
+            <Button variant="destructive" onClick={handleConfirmDelete} disabled={deleting}>
+              Delete
+            </Button>
+          </AlertDialogFooter>
+        </AlertDialogContent>
+      </AlertDialog>
     </div>
   );
 }
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx
index 2bf645fced2..4b90f693f2a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/RealtimePlayground.tsx
@@ -1,13 +1,13 @@
 "use client";
 
-import { AudioMutedOutlined, AudioOutlined, CloseCircleOutlined, SendOutlined, SoundOutlined } from "@ant-design/icons";
-import { Button, Input, Select, Typography } from "antd";
+import { CircleX, Mic, MicOff, Send, Volume2 } from "lucide-react";
 import React, { useCallback, useEffect, useRef, useState } from "react";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
 import { getProxyBaseUrl } from "@/components/networking";
 import { OPEN_AI_VOICE_SELECT_OPTIONS } from "./chatConstants";
 
-const { Text } = Typography;
-
 interface RealtimeMessage {
   role: "user" | "assistant" | "system" | "status";
   content: string;
@@ -364,28 +364,37 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
       {/* Header */}
       <div className="flex items-center justify-between px-4 py-3 border-b border-gray-200 bg-gray-50">
         <div className="flex items-center gap-3">
-          <SoundOutlined className="text-lg text-blue-500" />
-          <Text className="font-semibold text-gray-800">Realtime Voice Chat</Text>
+          <Volume2 className="size-5 text-blue-500" />
+          <span className="font-semibold text-gray-800">Realtime Voice Chat</span>
           <span className={`inline-block w-2 h-2 rounded-full ${isConnected ? "bg-green-500" : "bg-gray-300"}`} />
-          <Text className="text-xs text-gray-500">
+          <span className="text-xs text-gray-500">
             {isConnected ? "Connected" : isConnecting ? "Connecting..." : "Disconnected"}
-          </Text>
+          </span>
         </div>
         <div className="flex items-center gap-2">
           <Select
-            size="small"
             value={selectedVoice}
-            onChange={setSelectedVoice}
-            options={OPEN_AI_VOICE_SELECT_OPTIONS}
-            style={{ width: 220 }}
+            onValueChange={(voice) => setSelectedVoice(voice ?? selectedVoice)}
             disabled={isConnected}
-          />
+          >
+            <SelectTrigger size="sm" className="w-[220px]" aria-label="Voice">
+              <SelectValue>{OPEN_AI_VOICE_SELECT_OPTIONS.find((v) => v.value === selectedVoice)?.label}</SelectValue>
+            </SelectTrigger>
+            <SelectContent>
+              {OPEN_AI_VOICE_SELECT_OPTIONS.map((voice) => (
+                <SelectItem key={voice.value} value={voice.value}>
+                  {voice.label}
+                </SelectItem>
+              ))}
+            </SelectContent>
+          </Select>
           {!isConnected ? (
-            <Button type="primary" onClick={connect} loading={isConnecting} size="small">
+            <Button onClick={connect} disabled={isConnecting} size="sm">
               Connect
             </Button>
           ) : (
-            <Button danger onClick={disconnect} size="small" icon={<CloseCircleOutlined />}>
+            <Button variant="destructive" onClick={disconnect} size="sm">
+              <CircleX />
               Disconnect
             </Button>
           )}
@@ -396,12 +405,12 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
       <div className="flex-1 overflow-y-auto p-4 space-y-3">
         {messages.length === 0 && !isConnected && (
           <div className="flex flex-col items-center justify-center h-full text-gray-400 gap-3">
-            <SoundOutlined style={{ fontSize: 48 }} />
-            <Text className="text-lg text-gray-500">Realtime Voice Playground</Text>
-            <Text className="text-sm text-gray-400 text-center max-w-md">
+            <Volume2 className="size-12" />
+            <span className="text-lg text-gray-500">Realtime Voice Playground</span>
+            <p className="text-sm text-gray-400 text-center max-w-md">
               Click <b>Connect</b> to start a realtime session. You can speak using your microphone or type messages.
               The AI will respond with voice and text.
-            </Text>
+            </p>
           </div>
         )}
         {messages.map((msg, i) => (
@@ -433,30 +442,26 @@ const RealtimePlayground: React.FC<RealtimePlaygroundProps> = ({
         <div className="border-t border-gray-200 p-3 bg-white">
           <div className="flex items-center gap-2">
             <Button
-              shape="circle"
-              size="large"
-              type={isRecording ? "primary" : "default"}
-              danger={isRecording}
-              icon={isRecording ? <AudioMutedOutlined /> : <AudioOutlined />}
+              size="icon-lg"
+              variant={isRecording ? "destructive" : "outline"}
               onClick={isRecording ? stopRecording : startRecording}
               title={isRecording ? "Stop recording" : "Start recording"}
-              className={isRecording ? "animate-pulse" : ""}
-            />
+              className={`rounded-full ${isRecording ? "animate-pulse" : ""}`}
+            >
+              {isRecording ? <MicOff /> : <Mic />}
+            </Button>
             <Input
               placeholder="Type a message or use the mic..."
               value={inputText}
               onChange={(e) => setInputText(e.target.value)}
-              onPressEnter={sendTextMessage}
-              className="flex-1"
-              size="large"
-            />
-            <Button
-              type="primary"
-              icon={<SendOutlined />}
-              onClick={sendTextMessage}
-              disabled={!inputText.trim()}
-              size="large"
+              onKeyDown={(e) => {
+                if (e.key === "Enter") sendTextMessage();
+              }}
+              className="h-10 flex-1"
             />
+            <Button size="icon-lg" onClick={sendTextMessage} disabled={!inputText.trim()} aria-label="Send">
+              <Send />
+            </Button>
           </div>
           {isRecording && (
             <div className="mt-2 flex items-center gap-2 text-red-500 text-xs">
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
index 3b396924aba..efb8e809352 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
@@ -2,9 +2,12 @@
 
 import NotificationsManager from "@/components/molecules/notifications_manager";
 import { DEBOUNCE_WAIT_MS } from "@/utils/debounceConstants";
-import { ClearOutlined, DeleteOutlined, FilePdfOutlined, PlusOutlined } from "@ant-design/icons";
+import { Eraser, FileText, Plus, Trash2 } from "lucide-react";
 import { useDebouncedValue } from "@tanstack/react-pacer/debouncer";
-import { Button, Input, Select, Tooltip } from "antd";
+import { Button } from "@/components/ui/button";
+import { Input } from "@/components/ui/input";
+import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select";
+import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip";
 import { useEffect, useMemo, useState } from "react";
 import { v4 as uuidv4 } from "uuid";
 import ChatImageUpload from "../chat_ui/ChatImageUpload";
@@ -692,17 +695,22 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
               <span className="text-sm font-medium text-gray-600">Virtual Key Source</span>
               <Select
                 value={apiKeySource}
-                onChange={(value) => setApiKeySource(value as "session" | "custom")}
+                onValueChange={(value) => setApiKeySource(value as "session" | "custom")}
                 disabled={disabledPersonalKeyCreation}
-                className="w-48"
               >
-                <Select.Option value="session" disabled={!canUseSessionKey}>
-                  Current UI Session
-                </Select.Option>
-                <Select.Option value="custom">Virtual Key</Select.Option>
+                <SelectTrigger className="w-48" aria-label="Virtual Key Source">
+                  <SelectValue>{apiKeySource === "custom" ? "Virtual Key" : "Current UI Session"}</SelectValue>
+                </SelectTrigger>
+                <SelectContent>
+                  <SelectItem value="session" disabled={!canUseSessionKey}>
+                    Current UI Session
+                  </SelectItem>
+                  <SelectItem value="custom">Virtual Key</SelectItem>
+                </SelectContent>
               </Select>
               {apiKeySource === "custom" && (
-                <Input.Password
+                <Input
+                  type="password"
                   value={customApiKey}
                   onChange={(event) => setCustomApiKey(event.target.value)}
                   placeholder="Enter Virtual Key"
@@ -712,30 +720,35 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
             </div>
             <div className="flex items-center gap-2">
               <span className="text-sm font-medium text-gray-600">Endpoint</span>
-              <Select
-                value={selectedEndpoint}
-                onChange={(value) => setSelectedEndpoint(value as EndpointIdType)}
-                className="w-56"
-              >
-                {getAvailableEndpoints().map((endpoint) => (
-                  <Select.Option key={endpoint.value} value={endpoint.value}>
-                    {endpoint.label}
-                  </Select.Option>
-                ))}
+              <Select value={selectedEndpoint} onValueChange={(value) => setSelectedEndpoint(value as EndpointIdType)}>
+                <SelectTrigger className="w-56" aria-label="Endpoint">
+                  <SelectValue>{endpointConfig.label}</SelectValue>
+                </SelectTrigger>
+                <SelectContent>
+                  {getAvailableEndpoints().map((endpoint) => (
+                    <SelectItem key={endpoint.value} value={endpoint.value}>
+                      {endpoint.label}
+                    </SelectItem>
+                  ))}
+                </SelectContent>
               </Select>
             </div>
             <div className="flex items-center gap-3">
-              <Button onClick={clearAllChats} disabled={!hasMessages} icon={<ClearOutlined />}>
+              <Button variant="outline" onClick={clearAllChats} disabled={!hasMessages}>
+                <Eraser />
                 Clear All Chats
               </Button>
-              <Tooltip
-                title={
-                  comparisons.length >= maxComparisons ? "Compare up to 3 models at a time" : "Add another comparison"
-                }
-              >
-                <Button onClick={addComparison} disabled={comparisons.length >= maxComparisons} icon={<PlusOutlined />}>
-                  Add Comparison
-                </Button>
+              <Tooltip>
+                {/* A disabled button emits no pointer events, so the hover target has to be a wrapper */}
+                <TooltipTrigger render={<span className="inline-flex" />}>
+                  <Button variant="outline" onClick={addComparison} disabled={comparisons.length >= maxComparisons}>
+                    <Plus />
+                    Add Comparison
+                  </Button>
+                </TooltipTrigger>
+                <TooltipContent>
+                  {comparisons.length >= maxComparisons ? "Compare up to 3 models at a time" : "Add another comparison"}
+                </TooltipContent>
               </Tooltip>
             </div>
           </div>
@@ -807,8 +820,8 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
                   <div className="flex items-center gap-3 p-3 bg-gray-50 rounded-lg border border-gray-200">
                     <div className="relative inline-block">
                       {isUploadedFilePdf ? (
-                        <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center">
-                          <FilePdfOutlined style={{ fontSize: "16px", color: "white" }} />
+                        <div className="w-10 h-10 rounded-md bg-red-500 flex items-center justify-center text-white">
+                          <FileText className="size-4" aria-label="file-pdf" />
                         </div>
                       ) : (
                         <img
@@ -825,8 +838,9 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
                     <button
                       className="flex items-center justify-center w-6 h-6 text-gray-400 hover:text-gray-600 hover:bg-gray-200 rounded-full transition-colors"
                       onClick={handleRemoveFile}
+                      aria-label="Remove attachment"
                     >
-                      <DeleteOutlined style={{ fontSize: "12px" }} />
+                      <Trash2 className="size-3" />
                     </button>
                   </div>
                 </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx
index 3f853f42dc7..5ea77b2828a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.tsx
@@ -1,12 +1,15 @@
 import { Settings, X } from "lucide-react";
-import { useState } from "react";
+import { useId, useState } from "react";
 import { ComparisonInstance } from "../CompareUI";
 import { MessageDisplay } from "./MessageDisplay";
 import { UnifiedSelector } from "./UnifiedSelector";
 import TagSelector from "@/components/tag_management/TagSelector";
 import VectorStoreSelector from "@/components/vector_store_management/VectorStoreSelector";
 import GuardrailSelector from "@/components/guardrails/GuardrailSelector";
-import { Checkbox, Divider, Popover, Slider } from "antd";
+import { Checkbox } from "@/components/ui/checkbox";
+import { Popover, PopoverContent, PopoverTrigger } from "@/components/ui/popover";
+import { Separator } from "@/components/ui/separator";
+import { Slider } from "@/components/ui/slider";
 import { SelectorOption, EndpointConfig, isAgentEndpoint, getComparisonSelection } from "../endpoint_config";
 
 interface ComparisonPanelProps {
@@ -35,6 +38,8 @@ export function ComparisonPanel({
   const isA2AMode = isAgentEndpoint(endpointConfig.id);
   const currentSelection = getComparisonSelection(comparison, endpointConfig.id);
   const [popoverVisible, setPopoverVisible] = useState(false);
+  const syncId = useId();
+  const advancedParamsId = useId();
 
   const handleSyncChange = (checked: boolean) => {
     if (checked) {
@@ -103,12 +108,18 @@ export function ComparisonPanel({
       <div className="space-y-2">
         {/* Sync Checkbox */}
         <div className="flex items-center gap-2">
-          <Checkbox checked={comparison.applyAcrossModels} onChange={(e) => handleSyncChange(e.target.checked)}>
-            <span className="text-xs font-medium">Sync Settings Across Models</span>
-          </Checkbox>
+          <Checkbox
+            id={syncId}
+            checked={comparison.applyAcrossModels}
+            onCheckedChange={handleSyncChange}
+            aria-label="Sync Settings Across Models"
+          />
+          <label htmlFor={syncId} className="cursor-pointer text-xs font-medium">
+            Sync Settings Across Models
+          </label>
         </div>
 
-        <Divider className="border-gray-200" />
+        <Separator className="my-3" />
 
         {/* General Settings */}
         <div>
@@ -146,11 +157,14 @@ export function ComparisonPanel({
           <div className="space-y-2">
             <div className="flex items-center gap-2 pb-1">
               <Checkbox
+                id={advancedParamsId}
                 checked={comparison.useAdvancedParams}
-                onChange={(e) => handleAdvancedParamsChange(e.target.checked)}
-              >
-                <span className="text-sm font-medium">Use Advanced Parameters</span>
-              </Checkbox>
+                onCheckedChange={handleAdvancedParamsChange}
+                aria-label="Use Advanced Parameters"
+              />
+              <label htmlFor={advancedParamsId} className="cursor-pointer text-sm font-medium">
+                Use Advanced Parameters
+              </label>
             </div>
             <div className="space-y-2 transition-opacity duration-200" style={{ opacity: disabledOpacity }}>
               <div>
@@ -162,8 +176,8 @@ export function ComparisonPanel({
                   min={0}
                   max={2}
                   step={0.01}
-                  value={comparison.temperature}
-                  onChange={(value) => {
+                  value={[comparison.temperature]}
+                  onValueChange={(value) => {
                     const nextValue = Array.isArray(value) ? value[0] : value;
                     const clamped = Math.min(2, Math.max(0, Number(nextValue.toFixed(2))));
                     handleSettingChange("temperature", clamped);
@@ -180,8 +194,8 @@ export function ComparisonPanel({
                   min={1}
                   max={32768}
                   step={1}
-                  value={comparison.maxTokens}
-                  onChange={(value) => {
+                  value={[comparison.maxTokens]}
+                  onValueChange={(value) => {
                     const nextValue = Array.isArray(value) ? value[0] : value;
                     const clamped = Math.min(32768, Math.max(1, Math.round(nextValue)));
                     handleSettingChange("maxTokens", clamped);
@@ -209,26 +223,29 @@ export function ComparisonPanel({
           />
           <div className="flex items-center gap-2">
             <Popover
-              content={settingsContent}
-              trigger={[]}
               open={popoverVisible}
               onOpenChange={() => {
                 // Prevent automatic closing - we control it manually
               }}
-              placement="bottomRight"
-              destroyTooltipOnHide={false}
             >
-              <button
-                onClick={(event) => {
-                  event.stopPropagation();
-                  handleTogglePopover();
-                }}
-                className={`p-2 rounded-lg transition-colors ${
-                  popoverVisible ? "bg-gray-200 text-gray-700" : "hover:bg-gray-100 text-gray-600"
-                }`}
-              >
-                <Settings size={18} />
-              </button>
+              <PopoverTrigger
+                render={
+                  <button
+                    onClick={(event) => {
+                      event.stopPropagation();
+                      handleTogglePopover();
+                    }}
+                    className={`p-2 rounded-lg transition-colors ${
+                      popoverVisible ? "bg-gray-200 text-gray-700" : "hover:bg-gray-100 text-gray-600"
+                    }`}
+                  >
+                    <Settings size={18} />
+                  </button>
+                }
+              />
+              <PopoverContent side="bottom" align="end" className="w-auto">
+                {settingsContent}
+              </PopoverContent>
             </Popover>
           </div>
         </div>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx
index fae964b372a..b27341ff7a5 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/MessageInput.tsx
@@ -1,8 +1,7 @@
 import React from "react";
-import { Input, Button } from "antd";
-import { ArrowUpOutlined } from "@ant-design/icons";
-
-const { TextArea } = Input;
+import { ArrowUp } from "lucide-react";
+import { Button } from "@/components/ui/button";
+import { Textarea } from "@/components/ui/textarea";
 
 interface MessageInputProps {
   value: string;
@@ -29,25 +28,25 @@ export function MessageInput({ value, onChange, onSend, disabled, hasAttachment,
     <div className="flex items-center gap-2">
       <div className="flex items-center flex-1 bg-white border border-gray-300 rounded-xl px-3 py-1 min-h-[44px]">
         {uploadComponent && <div className="shrink-0 mr-2">{uploadComponent}</div>}
-        <TextArea
+        <Textarea
           value={value}
           onChange={(e) => onChange(e.target.value)}
           onKeyDown={handleKeyDown}
           placeholder="Type your message... (Shift+Enter for new line)"
           disabled={disabled}
-          className="flex-1"
-          autoSize={{ minRows: 1, maxRows: 4 }}
-          style={{
-            resize: "none",
-            border: "none",
-            boxShadow: "none",
-            background: "transparent",
-            padding: "4px 0",
-            fontSize: "14px",
-            lineHeight: "20px",
-          }}
+          rows={1}
+          className="max-h-20 min-h-0 flex-1 resize-none overflow-y-auto border-0 bg-transparent px-0 py-1 text-sm leading-5 shadow-none focus-visible:ring-0"
         />
-        <Button onClick={onSend} disabled={!canSend} icon={<ArrowUpOutlined />} shape="circle" />
+        <Button
+          onClick={onSend}
+          disabled={!canSend}
+          size="icon-sm"
+          variant="outline"
+          className="rounded-full"
+          aria-label="Send message"
+        >
+          <ArrowUp />
+        </Button>
       </div>
     </div>
   );
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx
index c37580efd39..915f0dbb320 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.tsx
@@ -3,7 +3,15 @@
  * based on the current endpoint configuration.
  */
 
-import { Select, Spin } from "antd";
+import {
+  Combobox,
+  ComboboxContent,
+  ComboboxEmpty,
+  ComboboxInput,
+  ComboboxItem,
+  ComboboxList,
+} from "@/components/ui/combobox";
+import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner";
 import { SelectorOption, EndpointConfig } from "../endpoint_config";
 
 interface UnifiedSelectorProps {
@@ -14,26 +22,44 @@ interface UnifiedSelectorProps {
   onChange: (value: string) => void;
 }
 
+const matchesQuery = (option: SelectorOption, query: string): boolean =>
+  option.label.toLowerCase().includes(query.trim().toLowerCase());
+
 export function UnifiedSelector({ value, options, loading, config, onChange }: UnifiedSelectorProps) {
+  const selected = options.find((option) => option.value === value) ?? null;
+  const noun = config.selectorLabel.toLowerCase();
+
   return (
-    <Select
-      value={value || undefined}
-      placeholder={loading ? `Loading ${config.selectorLabel.toLowerCase()}s...` : config.selectorPlaceholder}
-      onChange={onChange}
-      loading={loading}
-      showSearch
-      filterOption={(input, option) => (option?.label ?? "").toLowerCase().includes(input.toLowerCase())}
-      options={options}
-      className="w-48 md:w-64 lg:w-72"
-      notFoundContent={
-        loading ? (
-          <div className="flex items-center justify-center py-2">
-            <Spin size="small" />
-          </div>
-        ) : (
-          `No ${config.selectorLabel.toLowerCase()}s available`
-        )
-      }
-    />
+    <Combobox
+      items={options}
+      value={selected}
+      onValueChange={(option: SelectorOption | null) => onChange(option?.value ?? "")}
+      isItemEqualToValue={(a: SelectorOption, b: SelectorOption) => a.value === b.value}
+      itemToStringLabel={(option: SelectorOption) => option.label}
+      filter={matchesQuery}
+    >
+      <ComboboxInput
+        placeholder={loading ? `Loading ${noun}s...` : config.selectorPlaceholder}
+        className="w-48 md:w-64 lg:w-72"
+      />
+      <ComboboxContent>
+        <ComboboxEmpty>
+          {loading ? (
+            <span aria-busy="true" className="flex items-center justify-center py-2">
+              <UiLoadingSpinner className="size-4" />
+            </span>
+          ) : (
+            `No ${noun}s available`
+          )}
+        </ComboboxEmpty>
+        <ComboboxList>
+          {(option: SelectorOption) => (
+            <ComboboxItem key={option.value} value={option}>
+              {option.label}
+            </ComboboxItem>
+          )}
+        </ComboboxList>
+      </ComboboxContent>
+    </Combobox>
   );
 }
diff --git a/ui/litellm-dashboard/src/components/ui/slider.tsx b/ui/litellm-dashboard/src/components/ui/slider.tsx
new file mode 100644
index 00000000000..aaa7df73799
--- /dev/null
+++ b/ui/litellm-dashboard/src/components/ui/slider.tsx
@@ -0,0 +1,41 @@
+import { Slider as SliderPrimitive } from "@base-ui/react/slider";
+
+import { cn } from "@/lib/cva.config";
+
+function Slider({ className, defaultValue, value, min = 0, max = 100, ...props }: SliderPrimitive.Root.Props) {
+  const _values = Array.isArray(value) ? value : Array.isArray(defaultValue) ? defaultValue : [min, max];
+
+  return (
+    <SliderPrimitive.Root
+      className={cn("data-horizontal:w-full data-vertical:h-full", className)}
+      data-slot="slider"
+      defaultValue={defaultValue}
+      value={value}
+      min={min}
+      max={max}
+      thumbAlignment="edge"
+      {...props}
+    >
+      <SliderPrimitive.Control className="relative flex w-full touch-none items-center select-none data-disabled:opacity-50 data-vertical:h-full data-vertical:min-h-40 data-vertical:w-auto data-vertical:flex-col">
+        <SliderPrimitive.Track
+          data-slot="slider-track"
+          className="relative grow overflow-hidden rounded-full bg-muted select-none data-horizontal:h-1.5 data-horizontal:w-full data-vertical:h-full data-vertical:w-1.5"
+        >
+          <SliderPrimitive.Indicator
+            data-slot="slider-range"
+            className="bg-primary select-none data-horizontal:h-full data-vertical:w-full"
+          />
+        </SliderPrimitive.Track>
+        {Array.from({ length: _values.length }, (_, index) => (
+          <SliderPrimitive.Thumb
+            data-slot="slider-thumb"
+            key={index}
+            className="block size-4 shrink-0 rounded-full border border-primary bg-white shadow-sm ring-ring/50 transition-[color,box-shadow] select-none hover:ring-4 focus-visible:ring-4 focus-visible:outline-hidden disabled:pointer-events-none disabled:opacity-50"
+          />
+        ))}
+      </SliderPrimitive.Control>
+    </SliderPrimitive.Root>
+  );
+}
+
+export { Slider };

From 30b933dd5ebd577d6faa8943746173d9c8c471dc Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 16:21:01 -0700
Subject: [PATCH 16/31] test(ui): drop the metric grid class assertion

The assertion pinned the Tailwind track string, which jsdom can never
evaluate: it does no layout, so the test could not fail for the reason
that matters, and its pattern accepted any minmax minimum, staying green
if that minimum changed enough to break the layout outright.

Only a viewport-resize browser test can observe this, which belongs in
tests/e2e/ui rather than a route's unit tests. The responsive behaviour
itself is unchanged; the widths it reproduces are recorded in b341a22339
---
 .../_components/GuardrailsOverview.test.tsx              | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
index 2aff2083870..b616982d69b 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/GuardrailsOverview.test.tsx
@@ -152,15 +152,6 @@ describe("GuardrailsOverview", () => {
     await waitFor(() => expect(document.querySelector('[aria-busy="true"]')).toBeInTheDocument());
   });
 
-  it("wraps the metric cards on a flexible track instead of a fixed column count", async () => {
-    renderOverview();
-
-    const grid = (await screen.findByText("Total Evaluations")).closest("div.grid");
-
-    expect(grid).not.toBeNull();
-    expect(grid!.className).toMatch(/grid-cols-\[repeat\(auto-fit,minmax\([^)]+,1fr\)\)\]/);
-  });
-
   it("shows a failure message when the usage request rejects", async () => {
     mockGetGuardrailsUsageOverview.mockRejectedValue(new Error("network down"));
     renderOverview();

From 0ebbda92d7bbfd136565f8b741727293f080ace0 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 16:51:22 -0700
Subject: [PATCH 17/31] style(ui): drop the playground migration's explanatory
 comments

The repository's comment policy keeps reasoning out of the source, so the
notes on the delete dialog, the tooltip wrapper, and the test query choices
move to the pull request instead.
---
 .../playground/components/chat_ui/AgentBuilderView.test.tsx | 6 ------
 .../playground/components/chat_ui/AgentBuilderView.tsx      | 1 -
 .../playground/components/compareUI/CompareUI.tsx           | 1 -
 .../compareUI/components/ComparisonPanel.test.tsx           | 2 --
 .../compareUI/components/UnifiedSelector.test.tsx           | 4 ----
 5 files changed, 14 deletions(-)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
index 179b7b8557a..e30f87115cf 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.test.tsx
@@ -38,8 +38,6 @@ vi.mock("@/components/molecules/notifications_manager", () => ({
   default: { success: vi.fn(), fromBackend: vi.fn() },
 }));
 
-// Both tab bodies keep their own state, which is what the tab container has to
-// preserve when the user moves away and comes back.
 const StatefulPanel = ({ label }: { label: string }) => {
   const [draft, setDraft] = useState("");
   return <input aria-label={label} value={draft} onChange={(event) => setDraft(event.target.value)} />;
@@ -73,8 +71,6 @@ const props = {
   userRole: "Admin",
 };
 
-// The labels in the configure panel are not wired to their controls, so reach
-// the control through the field that the label heads.
 const controlUnder = (label: string): HTMLElement =>
   within(screen.getByText(label).parentElement!).getByRole("combobox");
 
@@ -212,8 +208,6 @@ describe("AgentBuilderView", () => {
     expect(await screen.findByText(/Are you sure you want to delete "support-agent"/)).toBeInTheDocument();
     expect(modelDeleteCall).not.toHaveBeenCalled();
 
-    // The trigger and the confirmation share a label; the confirmation is the
-    // one rendered last, in the overlay.
     const confirmations = screen.getAllByRole("button", { name: /Delete$/ });
     await user.click(confirmations[confirmations.length - 1]);
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
index 44995707b10..5a86b2d27f4 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/chat_ui/AgentBuilderView.tsx
@@ -752,7 +752,6 @@ export default function AgentBuilderView({
           </AlertDialogHeader>
           <AlertDialogFooter>
             <AlertDialogAction variant="outline">Cancel</AlertDialogAction>
-            {/* Not AlertDialogAction: that closes on click, and the dialog must stay up while the delete is in flight */}
             <Button variant="destructive" onClick={handleConfirmDelete} disabled={deleting}>
               Delete
             </Button>
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
index efb8e809352..984e04e93f1 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/CompareUI.tsx
@@ -739,7 +739,6 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
                 Clear All Chats
               </Button>
               <Tooltip>
-                {/* A disabled button emits no pointer events, so the hover target has to be a wrapper */}
                 <TooltipTrigger render={<span className="inline-flex" />}>
                   <Button variant="outline" onClick={addComparison} disabled={comparisons.length >= maxComparisons}>
                     <Plus />
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
index f18ea86b8d3..7b8c2a72055 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/ComparisonPanel.test.tsx
@@ -76,8 +76,6 @@ const mockProps = {
   apiKey: "test-api-key",
 };
 
-// Both the settings and the remove control are icon-only buttons whose lucide
-// icon survives the migration, so locate them by that icon.
 const buttonWithIcon = (icon: string): HTMLButtonElement => {
   const match = Array.from(document.querySelectorAll("button")).find((button) =>
     button.querySelector(`svg.lucide-${icon}`),
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
index 4686804afec..29a8092e2b0 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/components/compareUI/components/UnifiedSelector.test.tsx
@@ -7,8 +7,6 @@ import { EndpointId, ENDPOINT_CONFIGS } from "../endpoint_config";
 const CHAT = ENDPOINT_CONFIGS[EndpointId.CHAT_COMPLETIONS];
 const AGENTS = ENDPOINT_CONFIGS[EndpointId.A2A_AGENTS];
 
-// A select's prompt text is a rendered node in one library and an input
-// placeholder attribute in the other, so accept either.
 const promptIsVisible = (text: string) =>
   screen.queryByText(text) !== null || screen.queryByPlaceholderText(text) !== null;
 
@@ -67,8 +65,6 @@ describe("UnifiedSelector", () => {
 
     await openList(user);
 
-    // antd renders a hidden measurement copy of each option, so take the last
-    // match: it is the live one, and it is the only one when there is no copy.
     const matches = await screen.findAllByText("Option Two");
     await user.click(matches[matches.length - 1]);
 

From 62523c1734ef27d648f58614dff9562f970dfcd2 Mon Sep 17 00:00:00 2001
From: Yuneng Jiang <yuneng@berri.ai>
Date: Thu, 13 Aug 2026 17:27:26 -0700
Subject: [PATCH 18/31] fix(ui): restore pattern search and popup layering in
 the guardrail wizard

Two behaviour differences the port introduced, both found in browser QA

The antd Select matched a prebuilt pattern on its display_name and its
internal name; a Base UI Combobox only searches itemToStringLabel, so
queries like "amex" and "sg" stopped matching. Restore the second field
with a filter predicate on the Root, covered by a regression test that
searches on a token the visible label does not contain

Base UI portals a popup into a positioner whose "isolate z-50" is fixed
in the primitive, so inside an antd Modal at z-index 1000 the options
were visible but not clickable. antd hid this because its own dropdowns
and tooltips already sat above its Modal. The positioner is not reachable
from the call site, so this needs one app-wide rule keyed on an antd
modal being present, and it becomes deletable when the last one goes
---
 .../content_filter/PatternModal.test.tsx           | 14 ++++++++++++++
 .../_components/content_filter/PatternModal.tsx    |  6 ++++++
 ui/litellm-dashboard/src/app/globals.css           | 10 ++++++++++
 3 files changed, 30 insertions(+)

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
index 77cbcf6211d..5302ac3b7f7 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.test.tsx
@@ -84,6 +84,20 @@ describe("PatternModal", () => {
     expect(screen.queryByText("AWS access key")).not.toBeInTheDocument();
   });
 
+  it("should match the internal pattern name when it is absent from the display name", async () => {
+    const user = userEvent.setup();
+    renderModal();
+
+    expect(await screen.findByText("Add prebuilt pattern")).toBeInTheDocument();
+
+    await user.click(screen.getAllByRole("combobox")[0]);
+    await user.keyboard("ssn");
+
+    expect(await screen.findByText("US Social Security Number")).toBeInTheDocument();
+    expect(screen.queryByText("Email address")).not.toBeInTheDocument();
+    expect(screen.queryByText("Visa card")).not.toBeInTheDocument();
+  });
+
   it("should report the chosen action", async () => {
     const user = userEvent.setup();
     renderModal();
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
index df87c511cf4..2d98b67cb4a 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
+++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/content_filter/PatternModal.tsx
@@ -28,6 +28,11 @@ interface PatternGroup {
   items: PrebuiltPattern[];
 }
 
+const matchesPatternQuery = (pattern: PrebuiltPattern, query: string) => {
+  const needle = query.toLowerCase();
+  return pattern.display_name.toLowerCase().includes(needle) || pattern.name.toLowerCase().includes(needle);
+};
+
 interface PatternModalProps {
   visible: boolean;
   prebuiltPatterns: PrebuiltPattern[];
@@ -74,6 +79,7 @@ const PatternModal: React.FC<PatternModalProps> = ({
               value={selectedPattern}
               onValueChange={(pattern: PrebuiltPattern | null) => pattern && onPatternNameChange(pattern.name)}
               itemToStringLabel={(pattern: PrebuiltPattern) => pattern.display_name}
+              filter={matchesPatternQuery}
             >
               <ComboboxInput className="mt-2 w-full" placeholder="Choose pattern type" />
               <ComboboxContent>
diff --git a/ui/litellm-dashboard/src/app/globals.css b/ui/litellm-dashboard/src/app/globals.css
index 97f61a610d9..7013e502c22 100644
--- a/ui/litellm-dashboard/src/app/globals.css
+++ b/ui/litellm-dashboard/src/app/globals.css
@@ -242,3 +242,13 @@
 [data-slot="dialog-content"][data-nested-dialog-open] {
   visibility: hidden;
 }
+
+/* Base UI portals a popup into an `isolate z-50` positioner, which an antd Modal at z-index 1000
+   then paints over, so options on a half-migrated page are visible but not clickable. The
+   positioner is not reachable from the call site, hence the parent selector. Delete this once no
+   route renders an antd Modal. */
+body:has(.ant-modal-wrap) div:has(> [data-slot="select-content"]),
+body:has(.ant-modal-wrap) div:has(> [data-slot="combobox-content"]),
+body:has(.ant-modal-wrap) div:has(> [data-slot="tooltip-content"]) {
+  z-index: 1100;
+}

From cfbd43172aab1614f07193dc498b72b7cb02f670 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 13 Aug 2026 18:49:07 -0700
Subject: [PATCH 19/31] fix(scripts): end make check with a ran/skipped summary
 and verdict

---
 scripts/pre_commit_lint.sh                 | 26 ++++++++++++++
 tests/test_litellm/test_pre_commit_lint.py | 40 ++++++++++++++++++++++
 2 files changed, 66 insertions(+)

diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index afe55603466..e19012613cd 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -281,4 +281,30 @@ if [ -n "${gen_pid:-}" ]; then
     cat "$gen_log"; rm -f "$gen_log"
 fi
 
+summary_item() {
+    local check_name=$1 triggered=$2 skip_reason=$3
+    if [ -n "$triggered" ]; then
+        echo "    ran:     $check_name"
+    else
+        echo "    skipped: $check_name ($skip_reason)"
+    fi
+}
+
+echo "check: summary"
+summary_item "Python lint (make lint)" "$litellm_py_files" "no litellm/ Python files in scope"
+summary_item "tests/e2e checks (basedpyright + raw HTTP client ban)" "$e2e_py_files" "no tests/e2e Python files in scope"
+summary_item "dashboard lint (prettier + eslint + lint budgets)" "$ui_prettier_changed$ui_eslint_changed" "no dashboard files in scope"
+summary_item "dashboard API-type sync (npm run gen:api)" "$spec_files" "no litellm/proxy, litellm/types, or generator files in scope"
+
+if [ -z "$litellm_py_files$e2e_py_files$ui_prettier_changed$ui_eslint_changed$spec_files" ]; then
+    echo "check: NOTE - no gating lint check matches the files in scope, so nothing ran:" >&2
+    printf '%s\n' "$scope" | sed 's/^/    /' >&2
+    echo "  A pass here is a no-op, not a lint verdict." >&2
+fi
+
+if [ "$status" -eq 0 ]; then
+    echo "check: PASS"
+else
+    echo "check: FAIL"
+fi
 exit $status
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index 35d98903226..7ce16032641 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -384,3 +384,43 @@ def test_a_failing_block_fails_the_whole_run(tmp_path: Path, fail: str, message:
     proc = _run(repo, bin_dir, {"STUB_FAIL": fail})
     assert proc.returncode == 1
     assert message in proc.stdout + proc.stderr
+
+
+def test_run_ends_with_a_summary_of_ran_and_skipped_blocks(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    proc = _run(repo, bin_dir, {})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "check: summary" in proc.stdout
+    assert "ran:     Python lint (make lint)" in proc.stdout
+    assert "ran:     dashboard lint (prettier + eslint + lint budgets)" in proc.stdout
+    assert "ran:     dashboard API-type sync (npm run gen:api)" in proc.stdout
+    assert "skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope)" in proc.stdout
+    assert "check: PASS" in proc.stdout
+    assert "check: FAIL" not in proc.stdout
+
+
+def test_staged_files_matching_no_check_print_an_explicit_noop_note_and_nonempty_log(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    _commit_all(repo, "base")
+    tests_dir = repo / "tests" / "test_litellm"
+    tests_dir.mkdir(parents=True)
+    (tests_dir / "test_x.py").write_text("def test_x() -> None: ...\n")
+    subprocess.run(["git", "add", "tests"], cwd=repo, check=True)
+    proc = _run(repo, bin_dir, {})
+    assert proc.returncode == 0, proc.stdout + proc.stderr
+    assert "no gating lint check matches the files in scope, so nothing ran" in proc.stdout
+    assert "tests/test_litellm/test_x.py" in proc.stdout
+    assert "a no-op, not a lint verdict" in proc.stdout
+    assert "check: PASS" in proc.stdout
+    assert "linting Python" not in proc.stdout
+    log = (repo / ".git" / "pre_commit_lint.log").read_text()
+    assert "check: summary" in log
+    assert "skipped: Python lint (make lint) (no litellm/ Python files in scope)" in log
+
+
+def test_failing_run_ends_with_a_fail_verdict(tmp_path: Path) -> None:
+    repo, bin_dir = _sandbox(tmp_path)
+    proc = _run(repo, bin_dir, {"STUB_FAIL": "make-lint"})
+    assert proc.returncode == 1
+    assert "check: FAIL" in proc.stdout
+    assert "check: PASS" not in proc.stdout

From 909a2e6232957b104cb2fd3ad18d14f3e2c1bbd8 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 13 Aug 2026 18:56:37 -0700
Subject: [PATCH 20/31] perf(spend-logs): bound retention cleanup so one run
 cannot saturate the database (#36594)

---
 litellm/constants.py                          |   3 +
 litellm/proxy/_types.py                       |  16 +
 .../db_transaction_queue/spend_log_cleanup.py | 453 ++++++++++--
 .../spend_log_cleanup_metrics.py              | 122 ++++
 .../spend_logs_partition_manager.py           | 128 +++-
 litellm/proxy/proxy_server.py                 |  28 +-
 .../test_spend_logs_partition_manager.py      | 171 ++++-
 tests/test_litellm/proxy/test_proxy_server.py |  85 +++
 .../proxy/test_spend_log_cleanup.py           | 665 +++++++++++++++++-
 .../hooks/proxyConfig/useProxyConfig.ts       |   4 +
 .../useStoreRequestInSpendLogs.ts             |  12 +-
 .../LoggingSettings/LoggingSettings.test.tsx  | 319 ++++++++-
 .../LoggingSettings/LoggingSettings.tsx       | 229 ++++--
 ui/litellm-dashboard/src/lib/http/schema.d.ts |  20 +
 14 files changed, 2050 insertions(+), 205 deletions(-)
 create mode 100644 litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py

diff --git a/litellm/constants.py b/litellm/constants.py
index 554165f5d39..e9d2d719ae2 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1491,6 +1491,9 @@ SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEA
 SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float(
     os.getenv("SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS", 0.5)
 )
+SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300"))
+SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30"))
+SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000"))
 TOOL_SPEND_TOP_TOOLS: Final = 100
 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day")
 SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7))
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index 6c01575e271..385f39e02a4 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -2514,6 +2514,22 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase):
         None,
         description="If True and LiteLLM_SpendLogs has been converted to a range-partitioned table (db_scripts/partition_spend_logs.sql), retention cleanup drops expired partitions instead of deleting rows, and pre-creates upcoming partitions. Default is False.",
     )
+    maximum_spend_logs_cleanup_batch_size: int | None = Field(
+        None,
+        description="Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000.",
+    )
+    maximum_spend_logs_cleanup_max_batches: int | None = Field(
+        None,
+        description="Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500.",
+    )
+    maximum_spend_logs_cleanup_run_budget: str | None = Field(
+        None,
+        description="Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'.",
+    )
+    maximum_spend_logs_cleanup_batch_timeout: str | None = Field(
+        None,
+        description="Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'.",
+    )
     mcp_internal_ip_ranges: list[str] | None = Field(
         None,
         description="Custom CIDR ranges that define internal/private networks for MCP access control. When set, only these ranges are treated as internal. Defaults to RFC 1918 private ranges (10.0.0.0/8, 172.16.0.0/12, 192.168.0.0/16, 127.0.0.0/8).",
diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
index 9f01c719a5f..d19023862cb 100644
--- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
+++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
@@ -1,22 +1,60 @@
 import asyncio
+import time
+from dataclasses import dataclass
 from datetime import datetime, timedelta, timezone
-from typing import Final
+from typing import Final, Literal, TypeAlias
+
+from pydantic import BaseModel, TypeAdapter
 
 from litellm._logging import verbose_proxy_logger
 from litellm.caching import RedisCache
 from litellm.constants import (
     SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS,
     SPEND_LOG_CLEANUP_BATCH_SIZE,
+    SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS,
     SPEND_LOG_CLEANUP_JOB_NAME,
     SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES,
+    SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
+    SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
     SPEND_LOG_RUN_LOOPS,
 )
 from litellm.litellm_core_utils.duration_parser import duration_in_seconds
+from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
+    RunOutcome,
+    SpendLogCleanupMetrics,
+)
 from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
+    RemainingTimeoutMs,
     SpendLogsPartitionManager,
 )
 from litellm.proxy.utils import PrismaClient
 
+StopReason: TypeAlias = Literal["exhausted", "budget_exhausted", "batch_cap_reached", "aborted"]
+
+
+@dataclass(frozen=True, slots=True)
+class TableCleanupResult:
+    """Outcome of pruning one table, so the caller can report why a run ended."""
+
+    rows_deleted: int
+    stop_reason: StopReason
+
+
+class _RemainingRow(BaseModel):
+    """One row of the capped outstanding-rows probe, validated out of prisma's untyped result."""
+
+    remaining: int
+
+
+_REMAINING_ROWS: Final = TypeAdapter(list[_RemainingRow])
+
+SPEND_LOG_CLEANUP_BOUND_SETTINGS: Final = (
+    "maximum_spend_logs_cleanup_batch_size",
+    "maximum_spend_logs_cleanup_max_batches",
+    "maximum_spend_logs_cleanup_run_budget",
+    "maximum_spend_logs_cleanup_batch_timeout",
+)
+
 
 class SpendLogCleanup:
     """
@@ -26,6 +64,24 @@ class SpendLogCleanup:
     dropping whole partitions (instant, frees disk immediately). Otherwise it
     falls back to deleting logs in batches.
     Uses PodLockManager to ensure only one pod runs cleanup in multi-pod deployments.
+
+    Every run is bounded so it can never monopolise the database: a wall-clock
+    budget shared across all tables, a per-table batch cap, and a Postgres
+    statement/lock timeout on every statement the job issues, deletes and the
+    outstanding-rows probe alike. A run that hits a bound stops cleanly and the
+    next run resumes from where it left off, because the cutoff is recomputed
+    and deleted rows are gone.
+
+    The budget is a hard wall clock, not an advisory one. Every statement this
+    job issues, deletes, the outstanding-rows probe and partition DDL alike, is
+    issued with a timeout clamped to the budget that is still left, so one
+    started just under the deadline is cancelled by Postgres at the deadline
+    rather than running a further batch timeout past it. No statement is issued
+    at all once the budget is spent, which is why the probe is skipped on that
+    path. Partition DDL additionally carries a lock_timeout, because it takes an
+    ACCESS EXCLUSIVE lock and would otherwise queue behind a long-running reader
+    for as long as that reader lives; a partition this run cannot get is left
+    for the next one.
     """
 
     def __init__(
@@ -34,17 +90,88 @@ class SpendLogCleanup:
         redis_cache: RedisCache | None = None,
         partition_manager: SpendLogsPartitionManager | None = None,
     ):
-        self.batch_size = SPEND_LOG_CLEANUP_BATCH_SIZE
         self.retention_seconds: int | None = None
         self.partition_manager = partition_manager or SpendLogsPartitionManager()
         from litellm.proxy.proxy_server import general_settings as default_settings
 
         self.general_settings = general_settings or default_settings
+        self._refresh_bounds()
         from litellm.proxy.proxy_server import proxy_logging_obj
 
         pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager
         self.pod_lock_manager = pod_lock_manager
-        verbose_proxy_logger.info("SpendLogCleanup initialized with batch size: %s", self.batch_size)
+        verbose_proxy_logger.info(
+            "SpendLogCleanup initialized: batch_size=%s max_batches=%s run_budget=%ss batch_timeout=%ss",
+            self.batch_size,
+            self.max_batches,
+            self.run_budget_seconds,
+            self.batch_timeout_seconds,
+        )
+
+    def _refresh_bounds(self) -> None:
+        """
+        Re-read every bound in SPEND_LOG_CLEANUP_BOUND_SETTINGS from settings.
+
+        The scheduler holds one long-lived instance, so a bound captured at
+        construction would never reflect a dashboard change. general_settings is
+        the same dict the periodic config reload mutates in place, so reading it
+        per run is what makes these knobs live. Every bound falls back to its
+        shipped default, so clearing a field restores that default.
+        """
+        self.batch_size: int = self._positive_int_setting(
+            "maximum_spend_logs_cleanup_batch_size", SPEND_LOG_CLEANUP_BATCH_SIZE
+        )
+        self.max_batches: int = self._positive_int_setting(
+            "maximum_spend_logs_cleanup_max_batches", SPEND_LOG_RUN_LOOPS
+        )
+        self.run_budget_seconds: float = self._duration_setting(
+            "maximum_spend_logs_cleanup_run_budget", SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
+        )
+        self.batch_timeout_seconds: float = self._duration_setting(
+            "maximum_spend_logs_cleanup_batch_timeout", SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS
+        )
+
+    def _positive_int_setting(self, setting_name: str, default: int) -> int:
+        """
+        Read a positive-integer knob, falling back to the default when unset or unusable.
+        """
+        raw: Final = self.general_settings.get(setting_name)
+        if raw is None:
+            return default
+        try:
+            parsed: Final = int(raw)
+        except (TypeError, ValueError):
+            verbose_proxy_logger.warning("Invalid %s value: %s, using default %s", setting_name, raw, default)
+            return default
+        if parsed <= 0:
+            verbose_proxy_logger.warning("%s must be positive, got %s, using default %s", setting_name, parsed, default)
+            return default
+        return parsed
+
+    def _duration_setting(self, setting_name: str, default_seconds: float) -> float:
+        """
+        Read a duration knob (e.g. '5m'), falling back to the default when unset or unusable.
+
+        The knob must never be able to remove the bound it exists to enforce, so
+        anything the parser rejects (including the non-finite spellings 'inf' and
+        'nan') and anything non-positive falls back rather than being honoured.
+        """
+        raw: Final = self.general_settings.get(setting_name)
+        if raw is None:
+            return default_seconds
+        try:
+            parsed: Final = float(duration_in_seconds(str(raw)))
+        except (ValueError, TypeError) as e:
+            verbose_proxy_logger.warning(
+                "Invalid %s value: %s (%s), using default %ss", setting_name, raw, e, default_seconds
+            )
+            return default_seconds
+        if parsed <= 0:
+            verbose_proxy_logger.warning(
+                "%s must be a positive duration, got %s, using default %ss", setting_name, raw, default_seconds
+            )
+            return default_seconds
+        return parsed
 
     def _retention_seconds_for(self, setting_name: str) -> int | None:
         """
@@ -78,6 +205,91 @@ class SpendLogCleanup:
         self.retention_seconds = self._retention_seconds_for("maximum_spend_logs_retention_period")
         return self.retention_seconds is not None
 
+    def _timeout_ms(self, deadline: float) -> int:
+        """
+        The per-statement bound in milliseconds: the batch timeout, or whatever
+        is left of the run budget, whichever is smaller.
+
+        Clamping to the remaining budget is what makes the budget a real
+        wall-clock bound rather than an advisory one. Postgres offers no "stop
+        at time T", only a per-statement duration, so a statement issued just
+        under the deadline would otherwise run a full batch timeout past it, and
+        with several tables those overruns stack.
+
+        Interpolating this into SQL is safe by construction: an int cannot carry
+        SQL, and SET does not accept a bind parameter.
+        """
+        remaining_ms: Final = int((deadline - time.monotonic()) * 1000)
+        return max(1, min(int(self.batch_timeout_seconds * 1000), remaining_ms))
+
+    def _remaining_timeout_ms(self, deadline: float) -> RemainingTimeoutMs:
+        """
+        The per-statement bound for work this job delegates, as a callable.
+
+        Partition maintenance issues one statement per partition, so handing it a
+        number would bound each statement by the budget that was left before the
+        FIRST one and never by what remains. Re-evaluating per statement is what
+        makes the loop itself bounded, and None tells the callee to stop rather
+        than issue a statement it has no budget for.
+        """
+
+        def remaining() -> int | None:
+            return None if time.monotonic() >= deadline else self._timeout_ms(deadline)
+
+        return remaining
+
+    async def _execute_delete_batch(
+        self, prisma_client: PrismaClient, delete_sql: str, cutoff_date: datetime, deadline: float
+    ) -> int | None:
+        """
+        Run one delete batch under a Postgres statement and lock timeout.
+
+        The timeouts are what actually bound the work: a Prisma transaction
+        timeout cannot interrupt a statement that is already executing, so
+        without these a single batch blocked behind a lock would hold its
+        connection, and the row locks it already took, indefinitely. SET LOCAL
+        scopes both to this transaction so the pooled connection is unaffected.
+
+        Returns the row count, or None when the driver returned something that
+        is not a row count. That is a contract violation rather than a transient
+        fault, so the caller stops instead of retrying.
+        """
+        timeout_ms: Final = self._timeout_ms(deadline)
+        async with prisma_client.db.tx() as tx:
+            await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
+            await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
+            deleted_result: Final = await tx.execute_raw(delete_sql, cutoff_date, self.batch_size)
+        return deleted_result if isinstance(deleted_result, int) else None
+
+    async def _count_remaining(
+        self, prisma_client: PrismaClient, cutoff_date: datetime, table_name: str, time_column: str, deadline: float
+    ) -> int | None:
+        """
+        Count expired rows still outstanding, stopping at a cap.
+
+        An uncapped COUNT(*) over an expired backlog would itself be the kind of
+        long scan this job exists to avoid, so the probe reads at most
+        SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP index entries. A result equal to
+        the cap means "at least this many".
+        """
+        count_sql: Final = f"""
+            SELECT count(*)::int AS remaining FROM (
+                SELECT 1 FROM "{table_name}"
+                WHERE "{time_column}" < $1::timestamptz
+                LIMIT $2
+            ) capped
+            """
+        try:
+            async with prisma_client.db.tx() as tx:
+                await tx.execute_raw(f"SET LOCAL statement_timeout = {self._timeout_ms(deadline)}")
+                rows: Final = _REMAINING_ROWS.validate_python(
+                    await tx.query_raw(count_sql, cutoff_date, SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP)
+                )
+        except Exception as e:  # noqa: BLE001 - an observability probe must never fail the cleanup run
+            verbose_proxy_logger.warning("Could not count remaining %s rows: %s", table_name, e)
+            return None
+        return rows[0].remaining if rows else None
+
     async def _delete_old_rows_batched(
         self,
         prisma_client: PrismaClient,
@@ -85,10 +297,14 @@ class SpendLogCleanup:
         table_name: str,
         key_columns: tuple[str, ...],
         time_column: str,
-    ) -> int:
+        deadline: float,
+    ) -> TableCleanupResult:
         """
-        Helper method to delete a table's rows older than the cutoff in batches.
-        Returns the total number of rows deleted.
+        Delete a table's rows older than the cutoff in batches.
+
+        Stops at whichever bound is reached first: the backlog running out, the
+        shared wall-clock deadline, the per-table batch cap, or too many
+        consecutive batch failures.
         """
         key_list: Final = ", ".join(f'"{col}"' for col in key_columns)
         delete_sql: Final = f"""
@@ -103,23 +319,46 @@ class SpendLogCleanup:
         run_count = 0
         consecutive_failures = 0
         while True:
-            if run_count > SPEND_LOG_RUN_LOOPS:
+            if time.monotonic() >= deadline:
+                verbose_proxy_logger.info(
+                    "Run budget exhausted during %s cleanup after %d rows; the next run resumes from here",
+                    table_name,
+                    total_deleted,
+                )
+                return await self._finish_table(
+                    prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
+                )
+            if run_count >= self.max_batches:
                 verbose_proxy_logger.info(
                     "Max batches reached for %s cleanup, remaining rows will be deleted in next run", table_name
                 )
-                break
-            # Step 1: Find rows and delete them in one go without fetching to application
-            # Delete in batches, limited by self.batch_size
-            try:
-                deleted_result = await prisma_client.db.execute_raw(
-                    delete_sql,
-                    cutoff_date,
-                    self.batch_size,
+                return await self._finish_table(
+                    prisma_client, cutoff_date, table_name, time_column, total_deleted, "batch_cap_reached", deadline
                 )
+            # Find rows and delete them in one go without fetching to application
+            batch_started_at = time.monotonic()
+            try:
+                batch_result = await self._execute_delete_batch(prisma_client, delete_sql, cutoff_date, deadline)
             except Exception as batch_exc:
+                if time.monotonic() >= deadline:
+                    # The statement timeout was clamped to the budget that was
+                    # left, so this batch was cancelled by the deadline itself.
+                    # That is the bound working, not a database fault, and
+                    # counting it would both inflate the failure metric and push
+                    # every budget-exhausted run toward the abort threshold.
+                    verbose_proxy_logger.info(
+                        "Run budget exhausted mid-batch during %s cleanup after %d rows; "
+                        "the next run resumes from here",
+                        table_name,
+                        total_deleted,
+                    )
+                    return await self._finish_table(
+                        prisma_client, cutoff_date, table_name, time_column, total_deleted, "budget_exhausted", deadline
+                    )
                 # A single batch failure (e.g. Prisma/DB timeout) must not abort
                 # the whole run — subsequent batches may still succeed.
                 consecutive_failures += 1
+                SpendLogCleanupMetrics.record_batch_failure(table_name)
                 verbose_proxy_logger.exception(
                     "%s cleanup batch failed "
                     "(run_count=%d, consecutive_failures=%d, batch_size=%d, "
@@ -140,28 +379,31 @@ class SpendLogCleanup:
                         consecutive_failures,
                         total_deleted,
                     )
-                    break
+                    return await self._finish_table(
+                        prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
+                    )
                 await asyncio.sleep(SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS)
                 continue
 
-            consecutive_failures = 0
-
-            deleted_count = 0
-            if isinstance(deleted_result, int):
-                deleted_count = deleted_result
-            else:
+            if batch_result is None:
                 verbose_proxy_logger.error(
-                    "Unexpected execute_raw return type for %s cleanup: %s; aborting cleanup to avoid infinite loop",
+                    "Unexpected execute_raw return type for %s cleanup; aborting cleanup to avoid infinite loop",
                     table_name,
-                    type(deleted_result),
                 )
-                break
+                return await self._finish_table(
+                    prisma_client, cutoff_date, table_name, time_column, total_deleted, "aborted", deadline
+                )
 
+            consecutive_failures = 0
+            deleted_count = batch_result
+            SpendLogCleanupMetrics.record_batch(table_name, deleted_count, time.monotonic() - batch_started_at)
             verbose_proxy_logger.info("Deleted %s %s rows in this batch", deleted_count, table_name)
 
             if deleted_count == 0:
                 verbose_proxy_logger.info("No more %s rows to delete. Total deleted: %s", table_name, total_deleted)
-                break
+                return await self._finish_table(
+                    prisma_client, cutoff_date, table_name, time_column, total_deleted, "exhausted", deadline
+                )
 
             total_deleted += deleted_count
             run_count += 1
@@ -169,18 +411,49 @@ class SpendLogCleanup:
             # Add a small sleep to prevent overwhelming the database
             await asyncio.sleep(0.1)
 
-        return total_deleted
+    async def _finish_table(
+        self,
+        prisma_client: PrismaClient,
+        cutoff_date: datetime,
+        table_name: str,
+        time_column: str,
+        rows_deleted: int,
+        stop_reason: StopReason,
+        deadline: float,
+    ) -> TableCleanupResult:
+        """
+        Publish how much of this table is still outstanding, then report the run's result.
 
-    async def _delete_old_logs(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
+        The probe is skipped once the budget is spent. It is the one piece of
+        work that would otherwise be ISSUED after the deadline, and every table
+        exits through here, including the ones a spent run never started, so
+        keeping it would put one more statement per table past the bound. A run
+        that ends this way already reports "budget_exhausted", which tells an
+        operator the backlog was not drained; the gauge simply keeps its value
+        from the last run that finished inside its budget.
+        """
+        if time.monotonic() >= deadline:
+            return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
+        remaining: Final = await self._count_remaining(prisma_client, cutoff_date, table_name, time_column, deadline)
+        if remaining is not None:
+            SpendLogCleanupMetrics.set_rows_remaining(table_name, remaining)
+        return TableCleanupResult(rows_deleted=rows_deleted, stop_reason=stop_reason)
+
+    async def _delete_old_logs(
+        self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
+    ) -> TableCleanupResult:
         return await self._delete_old_rows_batched(
             prisma_client,
             cutoff_date,
             table_name="LiteLLM_SpendLogs",
             key_columns=("request_id", "startTime"),
             time_column="startTime",
+            deadline=deadline,
         )
 
-    async def _delete_old_tool_index_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
+    async def _delete_old_tool_index_rows(
+        self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
+    ) -> TableCleanupResult:
         # SpendLogToolIndex rows are derived from spend logs, so they expire on the
         # same cutoff; rows older than retention point at already-deleted logs.
         return await self._delete_old_rows_batched(
@@ -189,17 +462,87 @@ class SpendLogCleanup:
             table_name="LiteLLM_SpendLogToolIndex",
             key_columns=("request_id", "tool_name"),
             time_column="start_time",
+            deadline=deadline,
         )
 
-    async def _delete_old_autorouter_session_rows(self, prisma_client: PrismaClient, cutoff_date: datetime) -> int:
+    async def _delete_old_autorouter_session_rows(
+        self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float
+    ) -> TableCleanupResult:
         return await self._delete_old_rows_batched(
             prisma_client,
             cutoff_date,
             table_name="LiteLLM_AutoRouterSession",
             key_columns=("api_key", "session_id", "router_name"),
             time_column="last_turn_at",
+            deadline=deadline,
         )
 
+    async def _clean_spend_log_tables(
+        self, prisma_client: PrismaClient, deadline: float
+    ) -> tuple[TableCleanupResult, ...]:
+        """
+        Prune the spend logs and the tool index rows derived from them.
+
+        When the table is range-partitioned, whole expired partitions are dropped
+        first because that reclaims disk immediately. Expired rows can still sit in
+        the DEFAULT partition (backfill, coverage gaps) or in a partition that spans
+        the cutoff, so retention still deletes those stragglers row-wise.
+        """
+        cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds or 0))
+        verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
+
+        # Partition maintenance is DDL taking an ACCESS EXCLUSIVE lock, so it is
+        # only STARTED while the run still has budget, and each statement carries
+        # the same timeouts the batches do. Without those, a DROP would queue
+        # behind any long-running reader for as long as that reader lives, which
+        # is the one way this job could still outlast its budget without bound.
+        remaining_timeout_ms: Final = self._remaining_timeout_ms(deadline)
+        if time.monotonic() >= deadline:
+            verbose_proxy_logger.info("Run budget already spent, skipping partition maintenance this run")
+        elif self.general_settings.get(
+            "use_spend_logs_partitioning", False
+        ) and await self.partition_manager.is_partitioned(prisma_client, remaining_timeout_ms):
+            await self.partition_manager.ensure_partitions(prisma_client, remaining_timeout_ms)
+            dropped: Final = await self.partition_manager.drop_partitions_older_than(
+                prisma_client, cutoff_date, remaining_timeout_ms
+            )
+            verbose_proxy_logger.info("Dropped %d expired spend-log partitions: %s", len(dropped), dropped)
+
+        logs_result: Final = await self._delete_old_logs(prisma_client, cutoff_date, deadline)
+        verbose_proxy_logger.info("Deleted %s logs", logs_result.rows_deleted)
+
+        index_result: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date, deadline)
+        verbose_proxy_logger.info("Deleted %s expired tool index rows", index_result.rows_deleted)
+        return (logs_result, index_result)
+
+    async def _clean_session_rollup(
+        self, prisma_client: PrismaClient, retention_seconds: int, deadline: float
+    ) -> tuple[TableCleanupResult, ...]:
+        """
+        Prune auto-router session rollup rows, which carry their own retention horizon.
+        """
+        session_cutoff: Final = datetime.now(timezone.utc) - timedelta(seconds=float(retention_seconds))
+        sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline)
+        verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted)
+        return (sessions_result,)
+
+    @staticmethod
+    def _run_outcome(results: tuple[TableCleanupResult, ...]) -> RunOutcome:
+        """
+        Report the most operationally significant reason the run stopped.
+
+        A bound that was hit matters more than a table that simply ran dry, so
+        those win over "completed", and an abort wins over everything.
+        """
+        reasons: Final = frozenset(result.stop_reason for result in results)
+        if "aborted" in reasons:
+            return "aborted"
+        if "budget_exhausted" in reasons:
+            return "budget_exhausted"
+        if "batch_cap_reached" in reasons:
+            return "batch_cap_reached"
+        return "completed"
+
     async def cleanup_old_spend_logs(self, prisma_client: PrismaClient) -> None:
         """
         Main cleanup function. Deletes old spend logs in batches.
@@ -209,16 +552,19 @@ class SpendLogCleanup:
         lock_acquired = False
         try:
             verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now())
+            self._refresh_bounds()
 
             delete_spend_logs: Final = self._should_delete_spend_logs()
             autorouter_retention_seconds: Final = self._retention_seconds_for(
                 "maximum_autorouter_session_retention_period"
             )
             if not delete_spend_logs and autorouter_retention_seconds is None:
+                SpendLogCleanupMetrics.record_run("skipped_disabled")
                 return
 
             if delete_spend_logs and self.retention_seconds is None:
                 verbose_proxy_logger.error("Retention seconds is None, cannot proceed with cleanup")
+                SpendLogCleanupMetrics.record_run("skipped_disabled")
                 return
 
             # If we have a pod lock manager, try to acquire the lock
@@ -235,43 +581,23 @@ class SpendLogCleanup:
 
                 if not lock_acquired:
                     verbose_proxy_logger.info("Another pod is already running cleanup")
+                    SpendLogCleanupMetrics.record_run("skipped_locked")
                     return
 
-            if delete_spend_logs and self.retention_seconds is not None:
-                cutoff_date: Final = datetime.now(timezone.utc) - timedelta(seconds=float(self.retention_seconds))
-                verbose_proxy_logger.info("Removing logs older than %s", cutoff_date.isoformat())
+            deadline: Final = time.monotonic() + self.run_budget_seconds
 
-                if self.general_settings.get(
-                    "use_spend_logs_partitioning", False
-                ) and await self.partition_manager.is_partitioned(prisma_client):
-                    await self.partition_manager.ensure_partitions(prisma_client)
-                    dropped: Final = await self.partition_manager.drop_partitions_older_than(prisma_client, cutoff_date)
-                    verbose_proxy_logger.info(
-                        "Dropped %d expired spend-log partitions: %s",
-                        len(dropped),
-                        dropped,
-                    )
-                    # DROP only reclaims whole expired partitions. Expired rows can
-                    # still sit in the DEFAULT partition (backfill, coverage gaps)
-                    # or in a partition that spans the cutoff, so retention must
-                    # also delete those stragglers row-wise.
-                    total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
-                    verbose_proxy_logger.info(
-                        "Deleted %s expired logs not covered by dropped partitions", total_deleted
-                    )
-                else:
-                    total_deleted = await self._delete_old_logs(prisma_client, cutoff_date)
-                    verbose_proxy_logger.info("Deleted %s logs", total_deleted)
+            spend_log_results: Final = (
+                await self._clean_spend_log_tables(prisma_client, deadline)
+                if delete_spend_logs and self.retention_seconds is not None
+                else ()
+            )
+            session_results: Final = (
+                await self._clean_session_rollup(prisma_client, autorouter_retention_seconds, deadline)
+                if autorouter_retention_seconds is not None
+                else ()
+            )
 
-                index_deleted: Final = await self._delete_old_tool_index_rows(prisma_client, cutoff_date)
-                verbose_proxy_logger.info("Deleted %s expired tool index rows", index_deleted)
-
-            if autorouter_retention_seconds is not None:
-                session_cutoff: Final = datetime.now(timezone.utc) - timedelta(
-                    seconds=float(autorouter_retention_seconds)
-                )
-                sessions_deleted: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff)
-                verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_deleted)
+            SpendLogCleanupMetrics.record_run(self._run_outcome(spend_log_results + session_results))
 
         except Exception as e:
             # .exception() captures the traceback; str(e) alone on a Prisma/DB
@@ -281,6 +607,7 @@ class SpendLogCleanup:
                 type(e).__name__,
                 e,
             )
+            SpendLogCleanupMetrics.record_run("aborted")
             return  # Return after error handling
         finally:
             # Only release the lock if it was actually acquired
diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py
new file mode 100644
index 00000000000..340aeab938c
--- /dev/null
+++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup_metrics.py
@@ -0,0 +1,122 @@
+"""
+Prometheus metrics for the spend-log retention cleanup job.
+
+The job runs in the background on a single elected pod, so its cost is invisible
+from request-path metrics. These instruments make a run's database footprint
+observable: how much it deleted, how long each batch took, how much work is
+still outstanding, and why a run stopped.
+
+``prometheus_client`` is an optional dependency, so every recorder degrades to a
+no-op when it is absent.
+"""
+
+from typing import TYPE_CHECKING, Final, Literal, TypeAlias
+
+from litellm._logging import verbose_proxy_logger
+
+if TYPE_CHECKING:
+    # aliased so the annotations below cannot be mistaken for collections.Counter
+    from prometheus_client import Counter as PrometheusCounter
+    from prometheus_client import Gauge as PrometheusGauge
+    from prometheus_client import Histogram as PrometheusHistogram
+
+RunOutcome: TypeAlias = Literal[
+    "completed",
+    "budget_exhausted",
+    "batch_cap_reached",
+    "skipped_locked",
+    "skipped_disabled",
+    "aborted",
+]
+
+_BATCH_DURATION_BUCKETS: Final = (0.005, 0.025, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0)
+_TABLE_LABEL: Final = ("table",)
+_OUTCOME_LABEL: Final = ("outcome",)
+
+
+class SpendLogCleanupMetrics:
+    """
+    Lazily-registered Prometheus instruments for the retention cleanup job.
+
+    Registration is deferred to first use so that importing this module never
+    touches the Prometheus registry, which keeps it safe to import from the
+    proxy regardless of whether Prometheus is a configured callback.
+    """
+
+    _initialized: bool = False
+    rows_deleted: "PrometheusCounter | None" = None
+    batch_duration: "PrometheusHistogram | None" = None
+    rows_remaining: "PrometheusGauge | None" = None
+    batch_failures: "PrometheusCounter | None" = None
+    runs: "PrometheusCounter | None" = None
+
+    @classmethod
+    def _ensure_initialized(cls) -> None:
+        if cls._initialized:
+            return
+        cls._initialized = True
+        try:
+            # prometheus_client is an optional extra, so it is resolved here rather
+            # than at module import: this module is reachable from proxy startup
+            # regardless of whether Prometheus is a configured callback.
+            from prometheus_client import Counter, Gauge, Histogram
+
+            cls.rows_deleted = Counter(
+                "litellm_spend_log_cleanup_rows_deleted_total",
+                "Rows deleted by the spend-log retention cleanup job",
+                labelnames=_TABLE_LABEL,
+            )
+            cls.batch_duration = Histogram(
+                "litellm_spend_log_cleanup_batch_duration_seconds",
+                "Wall-clock duration of one retention cleanup delete batch",
+                labelnames=_TABLE_LABEL,
+                buckets=_BATCH_DURATION_BUCKETS,
+            )
+            cls.rows_remaining = Gauge(
+                "litellm_spend_log_cleanup_rows_remaining",
+                "Expired rows still awaiting deletion, counted only up to "
+                "SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP so the probe itself cannot scan a "
+                "large table; a value equal to that cap means at least that many remain",
+                labelnames=_TABLE_LABEL,
+                multiprocess_mode="livemax",
+            )
+            cls.batch_failures = Counter(
+                "litellm_spend_log_cleanup_batch_failures_total",
+                "Retention cleanup delete batches that raised",
+                labelnames=_TABLE_LABEL,
+            )
+            cls.runs = Counter(
+                "litellm_spend_log_cleanup_runs_total",
+                "Retention cleanup runs, labelled by why the run ended",
+                labelnames=_OUTCOME_LABEL,
+            )
+        except Exception as e:  # noqa: BLE001 - a metrics problem must never fail the cleanup run
+            # Covers the extra being absent, a duplicate registration (repeated
+            # imports under a test runner), and registry misconfiguration alike.
+            verbose_proxy_logger.warning("Could not register spend-log cleanup metrics: %s", e)
+
+    @classmethod
+    def record_batch(cls, table_name: str, rows_deleted: int, duration_seconds: float) -> None:
+        cls._ensure_initialized()
+        if cls.rows_deleted is not None:
+            cls.rows_deleted.labels(table=table_name).inc(rows_deleted)
+        if cls.batch_duration is not None:
+            cls.batch_duration.labels(table=table_name).observe(duration_seconds)
+
+    @classmethod
+    def record_batch_failure(cls, table_name: str) -> None:
+        cls._ensure_initialized()
+        if cls.batch_failures is not None:
+            cls.batch_failures.labels(table=table_name).inc()
+
+    @classmethod
+    def set_rows_remaining(cls, table_name: str, remaining: int) -> None:
+        cls._ensure_initialized()
+        if cls.rows_remaining is not None:
+            cls.rows_remaining.labels(table=table_name).set(remaining)
+
+    @classmethod
+    def record_run(cls, outcome: RunOutcome) -> None:
+        cls._ensure_initialized()
+        if cls.runs is not None:
+            cls.runs.labels(outcome=outcome).inc()
diff --git a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
index df17721d8e5..221c142d9d3 100644
--- a/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
+++ b/litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py
@@ -14,8 +14,9 @@ keeps the batched-DELETE path, so existing deployments are untouched.
 """
 
 import re
+from collections.abc import Callable
 from datetime import date, datetime, timedelta, timezone
-from typing import Final
+from typing import TYPE_CHECKING, Final, TypeAlias
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
@@ -23,8 +24,23 @@ from litellm.constants import (
     SPEND_LOG_PARTITION_PRECREATE_AHEAD,
 )
 
+if TYPE_CHECKING:
+    from litellm.proxy.utils import PrismaClient
+
 SPEND_LOGS_TABLE: Final = "LiteLLM_SpendLogs"
 
+RemainingTimeoutMs: TypeAlias = Callable[[], "int | None"]
+"""
+The per-statement bound in milliseconds, or None once the caller's budget is
+spent.
+
+Injected rather than passed as a number so it is re-evaluated before EVERY
+statement: a value read once at entry would let a loop issue N statements each
+bounded by the budget that was left before the first of them, which is not a
+bound on the loop at all. The caller owns the policy; this module only asks how
+much time it may still use.
+"""
+
 PartitionInterval = str  # "day" | "week" | "month"
 
 VALID_PARTITION_INTERVALS: Final = {"day", "week", "month"}
@@ -116,21 +132,26 @@ class SpendLogsPartitionManager:
         self.interval = interval
         self.precreate_ahead = precreate_ahead
 
-    async def is_partitioned(self, prisma_client) -> bool:
+    async def is_partitioned(self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs) -> bool:
+        budget_ms: Final = remaining_timeout_ms()
+        if budget_ms is None:
+            return False
         try:
-            rows: Final = await prisma_client.db.query_raw(
-                """
-                SELECT EXISTS (
-                    SELECT 1
-                    FROM pg_partitioned_table pt
-                    JOIN pg_class c ON c.oid = pt.partrelid
-                    JOIN pg_namespace n ON n.oid = c.relnamespace
-                    WHERE c.relname = $1
-                      AND n.nspname = current_schema()
-                ) AS partitioned
-                """,
-                SPEND_LOGS_TABLE,
-            )
+            async with prisma_client.db.tx() as tx:
+                await tx.execute_raw(f"SET LOCAL statement_timeout = {budget_ms}")
+                rows: Final = await tx.query_raw(
+                    """
+                    SELECT EXISTS (
+                        SELECT 1
+                        FROM pg_partitioned_table pt
+                        JOIN pg_class c ON c.oid = pt.partrelid
+                        JOIN pg_namespace n ON n.oid = c.relnamespace
+                        WHERE c.relname = $1
+                          AND n.nspname = current_schema()
+                    ) AS partitioned
+                    """,
+                    SPEND_LOGS_TABLE,
+                )
         except Exception as e:
             verbose_proxy_logger.warning(
                 "Could not determine if %s is partitioned, assuming it is not: %s",
@@ -140,7 +161,25 @@ class SpendLogsPartitionManager:
             return False
         return bool(rows and rows[0].get("partitioned"))
 
-    async def ensure_partitions(self, prisma_client) -> list[str]:
+    @staticmethod
+    async def _execute_bounded_ddl(prisma_client: "PrismaClient", statement: str, timeout_ms: int) -> None:
+        """
+        Run one DDL statement under a Postgres statement and lock timeout.
+
+        Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded statement
+        queues behind any long-running reader for as long as that reader lives,
+        and the caller's run budget cannot cut it short. lock_timeout bounds the
+        wait for the lock and statement_timeout bounds the work itself, so a
+        partition this run cannot get is simply left for the next one.
+        """
+        async with prisma_client.db.tx() as tx:
+            await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
+            await tx.execute_raw(f"SET LOCAL lock_timeout = {timeout_ms}")
+            await tx.execute_raw(statement)
+
+    async def ensure_partitions(
+        self, prisma_client: "PrismaClient", remaining_timeout_ms: RemainingTimeoutMs
+    ) -> list[str]:
         """
         Ensure the current and upcoming partitions exist, returning the names
         now present. CREATE TABLE IF NOT EXISTS is a no-op for partitions that
@@ -150,42 +189,61 @@ class SpendLogsPartitionManager:
         for name, lower, upper in upcoming_partitions(
             datetime.now(timezone.utc).date(), self.interval, self.precreate_ahead
         ):
+            budget_ms = remaining_timeout_ms()
+            if budget_ms is None:
+                verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
+                break
             try:
-                await prisma_client.db.execute_raw(
+                await self._execute_bounded_ddl(
+                    prisma_client,
                     f'CREATE TABLE IF NOT EXISTS "{name}" '
                     f'PARTITION OF "{SPEND_LOGS_TABLE}" '
-                    f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')"
+                    f"FOR VALUES FROM ('{lower.isoformat()}') TO ('{upper.isoformat()}')",
+                    budget_ms,
                 )
                 ensured.append(name)
             except Exception as e:
                 verbose_proxy_logger.warning("Failed to ensure spend-log partition %s: %s", name, e)
         return ensured
 
-    async def _list_partitions(self, prisma_client) -> list[tuple[str, datetime | None]]:
-        rows: Final = await prisma_client.db.query_raw(
-            """
-            SELECT c.relname AS name,
-                   pg_get_expr(c.relpartbound, c.oid) AS bound
-            FROM pg_inherits i
-            JOIN pg_class c ON c.oid = i.inhrelid
-            JOIN pg_class p ON p.oid = i.inhparent
-            JOIN pg_namespace n ON n.oid = p.relnamespace
-            WHERE p.relname = $1
-              AND n.nspname = current_schema()
-            """,
-            SPEND_LOGS_TABLE,
-        )
+    async def _list_partitions(
+        self, prisma_client: "PrismaClient", timeout_ms: int
+    ) -> list[tuple[str, datetime | None]]:
+        async with prisma_client.db.tx() as tx:
+            await tx.execute_raw(f"SET LOCAL statement_timeout = {timeout_ms}")
+            rows: Final = await tx.query_raw(
+                """
+                SELECT c.relname AS name,
+                       pg_get_expr(c.relpartbound, c.oid) AS bound
+                FROM pg_inherits i
+                JOIN pg_class c ON c.oid = i.inhrelid
+                JOIN pg_class p ON p.oid = i.inhparent
+                JOIN pg_namespace n ON n.oid = p.relnamespace
+                WHERE p.relname = $1
+                  AND n.nspname = current_schema()
+                """,
+                SPEND_LOGS_TABLE,
+            )
         return [(row["name"], parse_partition_upper_bound(row.get("bound") or "")) for row in rows]
 
-    async def drop_partitions_older_than(self, prisma_client, cutoff: datetime) -> list[str]:
+    async def drop_partitions_older_than(
+        self, prisma_client: "PrismaClient", cutoff: datetime, remaining_timeout_ms: RemainingTimeoutMs
+    ) -> list[str]:
         """DROP every partition whose whole range is older than `cutoff`."""
+        list_budget_ms: Final = remaining_timeout_ms()
+        if list_budget_ms is None:
+            return []
         cutoff_naive: Final = cutoff.astimezone(timezone.utc).replace(tzinfo=None)
-        partitions: Final = await self._list_partitions(prisma_client)
+        partitions: Final = await self._list_partitions(prisma_client, list_budget_ms)
         to_drop: Final = select_partitions_to_drop(partitions, cutoff_naive)
         dropped: Final[list[str]] = []
         for name in to_drop:
+            budget_ms = remaining_timeout_ms()
+            if budget_ms is None:
+                verbose_proxy_logger.info("Run budget spent, leaving the remaining partitions for the next run")
+                break
             try:
-                await prisma_client.db.execute_raw(f'DROP TABLE IF EXISTS "{name}"')
+                await self._execute_bounded_ddl(prisma_client, f'DROP TABLE IF EXISTS "{name}"', budget_ms)
                 dropped.append(name)
             except Exception as e:
                 verbose_proxy_logger.warning("Failed to drop spend-log partition %s: %s", name, e)
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 968c645ffcd..b1b5a7ffbe5 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -367,7 +367,10 @@ from litellm.proxy.config_resolvers.alerting import (
 )
 from litellm.proxy.container_endpoints.endpoints import router as container_router
 from litellm.proxy.credential_endpoints.endpoints import router as credential_router
-from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
+from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
+    SPEND_LOG_CLEANUP_BOUND_SETTINGS,
+    SpendLogCleanup,
+)
 from litellm.proxy.db.exception_handler import (
     PrismaDBExceptionHandler,
     call_with_db_reconnect_retry,
@@ -4079,6 +4082,7 @@ class ProxyConfig:
         # precedence over stale DB-cached values for these specific keys
         # during periodic config reloads (_update_general_settings).
         self._yaml_general_settings_keys: set[str] = set()  # mutable-ok: populated once at startup, read-only thereafter  # fmt: skip
+        self._yaml_spend_log_cleanup_bounds: dict[str, object] = {}  # mutable-ok: snapshot of YAML bounds at load time  # fmt: skip
 
     def is_yaml(self, config_file_path: str) -> bool:
         if not os.path.isfile(config_file_path):
@@ -5015,6 +5019,12 @@ class ProxyConfig:
             # These keys take precedence over DB-cached values during periodic
             # reloads (see _update_general_settings).
             self._yaml_general_settings_keys = set(general_settings.keys())  # mutable-ok: snapshot of YAML keys at load time  # fmt: skip
+            # The VALUES matter for the cleanup bounds, not just which keys were
+            # set: clearing one from the dashboard has to fall back to what the
+            # YAML declared, and a set of names cannot answer that.
+            self._yaml_spend_log_cleanup_bounds = {  # mutable-ok: snapshot of YAML bounds at load time  # fmt: skip
+                key: general_settings[key] for key in SPEND_LOG_CLEANUP_BOUND_SETTINGS if key in general_settings
+            }
 
             ### LOAD KEY MANAGEMENT SETTINGS FIRST (needed for custom secret manager) ###
             key_management_settings: Final = general_settings.get("key_management_settings", None)
@@ -6299,6 +6309,18 @@ class ProxyConfig:
             if old_session_value != new_session_value:
                 await self._reschedule_spend_log_cleanup_job()
 
+        ## SPEND LOG CLEANUP BOUNDS ##
+        # The dashboard writes these straight to the DB, so without copying them
+        # here the running cleanup job never sees them. A key the DB no longer
+        # carries was cleared from the dashboard, and falls back to whatever
+        # config.yaml declared, or to None (the shipped default) when it declared
+        # nothing. Leaving the deleted DB value in memory would keep enforcing the
+        # bound the operator just removed.
+        for cleanup_key in SPEND_LOG_CLEANUP_BOUND_SETTINGS:
+            general_settings[cleanup_key] = _general_settings.get(
+                cleanup_key, self._yaml_spend_log_cleanup_bounds.get(cleanup_key)
+            )
+
         for key in (
             "user_url_allowed_hosts",
             "user_url_validation",
@@ -15529,6 +15551,10 @@ _GENERAL_SETTINGS_CONFIG_LIST_FIELD_TYPES: Final[Mapping[str, str]] = MappingPro
         "store_model_in_db": "Boolean",
         "store_prompts_in_spend_logs": "Boolean",
         "maximum_spend_logs_retention_period": "String",
+        "maximum_spend_logs_cleanup_batch_size": "Integer",
+        "maximum_spend_logs_cleanup_max_batches": "Integer",
+        "maximum_spend_logs_cleanup_run_budget": "String",
+        "maximum_spend_logs_cleanup_batch_timeout": "String",
         "mcp_internal_ip_ranges": "List",
         "mcp_trusted_proxy_ranges": "List",
         "mcp_xff_num_trusted_hops": "Integer",
diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
index 289de707387..e949afce57b 100644
--- a/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
+++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_spend_logs_partition_manager.py
@@ -3,6 +3,7 @@ Tests for SpendLogsPartitionManager: partition naming/bounds math, retention
 selection, the non-partitioned no-op safety path, and the drop/ensure SQL flow.
 """
 
+from contextlib import asynccontextmanager
 from datetime import date, datetime, timezone
 from unittest.mock import AsyncMock, MagicMock
 
@@ -19,6 +20,46 @@ from litellm.proxy.db.db_transaction_queue.spend_logs_partition_manager import (
 )
 
 
+DDL_TIMEOUT_MS = 30000
+
+
+def _budget(ms: "int | None" = DDL_TIMEOUT_MS):
+    """The injected per-statement bound: a callable re-read before each statement."""
+    return lambda: ms
+
+
+def _wire_tx(db) -> list[str]:
+    """
+    Model the prisma seam the partition DDL uses.
+
+    Every statement this manager issues, DDL and catalog query alike, runs inside
+    db.tx() so it can carry SET LOCAL timeouts. Those SET LOCAL statements are
+    collected in the returned list rather than forwarded, so assertions on
+    db.execute_raw and db.query_raw still see only the real statements.
+    """
+    session_settings: list[str] = []
+
+    @asynccontextmanager
+    async def _tx():
+        tx = MagicMock()
+
+        async def _execute_raw(sql, *args):
+            if sql.lstrip().upper().startswith("SET LOCAL"):
+                session_settings.append(sql.strip())
+                return 0
+            return await db.execute_raw(sql, *args)
+
+        async def _query_raw(sql, *args):
+            return await db.query_raw(sql, *args)
+
+        tx.execute_raw = _execute_raw
+        tx.query_raw = _query_raw
+        yield tx
+
+    db.tx = _tx
+    return session_settings
+
+
 def test_period_start_per_interval():
     d = date(2026, 6, 3)  # a Wednesday
     assert period_start(d, "day") == date(2026, 6, 3)
@@ -78,11 +119,13 @@ async def test_is_partitioned_true_and_false():
 
     client_true = MagicMock()
     client_true.db.query_raw = AsyncMock(return_value=[{"partitioned": True}])
-    assert await mgr.is_partitioned(client_true) is True
+    _wire_tx(client_true.db)
+    assert await mgr.is_partitioned(client_true, _budget()) is True
 
     client_false = MagicMock()
     client_false.db.query_raw = AsyncMock(return_value=[{"partitioned": False}])
-    assert await mgr.is_partitioned(client_false) is False
+    _wire_tx(client_false.db)
+    assert await mgr.is_partitioned(client_false, _budget()) is False
 
 
 @pytest.mark.asyncio
@@ -94,13 +137,14 @@ async def test_catalog_queries_are_scoped_to_current_schema():
     mgr = SpendLogsPartitionManager()
     client = MagicMock()
     client.db.query_raw = AsyncMock(return_value=[])
+    _wire_tx(client.db)
 
-    await mgr.is_partitioned(client)
+    await mgr.is_partitioned(client, _budget())
     is_partitioned_sql = client.db.query_raw.call_args.args[0]
     assert "pg_namespace" in is_partitioned_sql
     assert "current_schema()" in is_partitioned_sql
 
-    await mgr._list_partitions(client)
+    await mgr._list_partitions(client, DDL_TIMEOUT_MS)
     list_sql = client.db.query_raw.call_args.args[0]
     assert "pg_namespace" in list_sql
     assert "current_schema()" in list_sql
@@ -112,7 +156,10 @@ async def test_is_partitioned_swallows_errors_and_returns_false():
     mgr = SpendLogsPartitionManager()
     client = MagicMock()
     client.db.query_raw = AsyncMock(side_effect=Exception("db down"))
-    assert await mgr.is_partitioned(client) is False
+    # Wire the real seam: without it the async with itself raises, and the test
+    # would pass on the wrong exception.
+    _wire_tx(client.db)
+    assert await mgr.is_partitioned(client, _budget()) is False
 
 
 @pytest.mark.asyncio
@@ -133,9 +180,10 @@ async def test_drop_partitions_older_than_drops_expired_only():
         ]
     )
     client.db.execute_raw = AsyncMock(return_value=0)
+    _wire_tx(client.db)
 
     cutoff = datetime(2026, 6, 5, 0, 0, 0, tzinfo=timezone.utc)
-    dropped = await mgr.drop_partitions_older_than(client, cutoff)
+    dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget())
 
     assert dropped == ["LiteLLM_SpendLogs_p20260601"]
     executed = " ".join(call.args[0] for call in client.db.execute_raw.call_args_list)
@@ -149,8 +197,9 @@ async def test_ensure_partitions_issues_create_for_each_period():
     mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
     client = MagicMock()
     client.db.execute_raw = AsyncMock(return_value=0)
+    _wire_tx(client.db)
 
-    created = await mgr.ensure_partitions(client)
+    created = await mgr.ensure_partitions(client, _budget())
 
     assert len(created) == 3  # current + 2 ahead
     assert client.db.execute_raw.await_count == 3
@@ -159,6 +208,105 @@ async def test_ensure_partitions_issues_create_for_each_period():
     assert "CREATE TABLE IF NOT EXISTS" in first_sql
 
 
+@pytest.mark.asyncio
+async def test_partition_ddl_carries_a_statement_and_lock_timeout():
+    """
+    Partition DDL takes an ACCESS EXCLUSIVE lock, so an unbounded DROP queues
+    behind any long-running reader for as long as that reader lives. That is the
+    one path by which cleanup could outlast its run budget without bound, and
+    lock_timeout is what bounds the wait rather than only the work.
+    """
+    mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=0)
+    client = MagicMock()
+    client.db.execute_raw = AsyncMock(return_value=0)
+    client.db.query_raw = AsyncMock(
+        return_value=[
+            {
+                "name": "LiteLLM_SpendLogs_p20260601",
+                "bound": "FOR VALUES FROM ('2026-06-01 00:00:00') TO ('2026-06-02 00:00:00')",
+            }
+        ]
+    )
+    session_settings = _wire_tx(client.db)
+
+    await mgr.ensure_partitions(client, _budget(7000))
+    await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), _budget(7000))
+
+    # Three statements were issued: the CREATE, the catalog list the drop needs,
+    # and the DROP. All three carry a statement timeout; only the two that take
+    # a lock also carry a lock timeout, since the catalog read takes none.
+    assert session_settings.count("SET LOCAL statement_timeout = 7000") == 3
+    assert session_settings.count("SET LOCAL lock_timeout = 7000") == 2
+
+
+@pytest.mark.asyncio
+async def test_catalog_queries_carry_a_statement_timeout():
+    """
+    Bounding only the DDL leaves the two catalog lookups as statements this job
+    issues with no bound at all, so a run could still outlast its budget waiting
+    on one. Every statement the manager issues carries the caller's timeout.
+    """
+    mgr = SpendLogsPartitionManager()
+    client = MagicMock()
+    client.db.query_raw = AsyncMock(return_value=[])
+    session_settings = _wire_tx(client.db)
+
+    await mgr.is_partitioned(client, _budget(4000))
+    assert session_settings == ["SET LOCAL statement_timeout = 4000"], (
+        f"is_partitioned issued no statement timeout: {session_settings}"
+    )
+
+    session_settings.clear()
+    await mgr._list_partitions(client, 4000)
+    assert session_settings == ["SET LOCAL statement_timeout = 4000"], (
+        f"_list_partitions issued no statement timeout: {session_settings}"
+    )
+
+
+@pytest.mark.asyncio
+async def test_partition_loops_stop_when_the_budget_runs_out_mid_way():
+    """
+    Each loop issues one statement per partition, so a bound read once at entry
+    would let N statements each run for the budget that was left before the
+    first of them. The bound is re-read per statement and the loop stops.
+    """
+    mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=4)
+    client = MagicMock()
+    client.db.execute_raw = AsyncMock(return_value=0)
+    _wire_tx(client.db)
+
+    # Budget for two statements, then spent.
+    calls = {"n": 0}
+
+    def budget() -> "int | None":
+        calls["n"] += 1
+        return 5000 if calls["n"] <= 2 else None
+
+    created = await mgr.ensure_partitions(client, budget)
+
+    assert len(created) == 2, f"the loop ran past its budget and created {len(created)}"
+    assert client.db.execute_raw.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_partition_maintenance_issues_nothing_when_the_budget_is_already_spent():
+    """A run with no budget left must not issue even the catalog lookups."""
+    mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
+    client = MagicMock()
+    client.db.execute_raw = AsyncMock(return_value=0)
+    client.db.query_raw = AsyncMock(return_value=[])
+    _wire_tx(client.db)
+
+    spent = _budget(None)
+
+    assert await mgr.is_partitioned(client, spent) is False
+    assert await mgr.ensure_partitions(client, spent) == []
+    assert await mgr.drop_partitions_older_than(client, datetime(2026, 6, 5, tzinfo=timezone.utc), spent) == []
+
+    client.db.execute_raw.assert_not_awaited()
+    client.db.query_raw.assert_not_awaited()
+
+
 def test_unsupported_interval_raises():
     with pytest.raises(ValueError):
         period_start(date(2026, 6, 1), "year")
@@ -178,8 +326,9 @@ async def test_ensure_partitions_continues_when_one_create_fails():
     mgr = SpendLogsPartitionManager(interval="day", precreate_ahead=2)
     client = MagicMock()
     client.db.execute_raw = AsyncMock(side_effect=[0, Exception("overlap"), 0])
+    _wire_tx(client.db)
 
-    created = await mgr.ensure_partitions(client)
+    created = await mgr.ensure_partitions(client, _budget())
 
     # the failed partition is skipped, the others still created
     assert len(created) == 2
@@ -202,8 +351,9 @@ async def test_invalid_interval_does_not_abort_ensure_partitions():
     mgr = SpendLogsPartitionManager(interval="fortnight", precreate_ahead=1)
     client = MagicMock()
     client.db.execute_raw = AsyncMock(return_value=0)
+    _wire_tx(client.db)
 
-    created = await mgr.ensure_partitions(client)
+    created = await mgr.ensure_partitions(client, _budget())
 
     assert len(created) == 2  # current + 1 ahead, day-based fallback
 
@@ -225,9 +375,10 @@ async def test_drop_partitions_continues_when_one_drop_fails():
         ]
     )
     client.db.execute_raw = AsyncMock(side_effect=[Exception("locked"), 0])
+    _wire_tx(client.db)
 
     cutoff = datetime(2026, 6, 10, 0, 0, 0, tzinfo=timezone.utc)
-    dropped = await mgr.drop_partitions_older_than(client, cutoff)
+    dropped = await mgr.drop_partitions_older_than(client, cutoff, _budget())
 
     # both were eligible; the first drop failed so only the second is reported
     assert dropped == ["LiteLLM_SpendLogs_p20260602"]
diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py
index 57b2c874962..918d39646b0 100644
--- a/tests/test_litellm/proxy/test_proxy_server.py
+++ b/tests/test_litellm/proxy/test_proxy_server.py
@@ -6872,6 +6872,91 @@ async def test_update_general_settings_propagates_apply_user_budget_to_team_keys
         assert ps.general_settings["apply_user_budget_to_team_keys"] is True
 
 
+@pytest.mark.asyncio
+async def test_update_general_settings_propagates_spend_log_cleanup_bounds():
+    """The dashboard writes the cleanup bounds straight to the DB config, so
+    without runtime propagation the scheduled job never sees them and the knobs
+    do nothing until the process restarts."""
+    from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
+        SPEND_LOG_CLEANUP_BOUND_SETTINGS,
+    )
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    proxy_config = ProxyConfig()
+    db_settings = {
+        "maximum_spend_logs_cleanup_batch_size": 2000,
+        "maximum_spend_logs_cleanup_max_batches": 250,
+        "maximum_spend_logs_cleanup_run_budget": "90s",
+        "maximum_spend_logs_cleanup_batch_timeout": "10s",
+    }
+    assert set(db_settings) == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)
+
+    with patch("litellm.proxy.proxy_server.general_settings", {}):
+        await proxy_config._update_general_settings(db_general_settings=db_settings)
+
+        import litellm.proxy.proxy_server as ps
+
+        assert {key: ps.general_settings.get(key) for key in db_settings} == db_settings
+
+
+@pytest.mark.asyncio
+async def test_update_general_settings_clears_a_spend_log_cleanup_bound_dropped_from_the_db():
+    """Blanking the field in the dashboard deletes the key outright, so leaving
+    the last value in memory would keep a bound the operator just removed."""
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    proxy_config = ProxyConfig()
+
+    with patch(
+        "litellm.proxy.proxy_server.general_settings",
+        {"maximum_spend_logs_cleanup_run_budget": "90s", "maximum_spend_logs_cleanup_batch_timeout": "10s"},
+    ):
+        await proxy_config._update_general_settings(
+            db_general_settings={"maximum_spend_logs_cleanup_batch_timeout": "10s"}
+        )
+
+        import litellm.proxy.proxy_server as ps
+
+        assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] is None
+        assert ps.general_settings["maximum_spend_logs_cleanup_batch_timeout"] == "10s"
+
+
+@pytest.mark.asyncio
+async def test_update_general_settings_keeps_a_yaml_set_spend_log_cleanup_bound():
+    """A YAML-set bound never appears in the DB object, so treating its absence
+    as a dashboard clear would discard the deployed config on every reload."""
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    proxy_config = ProxyConfig()
+    proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"}
+
+    with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "90s"}):
+        await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
+
+        import litellm.proxy.proxy_server as ps
+
+        assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s"
+
+
+@pytest.mark.asyncio
+async def test_update_general_settings_clearing_a_db_override_falls_back_to_the_yaml_bound():
+    """Clearing a dashboard override of a YAML-declared bound must restore the
+    YAML value. Leaving the deleted override in memory would keep enforcing the
+    bound the operator just removed, until the process restarted."""
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    proxy_config = ProxyConfig()
+    proxy_config._yaml_spend_log_cleanup_bounds = {"maximum_spend_logs_cleanup_run_budget": "90s"}
+
+    # Memory currently holds the dashboard override, and the DB no longer carries it.
+    with patch("litellm.proxy.proxy_server.general_settings", {"maximum_spend_logs_cleanup_run_budget": "30s"}):
+        await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True})
+
+        import litellm.proxy.proxy_server as ps
+
+        assert ps.general_settings["maximum_spend_logs_cleanup_run_budget"] == "90s"
+
+
 @pytest.mark.asyncio
 async def test_update_general_settings_apply_user_budget_to_team_keys_yaml_wins():
     """A DB value must not silently override an explicit YAML setting on reload."""
diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py
index 03eef14dacb..87fbdd4c933 100644
--- a/tests/test_litellm/proxy/test_spend_log_cleanup.py
+++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py
@@ -2,12 +2,66 @@
 Test cases for spend log cleanup functionality
 """
 
+import asyncio
+import math
+import time
+from contextlib import asynccontextmanager
 from datetime import datetime, timedelta, timezone
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
-from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup
+from litellm.constants import (
+    SPEND_LOG_CLEANUP_BATCH_SIZE,
+    SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP,
+    SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS,
+)
+from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import (
+    SPEND_LOG_CLEANUP_BOUND_SETTINGS,
+    SpendLogCleanup,
+    TableCleanupResult,
+)
+from litellm.proxy.db.db_transaction_queue.spend_log_cleanup_metrics import (
+    SpendLogCleanupMetrics,
+)
+
+
+def _far_deadline() -> float:
+    """A run deadline far enough out that only the other bounds can stop a batch loop."""
+    return time.monotonic() + 3600
+
+
+def _wire_tx(db):
+    """
+    Model the prisma seam the cleanup job actually uses.
+
+    Every statement the job issues runs inside db.tx() so it can carry a SET
+    LOCAL statement_timeout. Batch and probe statements are forwarded to
+    db.execute_raw and db.query_raw, which is what tests configure and assert
+    on, while the SET LOCAL statements are answered here so they neither consume
+    a side_effect entry nor show up in the recorded call list. Lookup is
+    deferred to call time so this can be wired before a test assigns its own
+    execute_raw.
+    """
+
+    @asynccontextmanager
+    async def _tx():
+        tx = MagicMock()
+
+        async def _execute_raw(sql, *args):
+            if sql.lstrip().upper().startswith("SET LOCAL"):
+                return 0
+            return await db.execute_raw(sql, *args)
+
+        async def _query_raw(sql, *args):
+            return await db.query_raw(sql, *args)
+
+        tx.execute_raw = _execute_raw
+        tx.query_raw = _query_raw
+        yield tx
+
+    db.tx = _tx
+    db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
 
 
 def test_spend_log_cleanup_cron_scheduling():
@@ -49,6 +103,7 @@ def test_spend_log_cleanup_cron_scheduler_integration():
     # Mock scheduler
     mock_scheduler = MagicMock()
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_cleanup_instance = MagicMock()
 
     # Test Case 1: Cron-based scheduling
@@ -155,7 +210,9 @@ async def test_cleanup_old_spend_logs_batch_deletion():
 
     # Setup Prisma client
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
 
     # Mock execute_raw to return deleted counts (3 spend-log batches, then the
     # tool-index cleanup's first batch returning 0)
@@ -207,7 +264,9 @@ async def test_cleanup_old_spend_logs_retention_period_cutoff():
     """
     # Setup Prisma client
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(return_value=0)
     mock_prisma_client.db = mock_db
 
@@ -244,6 +303,7 @@ async def test_cleanup_drops_partitions_when_enabled_and_partitioned():
     from unittest.mock import AsyncMock, MagicMock
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_prisma_client.db.execute_raw = AsyncMock(return_value=0)
 
     partition_manager = MagicMock()
@@ -285,6 +345,7 @@ async def test_cleanup_uses_delete_when_partitioning_not_enabled():
     from unittest.mock import AsyncMock, MagicMock
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
 
     partition_manager = MagicMock()
@@ -316,6 +377,7 @@ async def test_cleanup_uses_delete_when_not_partitioned():
     from unittest.mock import AsyncMock, MagicMock
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[10, 0, 0])
 
     partition_manager = MagicMock()
@@ -346,6 +408,7 @@ async def test_cleanup_old_spend_logs_no_retention_period():
     Test that no logs are deleted when no retention period is set
     """
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_prisma_client.db.execute_raw = AsyncMock()
 
     cleaner = SpendLogCleanup(general_settings={})  # no retention
@@ -361,6 +424,7 @@ async def test_lock_not_released_when_not_acquired():
     before the lock is ever acquired.
     """
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_prisma_client.db.execute_raw = AsyncMock()
 
     mock_redis_cache = MagicMock()
@@ -418,7 +482,9 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
     """should abort deletion loop immediately when execute_raw returns a non-int
     (e.g. None or dict), preventing an infinite loop."""
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(return_value=None)
     mock_prisma_client.db = mock_db
 
@@ -427,17 +493,19 @@ async def test_delete_old_logs_aborts_on_non_int_execute_raw_return():
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
     assert mock_db.execute_raw.call_count == 1
-    assert total_deleted == 0
+    assert result.rows_deleted == 0
 
 
 @pytest.mark.asyncio
 async def test_delete_old_logs_continues_on_valid_int_return():
     """should continue deletion loop across batches when execute_raw returns valid int counts."""
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(side_effect=[500, 300, 0])
     mock_prisma_client.db = mock_db
 
@@ -446,35 +514,37 @@ async def test_delete_old_logs_continues_on_valid_int_return():
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
     assert mock_db.execute_raw.call_count == 3
-    assert total_deleted == 800
+    assert result.rows_deleted == 800
 
 
 @pytest.mark.asyncio
-async def test_delete_old_rows_stops_at_max_batches(monkeypatch):
-    """The run-loop backstop must halt a cleanup that keeps finding rows, so a
-    huge backlog is spread across scheduled runs instead of one unbounded loop."""
-    import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module
-
-    monkeypatch.setattr(cleanup_module, "SPEND_LOG_RUN_LOOPS", 2)
-
+async def test_delete_old_rows_stops_at_max_batches():
+    """The batch cap must halt a cleanup that keeps finding rows, so a huge
+    backlog is spread across scheduled runs instead of one unbounded loop, and
+    the operator-facing knob must mean exactly the number of statements it names."""
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(return_value=1000)
     mock_prisma_client.db = mock_db
 
     cleaner = SpendLogCleanup(
-        general_settings={"maximum_spend_logs_retention_period": "7d"}
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_max_batches": 2,
+        }
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
-    # run_count exceeds the cap only after 3 full batches (0, 1, 2)
-    assert mock_db.execute_raw.call_count == 3
-    assert total_deleted == 3000
+    assert mock_db.execute_raw.call_count == 2
+    assert result.rows_deleted == 2000
+    assert result.stop_reason == "batch_cap_reached"
 
 
 @pytest.mark.asyncio
@@ -482,7 +552,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key():
     """Tool index rows are derived from spend logs and expire on the same cutoff;
     the delete must match on the table's composite primary key."""
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(side_effect=[5, 0])
     mock_prisma_client.db = mock_db
 
@@ -491,9 +563,9 @@ async def test_delete_old_tool_index_rows_deletes_on_composite_key():
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_tool_index_rows(mock_prisma_client, cutoff_date, _far_deadline())
 
-    assert total_deleted == 5
+    assert result.rows_deleted == 5
     delete_sql = mock_db.execute_raw.call_args_list[0][0][0]
     assert 'DELETE FROM "LiteLLM_SpendLogToolIndex"' in delete_sql
     assert 'WHERE ("request_id", "tool_name") IN' in delete_sql
@@ -513,7 +585,9 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch)
     )
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     # batch 1 succeeds, batch 2 raises (one-off DB timeout), batches 3-4 succeed,
     # batch 5 returns 0 → loop exits naturally.
     mock_db.execute_raw = AsyncMock(
@@ -526,11 +600,11 @@ async def test_delete_old_logs_continues_after_single_batch_failure(monkeypatch)
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
     # All 5 batches should have been attempted; 100 + 200 + 50 = 350 deleted.
     assert mock_db.execute_raw.call_count == 5
-    assert total_deleted == 350
+    assert result.rows_deleted == 350
 
 
 @pytest.mark.asyncio
@@ -548,7 +622,9 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
     )
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     # Every batch raises — must abort after exactly 3 attempts, not loop forever.
     mock_db.execute_raw = AsyncMock(
         side_effect=ConnectionError("simulated persistent DB outage")
@@ -560,10 +636,10 @@ async def test_delete_old_logs_aborts_after_consecutive_failures(monkeypatch):
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
     assert mock_db.execute_raw.call_count == 3
-    assert total_deleted == 0
+    assert result.rows_deleted == 0
 
 
 @pytest.mark.asyncio
@@ -580,7 +656,9 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc
     )
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     # Pattern: fail, fail, success (resets counter), fail, fail, success, done.
     # Without reset, three of these would trip abort; with reset, they don't.
     mock_db.execute_raw = AsyncMock(
@@ -601,10 +679,10 @@ async def test_delete_old_logs_resets_consecutive_failures_on_success(monkeypatc
     )
 
     cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
-    total_deleted = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date)
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, _far_deadline())
 
     assert mock_db.execute_raw.call_count == 7
-    assert total_deleted == 150
+    assert result.rows_deleted == 150
 
 
 @pytest.mark.asyncio
@@ -617,6 +695,7 @@ async def test_cleanup_uses_logger_exception_for_full_traceback(monkeypatch):
     monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger)
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     # Force the outer try/except to fire by making _should_delete_spend_logs raise.
     cleaner = cleanup_module.SpendLogCleanup(
         general_settings={"maximum_spend_logs_retention_period": "7d"}
@@ -653,7 +732,9 @@ async def test_cleanup_releases_lock_after_persistent_batch_failures(monkeypatch
     )
 
     mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
     mock_db = MagicMock()
+    _wire_tx(mock_db)
     mock_db.execute_raw = AsyncMock(side_effect=TimeoutError("DB down"))
     mock_prisma_client.db = mock_db
 
@@ -698,6 +779,7 @@ def _mock_prisma_for_retention(side_effect: list) -> "MagicMock":
     from unittest.mock import AsyncMock, MagicMock
 
     client = MagicMock()
+    _wire_tx(client.db)
     client.db.execute_raw = AsyncMock(side_effect=side_effect)
     return client
 
@@ -753,3 +835,536 @@ async def test_no_retention_keys_means_no_cleanup_at_all():
     cleaner.pod_lock_manager = None
     await cleaner.cleanup_old_spend_logs(client)
     assert client.db.execute_raw.await_count == 0
+
+
+@pytest.mark.asyncio
+async def test_run_budget_stops_the_loop_and_leaves_the_backlog_for_the_next_run():
+    """
+    The wall-clock budget is the bound that keeps a large backlog from turning
+    into one multi-hour run. With rows always available, the loop must stop on
+    the deadline rather than on the batch cap, and must report that reason so
+    operators can tell a budgeted stop from a drained table.
+    """
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=1000)
+    mock_prisma_client.db = mock_db
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            # Comfortably more batches than a sub-second budget can reach (each
+            # batch sleeps 0.1s), but small enough that a broken deadline fails
+            # this test in seconds instead of hanging it
+            "maximum_spend_logs_cleanup_max_batches": 50,
+        }
+    )
+
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=7)
+    started_at = time.monotonic()
+    result = await cleaner._delete_old_logs(mock_prisma_client, cutoff_date, time.monotonic() + 0.25)
+    elapsed = time.monotonic() - started_at
+
+    assert result.stop_reason == "budget_exhausted"
+    assert elapsed < 3, f"budgeted run overran its deadline: {elapsed}s"
+    assert mock_db.execute_raw.call_count < 50
+    assert result.rows_deleted > 0
+
+
+@pytest.mark.asyncio
+async def test_run_budget_is_shared_across_tables_not_granted_per_table():
+    """
+    A per-table budget would let a run take N times the configured bound. The
+    deadline is computed once per run, so once it is spent on the first table
+    the later tables must stop immediately rather than each getting a fresh one.
+    """
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=1000)
+    mock_prisma_client.db = mock_db
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_autorouter_session_retention_period": "365d",
+            # Comfortably more batches than a sub-second budget can reach (each
+            # batch sleeps 0.1s), but small enough that a broken deadline fails
+            # this test in seconds instead of hanging it
+            "maximum_spend_logs_cleanup_max_batches": 50,
+            "maximum_spend_logs_cleanup_run_budget": "1s",
+        }
+    )
+    cleaner.pod_lock_manager = None
+
+    started_at = time.monotonic()
+    await cleaner.cleanup_old_spend_logs(mock_prisma_client)
+    elapsed = time.monotonic() - started_at
+
+    # three tables are eligible; a per-table budget would push this past 3s
+    assert elapsed < 2.5, f"budget was granted per table, not per run: {elapsed}s"
+    tables_touched = {call[0][0].split('"')[1] for call in mock_db.execute_raw.call_args_list}
+    assert "LiteLLM_SpendLogs" in tables_touched
+
+
+@pytest.mark.asyncio
+async def test_each_batch_carries_a_statement_and_lock_timeout():
+    """
+    A Prisma transaction timeout cannot interrupt a statement already running,
+    so the Postgres statement_timeout and lock_timeout are the only things
+    stopping one batch from holding row locks and a pooled connection
+    indefinitely. Both must be set, inside the batch's own transaction, and
+    scoped with SET LOCAL so the pooled connection is left unchanged.
+    """
+    recorded: list[str] = []
+
+    mock_prisma_client = MagicMock()
+    mock_db = MagicMock()
+
+    @asynccontextmanager
+    async def _tx():
+        tx = MagicMock()
+
+        async def _execute_raw(sql, *args):
+            recorded.append(sql.strip())
+            return 0
+
+        tx.execute_raw = _execute_raw
+        yield tx
+
+    mock_db.tx = _tx
+    mock_db.query_raw = AsyncMock(return_value=[{"remaining": 0}])
+    mock_prisma_client.db = mock_db
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_batch_timeout": "12s",
+        }
+    )
+
+    await cleaner._delete_old_logs(
+        mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()
+    )
+
+    assert "SET LOCAL statement_timeout = 12000" in recorded
+    assert "SET LOCAL lock_timeout = 12000" in recorded
+    # the timeouts must precede the delete they are meant to bound
+    assert recorded.index("SET LOCAL statement_timeout = 12000") < next(
+        i for i, sql in enumerate(recorded) if sql.startswith("DELETE")
+    )
+
+
+@pytest.mark.parametrize(
+    "setting_value",
+    ["inf", "-inf", "nan", "1e400", "0s", "-5m", "not-a-duration"],
+)
+def test_a_non_finite_or_non_positive_budget_falls_back_to_the_default(setting_value):
+    """
+    The knob must not be able to remove the bound it exists to enforce.
+
+    'inf', 'nan' and '1e400' are the spellings that would turn the deadline
+    into no deadline at all, and '0s' and '-5m' would make every run stop before
+    deleting anything. All of them must land on the default rather than being
+    honoured, and the resulting budget must be usable arithmetic.
+    """
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_run_budget": setting_value,
+        }
+    )
+
+    assert cleaner.run_budget_seconds == SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS
+    assert math.isfinite(cleaner.run_budget_seconds)
+    assert cleaner.run_budget_seconds > 0
+
+
+@pytest.mark.parametrize("setting_value", [0, -1, "abc", "", 2.9])
+def test_a_bad_batch_size_falls_back_to_the_default(setting_value):
+    """A zero or negative batch size would make every DELETE a no-op and the
+    loop spin, so unusable values must fall back rather than be honoured."""
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_batch_size": setting_value,
+        }
+    )
+
+    assert cleaner.batch_size >= 1
+
+
+def test_operator_knobs_override_the_env_defaults():
+    """The knobs are meant to be reachable from general_settings (and therefore
+    from the admin UI), not only from environment variables."""
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_batch_size": 250,
+            "maximum_spend_logs_cleanup_max_batches": 7,
+            "maximum_spend_logs_cleanup_run_budget": "90s",
+            "maximum_spend_logs_cleanup_batch_timeout": "2m",
+        }
+    )
+
+    assert cleaner.batch_size == 250
+    assert cleaner.max_batches == 7
+    assert cleaner.run_budget_seconds == 90
+    assert cleaner.batch_timeout_seconds == 120
+
+
+_BOUND_SETTING_CASES = (
+    ("maximum_spend_logs_cleanup_batch_size", 137, "batch_size", 137),
+    ("maximum_spend_logs_cleanup_max_batches", 9, "max_batches", 9),
+    ("maximum_spend_logs_cleanup_run_budget", "45s", "run_budget_seconds", 45.0),
+    ("maximum_spend_logs_cleanup_batch_timeout", "8s", "batch_timeout_seconds", 8.0),
+)
+
+
+@pytest.mark.parametrize("setting_name, setting_value, attribute, expected", _BOUND_SETTING_CASES)
+@pytest.mark.asyncio
+async def test_a_bound_changed_after_construction_reaches_the_next_run(
+    setting_name, setting_value, attribute, expected
+):
+    """The scheduler holds one long-lived instance and the config reload mutates
+    general_settings in place, so a bound captured at construction would leave
+    every dashboard change inert until the process restarts."""
+    settings = {"maximum_spend_logs_retention_period": "7d"}
+    cleaner = SpendLogCleanup(general_settings=settings)
+    cleaner.pod_lock_manager = None
+    assert getattr(cleaner, attribute) != expected
+
+    settings[setting_name] = setting_value
+
+    await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
+
+    assert getattr(cleaner, attribute) == expected
+
+
+@pytest.mark.parametrize("cleared_to_none", [True, False])
+@pytest.mark.asyncio
+async def test_a_bound_cleared_after_construction_falls_back_to_its_default(cleared_to_none):
+    """Blanking the field in the dashboard has to restore the shipped default
+    rather than leave the operator's old bound in force, whether the reload
+    spells the clear as an explicit None or as an absent key."""
+    settings = {"maximum_spend_logs_retention_period": "7d", "maximum_spend_logs_cleanup_batch_size": 137}
+    cleaner = SpendLogCleanup(general_settings=settings)
+    cleaner.pod_lock_manager = None
+    assert cleaner.batch_size == 137
+
+    if cleared_to_none:
+        settings["maximum_spend_logs_cleanup_batch_size"] = None
+    else:
+        del settings["maximum_spend_logs_cleanup_batch_size"]
+
+    await cleaner.cleanup_old_spend_logs(_mock_prisma_for_retention([0, 0]))
+
+    assert cleaner.batch_size == SPEND_LOG_CLEANUP_BATCH_SIZE
+
+
+def test_every_declared_bound_setting_is_covered_by_a_live_reread_case():
+    """A bound added to the declared set without a live-reread case would be
+    propagated by the proxy and then ignored by the running job."""
+    assert {case[0] for case in _BOUND_SETTING_CASES} == set(SPEND_LOG_CLEANUP_BOUND_SETTINGS)
+
+
+@pytest.mark.asyncio
+async def test_remaining_rows_probe_is_capped_so_it_cannot_scan_the_table():
+    """The remaining-eligible-rows metric must never itself become the long
+    scan this job exists to avoid, so its probe carries a LIMIT."""
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=0)
+    mock_prisma_client.db = mock_db
+
+    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
+
+    await cleaner._delete_old_logs(
+        mock_prisma_client, datetime.now(timezone.utc) - timedelta(days=7), _far_deadline()
+    )
+
+    count_sql = mock_db.query_raw.call_args[0][0]
+    assert "count(*)" in count_sql
+    assert "LIMIT $2" in count_sql
+    assert mock_db.query_raw.call_args[0][2] == SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP
+
+
+@pytest.mark.asyncio
+async def test_a_run_skipped_because_another_pod_holds_the_lock_is_reported():
+    """Operators need to tell "nothing to do" apart from "someone else is doing
+    it", so a lock-skipped run is recorded under its own outcome."""
+    recorded: list[str] = []
+    original_record_run = SpendLogCleanupMetrics.record_run
+
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+
+    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
+    cleaner.pod_lock_manager = MagicMock()
+    cleaner.pod_lock_manager.redis_cache = MagicMock()
+    cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=False)
+    cleaner.pod_lock_manager.release_lock = AsyncMock()
+
+    SpendLogCleanupMetrics.record_run = classmethod(lambda cls, outcome: recorded.append(outcome))
+    try:
+        await cleaner.cleanup_old_spend_logs(mock_prisma_client)
+    finally:
+        SpendLogCleanupMetrics.record_run = original_record_run
+
+    assert recorded == ["skipped_locked"]
+    cleaner.pod_lock_manager.release_lock.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_the_outstanding_rows_probe_carries_a_statement_timeout():
+    """
+    The probe is a statement like any other, so if it were issued bare a slow one
+    would hold a connection past the budget the job advertises, which is exactly
+    what the bounds exist to prevent. With budget to spare it carries the same
+    per-statement timeout the delete batches do.
+    """
+    recorded: list[str] = []
+
+    mock_prisma_client = MagicMock()
+    mock_db = MagicMock()
+
+    @asynccontextmanager
+    async def _tx():
+        tx = MagicMock()
+
+        async def _execute_raw(sql, *args):
+            recorded.append(sql.strip())
+            return 0
+
+        async def _query_raw(sql, *args):
+            recorded.append(sql.strip())
+            return [{"remaining": 7}]
+
+        tx.execute_raw = _execute_raw
+        tx.query_raw = _query_raw
+        yield tx
+
+    mock_db.tx = _tx
+    mock_prisma_client.db = mock_db
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_batch_timeout": "8s",
+        }
+    )
+
+    remaining = await cleaner._count_remaining(
+        mock_prisma_client,
+        datetime.now(timezone.utc) - timedelta(days=7),
+        "LiteLLM_SpendLogs",
+        "startTime",
+        _far_deadline(),
+    )
+
+    assert remaining == 7
+    count_index = next(i for i, sql in enumerate(recorded) if sql.startswith("SELECT count(*)"))
+    assert "SET LOCAL statement_timeout = 8000" in recorded[:count_index], (
+        f"the probe ran without a statement timeout: {recorded}"
+    )
+
+
+@pytest.mark.asyncio
+async def test_a_statement_timeout_is_clamped_to_the_budget_that_is_left():
+    """
+    Postgres has no 'stop at time T', only a per-statement duration, so a batch
+    issued just under the deadline would run a whole batch timeout past it and
+    the run budget would be advisory. Clamping the timeout to the remaining
+    budget is what makes the budget a real wall clock.
+    """
+    recorded: list[str] = []
+    client = MagicMock()
+
+    @asynccontextmanager
+    async def _tx():
+        tx = MagicMock()
+
+        async def _execute_raw(sql, *args):
+            recorded.append(sql.strip())
+            return 0
+
+        tx.execute_raw = _execute_raw
+        tx.query_raw = AsyncMock(return_value=[{"remaining": 0}])
+        yield tx
+
+    client.db.tx = _tx
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "maximum_spend_logs_cleanup_batch_timeout": "30s",
+        }
+    )
+
+    # Only 2s of budget left against a 30s batch timeout.
+    await cleaner._execute_delete_batch(client, "DELETE FROM x", datetime.now(timezone.utc), time.monotonic() + 2)
+
+    timeouts = [sql for sql in recorded if "statement_timeout" in sql]
+    assert timeouts, f"no statement timeout was issued: {recorded}"
+    issued_ms = int(timeouts[0].split("=")[1].strip())
+    assert issued_ms <= 2000, f"the batch was given {issued_ms}ms with only 2000ms of budget left"
+
+
+@pytest.mark.asyncio
+async def test_no_statement_is_issued_once_the_budget_is_spent():
+    """
+    Every table exits through _finish_table, including the ones a spent run never
+    started, so an unconditional probe there would put one more statement per
+    table past the bound.
+    """
+    client = _mock_prisma_for_retention([0, 0])
+    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
+
+    result = await cleaner._finish_table(
+        client,
+        datetime.now(timezone.utc) - timedelta(days=7),
+        "LiteLLM_SpendLogs",
+        "startTime",
+        123,
+        "budget_exhausted",
+        time.monotonic() - 1,
+    )
+
+    assert result.rows_deleted == 123
+    assert result.stop_reason == "budget_exhausted"
+    client.db.query_raw.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_a_batch_cancelled_by_the_deadline_is_budget_exhaustion_not_a_failure(monkeypatch):
+    """
+    Clamping the timeout means the last batch of a budget-exhausted run is
+    cancelled by the deadline itself. Counting that as a batch failure would
+    inflate the failure metric on every such run and walk it toward the abort
+    threshold, so it has to be classified as the bound working.
+    """
+    failures: list[str] = []
+    client = MagicMock()
+    _wire_tx(client.db)
+
+    # The deadline has to pass DURING the batch, not before it: a deadline
+    # already spent is caught by the loop's own check and no batch is ever
+    # issued, which would exercise none of the classification under test.
+    async def _cancelled_after_the_deadline(sql, *args):
+        await asyncio.sleep(0.05)
+        raise Exception("canceling statement due to statement timeout")
+
+    client.db.execute_raw = _cancelled_after_the_deadline
+
+    cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"})
+    monkeypatch.setattr(SpendLogCleanupMetrics, "record_batch_failure", lambda table: failures.append(table))
+
+    result = await cleaner._delete_old_logs(
+        client, datetime.now(timezone.utc) - timedelta(days=7), time.monotonic() + 0.02
+    )
+
+    assert result.stop_reason == "budget_exhausted"
+    assert failures == [], f"a deadline cancellation was recorded as a batch failure: {failures}"
+
+
+@pytest.mark.asyncio
+async def test_partition_maintenance_is_skipped_once_the_run_budget_is_spent():
+    """
+    Dropping a partition is DDL holding an ACCESS EXCLUSIVE lock, and unlike a
+    delete batch it cannot be cut short once it has started. A run whose budget is
+    already gone must therefore not start it at all; the next tick picks it up.
+    """
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=0)
+    mock_prisma_client.db = mock_db
+
+    partition_manager = MagicMock()
+    partition_manager.is_partitioned = AsyncMock(return_value=True)
+    partition_manager.ensure_partitions = AsyncMock()
+    partition_manager.drop_partitions_older_than = AsyncMock(return_value=[])
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "use_spend_logs_partitioning": True,
+        },
+        partition_manager=partition_manager,
+    )
+    cleaner._should_delete_spend_logs()
+
+    # a deadline already in the past is what a run that spent its budget on an
+    # earlier table looks like
+    await cleaner._clean_spend_log_tables(mock_prisma_client, time.monotonic() - 1)
+
+    partition_manager.ensure_partitions.assert_not_awaited()
+    partition_manager.drop_partitions_older_than.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_partition_maintenance_still_runs_while_the_run_has_budget():
+    """The skip above must be caused by the spent budget, not by breaking the
+    partition path outright."""
+    mock_prisma_client = MagicMock()
+    _wire_tx(mock_prisma_client.db)
+    mock_db = MagicMock()
+    _wire_tx(mock_db)
+    mock_db.execute_raw = AsyncMock(return_value=0)
+    mock_prisma_client.db = mock_db
+
+    partition_manager = MagicMock()
+    partition_manager.is_partitioned = AsyncMock(return_value=True)
+    partition_manager.ensure_partitions = AsyncMock()
+    partition_manager.drop_partitions_older_than = AsyncMock(return_value=["LiteLLM_SpendLogs_p20260601"])
+
+    cleaner = SpendLogCleanup(
+        general_settings={
+            "maximum_spend_logs_retention_period": "7d",
+            "use_spend_logs_partitioning": True,
+        },
+        partition_manager=partition_manager,
+    )
+    cleaner._should_delete_spend_logs()
+
+    await cleaner._clean_spend_log_tables(mock_prisma_client, _far_deadline())
+
+    partition_manager.ensure_partitions.assert_awaited_once()
+    partition_manager.drop_partitions_older_than.assert_awaited_once()
+
+
+@pytest.mark.parametrize(
+    "stop_reasons, expected",
+    [
+        (("exhausted",), "completed"),
+        (("exhausted", "exhausted"), "completed"),
+        (("exhausted", "batch_cap_reached"), "batch_cap_reached"),
+        (("batch_cap_reached", "exhausted"), "batch_cap_reached"),
+        (("exhausted", "budget_exhausted"), "budget_exhausted"),
+        (("budget_exhausted", "exhausted"), "budget_exhausted"),
+        (("batch_cap_reached", "budget_exhausted"), "budget_exhausted"),
+        (("budget_exhausted", "batch_cap_reached"), "budget_exhausted"),
+        (("exhausted", "aborted"), "aborted"),
+        (("aborted", "exhausted"), "aborted"),
+        (("budget_exhausted", "aborted"), "aborted"),
+        (("aborted", "budget_exhausted"), "aborted"),
+        (("aborted", "budget_exhausted", "batch_cap_reached"), "aborted"),
+    ],
+)
+def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(stop_reasons, expected):
+    """
+    The run outcome answers "why did this run stop", so a table that merely ran
+    dry must never mask one that hit a bound, and an abort must outrank both.
+
+    Both orders of every pair are covered because this folds several per-table
+    results into one answer: a first-match-wins implementation would pass on
+    whichever order happened to be written and fail on its mirror.
+    """
+    results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons)
+    assert SpendLogCleanup._run_outcome(results) == expected
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts
index 485c7cc1f92..02c82611848 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/proxyConfig/useProxyConfig.ts
@@ -17,6 +17,10 @@ export enum ConfigType {
  */
 export enum GeneralSettingsFieldName {
   MAXIMUM_SPEND_LOGS_RETENTION_PERIOD = "maximum_spend_logs_retention_period",
+  MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_SIZE = "maximum_spend_logs_cleanup_batch_size",
+  MAXIMUM_SPEND_LOGS_CLEANUP_MAX_BATCHES = "maximum_spend_logs_cleanup_max_batches",
+  MAXIMUM_SPEND_LOGS_CLEANUP_RUN_BUDGET = "maximum_spend_logs_cleanup_run_budget",
+  MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_TIMEOUT = "maximum_spend_logs_cleanup_batch_timeout",
   // Add more field names here as needed
 }
 
diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts
index 88a37b30291..ca57c01679d 100644
--- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts
+++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs.ts
@@ -6,6 +6,10 @@ import { proxyConfigKeys } from "../proxyConfig/useProxyConfig";
 export interface StoreRequestInSpendLogsParams {
   store_prompts_in_spend_logs: boolean;
   maximum_spend_logs_retention_period?: string;
+  maximum_spend_logs_cleanup_batch_size?: number;
+  maximum_spend_logs_cleanup_max_batches?: number;
+  maximum_spend_logs_cleanup_run_budget?: string;
+  maximum_spend_logs_cleanup_batch_timeout?: string;
 }
 
 export interface StoreRequestInSpendLogsResponse {
@@ -19,6 +23,8 @@ const performStoreRequestInSpendLogs = async (
   const proxyBaseUrl = getProxyBaseUrl();
   const url = proxyBaseUrl ? `${proxyBaseUrl}/config/update` : `/config/update`;
 
+  const { store_prompts_in_spend_logs, ...optionalSettings } = params;
+
   const response = await fetch(url, {
     method: "POST",
     headers: {
@@ -27,10 +33,8 @@ const performStoreRequestInSpendLogs = async (
     },
     body: JSON.stringify({
       general_settings: {
-        store_prompts_in_spend_logs: params.store_prompts_in_spend_logs,
-        ...(params.maximum_spend_logs_retention_period && {
-          maximum_spend_logs_retention_period: params.maximum_spend_logs_retention_period,
-        }),
+        store_prompts_in_spend_logs,
+        ...optionalSettings,
       },
     }),
   });
diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx
index d90720df3c2..a5ef93d2151 100644
--- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.test.tsx
@@ -1,4 +1,8 @@
-import { useDeleteProxyConfigField, useProxyConfig } from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig";
+import {
+  DeleteProxyConfigFieldRequest,
+  useDeleteProxyConfigField,
+  useProxyConfig,
+} from "@/app/(dashboard)/hooks/proxyConfig/useProxyConfig";
 import { useStoreRequestInSpendLogs } from "@/app/(dashboard)/hooks/storeRequestInSpendLogs/useStoreRequestInSpendLogs";
 import NotificationsManager from "@/components/molecules/notifications_manager";
 import { parseErrorMessage } from "@/components/shared/errorUtils";
@@ -40,6 +44,64 @@ describe("LoggingSettings", () => {
   const mockDeleteField = vi.fn();
   const mockRefetch = vi.fn();
 
+  const clearedFieldNames = (): string[] =>
+    mockDeleteField.mock.calls.map((call) => (call[0] as DeleteProxyConfigFieldRequest).field_name);
+
+  // Every optional knob already persisted. Clearing is only ever issued for a
+  // field that has a stored value, so any test about the clear path has to say
+  // so; the default mock below is an empty config, which is a proxy that has
+  // never saved these settings and therefore has nothing to clear.
+  const withEveryOptionalFieldStored = () =>
+    mockUseProxyConfig.mockReturnValue({
+      data: [
+        {
+          field_name: "maximum_spend_logs_retention_period",
+          field_type: "string",
+          field_description: "Maximum retention period",
+          field_value: "30d",
+          stored_in_db: true,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_batch_size",
+          field_type: "Integer",
+          field_description: "Rows per delete",
+          field_value: 2000,
+          stored_in_db: true,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_max_batches",
+          field_type: "Integer",
+          field_description: "Deletes per table per run",
+          field_value: 50,
+          stored_in_db: true,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_run_budget",
+          field_type: "string",
+          field_description: "Wall clock budget per run",
+          field_value: "90s",
+          stored_in_db: true,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_batch_timeout",
+          field_type: "string",
+          field_description: "Statement and lock timeout per batch",
+          field_value: "10s",
+          stored_in_db: true,
+        },
+      ],
+      isLoading: false,
+      refetch: mockRefetch,
+    } as unknown as ReturnType<typeof useProxyConfig>);
+
+  // Blank every optional input the form rendered from stored values, which is
+  // what an admin does to reset a knob to its default.
+  const blankEveryOptionalField = async (user: ReturnType<typeof userEvent.setup>) => {
+    for (const placeholder of ["e.g., 7d, 30d", "e.g., 1000", "e.g., 500", "e.g., 5m", "e.g., 30s"]) {
+      await user.clear(screen.getByPlaceholderText(placeholder));
+    }
+  };
+
   beforeEach(() => {
     vi.resetAllMocks();
     mockUseStoreRequestInSpendLogs.mockReturnValue({
@@ -68,6 +130,19 @@ describe("LoggingSettings", () => {
     expect(screen.getByRole("button", { name: "Save Settings" })).toBeInTheDocument();
   });
 
+  it("should render a control for every spend logs cleanup knob", () => {
+    renderWithProviders(<LoggingSettings />);
+
+    expect(screen.getByLabelText("Spend Logs Cleanup Batch Size (Optional)")).toBeInTheDocument();
+    expect(screen.getByLabelText("Spend Logs Cleanup Max Batches (Optional)")).toBeInTheDocument();
+    expect(screen.getByLabelText("Spend Logs Cleanup Run Budget (Optional)")).toBeInTheDocument();
+    expect(screen.getByLabelText("Spend Logs Cleanup Batch Timeout (Optional)")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., 1000")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., 500")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., 5m")).toBeInTheDocument();
+    expect(screen.getByPlaceholderText("e.g., 30s")).toBeInTheDocument();
+  });
+
   it("should toggle store prompts switch", async () => {
     const user = userEvent.setup();
     renderWithProviders(<LoggingSettings />);
@@ -94,6 +169,9 @@ describe("LoggingSettings", () => {
 
   it("should submit form with store prompts enabled and retention period", async () => {
     const user = userEvent.setup();
+    mockDeleteField.mockImplementation((_params, options) => {
+      options?.onSettled?.();
+    });
     mockMutate.mockImplementation((_params, options) => {
       options?.onSuccess?.();
     });
@@ -110,7 +188,6 @@ describe("LoggingSettings", () => {
     await user.click(saveButton);
 
     await waitFor(() => {
-      expect(mockDeleteField).not.toHaveBeenCalled();
       expect(mockMutate).toHaveBeenCalledWith(
         {
           store_prompts_in_spend_logs: true,
@@ -119,9 +196,42 @@ describe("LoggingSettings", () => {
         expect.any(Object),
       );
     });
+    expect(clearedFieldNames()).not.toContain("maximum_spend_logs_retention_period");
   });
 
-  it("should delete retention period field when left empty on submit", async () => {
+  it("should submit every spend logs cleanup setting that has a value", async () => {
+    const user = userEvent.setup();
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    await user.click(screen.getByRole("switch"));
+    await user.type(screen.getByPlaceholderText("e.g., 7d, 30d"), "30d");
+    await user.type(screen.getByPlaceholderText("e.g., 1000"), "2000");
+    await user.type(screen.getByPlaceholderText("e.g., 500"), "50");
+    await user.type(screen.getByPlaceholderText("e.g., 5m"), "90s");
+    await user.type(screen.getByPlaceholderText("e.g., 30s"), "10s");
+
+    await user.click(screen.getByRole("button", { name: "Save Settings" }));
+
+    const expectedParams = {
+      store_prompts_in_spend_logs: true,
+      maximum_spend_logs_retention_period: "30d",
+      maximum_spend_logs_cleanup_batch_size: 2000,
+      maximum_spend_logs_cleanup_max_batches: 50,
+      maximum_spend_logs_cleanup_run_budget: "90s",
+      maximum_spend_logs_cleanup_batch_timeout: "10s",
+    };
+
+    await waitFor(() => {
+      expect(mockMutate).toHaveBeenCalledWith(expectedParams, expect.any(Object));
+    });
+    expect(mockDeleteField).not.toHaveBeenCalled();
+  });
+
+  it("should omit blank cleanup settings from the save payload instead of sending empty values", async () => {
     const user = userEvent.setup();
     mockDeleteField.mockImplementation((_params, options) => {
       options?.onSettled?.();
@@ -132,11 +242,70 @@ describe("LoggingSettings", () => {
 
     renderWithProviders(<LoggingSettings />);
 
+    await user.type(screen.getByPlaceholderText("e.g., 1000"), "2000");
+    await user.click(screen.getByRole("button", { name: "Save Settings" }));
+
+    await waitFor(() => {
+      expect(mockMutate).toHaveBeenCalled();
+    });
+
+    const submittedParams = mockMutate.mock.calls[0][0];
+    expect(submittedParams).not.toHaveProperty("maximum_spend_logs_retention_period");
+    expect(submittedParams).not.toHaveProperty("maximum_spend_logs_cleanup_max_batches");
+    expect(submittedParams).not.toHaveProperty("maximum_spend_logs_cleanup_run_budget");
+    expect(submittedParams).not.toHaveProperty("maximum_spend_logs_cleanup_batch_timeout");
+    expect(submittedParams).toEqual({
+      store_prompts_in_spend_logs: false,
+      maximum_spend_logs_cleanup_batch_size: 2000,
+    });
+  });
+
+  it("should clear the stored value of every cleanup setting left blank", async () => {
+    const user = userEvent.setup();
+    withEveryOptionalFieldStored();
+    mockDeleteField.mockImplementation((_params, options) => {
+      options?.onSettled?.();
+    });
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    await blankEveryOptionalField(user);
+    await user.type(screen.getByPlaceholderText("e.g., 5m"), "10m");
+    await user.click(screen.getByRole("button", { name: "Save Settings" }));
+
+    await waitFor(() => {
+      expect(mockMutate).toHaveBeenCalled();
+    });
+
+    expect(clearedFieldNames().sort()).toEqual([
+      "maximum_spend_logs_cleanup_batch_size",
+      "maximum_spend_logs_cleanup_batch_timeout",
+      "maximum_spend_logs_cleanup_max_batches",
+      "maximum_spend_logs_retention_period",
+    ]);
+  });
+
+  it("should delete retention period field when left empty on submit", async () => {
+    const user = userEvent.setup();
+    withEveryOptionalFieldStored();
+    mockDeleteField.mockImplementation((_params, options) => {
+      options?.onSettled?.();
+    });
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    await blankEveryOptionalField(user);
     const saveButton = screen.getByRole("button", { name: "Save Settings" });
     await user.click(saveButton);
 
     await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
+      expect(clearedFieldNames()).toContain("maximum_spend_logs_retention_period");
       expect(mockMutate).toHaveBeenCalledWith(
         {
           store_prompts_in_spend_logs: false,
@@ -234,6 +403,38 @@ describe("LoggingSettings", () => {
           stored_in_db: true,
           field_default_value: undefined,
         },
+        {
+          field_name: "maximum_spend_logs_cleanup_batch_size",
+          field_type: "Integer",
+          field_description: "Rows per delete",
+          field_value: 2000,
+          stored_in_db: true,
+          field_default_value: 1000,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_max_batches",
+          field_type: "Integer",
+          field_description: "Deletes per table per run",
+          field_value: 50,
+          stored_in_db: true,
+          field_default_value: 500,
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_run_budget",
+          field_type: "string",
+          field_description: "Wall clock budget per run",
+          field_value: "90s",
+          stored_in_db: true,
+          field_default_value: "5m",
+        },
+        {
+          field_name: "maximum_spend_logs_cleanup_batch_timeout",
+          field_type: "string",
+          field_description: "Statement and lock timeout per batch",
+          field_value: "10s",
+          stored_in_db: true,
+          field_default_value: "30s",
+        },
       ],
       isLoading: false,
       refetch: mockRefetch,
@@ -246,6 +447,10 @@ describe("LoggingSettings", () => {
 
     expect(switchElement).toBeChecked();
     expect(retentionInput).toHaveValue("30d");
+    expect(screen.getByPlaceholderText("e.g., 1000")).toHaveDisplayValue("2000");
+    expect(screen.getByPlaceholderText("e.g., 500")).toHaveDisplayValue("50");
+    expect(screen.getByPlaceholderText("e.g., 5m")).toHaveValue("90s");
+    expect(screen.getByPlaceholderText("e.g., 30s")).toHaveValue("10s");
   });
 
   it("should reflect persisted values that arrive after the initial loading render", async () => {
@@ -307,11 +512,11 @@ describe("LoggingSettings", () => {
     expect(skeletons.length).toBeGreaterThan(0);
   });
 
-  it("should continue with update even if deleteField fails", async () => {
+  it("should report an error and not claim success when clearing a field fails", async () => {
     const user = userEvent.setup();
-    const deleteError = new Error("Field does not exist");
+    withEveryOptionalFieldStored();
     mockDeleteField.mockImplementation((_params, options) => {
-      options?.onError?.(deleteError);
+      options?.onError?.(new Error("Field does not exist"));
       options?.onSettled?.();
     });
     mockMutate.mockImplementation((_params, options) => {
@@ -320,19 +525,52 @@ describe("LoggingSettings", () => {
 
     renderWithProviders(<LoggingSettings />);
 
+    await blankEveryOptionalField(user);
+    const saveButton = screen.getByRole("button", { name: "Save Settings" });
+    await user.click(saveButton);
+
+    await waitFor(() => {
+      expect(mockNotificationsManager.fromBackend).toHaveBeenCalled();
+    });
+    // the old value is still in force server side, so an unqualified success
+    // notification would tell the admin the opposite of what happened
+    expect(mockNotificationsManager.success).not.toHaveBeenCalled();
+  });
+
+  it("should clear fields one at a time, never concurrently", async () => {
+    const user = userEvent.setup();
+    withEveryOptionalFieldStored();
+    let inFlight = 0;
+    let maxInFlight = 0;
+
+    mockDeleteField.mockImplementation((_params, options) => {
+      inFlight += 1;
+      maxInFlight = Math.max(maxInFlight, inFlight);
+      // Settle on a microtask rather than synchronously, so a parallel
+      // implementation genuinely overlaps: Promise.all would issue every call
+      // before any of them settles, driving inFlight to the number of fields.
+      void Promise.resolve().then(() => {
+        inFlight -= 1;
+        options?.onSettled?.();
+      });
+    });
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    await blankEveryOptionalField(user);
     const saveButton = screen.getByRole("button", { name: "Save Settings" });
     await user.click(saveButton);
 
     await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
-      expect(mockMutate).toHaveBeenCalledWith(
-        {
-          store_prompts_in_spend_logs: false,
-        },
-        expect.any(Object),
-      );
       expect(mockNotificationsManager.success).toHaveBeenCalled();
     });
+    // /config/field/delete rewrites the whole general_settings object, so two of
+    // them in flight at once means the later write restores what the earlier cleared
+    expect(maxInFlight).toBe(1);
+    expect(mockDeleteField.mock.calls.length).toBeGreaterThan(1);
   });
 
   it("should submit with only store prompts enabled when retention is empty", async () => {
@@ -353,7 +591,6 @@ describe("LoggingSettings", () => {
     await user.click(saveButton);
 
     await waitFor(() => {
-      expect(mockDeleteField).toHaveBeenCalled();
       expect(mockMutate).toHaveBeenCalledWith(
         {
           store_prompts_in_spend_logs: true,
@@ -361,5 +598,57 @@ describe("LoggingSettings", () => {
         expect.any(Object),
       );
     });
+    // nothing is stored for the blank fields, so there is nothing to clear
+    expect(mockDeleteField).not.toHaveBeenCalled();
+  });
+
+  it("should save on a proxy that has never stored these settings, without clearing anything", async () => {
+    // The first save on a new deployment: no general_settings row exists, so
+    // /config/field/delete answers 400 for every blank field. Issuing those
+    // clears anyway failed the whole save and persisted nothing.
+    const user = userEvent.setup();
+    mockDeleteField.mockImplementation((_params, options) => {
+      options?.onError?.(new Error("Field name=... not in config"));
+      options?.onSettled?.();
+    });
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    const switchElement = screen.getByRole("switch");
+    await user.click(switchElement);
+    await user.click(screen.getByRole("button", { name: "Save Settings" }));
+
+    await waitFor(() => {
+      expect(mockNotificationsManager.success).toHaveBeenCalled();
+    });
+    expect(mockDeleteField).not.toHaveBeenCalled();
+    expect(mockMutate).toHaveBeenCalledWith({ store_prompts_in_spend_logs: true }, expect.any(Object));
+    expect(mockNotificationsManager.fromBackend).not.toHaveBeenCalled();
+  });
+
+  it("should still clear a field that does have a stored value", async () => {
+    // The guard above must not turn into "never clear anything": a field the
+    // admin blanks out that IS stored still has to be deleted server side.
+    const user = userEvent.setup();
+    withEveryOptionalFieldStored();
+    mockDeleteField.mockImplementation((_params, options) => {
+      options?.onSettled?.();
+    });
+    mockMutate.mockImplementation((_params, options) => {
+      options?.onSuccess?.();
+    });
+
+    renderWithProviders(<LoggingSettings />);
+
+    await user.clear(screen.getByPlaceholderText("e.g., 5m"));
+    await user.click(screen.getByRole("button", { name: "Save Settings" }));
+
+    await waitFor(() => {
+      expect(mockNotificationsManager.success).toHaveBeenCalled();
+    });
+    expect(clearedFieldNames()).toContain("maximum_spend_logs_cleanup_run_budget");
   });
 });
diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx
index 0d156a8dc10..07928768cd4 100644
--- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx
+++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/LoggingSettings/LoggingSettings.tsx
@@ -13,43 +13,165 @@ import {
 import NotificationsManager from "@/components/molecules/notifications_manager";
 import { parseErrorMessage } from "@/components/shared/errorUtils";
 import { ClockCircleOutlined } from "@ant-design/icons";
-import { Button, Card, Form, Input, Skeleton, Space, Switch, Typography } from "antd";
-import React, { useMemo } from "react";
+import { Button, Card, Form, Input, InputNumber, Skeleton, Space, Switch, Typography } from "antd";
+import React, { useCallback, useMemo } from "react";
+
+const STORE_PROMPTS_FIELD_NAME = "store_prompts_in_spend_logs";
+
+interface OptionalField {
+  readonly name: GeneralSettingsFieldName;
+  readonly kind: "duration" | "count";
+  readonly label: string;
+  readonly placeholder: string;
+  readonly fallbackTooltip: string;
+}
+
+const OPTIONAL_FIELDS: readonly OptionalField[] = [
+  {
+    name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD,
+    kind: "duration",
+    label: "Maximum Spend Logs Retention Period (Optional)",
+    placeholder: "e.g., 7d, 30d",
+    fallbackTooltip:
+      "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit.",
+  },
+  {
+    name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_SIZE,
+    kind: "count",
+    label: "Spend Logs Cleanup Batch Size (Optional)",
+    placeholder: "e.g., 1000",
+    fallbackTooltip: "Rows deleted per DELETE statement during cleanup. Leave empty to use the default of 1000.",
+  },
+  {
+    name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_CLEANUP_MAX_BATCHES,
+    kind: "count",
+    label: "Spend Logs Cleanup Max Batches (Optional)",
+    placeholder: "e.g., 500",
+    fallbackTooltip:
+      "Maximum number of DELETE statements run per table per cleanup run. Leave empty to use the default of 500.",
+  },
+  {
+    name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_CLEANUP_RUN_BUDGET,
+    kind: "duration",
+    label: "Spend Logs Cleanup Run Budget (Optional)",
+    placeholder: "e.g., 5m",
+    fallbackTooltip:
+      "Wall-clock budget for a whole cleanup run, shared across every table it cleans (e.g., '5m'). Leave empty to use the default of 5m.",
+  },
+  {
+    name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_CLEANUP_BATCH_TIMEOUT,
+    kind: "duration",
+    label: "Spend Logs Cleanup Batch Timeout (Optional)",
+    placeholder: "e.g., 30s",
+    fallbackTooltip:
+      "Postgres statement and lock timeout applied to each cleanup batch, so cleanup never monopolizes a connection (e.g., '30s'). Leave empty to use the default of 30s.",
+  },
+];
+
+interface LoggingSettingsFormValues {
+  store_prompts_in_spend_logs: boolean;
+  maximum_spend_logs_retention_period?: string | null;
+  maximum_spend_logs_cleanup_batch_size?: number | null;
+  maximum_spend_logs_cleanup_max_batches?: number | null;
+  maximum_spend_logs_cleanup_run_budget?: string | null;
+  maximum_spend_logs_cleanup_batch_timeout?: string | null;
+}
+
+const hasDuration = (value: string | null | undefined): value is string =>
+  typeof value === "string" && value.trim() !== "";
+
+const hasCount = (value: number | null | undefined): value is number =>
+  typeof value === "number" && Number.isFinite(value);
+
+const buildUpdateParams = (formValues: LoggingSettingsFormValues): StoreRequestInSpendLogsParams => ({
+  store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs,
+  ...(hasDuration(formValues.maximum_spend_logs_retention_period) && {
+    maximum_spend_logs_retention_period: formValues.maximum_spend_logs_retention_period,
+  }),
+  ...(hasCount(formValues.maximum_spend_logs_cleanup_batch_size) && {
+    maximum_spend_logs_cleanup_batch_size: formValues.maximum_spend_logs_cleanup_batch_size,
+  }),
+  ...(hasCount(formValues.maximum_spend_logs_cleanup_max_batches) && {
+    maximum_spend_logs_cleanup_max_batches: formValues.maximum_spend_logs_cleanup_max_batches,
+  }),
+  ...(hasDuration(formValues.maximum_spend_logs_cleanup_run_budget) && {
+    maximum_spend_logs_cleanup_run_budget: formValues.maximum_spend_logs_cleanup_run_budget,
+  }),
+  ...(hasDuration(formValues.maximum_spend_logs_cleanup_batch_timeout) && {
+    maximum_spend_logs_cleanup_batch_timeout: formValues.maximum_spend_logs_cleanup_batch_timeout,
+  }),
+});
+
+// A blank field only needs clearing when something is actually stored for it.
+// Asking the proxy to clear a field it has no value for is a 400 whenever no
+// general_settings row exists at all, which is the state of every deployment
+// that has never saved one, so clearing unconditionally would fail the first
+// save on a new proxy and take the rest of the form down with it.
+const omittedFieldNames = (
+  updateParams: StoreRequestInSpendLogsParams,
+  isStored: (name: GeneralSettingsFieldName) => boolean,
+): readonly GeneralSettingsFieldName[] =>
+  OPTIONAL_FIELDS.map((field) => field.name).filter((name) => !(name in updateParams) && isStored(name));
 
 const LoggingSettings: React.FC = () => {
-  const [form] = Form.useForm();
+  const [form] = Form.useForm<LoggingSettingsFormValues>();
   const { mutate, isPending } = useStoreRequestInSpendLogs();
   const { mutate: deleteField, isPending: isDeletingField } = useDeleteProxyConfigField();
   const { data: proxyConfigData, isLoading: isLoadingConfig } = useProxyConfig(ConfigType.GENERAL_SETTINGS);
 
+  const describeField = (name: string, fallback: string) =>
+    proxyConfigData?.find((field) => field.field_name === name)?.field_description || fallback;
+
+  const storedValue = useCallback(
+    (name: string) => proxyConfigData?.find((field) => field.field_name === name)?.field_value,
+    [proxyConfigData],
+  );
+
+  const isStored = (name: GeneralSettingsFieldName) => {
+    const value = storedValue(name);
+    return value !== null && value !== undefined;
+  };
+
   const initialValues = useMemo(() => {
-    if (!proxyConfigData) {
-      return {
-        store_prompts_in_spend_logs: false,
-        maximum_spend_logs_retention_period: undefined,
-      };
-    }
-
-    const storePromptsField = proxyConfigData.find((field) => field.field_name === "store_prompts_in_spend_logs");
-    const retentionPeriodField = proxyConfigData.find(
-      (field) => field.field_name === "maximum_spend_logs_retention_period",
-    );
-
     return {
-      store_prompts_in_spend_logs: storePromptsField?.field_value ?? false,
-      maximum_spend_logs_retention_period: retentionPeriodField?.field_value ?? undefined,
+      store_prompts_in_spend_logs: storedValue(STORE_PROMPTS_FIELD_NAME) ?? false,
+      ...Object.fromEntries(OPTIONAL_FIELDS.map((field) => [field.name, storedValue(field.name)])),
     };
-  }, [proxyConfigData]);
+  }, [storedValue]);
 
-  const handleFormSubmit = (formValues: StoreRequestInSpendLogsParams) => {
-    const retentionPeriodValue = formValues.maximum_spend_logs_retention_period;
-    const hasRetentionPeriod = typeof retentionPeriodValue === "string" && retentionPeriodValue.trim() !== "";
+  // Resolves to the field name when clearing it failed, or null when it worked.
+  const clearStoredField = (fieldName: GeneralSettingsFieldName) =>
+    new Promise<GeneralSettingsFieldName | null>((resolve) => {
+      let failed = false;
+      deleteField(
+        { config_type: ConfigType.GENERAL_SETTINGS, field_name: fieldName },
+        {
+          onError: () => {
+            failed = true;
+          },
+          onSettled: () => resolve(failed ? fieldName : null),
+        },
+      );
+    });
 
-    const updateParams: StoreRequestInSpendLogsParams = {
-      store_prompts_in_spend_logs: formValues.store_prompts_in_spend_logs,
-      ...(hasRetentionPeriod && { maximum_spend_logs_retention_period: retentionPeriodValue }),
-    };
+  // Clearing a field rewrites the whole stored general_settings object server
+  // side, so these must run one at a time: in parallel the last write back wins
+  // and silently restores the fields the earlier ones just cleared.
+  const clearStoredFieldsInSequence = async (
+    fieldNames: readonly GeneralSettingsFieldName[],
+  ): Promise<readonly GeneralSettingsFieldName[]> => {
+    const failed: GeneralSettingsFieldName[] = [];
+    for (const fieldName of fieldNames) {
+      const failure = await clearStoredField(fieldName);
+      if (failure !== null) {
+        failed.push(failure);
+      }
+    }
+    return failed;
+  };
 
+  const handleFormSubmit = (formValues: LoggingSettingsFormValues) => {
+    const updateParams = buildUpdateParams(formValues);
     const submitUpdate = () =>
       mutate(updateParams, {
         onSuccess: () => NotificationsManager.success("Spend logs settings updated successfully"),
@@ -57,21 +179,21 @@ const LoggingSettings: React.FC = () => {
           NotificationsManager.fromBackend("Failed to save spend logs settings: " + parseErrorMessage(error)),
       });
 
-    if (hasRetentionPeriod) {
+    const fieldsToClear = omittedFieldNames(updateParams, isStored);
+    if (fieldsToClear.length === 0) {
       submitUpdate();
       return;
     }
 
-    deleteField(
-      {
-        config_type: ConfigType.GENERAL_SETTINGS,
-        field_name: GeneralSettingsFieldName.MAXIMUM_SPEND_LOGS_RETENTION_PERIOD,
-      },
-      {
-        onError: (deleteError) => console.warn("Failed to delete retention period field (may not exist):", deleteError),
-        onSettled: submitUpdate,
-      },
-    );
+    void clearStoredFieldsInSequence(fieldsToClear).then((failed) => {
+      if (failed.length > 0) {
+        // Reporting an unqualified success here would tell the admin a setting
+        // was reset to its default while the old value is still in force.
+        NotificationsManager.fromBackend(`Failed to clear saved value for: ${failed.join(", ")}`);
+        return;
+      }
+      submitUpdate();
+    });
   };
 
   return (
@@ -87,27 +209,30 @@ const LoggingSettings: React.FC = () => {
           <Form form={form} layout="vertical" onFinish={handleFormSubmit} initialValues={initialValues}>
             <Form.Item
               label="Store Prompts in Spend Logs"
-              name="store_prompts_in_spend_logs"
-              tooltip={
-                proxyConfigData?.find((f) => f.field_name === "store_prompts_in_spend_logs")?.field_description ||
-                "When enabled, prompts will be stored in spend logs for tracking and analysis purposes."
-              }
+              name={STORE_PROMPTS_FIELD_NAME}
+              tooltip={describeField(
+                STORE_PROMPTS_FIELD_NAME,
+                "When enabled, prompts will be stored in spend logs for tracking and analysis purposes.",
+              )}
               valuePropName="checked"
             >
               <Switch />
             </Form.Item>
 
-            <Form.Item
-              label="Maximum Spend Logs Retention Period (Optional)"
-              name="maximum_spend_logs_retention_period"
-              tooltip={
-                proxyConfigData?.find((f) => f.field_name === "maximum_spend_logs_retention_period")
-                  ?.field_description ||
-                "Set the maximum retention period for spend logs (e.g., '7d' for 7 days, '30d' for 30 days). Leave empty for no limit."
-              }
-            >
-              <Input placeholder="e.g., 7d, 30d" prefix={<ClockCircleOutlined />} />
-            </Form.Item>
+            {OPTIONAL_FIELDS.map((field) => (
+              <Form.Item
+                key={field.name}
+                label={field.label}
+                name={field.name}
+                tooltip={describeField(field.name, field.fallbackTooltip)}
+              >
+                {field.kind === "duration" ? (
+                  <Input placeholder={field.placeholder} prefix={<ClockCircleOutlined />} />
+                ) : (
+                  <InputNumber min={1} precision={0} placeholder={field.placeholder} style={{ width: "100%" }} />
+                )}
+              </Form.Item>
+            ))}
 
             <Form.Item>
               <Button type="primary" htmlType="submit" loading={isPending || isDeletingField}>
diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts
index b6dbe0dfc6d..1c5fe408203 100644
--- a/ui/litellm-dashboard/src/lib/http/schema.d.ts
+++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts
@@ -23675,6 +23675,26 @@ export interface components {
              * @description Maximum retention period for auto-router benchmark session rollup rows (e.g., '365d'). Rows whose last turn is older than this are deleted by the spend log cleanup job, on that job's schedule. Unset means rollup rows are never deleted.
              */
             maximum_autorouter_session_retention_period?: string | null;
+            /**
+             * Maximum Spend Logs Cleanup Batch Size
+             * @description Rows deleted per DELETE statement by the spend log cleanup job. Defaults to 1000.
+             */
+            maximum_spend_logs_cleanup_batch_size?: number | null;
+            /**
+             * Maximum Spend Logs Cleanup Batch Timeout
+             * @description Postgres statement_timeout and lock_timeout applied to each spend log cleanup delete batch (e.g. '30s'), so cleanup cannot hold row locks or a connection indefinitely. Defaults to '30s'.
+             */
+            maximum_spend_logs_cleanup_batch_timeout?: string | null;
+            /**
+             * Maximum Spend Logs Cleanup Max Batches
+             * @description Maximum DELETE statements the spend log cleanup job issues per table per run. Defaults to 500.
+             */
+            maximum_spend_logs_cleanup_max_batches?: number | null;
+            /**
+             * Maximum Spend Logs Cleanup Run Budget
+             * @description Wall-clock budget for one spend log cleanup run (e.g. '5m'), shared across every table it prunes. A run that hits the budget stops and the next run resumes from where it left off. Defaults to '5m'.
+             */
+            maximum_spend_logs_cleanup_run_budget?: string | null;
             /**
              * Maximum Spend Logs Retention Period
              * @description Maximum retention period for spend logs (e.g., '7d' for 7 days). Logs older than this will be deleted.

From 3fbe40c9b10ff8fcd47a06a8101526c32e5a59ea Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 13 Aug 2026 19:13:44 -0700
Subject: [PATCH 21/31] fix(scripts): print the make check verdict on early
 informational exits

---
 scripts/pre_commit_lint.sh                 | 2 ++
 tests/test_litellm/test_pre_commit_lint.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/scripts/pre_commit_lint.sh b/scripts/pre_commit_lint.sh
index e19012613cd..82498ec10cd 100755
--- a/scripts/pre_commit_lint.sh
+++ b/scripts/pre_commit_lint.sh
@@ -55,11 +55,13 @@ else
     merge_base=$(git merge-base origin/litellm_internal_staging HEAD 2>/dev/null) || {
         echo "check: cannot resolve the merge base with origin/litellm_internal_staging." >&2
         echo "  Fix: git fetch origin litellm_internal_staging" >&2
+        echo "check: FAIL"
         exit 1
     }
     scope=$(printf '%s\n' "$(git diff --name-only --diff-filter=ACMRD "$merge_base")" "$untracked" | sed '/^$/d' | sort -u)
     if [ -z "$scope" ]; then
         echo "check: nothing to check (no staged files, no working-tree changes, no branch changes vs origin/litellm_internal_staging)"
+        echo "check: PASS"
         exit 0
     fi
     echo "check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging:"
diff --git a/tests/test_litellm/test_pre_commit_lint.py b/tests/test_litellm/test_pre_commit_lint.py
index 7ce16032641..33baf7474ce 100644
--- a/tests/test_litellm/test_pre_commit_lint.py
+++ b/tests/test_litellm/test_pre_commit_lint.py
@@ -224,6 +224,7 @@ def test_nothing_staged_and_no_changes_is_an_explicit_no_op(tmp_path: Path) -> N
     proc = _run(repo, bin_dir, {})
     assert proc.returncode == 0, proc.stdout + proc.stderr
     assert "nothing to check" in proc.stdout
+    assert "check: PASS" in proc.stdout
     assert "linting Python" not in proc.stdout
 
 
@@ -234,6 +235,7 @@ def test_nothing_staged_without_a_base_ref_fails_with_a_fetch_hint(tmp_path: Pat
     assert proc.returncode == 1
     assert "cannot resolve the merge base" in proc.stdout
     assert "git fetch origin litellm_internal_staging" in proc.stdout
+    assert "check: FAIL" in proc.stdout
 
 
 def test_partial_staging_warns_which_checks_were_skipped(tmp_path: Path) -> None:

From 60c6c2c703a0477b0acd5f482a9e16d481f7ef3e Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 13 Aug 2026 19:24:45 -0700
Subject: [PATCH 22/31] test: drop redundant explanatory comments from riva
 provider tests

---
 .../proxy/public_endpoints/test_public_endpoints.py          | 5 -----
 .../src/components/provider_info_helpers.test.tsx            | 3 ---
 2 files changed, 8 deletions(-)

diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
index ead5f4ab5cc..e99bdfb5c35 100644
--- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
+++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py
@@ -244,11 +244,6 @@ def test_bedrock_mantle_provider_fields():
 
 
 def test_nvidia_riva_provider_fields():
-    """The Add Model provider dropdown is populated from /public/providers/fields, so a
-    missing entry meant Riva could not be added through the UI. Riva is gRPC only with no
-    public default endpoint, hence the required api_base, and NVCF hosted Riva cannot be
-    called without nvcf_function_id.
-    """
     app_instance = FastAPI()
     app_instance.include_router(router)
     test_client = TestClient(app_instance)
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
index f43200570b0..b9516945faa 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
@@ -95,9 +95,6 @@ describe("provider_info_helpers", () => {
     });
 
     it("should resolve the nvidia_riva provider value to the Nvidia Riva display name and logo", () => {
-      // The backend registers nvidia_riva and it has a docs page, but the UI
-      // registry had no entry, so it could not be picked in Add Model and the
-      // slug rendered raw with no logo.
       const result = getProviderLogoAndName("nvidia_riva");
       expect(result.displayName).toBe(Providers.NVIDIA_RIVA);
       expect(provider_map.NVIDIA_RIVA).toBe("nvidia_riva");

From d9530bf3d158c3de9b5f00a92f748af8f3626c1b Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 13 Aug 2026 19:26:48 -0700
Subject: [PATCH 23/31] fix(proxy): fail config load when a callbacks entry is
 not dispatchable (#36858)

---
 litellm/proxy/common_utils/callback_utils.py  |  83 ++++++++-
 .../proxy/common_utils/test_callback_utils.py | 164 ++++++++++++++++++
 2 files changed, 240 insertions(+), 7 deletions(-)

diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py
index 1869328c039..60a03689804 100644
--- a/litellm/proxy/common_utils/callback_utils.py
+++ b/litellm/proxy/common_utils/callback_utils.py
@@ -1,7 +1,10 @@
 import copy
 import os
 from collections.abc import Callable, Iterable
-from typing import TYPE_CHECKING, Any, Final, Optional
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
+
+from typing_extensions import assert_never
 
 import litellm
 from litellm import get_secret
@@ -50,6 +53,66 @@ if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
 
 
+@dataclass(frozen=True, slots=True)
+class _CallbackResolvedToClass:
+    entry: str
+    loaded: type
+    tag: Literal["resolved_to_class"] = "resolved_to_class"
+
+
+@dataclass(frozen=True, slots=True)
+class _CallbackNotDispatchable:
+    entry: str
+    loaded: object
+    tag: Literal["not_dispatchable"] = "not_dispatchable"
+
+
+_CallbackLoadError: TypeAlias = _CallbackResolvedToClass | _CallbackNotDispatchable
+
+
+def _classify_loaded_callback(entry: str, loaded: object) -> CustomLogger | Callable[..., object] | _CallbackLoadError:
+    """
+    Decide whether what a ``litellm_settings.callbacks`` dotted path resolved to can be dispatched.
+
+    A dotted path only ever runs as a ``CustomLogger`` instance or as a callback function. Anything
+    else (most commonly a class instead of an instance) used to load without complaint and then be
+    skipped on every request, with no log line and no error.
+    """
+    if isinstance(loaded, CustomLogger) or (callable(loaded) and not isinstance(loaded, type)):
+        return loaded
+    if isinstance(loaded, type):
+        return _CallbackResolvedToClass(entry=entry, loaded=loaded)
+    return _CallbackNotDispatchable(entry=entry, loaded=loaded)
+
+
+def _raise_callback_load_error(error: _CallbackLoadError) -> NoReturn:
+    """The one edge that raises: map a load error onto config load's failure contract."""
+    match error:
+        case _CallbackResolvedToClass():
+            module_path: Final = error.entry.rsplit(".", 1)[0] if "." in error.entry else error.entry
+            raise ValueError(
+                f"litellm_settings.callbacks entry '{error.entry}' resolved to the class "
+                f"{error.loaded.__module__}.{error.loaded.__qualname__}, which is neither a "
+                "CustomLogger instance nor a callable, so the proxy would never run it."
+                f" Point it at an instance instead, e.g. add `proxy_handler_instance = {error.loaded.__name__}()` to "
+                f'{module_path} and set `callbacks: ["{module_path}.proxy_handler_instance"]`.'
+            )
+        case _CallbackNotDispatchable():
+            raise ValueError(
+                f"litellm_settings.callbacks entry '{error.entry}' resolved to "
+                f"{type(error.loaded).__name__} {error.loaded!r}, which is neither a "
+                "CustomLogger instance nor a callable, so the proxy would never run it."
+            )
+    assert_never(error)
+
+
+def _loaded_callback_or_raise(entry: str, loaded: object) -> CustomLogger | Callable[..., object]:
+    resolved: Final = _classify_loaded_callback(entry=entry, loaded=loaded)
+    if isinstance(resolved, _CallbackResolvedToClass | _CallbackNotDispatchable):
+        _raise_callback_load_error(resolved)
+    return resolved
+
+
 def initialize_callbacks_on_proxy(
     value: Any,
     premium_user: bool,
@@ -305,9 +368,12 @@ def initialize_callbacks_on_proxy(
                     "%s attempting to import custom calback=%s %s", blue_color_code, callback, reset_color_code
                 )
                 imported_list.append(
-                    get_instance_fn(
-                        value=callback,
-                        config_file_path=config_file_path,
+                    _loaded_callback_or_raise(
+                        entry=callback,
+                        loaded=get_instance_fn(
+                            value=callback,
+                            config_file_path=config_file_path,
+                        ),
                     )
                 )
         if isinstance(litellm.callbacks, list):
@@ -321,9 +387,12 @@ def initialize_callbacks_on_proxy(
             PrometheusLogger._mount_metrics_endpoint()
     else:
         litellm.callbacks = [
-            get_instance_fn(
-                value=value,
-                config_file_path=config_file_path,
+            _loaded_callback_or_raise(
+                entry=value,
+                loaded=get_instance_fn(
+                    value=value,
+                    config_file_path=config_file_path,
+                ),
             )
         ]
     verbose_proxy_logger.debug("%s Initialized Callbacks - %s %s", blue_color_code, litellm.callbacks, reset_color_code)
diff --git a/tests/test_litellm/proxy/common_utils/test_callback_utils.py b/tests/test_litellm/proxy/common_utils/test_callback_utils.py
index bfd4ffe1593..0c73c3fcf22 100644
--- a/tests/test_litellm/proxy/common_utils/test_callback_utils.py
+++ b/tests/test_litellm/proxy/common_utils/test_callback_utils.py
@@ -21,6 +21,10 @@ from litellm.proxy.common_utils.callback_utils import (
     strip_callback_config,
 )
 import litellm
+from litellm.caching.caching import DualCache
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.proxy.utils import ProxyLogging
 
 from unittest.mock import patch
 from litellm.proxy.common_utils.callback_utils import process_callback
@@ -491,3 +495,163 @@ def test_strip_callback_config_drops_credential_bearing_slots():
 @pytest.mark.parametrize("value", [None, "not-a-dict", 42])
 def test_strip_callback_config_passes_through_non_dicts(value):
     assert strip_callback_config(value) is value
+
+
+# ---------------------------------------------------------------------------
+# initialize_callbacks_on_proxy: dotted-path entries must resolve to something
+# the request path can actually dispatch
+# ---------------------------------------------------------------------------
+
+_PROBE_MODULE_NAME = "custom_callback_probe"
+
+_PROBE_MODULE_SOURCE = '''
+from litellm.integrations.custom_logger import CustomLogger
+
+
+class FloorMaxTokens(CustomLogger):
+    async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
+        data["max_tokens"] = 16
+        return data
+
+
+class NotALogger:
+    pass
+
+
+def log_event_fn(kwargs, response_obj, start_time, end_time):
+    return None
+
+
+NOT_A_CALLBACK = "some-plain-string"
+
+proxy_handler_instance = FloorMaxTokens()
+'''
+
+
+@pytest.fixture
+def probe_config_path(tmp_path):
+    """Write a callback module next to a config.yaml, the layout get_instance_fn's file
+    branch expects, and restore every global the load + dispatch path touches.
+
+    ``ProxyLogging._callback_capabilities_cache`` is keyed on the id()s of the
+    litellm.callbacks members, so an entry left behind here can be read back by an
+    unrelated test whose (len, ids) signature happens to collide.
+    """
+    (tmp_path / f"{_PROBE_MODULE_NAME}.py").write_text(_PROBE_MODULE_SOURCE)
+
+    original_callbacks = (
+        list(litellm.callbacks) if isinstance(litellm.callbacks, list) else litellm.callbacks
+    )
+    litellm.callbacks = []
+    ProxyLogging._callback_capabilities_cache.clear()
+    try:
+        yield str(tmp_path / "config.yaml")
+    finally:
+        litellm.callbacks = original_callbacks
+        ProxyLogging._callback_capabilities_cache.clear()
+
+
+def _load_callbacks(value, config_file_path):
+    initialize_callbacks_on_proxy(
+        value=value,
+        premium_user=False,
+        config_file_path=config_file_path,
+        litellm_settings={},
+        callback_specific_params={},
+    )
+
+
+def test_initialize_callbacks_on_proxy_rejects_class_valued_entry(probe_config_path):
+    """A class path loads an object that fails the `isinstance(_callback, CustomLogger)`
+    dispatch gate in ProxyLogging.pre_call_hook, so the proxy used to boot clean and
+    silently never run the hook. Config load must fail instead."""
+    entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens"
+
+    with pytest.raises(ValueError) as exc_info:
+        _load_callbacks([entry], probe_config_path)
+
+    message = str(exc_info.value)
+    assert entry in message
+    assert "the class" in message
+    assert "FloorMaxTokens" in message
+    assert f"{_PROBE_MODULE_NAME}.proxy_handler_instance" in message
+    assert litellm.callbacks == []
+
+
+@pytest.mark.parametrize(
+    "attribute, expected_fragment",
+    [
+        ("NotALogger", "the class"),
+        ("NOT_A_CALLBACK", "str 'some-plain-string'"),
+    ],
+)
+def test_initialize_callbacks_on_proxy_rejects_non_dispatchable_values(
+    probe_config_path, attribute, expected_fragment
+):
+    entry = f"{_PROBE_MODULE_NAME}.{attribute}"
+
+    with pytest.raises(ValueError) as exc_info:
+        _load_callbacks([entry], probe_config_path)
+
+    message = str(exc_info.value)
+    assert entry in message
+    assert expected_fragment in message
+    assert litellm.callbacks == []
+
+
+def test_initialize_callbacks_on_proxy_rejects_class_valued_non_list_value(probe_config_path):
+    entry = f"{_PROBE_MODULE_NAME}.FloorMaxTokens"
+
+    with pytest.raises(ValueError) as exc_info:
+        _load_callbacks(entry, probe_config_path)
+
+    assert entry in str(exc_info.value)
+
+
+@pytest.mark.asyncio
+async def test_initialize_callbacks_on_proxy_instance_entry_runs_pre_call_hook(probe_config_path):
+    """Positive control: the supported shape must still load AND still run. Drives the
+    real ProxyLogging.pre_call_hook, which is where a class-valued entry goes silent."""
+    _load_callbacks([f"{_PROBE_MODULE_NAME}.proxy_handler_instance"], probe_config_path)
+
+    assert len(litellm.callbacks) == 1
+    assert isinstance(litellm.callbacks[0], CustomLogger)
+
+    ProxyLogging._callback_capabilities_cache.clear()
+    proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
+    data = await proxy_logging.pre_call_hook(
+        user_api_key_dict=UserAPIKeyAuth(api_key="sk-probe"),
+        data={
+            "model": "gpt-4",
+            "messages": [{"role": "user", "content": "hello"}],
+            "max_tokens": 1,
+            "metadata": {},
+        },
+        call_type="acompletion",
+    )
+
+    assert data["max_tokens"] == 16
+
+
+def test_initialize_callbacks_on_proxy_keeps_known_string_callback(probe_config_path):
+    """Non-narrowing control: a known callback name never reaches get_instance_fn and
+    stays a plain string in litellm.callbacks."""
+    _load_callbacks(["langfuse"], probe_config_path)
+
+    assert litellm.callbacks == ["langfuse"]
+
+
+def test_initialize_callbacks_on_proxy_accepts_plain_function_callback(probe_config_path):
+    """Non-narrowing control: litellm.callbacks is typed
+    `Callable | <known name> | CustomLogger`, so a dotted path resolving to a plain
+    function is a supported shape and must keep loading."""
+    _load_callbacks([f"{_PROBE_MODULE_NAME}.log_event_fn"], probe_config_path)
+
+    assert [getattr(cb, "__name__", None) for cb in litellm.callbacks] == ["log_event_fn"]
+
+
+def test_initialize_callbacks_on_proxy_accepts_instance_non_list_value(probe_config_path):
+    _load_callbacks(f"{_PROBE_MODULE_NAME}.proxy_handler_instance", probe_config_path)
+
+    assert len(litellm.callbacks) == 1
+    assert isinstance(litellm.callbacks[0], CustomLogger)

From 86f2f16fd4cdb20ef7746110b3d501742d011430 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 13 Aug 2026 19:28:18 -0700
Subject: [PATCH 24/31] fix(bedrock): hoist custom.defer_loading before
 dropping custom on invoke tools (#36855)

---
 .../anthropic_claude3_transformation.py       |  6 +-
 litellm/llms/bedrock/common_utils.py          | 21 +++--
 .../anthropic_claude3_transformation.py       |  8 +-
 .../test_anthropic_claude3_transformation.py  | 84 +++++++++++++++++--
 4 files changed, 95 insertions(+), 24 deletions(-)

diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
index 647ccc33a44..b8b07af59c6 100644
--- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py
@@ -19,9 +19,9 @@ from litellm.llms.bedrock.common_utils import (
     convert_bedrock_invoke_output_format_to_inline_schema,
     get_anthropic_beta_from_headers,
     normalize_bedrock_opus_output_config_effort,
+    normalize_custom_field_on_tools,
     normalize_tool_input_schema_types_for_bedrock_invoke,
     pop_bedrock_invoke_output_config_format,
-    remove_custom_field_from_tools,
 )
 from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
 from litellm.types.llms.openai import AllMessageValues
@@ -243,8 +243,8 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
         if "anthropic_version" not in anthropic_request:
             anthropic_request["anthropic_version"] = self.anthropic_version
 
-        # Remove `custom` field from tools (Bedrock doesn't support it)
-        remove_custom_field_from_tools(anthropic_request)
+        # Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
+        normalize_custom_field_on_tools(anthropic_request)
         normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_request)
         return anthropic_request
 
diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py
index 48bc60a07e5..4ad20772ed0 100644
--- a/litellm/llms/bedrock/common_utils.py
+++ b/litellm/llms/bedrock/common_utils.py
@@ -176,13 +176,14 @@ def convert_bedrock_invoke_output_format_to_inline_schema(
     request_body["messages"] = new_messages
 
 
-def remove_custom_field_from_tools(request_body: dict) -> None:
+def normalize_custom_field_on_tools(request_body: dict) -> None:
     """
-    Remove ``custom`` field from each tool in the request body.
+    Drop the ``custom`` field from each tool, first hoisting a boolean
+    ``custom.defer_loading`` onto the top-level ``defer_loading`` flag that
+    Bedrock and Anthropic actually document, unless the tool already carries one.
 
-    Claude Code (v2.1.69+) sends ``custom: {defer_loading: true}`` on tool
-    definitions, which Anthropic's API accepts but Bedrock rejects with
-    ``"Extra inputs are not permitted"``.
+    Claude Code (v2.1.69+) is reported to send ``custom: {defer_loading: true}`` on
+    tool definitions, which Bedrock rejects with ``"Extra inputs are not permitted"``.
 
     Args:
         request_body: The request dictionary to modify in-place.
@@ -193,8 +194,14 @@ def remove_custom_field_from_tools(request_body: dict) -> None:
     if not tools or not isinstance(tools, list):
         return
     for tool in tools:
-        if isinstance(tool, dict):
-            tool.pop("custom", None)
+        if not isinstance(tool, dict):
+            continue
+        custom: dict[str, object] | None = tool.pop("custom", None)
+        if not isinstance(custom, dict) or "defer_loading" in tool:
+            continue
+        deferred: object = custom.get("defer_loading")
+        if isinstance(deferred, bool):
+            tool["defer_loading"] = deferred
 
 
 def normalize_json_schema_custom_types_to_object(schema: dict) -> None:
diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
index 8d039d95bb1..372cf110f7c 100644
--- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
+++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
@@ -33,9 +33,9 @@ from litellm.llms.bedrock.common_utils import (
     get_anthropic_beta_from_headers,
     is_claude_4_5_on_bedrock,
     normalize_bedrock_opus_output_config_effort,
+    normalize_custom_field_on_tools,
     normalize_tool_input_schema_types_for_bedrock_invoke,
     pop_bedrock_invoke_output_config_format,
-    remove_custom_field_from_tools,
 )
 from litellm.types.llms.anthropic import (
     ANTHROPIC_BETA_HEADER_VALUES,
@@ -749,11 +749,9 @@ class AmazonAnthropicClaudeMessagesConfig(
                     model,
                 )
 
-        # 5b. Remove `custom` field from tools (Bedrock doesn't support it)
-        # Claude Code sends `custom: {defer_loading: true}` on tool definitions,
-        # which causes Bedrock to reject the request with "Extra inputs are not permitted"
+        # 5b. Hoist `custom.defer_loading` then drop `custom` (Bedrock doesn't support it)
         # Ref: https://github.com/BerriAI/litellm/issues/22847
-        remove_custom_field_from_tools(anthropic_messages_request)
+        normalize_custom_field_on_tools(anthropic_messages_request)
         normalize_tool_input_schema_types_for_bedrock_invoke(anthropic_messages_request)
         ensure_bedrock_anthropic_messages_tool_names(anthropic_messages_request)
 
diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
index 02bd3535c0a..fd66667af64 100644
--- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
+++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py
@@ -16,8 +16,8 @@ sys.path.insert(0, os.path.abspath("../../../../../.."))
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.bedrock.common_utils import (
     ensure_bedrock_anthropic_messages_tool_names,
+    normalize_custom_field_on_tools,
     normalize_tool_input_schema_types_for_bedrock_invoke,
-    remove_custom_field_from_tools,
 )
 from litellm.constants import (
     BEDROCK_MIN_THINKING_BUDGET_TOKENS,
@@ -353,12 +353,13 @@ def test_remove_ttl_from_cache_control():
     assert request5 == {}
 
 
-def test_remove_custom_field_from_tools():
+def test_normalize_custom_field_on_tools():
     """
-    Ensure the `custom` field is stripped from every tool definition.
+    Ensure the `custom` field is stripped from every tool definition, and that a
+    boolean `custom.defer_loading` is hoisted onto the top-level `defer_loading`
+    flag Bedrock documents instead of being dropped with the wrapper.
 
-    Claude Code v2.1.69+ sends `custom: {defer_loading: true}` on tool
-    objects.  Bedrock does not accept this extra field and returns
+    Bedrock does not accept a `custom` object on a tool and returns
     "Extra inputs are not permitted".
 
     Ref: https://github.com/BerriAI/litellm/issues/22847
@@ -381,29 +382,94 @@ def test_remove_custom_field_from_tools():
         ]
     }
 
-    remove_custom_field_from_tools(request)
+    normalize_custom_field_on_tools(request)
 
     for tool in request["tools"]:
         assert "custom" not in tool, f"Tool {tool['name']} still has 'custom' field"
     # Other fields should be preserved
     assert request["tools"][0]["name"] == "Read"
     assert request["tools"][1]["name"] == "Write"
+    # `custom.defer_loading` is hoisted; the tool that never carried it is untouched
+    assert request["tools"][0]["defer_loading"] is True
+    assert "defer_loading" not in request["tools"][1]
 
     # Case 2: request without tools key (should not raise error)
     request2 = {"messages": [{"role": "user", "content": "hi"}]}
-    remove_custom_field_from_tools(request2)
+    normalize_custom_field_on_tools(request2)
     assert "tools" not in request2
 
     # Case 3: empty tools list (should not raise error)
     request3 = {"tools": []}
-    remove_custom_field_from_tools(request3)
+    normalize_custom_field_on_tools(request3)
     assert request3["tools"] == []
 
     # Case 4: tools with None value (should not raise error)
     request4 = {"tools": None}
-    remove_custom_field_from_tools(request4)
+    normalize_custom_field_on_tools(request4)
     assert request4["tools"] is None
 
+    # Case 5: an explicit top-level flag wins over a conflicting wrapped one
+    request5 = {
+        "tools": [
+            {"name": "Read", "defer_loading": False, "custom": {"defer_loading": True}}
+        ]
+    }
+    normalize_custom_field_on_tools(request5)
+    assert request5["tools"][0] == {"name": "Read", "defer_loading": False}
+
+    # Case 6: a non-boolean `custom.defer_loading` is dropped, never forwarded
+    for junk in ("true", 1, None, {"nested": True}):
+        request6 = {"tools": [{"name": "Read", "custom": {"defer_loading": junk}}]}
+        normalize_custom_field_on_tools(request6)
+        assert request6["tools"][0] == {"name": "Read"}, f"leaked defer_loading={junk!r}"
+
+    # Case 7: a `custom` that is not a dict is dropped without raising
+    request7 = {
+        "tools": [
+            {"name": "Read", "custom": "defer_loading"},
+            {"name": "Write", "custom": None},
+        ]
+    }
+    normalize_custom_field_on_tools(request7)
+    assert request7["tools"] == [{"name": "Read"}, {"name": "Write"}]
+
+
+@pytest.mark.parametrize(
+    "deferred_marker", [{"custom": {"defer_loading": True}}, {"defer_loading": True}]
+)
+def test_bedrock_invoke_messages_transform_emits_top_level_defer_loading(
+    deferred_marker,
+):
+    """A deferred tool must reach Bedrock as top-level ``defer_loading``, whether the
+    client wrapped the flag in ``custom`` or sent it top-level, and the outbound body
+    must still carry the Bedrock tool-search beta."""
+    from litellm.types.router import GenericLiteLLMParams
+
+    cfg = AmazonAnthropicClaudeMessagesConfig()
+    result = cfg.transform_anthropic_messages_request(
+        model="us.anthropic.claude-haiku-4-5-20251001-v1:0",
+        messages=[{"role": "user", "content": "hi"}],
+        anthropic_messages_optional_request_params={
+            "max_tokens": 128,
+            "stream": False,
+            "betas": ["advanced-tool-use-2025-11-20"],
+            "tools": [
+                {
+                    "name": "Read",
+                    "description": "Read a file",
+                    "input_schema": {"type": "object", "properties": {}},
+                    **deferred_marker,
+                },
+                {"type": "tool_search_tool_regex_20251119", "name": "tool_search"},
+            ],
+        },
+        litellm_params=GenericLiteLLMParams(),
+        headers={},
+    )
+    assert result["tools"][0]["defer_loading"] is True
+    assert "custom" not in result["tools"][0]
+    assert result["anthropic_beta"] == ["tool-search-tool-2025-10-19"]
+
 
 def test_normalize_tool_input_schema_types_for_bedrock_invoke():
     """

From efbdb6901a2cbace882705e13ea4e52ba7ef47e2 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 13 Aug 2026 19:47:02 -0700
Subject: [PATCH 25/31] fix(access groups): sync assigned_key_ids from the key
 write paths (#36843)

---
 .../key_management_endpoints.py               |  55 +-
 .../access_group_key_sync.py                  | 173 ++++
 .../test_key_management_endpoints.py          | 741 ++++++++++++++++++
 3 files changed, 968 insertions(+), 1 deletion(-)
 create mode 100644 litellm/proxy/management_helpers/access_group_key_sync.py

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
index 116dea464ff..7e190e8b19d 100644
--- a/litellm/proxy/management_endpoints/key_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -88,6 +88,11 @@ from litellm.proxy.management_endpoints.common_utils import (
 from litellm.proxy.management_endpoints.model_management_endpoints import (
     _add_model_to_db,
 )
+from litellm.proxy.management_helpers.access_group_key_sync import (
+    sync_key_access_group_membership,
+    sync_key_regeneration_access_group_membership,
+    sync_key_update_access_group_membership,
+)
 from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
 from litellm.proxy.management_helpers.object_permission_utils import (
     _set_object_permission,
@@ -2347,6 +2352,17 @@ async def _process_single_key_update(
         proxy_logging_obj=proxy_logging_obj,
     )
 
+    # After the key's own cache entry is dropped, so a failure here cannot leave the key
+    # authenticating against the access groups it just lost.
+    await sync_key_update_access_group_membership(
+        prisma_client=prisma_client,
+        key_token=_hash_token_if_needed(
+            _resolve_token_to_update(data=update_key_request, existing_key_row=existing_key_row)
+        ),
+        data=update_key_request,
+        existing_key_row=existing_key_row,
+    )
+
     # Trigger async hook
     asyncio.create_task(
         KeyManagementEventHooks.async_key_updated_hook(
@@ -2828,6 +2844,15 @@ async def update_key_fn(
             proxy_logging_obj=proxy_logging_obj,
         )
 
+        # After the key's own cache entry is dropped, so a failure here cannot leave the key
+        # authenticating against the access groups it just lost.
+        await sync_key_update_access_group_membership(
+            prisma_client=prisma_client,
+            key_token=_hash_token_if_needed(key),
+            data=data,
+            existing_key_row=existing_key_row,
+        )
+
         if data.spend is not None:
             from litellm.proxy.proxy_server import spend_counter_cache
 
@@ -3771,7 +3796,7 @@ async def generate_key_helper_fn(
     auto_rotate: bool | None = None,
     rotation_interval: str | None = None,
     router_settings: dict | None = None,
-    access_group_ids: list | None = None,
+    access_group_ids: list[str] | None = None,
     budget_limits: list | None = None,  # multiple concurrent budget windows
 ):
     from litellm.proxy.proxy_server import premium_user, prisma_client
@@ -3979,6 +4004,14 @@ async def generate_key_helper_fn(
             create_key_response: Final = await prisma_client.insert_data(data=key_data, table_name="key")
 
             key_data["token_id"] = getattr(create_key_response, "token", None)
+            created_token_hash: Final = getattr(create_key_response, "token", None)
+            if isinstance(created_token_hash, str):
+                await sync_key_access_group_membership(
+                    prisma_client=prisma_client,
+                    key_token=created_token_hash,
+                    previous_access_group_ids=None,
+                    updated_access_group_ids=access_group_ids,
+                )
             key_data["litellm_budget_table"] = getattr(create_key_response, "litellm_budget_table", None)
             key_data["created_at"] = getattr(create_key_response, "created_at", None)
             key_data["updated_at"] = getattr(create_key_response, "updated_at", None)
@@ -4196,6 +4229,7 @@ async def delete_verification_tokens(
                 deleted_tokens = [key.token for key in authorized_keys]
                 if len(deleted_tokens) != len(tokens):
                     failed_tokens = [token for token in tokens if token not in deleted_tokens]
+
         else:
             raise Exception("DB not connected. prisma_client is None")
     except Exception as e:
@@ -4211,6 +4245,16 @@ async def delete_verification_tokens(
         hashed_token = hash_token(cast(str, key))
         user_api_key_cache.delete_cache(hashed_token)
 
+    # After credential invalidation, so a failure here can never keep a deleted key alive.
+    for deleted_key in authorized_keys:
+        if deleted_key.token is not None:
+            await sync_key_access_group_membership(
+                prisma_client=prisma_client,
+                key_token=deleted_key.token,
+                previous_access_group_ids=deleted_key.access_group_ids,
+                updated_access_group_ids=None,
+            )
+
     return {
         "deleted_keys": deleted_tokens,
         "failed_tokens": failed_tokens,
@@ -4726,6 +4770,15 @@ async def _execute_virtual_key_regeneration(
             proxy_logging_obj=proxy_logging_obj,
         )
 
+    # After credential invalidation, so a failure here can never keep the old key alive.
+    await sync_key_regeneration_access_group_membership(
+        prisma_client=prisma_client,
+        previous_key_token=hashed_api_key,
+        new_key_token=new_token_hash,
+        data=data,
+        existing_key_row=key_in_db,
+    )
+
     response: Final = GenerateKeyResponse.model_validate(updated_token_dict)
     asyncio.create_task(
         KeyManagementEventHooks.async_key_rotated_hook(
diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py
new file mode 100644
index 00000000000..5d43cb29978
--- /dev/null
+++ b/litellm/proxy/management_helpers/access_group_key_sync.py
@@ -0,0 +1,173 @@
+"""
+Reverse sync for the key side of the key <-> access group relationship.
+
+`litellm_accessgrouptable.assigned_key_ids` and `litellm_verificationtoken.access_group_ids`
+are the two halves of one relationship and BOTH are read: the access group's
+attached-keys view reads the former, and so does the grant check in
+`auth_checks.get_authorized_resources_from_key_access_groups`, which authorizes a
+key only when the group lists the key's token (or the key's team). The access-group
+endpoints maintain both halves already; this module is what the key write paths call
+so an edit from that side is mirrored back.
+
+Every write is a single guarded statement rather than a read-modify-write. Prisma has no
+atomic scalar-list removal (see `TeamRepository.remove_member`), and the read-modify-write
+it otherwise forces is not safe here: a lost update would put an already revoked token back
+into a group and restore its grants, or drop a grant an admin just made. The guards also
+make each statement idempotent, so a retry cannot duplicate an entry. Each statement covers
+every group the request touches at once, so the size of the caller's id list does not turn
+into a matching number of round trips, and returns the ids it actually moved so only those
+groups are dropped from cache.
+
+It deliberately lives outside `access_group_endpoints`, which is a lazily
+registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
+module eagerly from `key_management_endpoints` would put it in `sys.modules`
+without its router ever being included, which drops its routes from the OpenAPI
+schema.
+"""
+
+from collections.abc import Sequence
+from typing import Final, Protocol
+
+from pydantic import BaseModel
+
+from litellm.proxy._types import (
+    LiteLLM_VerificationToken,
+    RegenerateKeyRequest,
+    UpdateKeyRequest,
+)
+from litellm.proxy.auth.auth_checks import (
+    _delete_cache_access_object,  # pyright: ignore[reportPrivateUsage]  # the access-group endpoints reach for this same cache primitive
+)
+from litellm.repositories.table_repositories import AccessGroupRepository
+
+
+class _MovedGroupRow(BaseModel):
+    access_group_id: str
+
+
+class _RawExecutor(Protocol):
+    async def query_raw(self, query: str, *args: str | Sequence[str]) -> Sequence[object]: ...
+
+
+_ATTACH_KEY_SQL: Final = (
+    'UPDATE "LiteLLM_AccessGroupTable" '
+    'SET "assigned_key_ids" = array_append("assigned_key_ids", $1) '
+    'WHERE "access_group_id" = ANY($2::text[]) AND NOT ($1 = ANY("assigned_key_ids")) '
+    'RETURNING "access_group_id"'
+)
+
+_DETACH_KEY_SQL: Final = (
+    'UPDATE "LiteLLM_AccessGroupTable" '
+    'SET "assigned_key_ids" = array_remove("assigned_key_ids", $1) '
+    'WHERE "access_group_id" = ANY($2::text[]) AND $1 = ANY("assigned_key_ids") '
+    'RETURNING "access_group_id"'
+)
+
+_REPOINT_KEY_SQL: Final = (
+    'UPDATE "LiteLLM_AccessGroupTable" '
+    'SET "assigned_key_ids" = array_append(array_remove(array_remove("assigned_key_ids", $1), $2), $2) '
+    'WHERE $1 = ANY("assigned_key_ids") '
+    'RETURNING "access_group_id"'
+)
+
+
+def _raw_executor(prisma_client: object) -> _RawExecutor:
+    """Narrow the untyped Prisma client down to the raw-query call this module makes."""
+    return AccessGroupRepository(prisma_client).prisma_client.db  # pyright: ignore[reportAny]  # untyped Prisma client
+
+
+async def _invalidate_access_group_cache(access_group_id: str) -> None:
+    """
+    Drop an access group entry from both the in-memory and Redis caches.
+
+    Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
+    to avoid circular imports, following the same pattern as key_management_endpoints.
+    """
+    from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
+
+    await _delete_cache_access_object(
+        access_group_id=access_group_id,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+
+
+async def _invalidate_moved_groups(moved_rows: Sequence[object]) -> None:
+    for row in moved_rows:
+        await _invalidate_access_group_cache(_MovedGroupRow.model_validate(row).access_group_id)
+
+
+async def _write_membership(prisma_client: object, sql: str, access_group_ids: frozenset[str], key_token: str) -> None:
+    """Run one guarded membership statement for every listed group, dropping the cache of those it moved."""
+    if not access_group_ids:
+        return
+    await _invalidate_moved_groups(
+        await _raw_executor(prisma_client).query_raw(sql, key_token, sorted(access_group_ids))
+    )
+
+
+async def sync_key_access_group_membership(
+    prisma_client: object,
+    key_token: str,
+    previous_access_group_ids: Sequence[str] | None,
+    updated_access_group_ids: Sequence[str] | None,
+) -> None:
+    """Mirror a key-side change to `access_group_ids` onto each access group's `assigned_key_ids`."""
+    previous: Final = frozenset(previous_access_group_ids or ())
+    updated: Final = frozenset(updated_access_group_ids or ())
+
+    await _write_membership(prisma_client, _ATTACH_KEY_SQL, updated - previous, key_token)
+    await _write_membership(prisma_client, _DETACH_KEY_SQL, previous - updated, key_token)
+
+
+async def sync_key_update_access_group_membership(
+    prisma_client: object,
+    key_token: str,
+    data: UpdateKeyRequest | RegenerateKeyRequest,
+    existing_key_row: LiteLLM_VerificationToken,
+) -> None:
+    """
+    Mirror a key UPDATE onto the group side, honouring `exclude_unset` semantics.
+
+    The key row is written from `model_dump(exclude_unset=True)`, so a request that never
+    mentions `access_group_ids` leaves the key's own list alone and must leave the group's
+    copy alone too. Reading the attribute instead of `model_fields_set` would see None on
+    every unrelated edit and withdraw the token from every group it belongs to.
+    """
+    if "access_group_ids" not in data.model_fields_set:
+        return
+    await sync_key_access_group_membership(
+        prisma_client=prisma_client,
+        key_token=key_token,
+        previous_access_group_ids=existing_key_row.access_group_ids,
+        updated_access_group_ids=data.access_group_ids,
+    )
+
+
+async def sync_key_regeneration_access_group_membership(
+    prisma_client: object,
+    previous_key_token: str,
+    new_key_token: str,
+    data: RegenerateKeyRequest | None,
+    existing_key_row: LiteLLM_VerificationToken,
+) -> None:
+    """
+    Re-point every group's copy from the old token to the regenerated one.
+
+    Regeneration replaces the token, which is the identity `assigned_key_ids` stores, so
+    leaving the old hash behind both points the group at a row that no longer exists and
+    denies the regenerated key the group's grants. The swap is driven by the groups that
+    hold the old token when the statement runs, not by the key row read earlier, so a group
+    edited in between is neither resurrected nor skipped. Removing the new token before
+    appending it keeps a re-run from duplicating it.
+    """
+    await _invalidate_moved_groups(
+        await _raw_executor(prisma_client).query_raw(_REPOINT_KEY_SQL, previous_key_token, new_key_token)
+    )
+    if data is not None:
+        await sync_key_update_access_group_membership(
+            prisma_client=prisma_client,
+            key_token=new_key_token,
+            data=data,
+            existing_key_row=existing_key_row,
+        )
diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
index 0ebdc07d282..bdf09a95e4b 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
@@ -791,6 +791,7 @@ async def test_generate_key_helper_fn_with_access_group_ids(monkeypatch):
     mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock(
         return_value=MagicMock(object_permission_id=None)
     )
+    mock_prisma_client.db.query_raw = AsyncMock(return_value=[])
 
     captured_key_data = {}
 
@@ -15703,3 +15704,743 @@ async def test_key_generate_omitted_budget_duration_still_filled_by_upperbound(m
 
     assert key_row["budget_duration"] == "30d"
     assert key_row["budget_reset_at"] is not None
+from litellm.proxy.management_helpers.access_group_key_sync import (
+    _ATTACH_KEY_SQL,
+    _DETACH_KEY_SQL,
+    _REPOINT_KEY_SQL,
+)
+
+ACCESS_GROUP_SYNC_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
+
+
+def _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups):
+    """
+    Back the access group table with an in-memory dict so the sync's writes are observable.
+
+    The sync writes through guarded set-based SQL statements, so this emulates exactly what
+    Postgres does with them, including the guards that make each one idempotent and the
+    `RETURNING` clause that reports which groups actually moved.
+    """
+
+    def _repoint(previous_token, new_token):
+        moved = [
+            group_id
+            for group_id, stored in access_groups.items()
+            if previous_token in stored["assigned_key_ids"]
+        ]
+        for group_id in moved:
+            current = access_groups[group_id]["assigned_key_ids"]
+            access_groups[group_id]["assigned_key_ids"] = [
+                *(t for t in current if t not in (previous_token, new_token)),
+                new_token,
+            ]
+        return moved
+
+    def _attach(key_token, access_group_ids):
+        moved = [
+            group_id
+            for group_id in access_group_ids
+            if group_id in access_groups
+            and key_token not in access_groups[group_id]["assigned_key_ids"]
+        ]
+        for group_id in moved:
+            stored = access_groups[group_id]
+            stored["assigned_key_ids"] = [*stored["assigned_key_ids"], key_token]
+        return moved
+
+    def _detach(key_token, access_group_ids):
+        moved = [
+            group_id
+            for group_id in access_group_ids
+            if group_id in access_groups
+            and key_token in access_groups[group_id]["assigned_key_ids"]
+        ]
+        for group_id in moved:
+            stored = access_groups[group_id]
+            stored["assigned_key_ids"] = [
+                t for t in stored["assigned_key_ids"] if t != key_token
+            ]
+        return moved
+
+    async def _query_raw(query, *args):
+        if query == _REPOINT_KEY_SQL:
+            moved = _repoint(*args)
+        elif query == _ATTACH_KEY_SQL:
+            moved = _attach(*args)
+        else:
+            assert query == _DETACH_KEY_SQL, f"unexpected statement: {query}"
+            moved = _detach(*args)
+        return [{"access_group_id": group_id} for group_id in moved]
+
+    raw_mock = AsyncMock(side_effect=_query_raw)
+    mock_prisma_client.db.query_raw = raw_mock
+    return raw_mock
+
+
+async def _authorized_models_for_key(access_groups, token, key_access_group_ids):
+    """Run the real auth-time reader against the post-sync access group rows."""
+    from litellm.proxy._types import LiteLLM_AccessGroupTable, LiteLLM_TeamTable
+    from litellm.proxy.auth.auth_checks import (
+        get_authorized_resources_from_key_access_groups,
+    )
+
+    async def _get_access_object(*, access_group_id, **_kwargs):
+        stored = access_groups[access_group_id]
+        return LiteLLM_AccessGroupTable(
+            access_group_id=access_group_id,
+            access_group_name=access_group_id,
+            access_model_names=list(stored["access_model_names"]),
+            assigned_team_ids=[],
+            assigned_key_ids=list(stored["assigned_key_ids"]),
+        )
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
+        patch(
+            "litellm.proxy.auth.auth_checks.get_access_object",
+            new_callable=AsyncMock,
+            side_effect=_get_access_object,
+        ),
+    ):
+        return await get_authorized_resources_from_key_access_groups(
+            valid_token=UserAPIKeyAuth(
+                token=token,
+                models=[],
+                team_id="team-a",
+                access_group_ids=list(key_access_group_ids),
+            ),
+            team_object=LiteLLM_TeamTable(team_id="team-a", models=[]),
+            resource_field="access_model_names",
+        )
+
+
+@pytest.mark.asyncio
+async def test_update_key_syncs_access_group_assigned_key_ids_in_both_directions(
+    monkeypatch,
+):
+    """
+    A key-side edit of `access_group_ids` must be mirrored onto every affected access
+    group's `assigned_key_ids`, in one operation, in both directions.
+
+    `assigned_key_ids` is not display-only. `get_authorized_resources_from_key_access_groups`
+    reads it as an authorization input and authorizes only when the group lists the key's
+    token, so a group the key just added must start granting its resources and a group the
+    key dropped must stop. A single-direction assertion would pass against a fix that only
+    ever adds (or only ever removes), so this covers add, remove, untouched, and the
+    authorization consequence of each.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=["ag-drop", "ag-keep"],
+    )
+    access_groups = {
+        "ag-drop": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+            "access_model_names": ["dropped-model"],
+        },
+        "ag-keep": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+            "access_model_names": ["kept-model"],
+        },
+        "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=key_in_db
+    )
+    mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
+        return_value=None
+    )
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
+    raw_mock = _access_group_table_mocks(
+        monkeypatch, mock_prisma_client, access_groups
+    )
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ) as invalidate_cache,
+    ):
+        await update_key_fn(
+            request=MagicMock(),
+            data=UpdateKeyRequest(
+                key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-keep", "ag-add"]
+            ),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by=None,
+        )
+
+    assert access_groups["ag-drop"]["assigned_key_ids"] == []
+    assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+    assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+
+    # Both halves go out as single guarded statements. A read-modify-write here lets two
+    # admins editing one group lose each other's change: an attach can vanish, and a detach
+    # can put an already revoked token back and restore its grants.
+    assert sorted(call.args for call in raw_mock.call_args_list) == sorted(
+        [
+            (_ATTACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"]),
+            (_DETACH_KEY_SQL, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop"]),
+        ]
+    )
+    assert {call.args[0] for call in invalidate_cache.call_args_list} == {
+        "ag-drop",
+        "ag-add",
+    }
+
+    authorized_models = await _authorized_models_for_key(
+        access_groups,
+        ACCESS_GROUP_SYNC_TOKEN,
+        ["ag-drop", "ag-keep", "ag-add"],
+    )
+    assert sorted(authorized_models) == ["added-model", "kept-model"]
+
+
+@pytest.mark.asyncio
+async def test_update_key_leaves_access_groups_alone_when_field_is_unset(monkeypatch):
+    """
+    An update that never mentions `access_group_ids` must not touch the group rows.
+
+    `prepare_key_update_data` writes from `model_dump(exclude_unset=True)`, so an omitted
+    field leaves the key row's own list intact. Reading the request attribute instead of
+    its `model_fields_set` would see None and wipe every group's copy of the token on any
+    unrelated edit, e.g. a max_budget change.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=["ag-keep"],
+    )
+    access_groups = {
+        "ag-keep": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+            "access_model_names": ["kept-model"],
+        },
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=key_in_db
+    )
+    mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
+        return_value=None
+    )
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
+    raw_mock = _access_group_table_mocks(
+        monkeypatch, mock_prisma_client, access_groups
+    )
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with patch(
+        "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+        new_callable=AsyncMock,
+    ):
+        await update_key_fn(
+            request=MagicMock(),
+            data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, max_budget=50.0),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by=None,
+        )
+
+    raw_mock.assert_not_called()
+    assert access_groups["ag-keep"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+    assert await _authorized_models_for_key(
+        access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-keep"]
+    ) == ["kept-model"]
+
+
+@pytest.mark.asyncio
+async def test_bulk_update_keys_syncs_access_group_assigned_key_ids(monkeypatch):
+    """
+    /key/bulk_update and /team/keys/bulk_update reach the DB through
+    `_process_single_key_update`, which is a separate write path from /key/update's own
+    inline one. Both have to maintain the group's copy or a bulk attach grants nothing.
+    """
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=["ag-drop"],
+    )
+    access_groups = {
+        "ag-drop": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+            "access_model_names": ["dropped-model"],
+        },
+        "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.TeamMemberPermissionChecks.can_team_member_execute_key_management_endpoint",
+            new_callable=AsyncMock,
+        ),
+    ):
+        await _process_single_key_update(
+            update_key_request=UpdateKeyRequest(
+                key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=["ag-add"]
+            ),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by=None,
+            prisma_client=mock_prisma_client,
+            user_api_key_cache=AsyncMock(),
+            proxy_logging_obj=MagicMock(),
+            llm_router=None,
+            existing_key_row=key_in_db,
+        )
+
+    assert access_groups["ag-drop"]["assigned_key_ids"] == []
+    assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+    assert await _authorized_models_for_key(
+        access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-drop", "ag-add"]
+    ) == ["added-model"]
+
+
+@pytest.mark.asyncio
+async def test_delete_key_withdraws_token_from_its_access_groups(monkeypatch):
+    """
+    Deleting a key must withdraw its token from every group that lists it.
+
+    Without the withdrawal the group keeps a token that no longer resolves to a row, so
+    the access group page lists a key that does not exist and the list grows without bound.
+    """
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=["ag-keep"],
+    )
+    access_groups = {
+        "ag-keep": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN, "other-key"],
+            "access_model_names": ["kept-model"],
+        },
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
+        return_value=[key_in_db]
+    )
+    mock_prisma_client.delete_data = AsyncMock(return_value={"deleted_keys": 1})
+    mock_prisma_client.db.litellm_deletedverificationtoken.create_many = AsyncMock()
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client", mock_prisma_client
+    )
+
+    mock_cache = MagicMock()
+    mock_cache.delete_cache = MagicMock()
+
+    with patch(
+        "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+        new_callable=AsyncMock,
+    ):
+        await delete_verification_tokens(
+            tokens=[ACCESS_GROUP_SYNC_TOKEN],
+            user_api_key_cache=mock_cache,
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by="admin-user",
+        )
+
+    assert access_groups["ag-keep"]["assigned_key_ids"] == ["other-key"]
+
+
+@pytest.mark.asyncio
+async def test_generate_key_records_token_in_its_access_groups(monkeypatch):
+    """
+    /key/generate with `access_group_ids` must record the new token on the group side.
+
+    The key row's own list alone does not authorize: the group has to list the token back
+    or `get_authorized_resources_from_key_access_groups` contributes nothing, so a key
+    created against a group silently gets none of its models.
+    """
+    access_groups = {
+        "ag-add": {"assigned_key_ids": [], "access_model_names": ["added-model"]},
+    }
+
+    created_key = MagicMock()
+    created_key.token = ACCESS_GROUP_SYNC_TOKEN
+    created_key.litellm_budget_table = None
+    created_key.created_at = None
+    created_key.updated_at = None
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.insert_data = AsyncMock(return_value=created_key)
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+    monkeypatch.setattr(
+        "litellm.proxy.proxy_server.prisma_client", mock_prisma_client
+    )
+    monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
+    monkeypatch.setattr("litellm.store_audit_logs", False)
+
+    with patch(
+        "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+        new_callable=AsyncMock,
+    ):
+        await generate_key_helper_fn(
+            request_type="key",
+            access_group_ids=["ag-add"],
+            table_name="key",
+            user_id="test-user",
+        )
+
+    assert access_groups["ag-add"]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+    assert await _authorized_models_for_key(
+        access_groups, ACCESS_GROUP_SYNC_TOKEN, ["ag-add"]
+    ) == ["added-model"]
+
+
+@pytest.mark.asyncio
+async def test_regenerate_key_repoints_access_group_assigned_key_ids(monkeypatch):
+    """
+    Regeneration replaces the key's token, which is the identity `assigned_key_ids` stores.
+
+    Leaving the old hash behind points the group at a token that no longer exists AND
+    denies the regenerated key the group's grants, so the group's copy has to be
+    re-pointed from the old hash to the new one in the same operation.
+    """
+    from litellm.proxy._types import RegenerateKeyRequest
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _execute_virtual_key_regeneration,
+    )
+
+    from litellm.proxy.utils import hash_token
+
+    new_token_hash = hash_token("sk-newtoken1234ab12")
+    existing_key = LiteLLM_VerificationToken(
+        token="abc123",
+        user_id="user-1",
+        models=["gpt-4"],
+        access_group_ids=["ag-keep"],
+    )
+    access_groups = {
+        "ag-keep": {
+            "assigned_key_ids": ["abc123"],
+            "access_model_names": ["kept-model"],
+        },
+    }
+
+    mock_prisma_client = _make_regenerate_mock_prisma()
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
+            new_callable=AsyncMock,
+            return_value="sk-newtoken1234ab12",
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ),
+    ):
+        await _execute_virtual_key_regeneration(
+            prisma_client=mock_prisma_client,
+            key_in_db=existing_key,
+            hashed_api_key="abc123",
+            key="abc123",
+            data=RegenerateKeyRequest(),
+            user_api_key_dict=_make_regenerate_user_api_key_dict(),
+            litellm_changed_by=None,
+            user_api_key_cache=MagicMock(),
+            proxy_logging_obj=MagicMock(),
+        )
+
+    assert access_groups["ag-keep"]["assigned_key_ids"] == [new_token_hash]
+    assert await _authorized_models_for_key(
+        access_groups, new_token_hash, ["ag-keep"]
+    ) == ["kept-model"]
+    assert (
+        await _authorized_models_for_key(access_groups, "abc123", ["ag-keep"]) == []
+    )
+
+
+@pytest.mark.asyncio
+async def test_key_write_paths_revoke_the_key_cache_before_syncing_access_groups(
+    monkeypatch,
+):
+    """
+    Credential invalidation must not sit behind the group sync on any key write path.
+
+    The cached auth object still carries the key's old `access_group_ids`, so if the sync
+    raises first, the request fails with the key still authenticating against groups it
+    just lost, until that entry expires. Ordering it last means a failed sync degrades to
+    the stale listing this PR fixes rather than to a stale grant.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    order = []
+
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=["ag-drop"],
+    )
+    access_groups = {
+        "ag-drop": {
+            "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+            "access_model_names": ["dropped-model"],
+        },
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=key_in_db
+    )
+    mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
+        return_value=None
+    )
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+    mock_prisma_client.db.query_raw = AsyncMock(
+        side_effect=lambda *a, **k: order.append("sync") or []
+    )
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+            side_effect=lambda **kwargs: order.append("revoke_key_cache"),
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ),
+    ):
+        await update_key_fn(
+            request=MagicMock(),
+            data=UpdateKeyRequest(key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=[]),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by=None,
+        )
+
+    assert order == ["revoke_key_cache", "sync"]
+
+
+@pytest.mark.asyncio
+async def test_update_key_syncs_many_access_groups_in_one_statement_per_direction(
+    monkeypatch,
+):
+    """
+    The number of groups on a request must not become a matching number of round trips.
+
+    Anyone allowed to assign access groups picks the size of `access_group_ids`, so a
+    per-group statement lets one /key/update hold a connection for hundreds of sequential
+    writes. Both halves are set-based, so the cost is two statements no matter the size.
+    """
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        update_key_fn,
+    )
+
+    dropped = [f"ag-drop-{i}" for i in range(60)]
+    added = [f"ag-add-{i}" for i in range(60)]
+    key_in_db = LiteLLM_VerificationToken(
+        token=ACCESS_GROUP_SYNC_TOKEN,
+        user_id="test-user",
+        access_group_ids=dropped,
+    )
+    access_groups = {
+        **{
+            group_id: {
+                "assigned_key_ids": [ACCESS_GROUP_SYNC_TOKEN],
+                "access_model_names": [f"{group_id}-model"],
+            }
+            for group_id in dropped
+        },
+        **{
+            group_id: {"assigned_key_ids": [], "access_model_names": [f"{group_id}-model"]}
+            for group_id in added
+        },
+    }
+
+    mock_prisma_client = AsyncMock()
+    mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
+        return_value=key_in_db
+    )
+    mock_prisma_client.db.litellm_verificationtoken.find_first = AsyncMock(
+        return_value=None
+    )
+    mock_prisma_client.update_data = AsyncMock(return_value={"data": {}})
+    raw_mock = _access_group_table_mocks(
+        monkeypatch, mock_prisma_client, access_groups
+    )
+    _setup_update_key_mocks(monkeypatch, mock_prisma_client)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ),
+    ):
+        await update_key_fn(
+            request=MagicMock(),
+            data=UpdateKeyRequest(
+                key=ACCESS_GROUP_SYNC_TOKEN, access_group_ids=added
+            ),
+            user_api_key_dict=UserAPIKeyAuth(
+                user_role=LitellmUserRoles.PROXY_ADMIN,
+                api_key="sk-admin",
+                user_id="admin-user",
+            ),
+            litellm_changed_by=None,
+        )
+
+    assert [call.args[0] for call in raw_mock.call_args_list] == [
+        _ATTACH_KEY_SQL,
+        _DETACH_KEY_SQL,
+    ]
+    assert sorted(raw_mock.call_args_list[0].args[2]) == sorted(added)
+    assert sorted(raw_mock.call_args_list[1].args[2]) == sorted(dropped)
+    assert all(
+        access_groups[group_id]["assigned_key_ids"] == [ACCESS_GROUP_SYNC_TOKEN]
+        for group_id in added
+    )
+    assert all(access_groups[group_id]["assigned_key_ids"] == [] for group_id in dropped)
+
+
+@pytest.mark.asyncio
+async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read(
+    monkeypatch,
+):
+    """
+    Regeneration must move whatever the groups hold when it writes, not the key row's list.
+
+    That list is read before the new token exists, so replaying it re-adds the key to a
+    group an admin revoked in between and leaves the dead hash in a group an admin attached
+    in between, which silently restores one grant and drops another.
+    """
+    from litellm.proxy._types import RegenerateKeyRequest
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        _execute_virtual_key_regeneration,
+    )
+    from litellm.proxy.utils import hash_token
+
+    new_token_hash = hash_token("sk-newtoken1234ab12")
+    existing_key = LiteLLM_VerificationToken(
+        token="abc123",
+        user_id="user-1",
+        models=["gpt-4"],
+        access_group_ids=["ag-revoked-since"],
+    )
+    access_groups = {
+        "ag-revoked-since": {
+            "assigned_key_ids": [],
+            "access_model_names": ["revoked-model"],
+        },
+        "ag-attached-since": {
+            "assigned_key_ids": ["abc123"],
+            "access_model_names": ["attached-model"],
+        },
+    }
+
+    mock_prisma_client = _make_regenerate_mock_prisma()
+    _access_group_table_mocks(monkeypatch, mock_prisma_client, access_groups)
+
+    with (
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.get_new_token",
+            new_callable=AsyncMock,
+            return_value="sk-newtoken1234ab12",
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._insert_deprecated_key",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_endpoints.key_management_endpoints.KeyManagementEventHooks.async_key_rotated_hook",
+            new_callable=AsyncMock,
+        ),
+        patch(
+            "litellm.proxy.management_helpers.access_group_key_sync._invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ),
+    ):
+        await _execute_virtual_key_regeneration(
+            prisma_client=mock_prisma_client,
+            key_in_db=existing_key,
+            hashed_api_key="abc123",
+            key="abc123",
+            data=RegenerateKeyRequest(),
+            user_api_key_dict=_make_regenerate_user_api_key_dict(),
+            litellm_changed_by=None,
+            user_api_key_cache=MagicMock(),
+            proxy_logging_obj=MagicMock(),
+        )
+
+    assert access_groups["ag-revoked-since"]["assigned_key_ids"] == []
+    assert access_groups["ag-attached-since"]["assigned_key_ids"] == [new_token_hash]
+    assert await _authorized_models_for_key(
+        access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"]
+    ) == ["attached-model"]

From ea4316b2452748cb6254032f6a5c9274bc1a6c96 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 13 Aug 2026 19:58:06 -0700
Subject: [PATCH 26/31] fix(ui): resolve provider dropdown enum keys in
 getPlaceholder

---
 .../src/components/provider_info_helpers.test.tsx           | 5 +++++
 .../src/components/provider_info_helpers.tsx                | 6 ++++--
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
index b9516945faa..b163a6b341c 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx
@@ -237,6 +237,11 @@ describe("provider_info_helpers", () => {
       expect(getPlaceholder(Providers.NVIDIA_RIVA)).toBe("nvidia_riva/nvidia/parakeet-ctc-1_1b-asr");
     });
 
+    it("should resolve enum keys from the provider dropdown, not just enum values", () => {
+      expect(getPlaceholder("NVIDIA_RIVA")).toBe("nvidia_riva/nvidia/parakeet-ctc-1_1b-asr");
+      expect(getPlaceholder("WATSONX")).toBe("watsonx/ibm/granite-3-3-8b-instruct");
+    });
+
     it("should return default gpt-3.5-turbo placeholder for unknown provider", () => {
       expect(getPlaceholder("UnknownProvider" as any)).toBe("gpt-3.5-turbo");
     });
diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
index 86d67867fcf..ae311070c07 100644
--- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
+++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx
@@ -426,8 +426,10 @@ const providerPlaceholderMap: Partial<Record<Providers, string>> = {
   [Providers.ZAI]: "zai/glm-4.5",
 };
 
-export const getPlaceholder = (selectedProvider: string): string =>
-  providerPlaceholderMap[selectedProvider as Providers] ?? "gpt-3.5-turbo";
+export const getPlaceholder = (selectedProvider: string): string => {
+  const resolvedProvider = Providers[selectedProvider as keyof typeof Providers] ?? (selectedProvider as Providers);
+  return providerPlaceholderMap[resolvedProvider] ?? "gpt-3.5-turbo";
+};
 
 export const getProviderModels = (provider: Providers, modelMap: any): Array<string> => {
   let providerKey = provider;

From 59eeae374ca6e8c2b9533506a4da84f62d7e860f Mon Sep 17 00:00:00 2001
From: "devin-ai-integration[bot]"
 <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 13 Aug 2026 20:07:16 -0700
Subject: [PATCH 27/31] fix(mcp): expose client HTTP headers to logging
 callbacks and hooks (#36724)

* fix(mcp): expose client HTTP headers to logging callbacks and hooks

MCP protocol tool calls built a synthetic Request with only content-type, so metadata.headers reaching logging callbacks and guardrails was empty while /mcp-rest/tools/call exposed the full set. Rebuild the synthetic request from the connection's raw headers (shared with the sampling path), and pass sanitized headers to the pre-call hook, the MCP to LLM guardrail bridge and the Responses API MCP bridge. Credential headers stay masked and proxy key headers stripped.

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

* fix(mcp): strip custom proxy key and upstream MCP credential headers from logging copies

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

* refactor(mcp): make client side auth header name accessor public

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

* fix(mcp): strip custom proxy key and client redaction opt-out from mcp headers

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

* fix(mcp): drop custom proxy key header in the synthetic request builder

Strips general_settings.litellm_key_header_name in build_synthetic_mcp_request so every caller, including sampling, is covered, and reverts passing general_settings into add_litellm_data_to_request on the tool call path since that also switches on enforced_params.

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

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Co-authored-by: shivam <shivam@berri.ai>
---
 .../mcp_server/auth/user_api_key_auth_mcp.py  |   4 +-
 .../mcp_server/mcp_server_manager.py          |   2 +
 .../mcp_server/sampling_handler.py            |  99 +-----------
 .../proxy/_experimental/mcp_server/server.py  |  36 ++---
 .../proxy/_experimental/mcp_server/utils.py   | 147 ++++++++++++++++++
 litellm/proxy/litellm_pre_call_utils.py       |   4 +-
 litellm/proxy/utils.py                        |   1 +
 .../mcp/litellm_proxy_mcp_handler.py          |   5 +-
 .../auth/test_user_api_key_auth_mcp.py        |   6 +-
 .../mcp_server/test_mcp_hook_extra_headers.py |  42 +++++
 .../mcp_server/test_mcp_server.py             | 107 ++++++++++++-
 .../_experimental/mcp_server/test_utils.py    | 100 ++++++++++++
 .../utils/proxy_logging/test_mcp_bridging.py  |  17 ++
 .../mcp/test_litellm_proxy_mcp_handler.py     |  31 ++++
 type-discipline-budget.json                   |   6 +-
 15 files changed, 483 insertions(+), 124 deletions(-)

diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
index 554b6ea952e..95a3806e8ad 100644
--- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
+++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
@@ -1190,7 +1190,7 @@ class MCPRequestHandler:
 
         DEPRECATED: This method is deprecated in favor of server-specific auth headers using the format x-mcp-{{server_alias}}-{{header_name}} instead.
         """
-        mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler._get_mcp_client_side_auth_header_name()
+        mcp_client_side_auth_header_name: Final[str] = MCPRequestHandler.get_mcp_client_side_auth_header_name()
         auth_header: Final = headers.get(mcp_client_side_auth_header_name)
         if auth_header:
             verbose_logger.warning(
@@ -1265,7 +1265,7 @@ class MCPRequestHandler:
         return oauth2_headers
 
     @staticmethod
-    def _get_mcp_client_side_auth_header_name() -> str:
+    def get_mcp_client_side_auth_header_name() -> str:
         """
         Get the header name used to pass the MCP auth header to the MCP server
 
diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
index c8ff6e262d2..a1adda2bc95 100644
--- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
+++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py
@@ -118,6 +118,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
     is_short_mcp_tool_prefix_enabled,
     iter_known_server_prefixes,
     iter_known_tool_name_spellings,
+    logging_safe_mcp_headers,
     match_known_server_prefix,
     match_known_tool_name,
     merge_mcp_headers,
@@ -4603,6 +4604,7 @@ class MCPServerManager:
             ),
             "user_api_key_hash": (getattr(user_api_key_auth, "api_key_hash", None) if user_api_key_auth else None),
             "incoming_bearer_token": incoming_bearer_token,
+            "headers": logging_safe_mcp_headers(raw_headers),
         }
 
         # Create MCP request object for processing
diff --git a/litellm/proxy/_experimental/mcp_server/sampling_handler.py b/litellm/proxy/_experimental/mcp_server/sampling_handler.py
index a9a3367cd93..125dc3d773d 100644
--- a/litellm/proxy/_experimental/mcp_server/sampling_handler.py
+++ b/litellm/proxy/_experimental/mcp_server/sampling_handler.py
@@ -1042,100 +1042,15 @@ def _build_sampling_request(
     raw_headers: dict[str, str] | None = None,
     client_ip: str | None = None,
 ) -> "Request":
-    """Build a synthetic FastAPI Request for sampling sub-calls.
+    """The synthetic FastAPI Request for sampling sub-calls, carrying the original
+    MCP connection's headers and client IP."""
+    from litellm.proxy._experimental.mcp_server.utils import build_synthetic_mcp_request
 
-    Converts the original MCP connection's HTTP headers into ASGI
-    scope format so that ``add_litellm_data_to_request`` can apply
-    header-dependent guardrails, tag-based routing, trace correlation,
-    and ``forward_llm_provider_auth_headers``.
-
-    Key fields populated:
-    - **headers**: All original HTTP headers are forwarded (except
-      hop-by-hop: content-length, transfer-encoding).  This ensures
-      ``traceparent``, ``authorization``, ``user-agent``, and
-      ``x-litellm-api-key`` are visible to pre-call utils.
-    - **client**: The ASGI ``(host, port)`` tuple so that
-      ``request.client.host`` returns the real client IP for
-      IP-based routing and guardrails.
-    - **server**: Derived from the running proxy's ``server_host``
-      / ``server_port`` when available, avoiding the misleading
-      ``127.0.0.1:0`` placeholder.
-    - **x-forwarded-for**: Injected from ``client_ip`` if the
-      original headers don't already carry it, as a fallback for
-      IP attribution.
-    """
-    from fastapi import Request
-
-    # --- Build ASGI headers ---
-    _scope_headers: Final[list[tuple[bytes, bytes]]] = [(b"content-type", b"application/json")]
-    # Hop-by-hop headers that must NOT be forwarded into the
-    # synthetic request (they describe the original HTTP framing,
-    # not the logical request).
-    _HOP_BY_HOP: Final = frozenset(
-        {
-            "content-length",
-            "transfer-encoding",
-            "connection",
-            "keep-alive",
-            "upgrade",
-            "te",
-            "trailer",
-        }
+    return build_synthetic_mcp_request(
+        path="/mcp/sampling/createMessage",
+        raw_headers=raw_headers,
+        client_ip=client_ip,
     )
-    if raw_headers:
-        for hdr_name, hdr_value in raw_headers.items():
-            _key = hdr_name.lower()
-            # Skip content-type (already set), x-forwarded-for (use resolved
-            # client_ip instead to prevent spoofing), and hop-by-hop headers
-            if _key in {"content-type", "x-forwarded-for"} or _key in _HOP_BY_HOP:
-                continue
-            _scope_headers.append(
-                (
-                    _key.encode("latin-1", errors="replace"),
-                    hdr_value.encode("utf-8"),
-                )
-            )
-
-    # Inject x-forwarded-for from captured client_ip if the
-    # original headers don't already carry it
-    if client_ip and not any(h[0] == b"x-forwarded-for" for h in _scope_headers):
-        _scope_headers.append((b"x-forwarded-for", client_ip.encode("utf-8")))
-
-    # --- Derive server (host, port) from the running proxy ---
-    _server_host = "127.0.0.1"
-    _server_port = 4000  # LiteLLM default
-    try:
-        from litellm.proxy import proxy_server
-
-        _proxy_host: Final[str | None] = getattr(proxy_server, "server_host", None)
-        _proxy_port: Final[str | int | None] = getattr(proxy_server, "server_port", None)
-
-        if _proxy_host:
-            _server_host = str(_proxy_host)
-        if _proxy_port:
-            _server_port = int(_proxy_port)
-    except (ImportError, AttributeError, TypeError, ValueError):
-        pass
-
-    # --- Build ASGI client tuple for request.client.host ---
-    _client_tuple = None
-    if client_ip:
-        _client_tuple = (client_ip, 0)
-
-    scope: Final[dict[str, object]] = {
-        "type": "http",
-        "method": "POST",
-        "path": "/mcp/sampling/createMessage",
-        "scheme": "http",
-        "server": (_server_host, _server_port),
-        "query_string": b"",
-        "root_path": "",
-        "headers": _scope_headers,
-    }
-    if _client_tuple is not None:
-        scope["client"] = _client_tuple
-
-    return Request(scope=scope)
 
 
 async def _build_completion_kwargs(
diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py
index 49a1f1314f0..f237529b319 100644
--- a/litellm/proxy/_experimental/mcp_server/server.py
+++ b/litellm/proxy/_experimental/mcp_server/server.py
@@ -58,9 +58,11 @@ from litellm.proxy._experimental.mcp_server.utils import (
     LITELLM_MCP_SERVER_VERSION,
     MCPMissingUserEnvVarsError,
     add_server_prefix_to_name,
+    build_synthetic_mcp_request,
     extract_mcp_tool_result_error_message,
     get_server_prefix,
     iter_known_server_prefixes,
+    logging_safe_mcp_headers,
     match_known_tool_name,
 )
 from litellm.proxy._types import (
@@ -860,11 +862,11 @@ if MCP_AVAILABLE:
         name: str,
         arguments: dict[str, object],
         user_api_key_auth: UserAPIKeyAuth,
+        raw_headers: Mapping[str, str] | None = None,
+        client_ip: str | None = None,
     ) -> LiteLLMLoggingObj | None:
         """Run the pre-call pipeline (guardrails + logging setup) for a virtual
         mcp_tool_call so the SSE path spend-logs like the REST path."""
-        from fastapi import Request
-
         from litellm.proxy.common_request_processing import (
             ProxyBaseLLMRequestProcessing,
         )
@@ -874,13 +876,10 @@ if MCP_AVAILABLE:
             proxy_logging_obj,
         )
 
-        request: Final = Request(
-            scope={
-                "type": "http",
-                "method": "POST",
-                "path": "/mcp/tools/call",
-                "headers": [(b"content-type", b"application/json")],
-            }
+        request: Final = build_synthetic_mcp_request(
+            path="/mcp/tools/call",
+            raw_headers=raw_headers,
+            client_ip=client_ip,
         )
         _, virtual_logging_obj = await ProxyBaseLLMRequestProcessing(
             data={"name": name, "arguments": arguments}
@@ -952,7 +951,11 @@ if MCP_AVAILABLE:
 
         assert user_api_key_auth is not None  # guaranteed by the flag check above
         virtual_logging_obj: Final = await _build_virtual_call_logging_obj(
-            name=name, arguments=args, user_api_key_auth=user_api_key_auth
+            name=name,
+            arguments=args,
+            user_api_key_auth=user_api_key_auth,
+            raw_headers=raw_headers,
+            client_ip=client_ip,
         )
         return await handle_mcp_tool_call(
             tool_name=args.get("tool_name", ""),
@@ -979,7 +982,6 @@ if MCP_AVAILABLE:
         Raises:
             HTTPException: If tool not found or arguments missing
         """
-        from fastapi import Request
         from mcp.server.lowlevel.server import request_ctx
         from mcp.types import CallToolResult
 
@@ -1041,13 +1043,10 @@ if MCP_AVAILABLE:
                     body_data["litellm_trace_id"] = chain_id
                     body_data["litellm_session_id"] = chain_id
 
-                request: Final = Request(
-                    scope={
-                        "type": "http",
-                        "method": "POST",
-                        "path": "/mcp/tools/call",
-                        "headers": [(b"content-type", b"application/json")],
-                    }
+                request: Final = build_synthetic_mcp_request(
+                    path="/mcp/tools/call",
+                    raw_headers=raw_headers,
+                    client_ip=_client_ip,
                 )
                 if user_api_key_auth is not None:
                     data = await add_litellm_data_to_request(
@@ -1905,6 +1904,7 @@ if MCP_AVAILABLE:
                 "litellm_trace_id": effective_litellm_trace_id,
                 "metadata": {
                     "spend_logs_metadata": spend_logs_metadata,
+                    "headers": logging_safe_mcp_headers(raw_headers),
                     **({"tags": request_tags} if request_tags else {}),
                 },
                 # Provide a small input payload for standard logging
diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py
index 5534b8b5a3c..4cf84dd0725 100644
--- a/litellm/proxy/_experimental/mcp_server/utils.py
+++ b/litellm/proxy/_experimental/mcp_server/utils.py
@@ -7,6 +7,7 @@ import importlib
 import json
 import os
 import re
+import typing
 from collections.abc import Iterable, Iterator, Mapping, MutableMapping, MutableSequence
 from collections.abc import Set as AbstractSet
 from typing import Any, Final, Protocol
@@ -14,6 +15,9 @@ from urllib.parse import quote
 
 from litellm.types.mcp_server.mcp_server_manager import MCPServer
 
+if typing.TYPE_CHECKING:
+    from fastapi import Request
+
 
 class _McpServerLike(Protocol):
     @property
@@ -862,3 +866,146 @@ def set_mcp_tool_result_structured_content(result: object, value: object) -> boo
         return True
     except (AttributeError, TypeError, ValueError):
         return False
+
+
+_HOP_BY_HOP_HEADERS: Final = frozenset(
+    {
+        "content-length",
+        "transfer-encoding",
+        "connection",
+        "keep-alive",
+        "upgrade",
+        "te",
+        "trailer",
+    }
+)
+
+_SYNTHETIC_REQUEST_EXCLUDED_HEADERS: Final = _HOP_BY_HOP_HEADERS | frozenset({"content-type", "x-forwarded-for"})
+
+_SYNTHETIC_REQUEST_SERVER: Final = ("127.0.0.1", 4000)
+
+_MCP_SERVER_AUTH_HEADER_PREFIX: Final = "x-mcp-"
+
+
+def _custom_litellm_key_header_name() -> str | None:
+    """``general_settings.litellm_key_header_name``, the deployment's custom header name for
+    the proxy virtual key, so it is stripped from observability copies like the standard ones."""
+    try:
+        from litellm.proxy.proxy_server import general_settings
+    except ImportError:
+        return None
+    return general_settings.get("litellm_key_header_name") if general_settings else None
+
+
+def _mcp_client_side_auth_header_name() -> str:
+    """The header name the client passes the upstream MCP credential in, falling back to the
+    default when ``general_settings`` is unavailable (the SDK, outside a running proxy)."""
+    from .auth.user_api_key_auth_mcp import MCPRequestHandler
+
+    try:
+        return MCPRequestHandler.get_mcp_client_side_auth_header_name()
+    except ImportError:
+        return MCPRequestHandler.LITELLM_MCP_AUTH_HEADER_NAME
+
+
+def _upstream_credential_headers(header_names: Iterable[str]) -> frozenset[str]:
+    """Lowercased names of the headers in ``header_names`` that carry an upstream MCP
+    credential rather than request context: the configured client side auth header and
+    the per-server ``x-mcp-{alias}-{header}`` family. ``clean_headers`` only knows the
+    credential headers of the chat completions path, so these are dropped on top of it.
+    """
+    from .auth.user_api_key_auth_mcp import MCPRequestHandler
+
+    non_credential: Final = frozenset(
+        {
+            MCPRequestHandler.LITELLM_MCP_SERVERS_HEADER_NAME.lower(),
+            MCPRequestHandler.LITELLM_MCP_ACCESS_GROUPS_HEADER_NAME.lower(),
+        }
+    )
+    client_side_auth: Final = _mcp_client_side_auth_header_name().lower()
+    return frozenset(
+        name
+        for name in (raw_name.lower() for raw_name in header_names)
+        if name == client_side_auth or (name.startswith(_MCP_SERVER_AUTH_HEADER_PREFIX) and name not in non_credential)
+    )
+
+
+def build_synthetic_mcp_request(
+    *,
+    path: str,
+    raw_headers: Mapping[str, str] | None = None,
+    client_ip: str | None = None,
+) -> "Request":
+    """A synthetic FastAPI ``Request`` carrying the MCP connection's HTTP headers.
+
+    The MCP protocol transports do not hand a per-call ``Request`` to the tool
+    handlers, so one is reconstructed from the connection's ``raw_headers``. That
+    lets ``add_litellm_data_to_request`` derive ``metadata.headers``,
+    ``proxy_server_request``, header-based tags, guardrails and trace correlation
+    exactly as on the chat completions path. Hop-by-hop headers describe the
+    original HTTP framing rather than the logical request, so they are dropped, and
+    ``x-forwarded-for`` comes from the resolved ``client_ip`` to avoid spoofing. Upstream
+    MCP credentials and the deployment's proxy key header, including a custom
+    ``litellm_key_header_name``, are dropped so they cannot reach a callback or a guardrail
+    through the derived metadata even when a caller omits ``general_settings``.
+    """
+    from fastapi import Request
+
+    custom_key_header: Final = _custom_litellm_key_header_name()
+    excluded: Final = (
+        _SYNTHETIC_REQUEST_EXCLUDED_HEADERS
+        | _upstream_credential_headers(raw_headers.keys() if raw_headers else ())
+        | (frozenset({custom_key_header.lower()}) if custom_key_header else frozenset())
+    )
+    forwarded: Final = tuple(
+        (
+            name.lower().encode("latin-1", errors="replace"),
+            value.encode("utf-8", errors="replace"),
+        )
+        for name, value in (raw_headers.items() if raw_headers else ())
+        if name.lower() not in excluded
+    )
+    xff: Final = ((b"x-forwarded-for", client_ip.encode("utf-8")),) if client_ip else ()
+    return Request(
+        scope={
+            "type": "http",
+            "method": "POST",
+            "path": path,
+            "scheme": "http",
+            "server": _SYNTHETIC_REQUEST_SERVER,
+            "query_string": b"",
+            "root_path": "",
+            "headers": ((b"content-type", b"application/json"), *forwarded, *xff),
+            **({"client": (client_ip, 0)} if client_ip else {}),
+        }
+    )
+
+
+def logging_safe_mcp_headers(raw_headers: Mapping[str, str] | None) -> Mapping[str, str]:
+    """The MCP request's client headers, sanitized the way the chat completions path
+    sanitizes them before they reach a logging callback or a guardrail: proxy key
+    headers stripped, including the custom key header name the deployment configured,
+    upstream MCP credentials dropped, and credential-bearing values masked.
+
+    Client-controlled behaviour flags (``litellm-disable-message-redaction``) are dropped
+    too: these headers are read back out of the metadata to change proxy behaviour, so
+    leaving one in place would let any MCP client turn off the redaction an admin
+    configured. This path carries no key or team object to authorize an opt-out with, so
+    it always strips them."""
+    from starlette.datastructures import Headers
+
+    from litellm.proxy.litellm_pre_call_utils import (
+        UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS,
+        clean_headers,
+        redact_credential_headers,
+    )
+
+    excluded: Final = (
+        _upstream_credential_headers(raw_headers.keys() if raw_headers else ())
+        | UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS
+    )
+    cleaned: Final = clean_headers(
+        Headers(raw_headers),
+        litellm_key_header_name=_custom_litellm_key_header_name(),
+    )
+    return redact_credential_headers({name: value for name, value in cleaned.items() if name.lower() not in excluded})
diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py
index 251ed1feb10..0a5626ba0a7 100644
--- a/litellm/proxy/litellm_pre_call_utils.py
+++ b/litellm/proxy/litellm_pre_call_utils.py
@@ -274,7 +274,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
     PRE_CALL_EXECUTED_GUARDRAILS_KEY,
 )
 
-_UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
+UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS: Final = frozenset(
     {
         "litellm-disable-message-redaction",
     }
@@ -355,7 +355,7 @@ def _strip_untrusted_request_header_controls(
         return
 
     for header_name in list(headers.keys()):
-        if isinstance(header_name, str) and header_name.lower() in _UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
+        if isinstance(header_name, str) and header_name.lower() in UNTRUSTED_REQUEST_HEADER_CONTROL_FIELDS:
             if allow_client_message_redaction_opt_out:
                 continue
             headers.pop(header_name, None)
diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py
index ef0376ade84..d3ca2fa64ed 100644
--- a/litellm/proxy/utils.py
+++ b/litellm/proxy/utils.py
@@ -679,6 +679,7 @@ class ProxyLogging:
             # (e.g. MCPJWTSigner) to independently verify the caller's identity
             # before re-signing an outbound token (FR-5 verify+re-sign).
             "incoming_bearer_token": kwargs.get("incoming_bearer_token"),
+            "metadata": {"headers": kwargs.get("headers") or {}},
         }
 
         return synthetic_data
diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
index c6e17502e5d..56818717c09 100644
--- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py
+++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py
@@ -11,6 +11,7 @@ from litellm._logging import verbose_logger
 from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy._experimental.mcp_server.utils import (
+    logging_safe_mcp_headers,
     split_server_prefix_from_name,
     strip_known_server_prefix,
 )
@@ -653,6 +654,7 @@ class LiteLLM_Proxy_MCP_Handler:
         tool_results: Final[list[MCPToolResult]] = []
         tool_call_id: str | None = None
         rules_obj: Final = Rules()
+        logging_safe_headers: Final = logging_safe_mcp_headers(raw_headers)
         for tool_call in tool_calls:
             logging_request_data: dict[str, object] = {}
             tool_name: str | None = None
@@ -697,6 +699,7 @@ class LiteLLM_Proxy_MCP_Handler:
                     "tool_call_id": tool_call_id,
                     "tool_name": sanitized_tool_name,
                     "server_name": server_name,
+                    "headers": logging_safe_headers,
                 }
                 logging_request_data = {
                     "model": f"MCP: {tool_name}",
@@ -708,7 +711,7 @@ class LiteLLM_Proxy_MCP_Handler:
                     "proxy_server_request": {
                         "url": "/mcp/tools/call",
                         "method": "POST",
-                        "headers": {},
+                        "headers": logging_safe_headers,
                         "body": {
                             "name": sanitized_tool_name,
                             "arguments": parsed_arguments,
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
index 0b95a497882..52dc91ce24d 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
@@ -2874,7 +2874,7 @@ class TestMCPCustomHeaderName:
                 mock_general_settings.get.return_value = general_setting
 
                 # Call the method
-                result = MCPRequestHandler._get_mcp_client_side_auth_header_name()
+                result = MCPRequestHandler.get_mcp_client_side_auth_header_name()
 
                 # Assert the result
                 assert result == expected_header_name
@@ -2938,7 +2938,7 @@ class TestMCPCustomHeaderName:
         # Mock the header name method
         with patch.object(
             MCPRequestHandler,
-            "_get_mcp_client_side_auth_header_name",
+            "get_mcp_client_side_auth_header_name",
             return_value=custom_header_name,
         ):
             # Create headers from the test data
@@ -2963,7 +2963,7 @@ class TestMCPCustomHeaderName:
         # Mock the custom header name
         with patch.object(
             MCPRequestHandler,
-            "_get_mcp_client_side_auth_header_name",
+            "get_mcp_client_side_auth_header_name",
             return_value="custom-auth-header",
         ):
             # Create ASGI scope with custom header
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
index b56a12db5b1..4081681daef 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_hook_extra_headers.py
@@ -1196,3 +1196,45 @@ class TestOpenApiResolvedUpstreamAuth:
             )
         assert resolved is None
         lookup.assert_not_awaited()
+
+
+class TestPreCallToolCheckExposesClientHeaders:
+    """The pre_mcp_call guardrail payload must carry the caller's sanitized HTTP headers."""
+
+    @pytest.mark.asyncio
+    async def test_sanitized_client_headers_reach_the_guardrail_payload(self):
+        manager = MCPServerManager()
+        server = MCPServer(
+            server_id="test-id",
+            name="test_server",
+            server_name="test_server",
+            url="https://example.com",
+            transport=MCPTransport.http,
+            auth_type=MCPAuth.none,
+        )
+
+        captured: Dict[str, Any] = {}
+
+        def capture(request_obj, kwargs):
+            captured.update(kwargs)
+            return {"model": "fake"}
+
+        proxy_logging = MagicMock(spec=ProxyLogging)
+        proxy_logging._create_mcp_request_object_from_kwargs = MagicMock(return_value=MagicMock())
+        proxy_logging._convert_mcp_to_llm_format = MagicMock(side_effect=capture)
+        proxy_logging.pre_call_hook = AsyncMock(return_value=None)
+
+        with patch.object(manager, "check_allowed_or_banned_tools", return_value=True):
+            with patch.object(manager, "check_tool_permission_for_key_team", new_callable=AsyncMock):
+                with patch.object(manager, "validate_allowed_params"):
+                    await manager.pre_call_tool_check(
+                        name="test_tool",
+                        arguments={"key": "val"},
+                        server_name="test_server",
+                        user_api_key_auth=None,
+                        proxy_logging_obj=proxy_logging,
+                        server=server,
+                        raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy"},
+                    )
+
+        assert captured["headers"] == {"x-nuid": "nuid-1"}
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
index 850d01c6e34..7df83065865 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
@@ -77,7 +77,7 @@ async def test_mcp_server_tool_call_body_contains_request_data():
     # Mock the add_litellm_data_to_request function to capture the data
     captured_data = {}
 
-    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
+    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
         captured_data.update(data)
         # Simulate the proxy_server_request creation
         captured_data["proxy_server_request"] = {
@@ -116,6 +116,107 @@ async def test_mcp_server_tool_call_body_contains_request_data():
     assert body["arguments"] == tool_arguments
 
 
+@pytest.mark.asyncio
+async def test_mcp_server_tool_call_forwards_client_headers_to_logging():
+    """The MCP protocol path must hand the connection's client headers to the pre-call
+    pipeline, so logging callbacks and guardrails see them the way the REST path does."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            mcp_server_tool_call,
+            set_auth_context,
+        )
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    set_auth_context(
+        UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+        raw_headers={
+            "x-nuid": "nuid-1",
+            "x-app-id": "app-1",
+            "content-length": "42",
+            "x-forwarded-for": "9.9.9.9",
+        },
+        client_ip="1.2.3.4",
+    )
+
+    captured_headers = {}
+
+    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
+        captured_headers.update(request.headers)
+        return data
+
+    async def mock_call_mcp_tool(*args, **kwargs):
+        return [{"type": "text", "text": "mocked response"}]
+
+    with patch(
+        "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+        mock_add_litellm_data_to_request,
+    ):
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
+            mock_call_mcp_tool,
+        ):
+            with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
+                await mcp_server_tool_call("test_tool", {"param": "value"})
+
+    assert captured_headers.get("x-nuid") == "nuid-1"
+    assert captured_headers.get("x-app-id") == "app-1"
+    assert "content-length" not in captured_headers
+    assert captured_headers.get("x-forwarded-for") == "1.2.3.4"
+
+
+@pytest.mark.asyncio
+async def test_mcp_server_tool_call_strips_custom_litellm_key_header():
+    """The deployment can rename the proxy key header via general_settings.litellm_key_header_name.
+    The pre-call pipeline only knows that name if it is passed in, so without it the virtual key
+    reaches metadata.headers and proxy_server_request.headers in plaintext."""
+    try:
+        from litellm.proxy._experimental.mcp_server.server import (
+            mcp_server_tool_call,
+            set_auth_context,
+        )
+        from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    set_auth_context(
+        UserAPIKeyAuth(api_key="test_key", user_id="test_user"),
+        raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"},
+        client_ip="1.2.3.4",
+    )
+
+    captured_data = {}
+
+    async def capturing_add_litellm_data_to_request(**kwargs):
+        data = await add_litellm_data_to_request(**kwargs)
+        captured_data.update(data)
+        return data
+
+    async def mock_call_mcp_tool(*args, **kwargs):
+        return [{"type": "text", "text": "mocked response"}]
+
+    with patch(
+        "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request",
+        capturing_add_litellm_data_to_request,
+    ):
+        with patch(
+            "litellm.proxy._experimental.mcp_server.server.call_mcp_tool",
+            mock_call_mcp_tool,
+        ):
+            with patch("litellm.proxy.proxy_server.proxy_config", MagicMock()):
+                with patch.dict(
+                    "litellm.proxy.proxy_server.general_settings",
+                    {"litellm_key_header_name": "x-company-key"},
+                    clear=False,
+                ):
+                    await mcp_server_tool_call("test_tool", {"param": "value"})
+
+    metadata_headers = captured_data["metadata"]["headers"]
+    assert metadata_headers.get("x-nuid") == "nuid-1"
+    assert "x-company-key" not in metadata_headers
+    assert "x-company-key" not in captured_data["proxy_server_request"]["headers"]
+
+
 @pytest.mark.asyncio
 async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
     """The MCP session manager serializes handler exceptions as JSON-RPC errors, so a mid-session
@@ -133,7 +234,7 @@ async def test_mcp_server_tool_call_relays_upstream_auth_error_as_iserror():
 
     set_auth_context(UserAPIKeyAuth(api_key="test_key", user_id="test_user"))
 
-    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
+    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
         return data
 
     async def mock_call_mcp_tool(*args, **kwargs):
@@ -1245,7 +1346,7 @@ async def test_mcp_server_tool_call_body_with_none_arguments():
     # Mock the add_litellm_data_to_request function to capture the data
     captured_data = {}
 
-    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config):
+    async def mock_add_litellm_data_to_request(data, request, user_api_key_dict, proxy_config, **kwargs):
         captured_data.update(data)
         captured_data["proxy_server_request"] = {
             "url": str(request.url),
diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py
index 73fdee9cde3..00ed4e91efa 100644
--- a/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py
+++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_utils.py
@@ -1,7 +1,11 @@
+from unittest.mock import patch
+
 import pytest
 from fastapi import HTTPException
 
 from litellm.proxy._experimental.mcp_server.utils import (
+    build_synthetic_mcp_request,
+    logging_safe_mcp_headers,
     validate_and_normalize_mcp_server_payload,
     validate_tool_display_names,
 )
@@ -47,3 +51,99 @@ class TestValidateAndNormalizeMcpServerPayload:
             tool_name_to_display_name={"read_wiki_structure": "browse_repo_docs"},
         )
         validate_and_normalize_mcp_server_payload(payload)
+
+
+class TestLoggingSafeMcpHeaders:
+    def test_returns_empty_for_missing_headers(self):
+        assert logging_safe_mcp_headers(None) == {}
+        assert logging_safe_mcp_headers({}) == {}
+
+    def test_exposes_custom_headers_and_masks_credentials(self):
+        safe = logging_safe_mcp_headers(
+            {
+                "x-nuid": "nuid-1",
+                "x-app-id": "app-1",
+                "x-litellm-api-key": "sk-proxy",
+                "cookie": "session=secret",
+            }
+        )
+        assert safe == {
+            "x-nuid": "nuid-1",
+            "x-app-id": "app-1",
+            "cookie": "***REDACTED***",
+        }
+
+    def test_strips_custom_litellm_key_header(self):
+        """general_settings.litellm_key_header_name carries the proxy virtual key, so it must
+        never reach a callback or a guardrail even though clean_headers cannot know its name."""
+        with patch.dict(
+            "litellm.proxy.proxy_server.general_settings",
+            {"litellm_key_header_name": "x-company-key"},
+            clear=False,
+        ):
+            safe = logging_safe_mcp_headers({"x-company-key": "sk-proxy", "x-nuid": "nuid-1"})
+
+        assert safe == {"x-nuid": "nuid-1"}
+
+    def test_strips_client_controlled_redaction_opt_out(self):
+        """litellm-disable-message-redaction is read back out of the logged metadata to turn off
+        redaction, so leaving it in place lets any MCP client undo what an admin configured."""
+        safe = logging_safe_mcp_headers({"litellm-disable-message-redaction": "true", "x-nuid": "nuid-1"})
+
+        assert safe == {"x-nuid": "nuid-1"}
+
+    def test_strips_upstream_mcp_credentials(self):
+        safe = logging_safe_mcp_headers(
+            {
+                "x-mcp-auth": "Bearer upstream",
+                "x-mcp-github-authorization": "Bearer gh_token",
+                "x-mcp-zapier-x-api-key": "zapier-key",
+                "x-nuid": "nuid-1",
+            }
+        )
+
+        assert safe == {"x-nuid": "nuid-1"}
+
+    def test_strips_custom_mcp_client_side_auth_header(self):
+        with patch.dict(
+            "litellm.proxy.proxy_server.general_settings",
+            {"mcp_client_side_auth_header_name": "x-upstream-token"},
+            clear=False,
+        ):
+            safe = logging_safe_mcp_headers({"x-upstream-token": "Bearer upstream", "x-nuid": "nuid-1"})
+
+        assert safe == {"x-nuid": "nuid-1"}
+
+
+class TestBuildSyntheticMcpRequest:
+    def test_forwards_client_headers_without_upstream_credentials(self):
+        """The synthetic request feeds add_litellm_data_to_request, which derives
+        metadata.headers, so upstream MCP credentials must not ride along."""
+        request = build_synthetic_mcp_request(
+            path="/mcp/tools/call",
+            raw_headers={
+                "x-nuid": "nuid-1",
+                "x-mcp-auth": "Bearer upstream",
+                "x-mcp-github-authorization": "Bearer gh_token",
+            },
+        )
+
+        assert request.headers.get("x-nuid") == "nuid-1"
+        assert "x-mcp-auth" not in request.headers
+        assert "x-mcp-github-authorization" not in request.headers
+
+    def test_drops_custom_litellm_key_header(self):
+        """Callers such as the sampling flow build metadata off this request, so the
+        deployment's custom proxy key header must never be forwarded on it."""
+        with patch.dict(
+            "litellm.proxy.proxy_server.general_settings",
+            {"litellm_key_header_name": "x-company-key"},
+            clear=False,
+        ):
+            request = build_synthetic_mcp_request(
+                path="/mcp/sampling/createMessage",
+                raw_headers={"x-company-key": "sk-proxy-secret", "x-nuid": "nuid-1"},
+            )
+
+        assert request.headers.get("x-nuid") == "nuid-1"
+        assert "x-company-key" not in request.headers
diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
index 9defb309863..56057dce7e0 100644
--- a/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
+++ b/tests/test_litellm/proxy/utils/proxy_logging/test_mcp_bridging.py
@@ -76,6 +76,23 @@ def test_convert_mcp_to_llm_format_defaults_model(proxy_logging, make_mcp_reques
     }
 
 
+def test_convert_mcp_to_llm_format_exposes_headers_on_metadata(proxy_logging, make_mcp_request_obj):
+    """Guardrails read the caller's HTTP headers off ``metadata.headers`` on the chat
+    completions path, so the MCP bridge has to put them in the same place."""
+    req = make_mcp_request_obj()
+    out = proxy_logging._convert_mcp_to_llm_format(
+        request_obj=req,
+        kwargs={"headers": {"x-nuid": "nuid-1"}},
+    )
+    assert out["metadata"]["headers"] == {"x-nuid": "nuid-1"}
+
+
+def test_convert_mcp_to_llm_format_defaults_headers_to_empty(proxy_logging, make_mcp_request_obj):
+    req = make_mcp_request_obj()
+    out = proxy_logging._convert_mcp_to_llm_format(request_obj=req, kwargs={})
+    assert out["metadata"]["headers"] == {}
+
+
 def test_convert_mcp_to_llm_format_missing_request_obj_raises(proxy_logging):
     with pytest.raises(AttributeError):
         proxy_logging._convert_mcp_to_llm_format(request_obj=None, kwargs={})
diff --git a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
index d60fff66c44..4981caa10c3 100644
--- a/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
+++ b/tests/test_litellm/responses/mcp/test_litellm_proxy_mcp_handler.py
@@ -536,6 +536,37 @@ async def test_get_mcp_tools_from_manager_forwards_request_tags(monkeypatch):
     assert mock_get_tools.await_args.kwargs["request_tags"] == ["team-a"]
 
 
+@pytest.mark.asyncio
+async def test_execute_tool_calls_exposes_sanitized_client_headers_to_logging(monkeypatch):
+    """The Responses API MCP bridge used to log an empty header dict, hiding the caller's
+    headers from logging callbacks and hooks."""
+    _setup_proxy_logging(monkeypatch)
+    _setup_mcp_call_environment(monkeypatch)
+
+    captured = {}
+
+    def fake_function_setup(*_args, **kwargs):
+        captured.update(kwargs)
+        return None, None
+
+    handler_module = importlib.import_module(
+        "litellm.responses.mcp.litellm_proxy_mcp_handler"
+    )
+    monkeypatch.setattr(handler_module, "function_setup", fake_function_setup)
+
+    tool_name = "deepwiki-read_wiki_structure"
+    await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
+        tool_server_map={tool_name: "deepwiki"},
+        tool_calls=[{"id": "call-1", "function": {"name": tool_name, "arguments": "{}"}}],
+        user_api_key_auth=None,
+        raw_headers={"x-nuid": "nuid-1", "x-litellm-api-key": "sk-proxy", "cookie": "s=1"},
+    )
+
+    expected = {"x-nuid": "nuid-1", "cookie": "***REDACTED***"}
+    assert captured["metadata"]["headers"] == expected
+    assert captured["proxy_server_request"]["headers"] == expected
+
+
 @pytest.mark.asyncio
 async def test_execute_tool_calls_propagates_request_tags_to_function_setup(monkeypatch):
     _setup_proxy_logging(monkeypatch)
diff --git a/type-discipline-budget.json b/type-discipline-budget.json
index 0104733bcf3..894d99c92e0 100644
--- a/type-discipline-budget.json
+++ b/type-discipline-budget.json
@@ -1,9 +1,9 @@
 {
   "LIT001": {
-    "limit": 22943
+    "limit": 22941
   },
   "LIT002": {
-    "limit": 27141
+    "limit": 27139
   },
   "LIT003": {
     "limit": 269
@@ -27,7 +27,7 @@
     "limit": 0
   },
   "LIT010": {
-    "limit": 16722
+    "limit": 16716
   },
   "LIT011": {
     "limit": 5596

From a36ba0588205ade0f6fe2246fbd8630ba09e5ed4 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Thu, 13 Aug 2026 20:11:06 -0700
Subject: [PATCH 28/31] test(proxy): stop monkeypatch.undo re-planting
 fixture-mocked prisma_client

---
 tests/test_litellm/proxy/conftest.py          | 55 +++++++++++++------
 .../test_mcp_management_endpoints.py          | 20 ++++++-
 tests/test_litellm/proxy/test_conftest.py     | 31 +++++++++++
 3 files changed, 85 insertions(+), 21 deletions(-)
 create mode 100644 tests/test_litellm/proxy/test_conftest.py

diff --git a/tests/test_litellm/proxy/conftest.py b/tests/test_litellm/proxy/conftest.py
index 230ccaf5fd4..61752997f0f 100644
--- a/tests/test_litellm/proxy/conftest.py
+++ b/tests/test_litellm/proxy/conftest.py
@@ -43,32 +43,51 @@ def disconnected_prisma() -> DisconnectedPrisma:
     return DisconnectedPrisma()
 
 
-@pytest.fixture(autouse=True)
-def _isolate_proxy_module_globals():
-    """
-    Snapshot and restore module-level globals on litellm.proxy.proxy_server
-    that tests sometimes mutate via raw setattr (not monkeypatch).
+_MODULE_GLOBAL_MISSING = object()
+_proxy_module_globals_snapshot = pytest.StashKey[Dict[str, object]]()
 
-    Without this, a leaked value — e.g. master_key set by a sibling test —
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_setup(item):
+    """
+    Snapshot module-level globals on litellm.proxy.proxy_server before any
+    fixture runs, and restore them in pytest_runtest_teardown after every
+    fixture finalizer has run.
+
+    Without this, a leaked value (e.g. master_key set by a sibling test)
     flips the auth short-circuit in user_api_key_auth and causes unrelated
     tests in the same xdist worker to return 401 instead of 200.
+
+    This must be a hook pair, not an autouse fixture: an autouse fixture in
+    the root conftest requests monkeypatch, so monkeypatch's undo stack
+    unwinds after every other fixture finalizer. A test that monkeypatches a
+    global while a fixture has it patched records the fixture's mock as the
+    "original", and monkeypatch.undo re-plants that mock after all restores
+    have run, poisoning the global for the rest of the xdist worker.
     """
     from litellm.proxy import proxy_server
 
-    sentinel = object()
-    snapshot = {
-        name: getattr(proxy_server, name, sentinel)
+    item.stash[_proxy_module_globals_snapshot] = {
+        name: getattr(proxy_server, name, _MODULE_GLOBAL_MISSING)
         for name in _PROXY_MODULE_GLOBALS_TO_ISOLATE
     }
-    try:
-        yield
-    finally:
-        for name, value in snapshot.items():
-            if value is sentinel:
-                if hasattr(proxy_server, name):
-                    delattr(proxy_server, name)
-            else:
-                setattr(proxy_server, name, value)
+    yield
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_teardown(item, nextitem):
+    yield
+    snapshot = item.stash.get(_proxy_module_globals_snapshot, None)
+    if snapshot is None:
+        return
+    from litellm.proxy import proxy_server
+
+    for name, value in snapshot.items():
+        if value is _MODULE_GLOBAL_MISSING:
+            if hasattr(proxy_server, name):
+                delattr(proxy_server, name)
+        else:
+            setattr(proxy_server, name, value)
 
 
 @pytest.fixture(autouse=True)
diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
index b6c2fd4c8e3..84dee5b05c5 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py
@@ -1688,9 +1688,15 @@ class TestTemporaryMCPSessionEndpoints:
             expires_at=datetime.utcnow() - timedelta(seconds=30),
         )
         cache = {"expired": expired_entry}
-        with patch(
-            "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
-            cache,
+        with (
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._temporary_mcp_servers",
+                cache,
+            ),
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
+                return_value=None,
+            ),
         ):
             result = await get_cached_temporary_mcp_server("expired")
 
@@ -2274,6 +2280,10 @@ class TestTemporaryMCPSessionEndpoints:
                 "litellm.proxy.management_endpoints.mcp_management_endpoints._cache_temporary_mcp_server_in_redis",
                 AsyncMock(),
             ) as redis_cache_mock,
+            patch(
+                "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
+                return_value=None,
+            ),
         ):
             response = await add_session_mcp_server(
                 payload=payload,
@@ -3419,6 +3429,10 @@ class TestTemporaryMCPSessionEndpoints:
                     "litellm.proxy.management_endpoints.mcp_management_endpoints.decrypt_value_helper",
                     return_value=serialized,
                 ),
+                patch(
+                    "litellm.proxy.management_endpoints.mcp_management_endpoints._get_prisma_client_or_none",
+                    return_value=None,
+                ),
             ):
                 result = await get_cached_temporary_mcp_server("from-redis")
         finally:
diff --git a/tests/test_litellm/proxy/test_conftest.py b/tests/test_litellm/proxy/test_conftest.py
new file mode 100644
index 00000000000..6df692a67c9
--- /dev/null
+++ b/tests/test_litellm/proxy/test_conftest.py
@@ -0,0 +1,31 @@
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+
+@pytest.fixture
+def fixture_planted_prisma_mock():
+    with patch("litellm.proxy.proxy_server.prisma_client", MagicMock()):
+        yield
+
+
+def test_monkeypatch_over_fixture_patched_prisma_client(
+    fixture_planted_prisma_mock, monkeypatch
+):
+    """
+    Mirrors the flake in test_team_endpoints.py: an autouse fixture patches
+    prisma_client, the test monkeypatches the same global, and monkeypatch
+    records the fixture's MagicMock as the value to restore. Its undo runs
+    after every other finalizer, so without hook-level isolation the mock
+    leaks and every later no-database test on the worker fails awaiting it.
+    """
+    import litellm.proxy.proxy_server as proxy_server
+
+    monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", AsyncMock())
+    assert isinstance(proxy_server.prisma_client, AsyncMock)
+
+
+def test_prisma_client_did_not_leak_from_previous_test():
+    import litellm.proxy.proxy_server as proxy_server
+
+    assert not isinstance(proxy_server.prisma_client, MagicMock)

From 0a25756e782b3ad52b9d9247b5a623bd62b523bf Mon Sep 17 00:00:00 2001
From: yucheng-berri <yucheng@berri.ai>
Date: Thu, 13 Aug 2026 20:16:12 -0700
Subject: [PATCH 29/31] fix(ptu): stop per-token billing on a PTU-configured
 deployment (#36829)

A deployment with PTU flat-cost attribution also billed every request per
token, so a team paid for reserved capacity and again for the traffic that
capacity serves. Nothing set the per-token price and an unset price falls
back to the public cost map, which made the double charge the default.

/model/new and /model/{id}/update now store zero for every pricing field the
cost map could otherwise fill, refuse a price the caller supplies alongside
PTU config with a 400 naming the field, zero a price already on the row
rather than rejecting later edits of unrelated fields, and drop the zeros
again when the PTU config goes.

A PTU deployment is no longer read as a free model by the budget checks,
which would have waived every budget for it.
---
 litellm/constants.py                          |   3 +
 litellm/proxy/auth/auth_checks.py             |  29 ++
 .../model_management_endpoints.py             | 171 +++++++-
 .../spend_tracking/ptu_flat_cost_rollup.py    |  38 ++
 .../test_ptu_model_settings.py                | 403 +++++++++++++++++-
 .../test_ptu_flat_cost_rollup.py              |  81 ++++
 .../add_model/advanced_settings.tsx           |  16 +-
 .../src/components/model_info_view.test.tsx   |  69 +++
 .../src/components/model_info_view.tsx        |  51 ++-
 .../src/utils/ptuValidation.test.ts           |  34 ++
 .../src/utils/ptuValidation.ts                |  20 +
 11 files changed, 889 insertions(+), 26 deletions(-)

diff --git a/litellm/constants.py b/litellm/constants.py
index e9d2d719ae2..6449834d6a4 100644
--- a/litellm/constants.py
+++ b/litellm/constants.py
@@ -1745,6 +1745,9 @@ PTU_ROLLUP_LOCK_TTL_SECONDS: Final[int] = 900
 # Furthest back the catch-up pass looks for unpriced PTU days when a deployment
 # declares no ptu_effective_from, bounding the scan for an open-ended window.
 PTU_ROLLUP_MAX_BACKFILL_DAYS: Final[int] = 90
+# Deployments named in the lapsed-window alert before it is truncated, so a fleet-wide
+# expiry cannot produce an alert too large for the channel delivering it.
+PTU_LAPSED_ALERT_LIMIT: Final[int] = 10
 # Slack allowed when deciding a sentinel row is stale. The row's updated_at and the
 # run's cutoff are stamped by different hosts, so clock skew between them must not let
 # one run delete a charge another just wrote. A stale row is hours old and a concurrent
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index 0239368ad0e..51050e62494 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -14,6 +14,7 @@ import math
 import re
 import time
 from collections.abc import Iterator, Mapping, Sequence
+from types import MappingProxyType
 from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast
 
 from fastapi import HTTPException, Request, status
@@ -376,6 +377,16 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
                     zero_cost_cache[model_name] = False
                 return False
 
+            if _has_ptu_flat_cost(model_name, llm_router):
+                verbose_proxy_logger.debug(
+                    "Model %s prices reserved PTU capacity as a flat cost, so its zero per-token "
+                    "rate is not a free model (enforce budget)",
+                    safe_name,
+                )
+                if zero_cost_cache is not None:
+                    zero_cost_cache[model_name] = False
+                return False
+
             verbose_proxy_logger.debug(
                 "Model %s has zero cost explicitly configured (input: %s, output: %s)",
                 safe_name,
@@ -394,6 +405,24 @@ def _is_model_cost_zero(model: str | list[str] | None, llm_router: Router | None
     return True
 
 
+_NO_MODEL_INFO: Final[Mapping[str, object]] = MappingProxyType({})
+
+
+def _has_ptu_flat_cost(model: str, llm_router: "Router") -> bool:
+    """Whether any deployment in the model group bills reserved PTU capacity as a flat cost.
+
+    Such a deployment carries an explicit zero per-token price so the flat cost is not charged
+    twice, which otherwise reads here as a free model and waives every budget check for it.
+    """
+    for deployment in llm_router.model_list:
+        if deployment.get("model_name") != model:
+            continue
+        model_info = deployment.get("model_info") or _NO_MODEL_INFO
+        if model_info.get("ptu_count") is not None and model_info.get("cost_per_ptu_per_hour") is not None:
+            return True
+    return False
+
+
 def _is_cost_explicitly_configured(model: str, llm_router: "Router") -> bool:
     """
     Check if any deployment in the model group has cost fields explicitly
diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py
index 4bc72f554b1..912e18150b3 100644
--- a/litellm/proxy/management_endpoints/model_management_endpoints.py
+++ b/litellm/proxy/management_endpoints/model_management_endpoints.py
@@ -89,6 +89,7 @@ from litellm.types.router import (
     ModelInfo,
     updateDeployment,
 )
+from litellm.types.utils import CustomPricingLiteLLMParams
 from litellm.utils import get_utc_datetime
 
 router: Final = APIRouter()
@@ -242,6 +243,7 @@ def _raise_on_strategy_router_write_violation(
 
 
 _PTU_MODEL_INFO_FIELDS: Final = ("ptu_count", "cost_per_ptu_per_hour", "ptu_effective_from", "ptu_effective_to")
+_PTU_PRICED_PAIR: Final = frozenset({"ptu_count", "cost_per_ptu_per_hour"})
 
 
 def _explicitly_cleared_ptu_fields(model_info: ModelInfo | None) -> frozenset[str]:
@@ -265,9 +267,10 @@ def _merged_ptu_model_info(*, db_model: Deployment, patch_data: updateDeployment
     A PTU invariant holds over the deployment as it will exist, not over whichever subset
     of fields a caller happened to send.
     """
-    empty: Final[Mapping[str, object]] = MappingProxyType({})
-    stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else empty
-    incoming: Final = patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else empty
+    stored: Final = db_model.model_info.model_dump(exclude_none=True) if db_model.model_info else _EMPTY_MODEL_INFO
+    incoming: Final = (
+        patch_data.model_info.model_dump(exclude_none=True) if patch_data.model_info else _EMPTY_MODEL_INFO
+    )
     cleared: Final = _explicitly_cleared_ptu_fields(patch_data.model_info)
     return MappingProxyType({k: v for k, v in {**stored, **incoming}.items() if k not in cleared})
 
@@ -339,6 +342,140 @@ def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None:
         )
 
 
+# The six mirrored pricing fields plus the three remaining fields
+# Router._inherit_builtin_cache_pricing back-fills from the public cost map. An unset field is
+# what that back-fill targets, so a field left out here is one a PTU deployment still bills.
+_PTU_ZEROED_PRICING_FIELDS: Final = SPECIAL_MODEL_INFO_PARAMS + (
+    "cache_creation_input_token_cost_above_1hr",
+    "cache_creation_input_token_cost_above_200k_tokens",
+    "cache_read_input_token_cost_above_200k_tokens",
+)
+_PTU_ZEROED_PRICING: Final[Mapping[str, float]] = MappingProxyType(dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0))
+_NO_PRICING_OVERRIDE: Final[Mapping[str, float]] = MappingProxyType({})
+_EMPTY_MODEL_INFO: Final[Mapping[str, object]] = _NO_PRICING_OVERRIDE
+# Rate fields only. CustomPricingLiteLLMParams also carries settings that are not charges
+# (an embedding's output_vector_size, the regional uplift multipliers), and zeroing one of
+# those would destroy the deployment's configuration rather than stop a charge.
+_CUSTOM_PRICING_FIELDS: Final = frozenset(f for f in CustomPricingLiteLLMParams.model_fields if "cost" in f)
+
+
+def _is_nonzero_price(value: object) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool) and value != 0
+
+
+def _is_zero_price(value: object) -> bool:
+    return isinstance(value, (int, float)) and not isinstance(value, bool) and value == 0
+
+
+def _raise_if_ptu_deployment_is_priced(*, model_info: Mapping[str, object], supplied: Mapping[str, object]) -> None:
+    """Refuse a rate the caller supplies for a deployment that bills reserved capacity.
+
+    Separate from the zeroing so the team-model path can run it before it touches the team, whose
+    ACL write autocommits: a refusal raised after it would leave the team changed and the
+    deployment row never written.
+    """
+    if not is_ptu_cost_attribution_enabled():
+        return
+    if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
+        return
+    priced: Final = tuple(sorted(field for field in _CUSTOM_PRICING_FIELDS if _is_nonzero_price(supplied.get(field))))
+    if not priced:
+        return
+    raise HTTPException(
+        status_code=400,
+        detail=(
+            f"A PTU deployment bills by reserved capacity, so {', '.join(priced)} cannot be charged on "
+            "top of it. Send 0 or no value, or remove ptu_count and cost_per_ptu_per_hour to bill per token."
+        ),
+    )
+
+
+def _ptu_zeroed_pricing(
+    *,
+    model_info: Mapping[str, object],
+    litellm_params: Mapping[str, object],
+    supplied: Mapping[str, object],
+) -> Mapping[str, float]:
+    """The pricing a PTU deployment must carry, empty unless one is being stored.
+
+    Reserved capacity is already billed by the flat cost the rollup writes, so charging the
+    traffic it serves bills the same tokens twice. Left unset the rate falls back to the public
+    cost map, which makes the double charge the default rather than an opt-in.
+
+    Only a price the caller supplies is refused. A non-zero price already on the row is zeroed
+    instead, so a deployment priced through a path this rule does not cover heals on its next
+    save rather than rejecting every later edit of a field that has nothing to do with pricing.
+
+    ``supplied`` is the caller's litellm_params alone, because that is the blob a price is
+    authored on. model_info's copy is written by the server, both by the mirror in
+    ``Deployment.__init__`` and by the cost-map defaults /model/info fills in, so a client that
+    round-trips a model_info blob sends back prices it never chose.
+    """
+    if not is_ptu_cost_attribution_enabled():
+        return _NO_PRICING_OVERRIDE
+    if model_info.get("ptu_count") is None or model_info.get("cost_per_ptu_per_hour") is None:
+        return _NO_PRICING_OVERRIDE
+    _raise_if_ptu_deployment_is_priced(model_info=model_info, supplied=supplied)
+    stored: Final = frozenset(
+        field
+        for field in _CUSTOM_PRICING_FIELDS
+        if _is_nonzero_price(model_info.get(field)) or _is_nonzero_price(litellm_params.get(field))
+    )
+    if not stored:
+        return _PTU_ZEROED_PRICING
+    return MappingProxyType({**_PTU_ZEROED_PRICING, **dict.fromkeys(stored, 0.0)})
+
+
+def _ptu_pricing_delta(
+    *,
+    stored_model_info: Mapping[str, object],
+    model_info: Mapping[str, object],
+    litellm_params: Mapping[str, object],
+    patch: updateDeployment,
+) -> tuple[Mapping[str, float], frozenset[str]]:
+    """The pricing a patch must write into both blobs, and the pricing it must drop from them.
+
+    A patch that takes the deployment off PTU takes the zeroed pricing with it, since the zeros
+    exist only to stop the double charge. Left behind they would serve the deployment for free.
+    Reading the stored row rather than the patch alone keeps that release off a deployment that
+    never carried PTU config, whose zero price is a rate its operator chose. A zero the patch
+    itself carries is released with the rest, because the dashboard echoes the whole stored
+    blob on every save, so a supplied zero cannot be told apart from the one this rule wrote.
+
+    The release spans every field the zeroing could have written, not just the mirrored ones, or
+    a rate zeroed on the way in (per-second, per-character tiers) would bill nothing forever.
+    """
+    supplied: Final = patch.litellm_params.model_dump(exclude_none=True) if patch.litellm_params else _EMPTY_MODEL_INFO
+    zeroed: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=supplied)
+    if zeroed:
+        return zeroed, frozenset()
+    was_ptu: Final = any(stored_model_info.get(field) is not None for field in _PTU_PRICED_PAIR)
+    if not was_ptu or not _explicitly_cleared_ptu_fields(patch.model_info) & _PTU_PRICED_PAIR:
+        return _NO_PRICING_OVERRIDE, frozenset()
+    return _NO_PRICING_OVERRIDE, frozenset(
+        field
+        for field in _CUSTOM_PRICING_FIELDS.union(_PTU_ZEROED_PRICING_FIELDS)
+        if _is_zero_price(model_info.get(field)) or _is_zero_price(litellm_params.get(field))
+    )
+
+
+def _ptu_priced_deployment(model_params: Deployment) -> Deployment:
+    """``model_params`` with PTU pricing applied, or itself when it configures no PTU."""
+    model_info: Final = model_params.model_info.model_dump(exclude_none=True)
+    litellm_params: Final = model_params.litellm_params.model_dump(exclude_none=True)
+    override: Final = _ptu_zeroed_pricing(model_info=model_info, litellm_params=litellm_params, supplied=litellm_params)
+    if not override:
+        return model_params
+    return model_params.model_copy(
+        update=MappingProxyType(
+            {
+                "litellm_params": model_params.litellm_params.model_copy(update=override),
+                "model_info": model_params.model_info.model_copy(update=override),
+            }
+        )
+    )
+
+
 def _parse_ptu_datetime(value: object) -> datetime.datetime | None:
     """``value`` as a datetime, parsing an ISO string, else None."""
     if isinstance(value, datetime.datetime):
@@ -404,6 +541,19 @@ def update_db_model(db_model: Deployment, updated_patch: updateDeployment) -> Pr
             merged_model_info.pop(field, None)
 
     _validate_ptu_model_info(merged_model_info)
+    ptu_pricing, ptu_released = _ptu_pricing_delta(
+        stored_model_info=db_model.model_info.model_dump(exclude_none=True)
+        if db_model.model_info
+        else _EMPTY_MODEL_INFO,
+        model_info=merged_model_info,
+        litellm_params=merged_litellm_params,
+        patch=updated_patch,
+    )
+    merged_model_info.update(ptu_pricing)
+    merged_litellm_params.update(ptu_pricing)
+    for field in ptu_released:
+        merged_model_info.pop(field, None)
+        merged_litellm_params.pop(field, None)
 
     # convert to prisma compatible format
 
@@ -863,6 +1013,12 @@ async def _update_team_model_in_db(
     if patch_data.model_info is not None:
         _raise_if_ptu_cost_attribution_disabled(patch_data.model_info.model_dump(exclude_none=True))
         _validate_ptu_model_info(_merged_ptu_model_info(db_model=db_model, patch_data=patch_data))
+    _raise_if_ptu_deployment_is_priced(
+        model_info=_merged_ptu_model_info(db_model=db_model, patch_data=patch_data),
+        supplied=(
+            patch_data.litellm_params.model_dump(exclude_none=True) if patch_data.litellm_params else _EMPTY_MODEL_INFO
+        ),
+    )
 
     patch_team_id: Final = patch_data.model_info.team_id if patch_data.model_info else None
 
@@ -1589,6 +1745,7 @@ async def add_new_model(
         incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
         _raise_if_ptu_cost_attribution_disabled(incoming_model_info)
         _validate_ptu_model_info(incoming_model_info)
+        priced_model_params: Final = _ptu_priced_deployment(model_params)
 
         if store_model_in_db is True:
             """
@@ -1602,13 +1759,13 @@ async def add_new_model(
                 _original_litellm_model_name: Final = model_params.model_name
                 if model_params.model_info.team_id is None:
                     model_response = await _add_model_to_db(
-                        model_params=model_params,
+                        model_params=priced_model_params,
                         user_api_key_dict=user_api_key_dict,
                         prisma_client=prisma_client,
                     )
                 else:
                     model_response = await _add_team_model_to_db(
-                        model_params=model_params,
+                        model_params=priced_model_params,
                         user_api_key_dict=user_api_key_dict,
                         prisma_client=prisma_client,
                     )
@@ -1620,9 +1777,9 @@ async def add_new_model(
                 if "slack" in _alerting:
                     # send notification - new model added
                     await proxy_logging_obj.slack_alerting_instance.model_added_alert(
-                        model_name=model_params.model_name,
+                        model_name=priced_model_params.model_name,
                         litellm_model_name=_original_litellm_model_name,
-                        passed_model_info=model_params.model_info,
+                        passed_model_info=priced_model_params.model_info,
                     )
             except Exception as e:
                 verbose_proxy_logger.exception("Exception in add_new_model: %s", e)
diff --git a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
index 029648f7901..efdbda47fdc 100644
--- a/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
+++ b/litellm/proxy/spend_tracking/ptu_flat_cost_rollup.py
@@ -21,6 +21,7 @@ from typing import TYPE_CHECKING, Final
 
 from litellm._logging import verbose_proxy_logger
 from litellm.constants import (
+    PTU_LAPSED_ALERT_LIMIT,
     PTU_PRUNE_SKEW_GRACE_SECONDS,
     PTU_ROLLUP_JOB_ID,
     PTU_ROLLUP_LOCK_TTL_SECONDS,
@@ -45,6 +46,7 @@ class RollupResult:
     models_processed: int
     rows_written: int
     rows_failed: int = 0
+    lapsed: tuple[str, ...] = ()
 
 
 @dataclass(frozen=True, slots=True)
@@ -387,6 +389,34 @@ async def run_ptu_flat_cost_rollup(
         models_processed=len(ptu_models),
         rows_written=rows_written,
         rows_failed=rows_failed,
+        lapsed=_lapsed_models(ptu_models, run_started),
+    )
+
+
+def _slack_safe(model_name: str) -> str:
+    """``model_name`` with the characters Slack reads as markup escaped.
+
+    A model name is operator-supplied and this alert is delivered to an operator channel, so an
+    unescaped name could post a channel-wide mention or a disguised link.
+    """
+    return model_name.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
+
+
+def _lapsed_models(ptu_models: tuple[PTUModel, ...], now: datetime) -> tuple[str, ...]:
+    """PTU deployments whose window has closed, newest bound first.
+
+    The provider bills reserved capacity until the deployment is deleted, so a closed window
+    stops this attribution without stopping the charge. The deployment is left alone: the
+    window is what the operator asked to be attributed, and per-token pricing would invent a
+    charge the provider does not make for reserved capacity.
+    """
+    return tuple(
+        _slack_safe(model.model_name)
+        for model in sorted(
+            (m for m in ptu_models if m.effective_to is not None and m.effective_to <= now),
+            key=lambda m: m.effective_to,
+            reverse=True,
+        )
     )
 
 
@@ -585,6 +615,14 @@ async def _run_and_alert(
             f"{result.rows_written + result.rows_failed} team charges failed to write. Those teams show no PTU "
             f"cost for that date until the rollup is rerun for it.",
         )
+    if result.lapsed:
+        await _deliver_alert(
+            alert,
+            f"PTU flat-cost attribution has stopped for {len(result.lapsed)} deployment(s) whose effective "
+            f"window has closed: {', '.join(result.lapsed[:PTU_LAPSED_ALERT_LIMIT])}. Reserved capacity is billed "
+            "until the deployment is deleted, so a deployment still serving traffic is still being charged for "
+            "by the provider with nothing attributing it here. Extend the window, or retire the deployment.",
+        )
     if target_date is None:
         await _backfill_and_alert(prisma_client, alert=alert)
     return result
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
index 178e9379ea7..6e670e48b6a 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ptu_model_settings.py
@@ -4,6 +4,7 @@ import datetime
 import json
 from contextlib import ExitStack
 from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import patch as patch_ctx
 
 import pytest
 from fastapi import HTTPException
@@ -14,15 +15,28 @@ from litellm.proxy._types import (
     ReconcileOutcome,
     UserAPIKeyAuth,
 )
+from litellm.proxy.auth.auth_checks import _is_model_cost_zero
 from litellm.proxy.management_endpoints.model_management_endpoints import (
+    _PTU_ZEROED_PRICING_FIELDS,
     _merged_ptu_model_info,
+    _update_team_model_in_db,
+    _ptu_priced_deployment,
+    _ptu_zeroed_pricing,
     _raise_if_ptu_cost_attribution_disabled,
     _validate_ptu_model_info,
     add_new_model,
     update_db_model,
 )
 from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR
-from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment
+from litellm.router import Router
+from litellm.types.router import (
+    SPECIAL_MODEL_INFO_PARAMS,
+    Deployment,
+    LiteLLM_Params,
+    ModelInfo,
+    updateDeployment,
+    updateLiteLLMParams,
+)
 
 
 def test_model_info_accepts_valid_ptu_fields():
@@ -717,3 +731,390 @@ class TestAddNewModelPtuGate:
 
         assert result.model_id == "ptu-gate-model"
         add_team_model_to_db.assert_called_once()
+
+
+
+class TestPtuDeploymentsAreNotBilledPerToken:
+    """Reserved capacity is billed by the flat cost the rollup writes, so a PTU deployment must
+    not also bill the traffic that capacity serves."""
+
+    PTU = {"ptu_count": 15, "cost_per_ptu_per_hour": 2.0}
+
+    @pytest.fixture(autouse=True)
+    def _flag_on(self, monkeypatch):
+        monkeypatch.setenv(PTU_COST_ATTRIBUTION_ENV_VAR, "true")
+        # update_db_model encrypts every litellm_params value it is handed, and the salt falls
+        # back to the master key the proxy sets at boot, which no unit test has.
+        monkeypatch.setenv("LITELLM_SALT_KEY", "test-salt-key")
+
+    @staticmethod
+    def _zeroed(model_info=None, litellm_params=None, supplied=None):
+        return _ptu_zeroed_pricing(
+            model_info=model_info if model_info is not None else {},
+            litellm_params=litellm_params if litellm_params is not None else {},
+            supplied=supplied if supplied is not None else {},
+        )
+
+    def test_a_deployment_without_ptu_config_keeps_its_pricing(self):
+        assert self._zeroed(model_info={"team_id": "t"}, litellm_params={"input_cost_per_token": 5e-07}) == {}
+
+    def test_a_half_set_pair_is_not_treated_as_ptu(self):
+        assert self._zeroed(model_info={"ptu_count": 15}) == {}
+
+    def test_every_field_the_cost_map_could_fill_is_zeroed(self):
+        assert self._zeroed(model_info=self.PTU) == dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
+
+    def test_nothing_is_zeroed_while_the_feature_is_disabled(self, monkeypatch):
+        monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False)
+        assert self._zeroed(model_info=self.PTU) == {}
+
+    @pytest.mark.parametrize("field", ["input_cost_per_token", "cache_read_input_token_cost", "input_cost_per_second"])
+    def test_a_price_the_caller_supplies_is_refused(self, field):
+        """Every custom-pricing field, not only the mirrored ones: per-second pricing bills a
+        PTU deployment just as surely as per-token pricing does."""
+        with pytest.raises(HTTPException) as exc:
+            self._zeroed(model_info=self.PTU, supplied={field: 5e-07})
+        assert exc.value.status_code == 400
+        assert field in str(exc.value.detail)
+
+    def test_a_price_the_caller_supplies_as_zero_is_accepted(self):
+        assert self._zeroed(model_info={**self.PTU, "input_cost_per_token": 0}, supplied={"input_cost_per_token": 0})[
+            "input_cost_per_token"
+        ] == 0
+
+    def test_a_price_already_on_the_row_is_zeroed_rather_than_refused(self):
+        """A row priced through a path this rule does not cover must heal on its next save. The
+        alternative refuses every later edit of a field that has nothing to do with pricing."""
+        zeroed = self._zeroed(model_info={**self.PTU, "input_cost_per_second": 3.0}, litellm_params={})
+        assert zeroed["input_cost_per_second"] == 0
+        assert zeroed["input_cost_per_token"] == 0
+
+    @pytest.mark.asyncio
+    async def test_a_refused_price_does_not_leave_the_team_changed(self):
+        """The team ACL write autocommits, so the refusal has to run before it. Otherwise a
+        rejected edit grants the team a model whose settings were never saved."""
+        db_model = Deployment(
+            model_name="gpt-4o",
+            litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
+            model_info=ModelInfo(
+                id="dep-0",
+                team_id="team-1",
+                ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                **self.PTU,
+            ),
+        )
+        patch = updateDeployment(
+            litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07),
+            model_info=ModelInfo(id="dep-0", team_id="team-2"),
+        )
+        endpoints = "litellm.proxy.management_endpoints.model_management_endpoints"
+        setup_new = AsyncMock()
+        update_existing = AsyncMock()
+        with ExitStack() as stack:
+            stack.enter_context(
+                patch_ctx(f"{endpoints}.ModelManagementAuthChecks.allow_team_model_action", AsyncMock(return_value=True))
+            )
+            stack.enter_context(patch_ctx(f"{endpoints}._setup_new_team_model_assignment", setup_new))
+            stack.enter_context(patch_ctx(f"{endpoints}._update_existing_team_model_assignment", update_existing))
+            stack.enter_context(patch_ctx("litellm.proxy.proxy_server.premium_user", True))
+            with pytest.raises(HTTPException) as exc:
+                await _update_team_model_in_db(
+                    db_model=db_model,
+                    patch_data=patch,
+                    user_api_key_dict=UserAPIKeyAuth(user_id="a", user_role=LitellmUserRoles.PROXY_ADMIN),
+                    prisma_client=MagicMock(),
+                )
+
+        assert exc.value.status_code == 400
+        setup_new.assert_not_called()
+        update_existing.assert_not_called()
+
+    def test_a_setting_that_is_not_a_charge_is_left_alone(self):
+        """CustomPricingLiteLLMParams also carries an embedding's output vector size and the
+        regional uplift multipliers. Zeroing one of those destroys the deployment's config, and
+        refusing it answers with a message calling a setting a charge."""
+        priced = _ptu_priced_deployment(
+            Deployment(
+                model_name="embeddings",
+                litellm_params=LiteLLM_Params(
+                    model="azure/text-embedding-3-large",
+                    output_vector_size=1536,
+                    regional_processing_uplift_multiplier_eu=1.15,
+                ),
+                model_info=ModelInfo(
+                    id="dep-emb",
+                    team_id="team-1",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                ),
+            )
+        )
+        assert priced.litellm_params.get("output_vector_size") == 1536
+        assert priced.litellm_params.get("regional_processing_uplift_multiplier_eu") == 1.15
+        assert priced.litellm_params.get("input_cost_per_token") == 0
+
+    def test_removing_ptu_config_releases_every_rate_it_zeroed(self):
+        """The zeroing covers any stored rate, so a release that only spans the mirrored fields
+        leaves a per-second deployment billing nothing for that dimension forever."""
+        on = update_db_model(
+            db_model=Deployment(
+                model_name="audio",
+                litellm_params=LiteLLM_Params(model="azure/whisper", input_cost_per_second=0.006),
+                model_info=ModelInfo(id="dep-audio", team_id="t"),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(
+                    id="dep-audio",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                )
+            ),
+        )
+        assert json.loads(on["litellm_params"])["input_cost_per_second"] == 0
+
+        off = update_db_model(
+            db_model=Deployment(
+                model_name="audio",
+                litellm_params=LiteLLM_Params(**json.loads(on["litellm_params"])),
+                model_info=ModelInfo(**json.loads(on["model_info"])),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(id="dep-audio", ptu_count=None, cost_per_ptu_per_hour=None)
+            ),
+        )
+        assert "input_cost_per_second" not in json.loads(off["litellm_params"])
+
+    @pytest.mark.parametrize(
+        "backend", ["azure/gpt-4o", "anthropic/claude-sonnet-4-5", "bedrock/anthropic.claude-sonnet-4-20250514-v1:0"]
+    )
+    def test_the_cost_map_contributes_no_price_to_a_priced_ptu_deployment(self, backend):
+        """The acceptance criterion, read off the entry the router registers for the deployment.
+
+        Zeroing only the per-token pair leaves the cache-tier fields unset, which is exactly what
+        Router._inherit_builtin_cache_pricing back-fills from the public cost map, so a cached
+        prompt would still be billed at the public rate."""
+        priced = _ptu_priced_deployment(
+            Deployment(
+                model_name="ptu-deployment",
+                litellm_params=LiteLLM_Params(model=backend, api_key="fake-key"),
+                model_info=ModelInfo(
+                    id="dep-ptu",
+                    team_id="team-1",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                ),
+            )
+        )
+        registered = Router._deployment_model_cost_payload(priced)
+        charged = {k: v for k, v in registered.items() if "cost" in k and k != "cost_per_ptu_per_hour" and v}
+        assert charged == {}
+
+    def test_the_zeroed_pricing_does_not_waive_budget_enforcement(self):
+        """A zero price otherwise tells auth the model is free and skips every budget check."""
+        priced = _ptu_priced_deployment(
+            Deployment(
+                model_name="model_name_team-1_dep-ptu",
+                litellm_params=LiteLLM_Params(model="gemini/gemini-2.5-flash", api_key="fake-key"),
+                model_info=ModelInfo(
+                    id="dep-ptu",
+                    team_id="team-1",
+                    team_public_model_name="ptu-model",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                ),
+            )
+        )
+        router = Router(model_list=[priced.to_json(exclude_none=True)])
+        assert _is_model_cost_zero(model="model_name_team-1_dep-ptu", llm_router=router) is False
+        assert _is_model_cost_zero(model="ptu-model", llm_router=router) is False
+
+    def test_an_unrelated_patch_heals_a_deployment_stored_before_this_rule(self):
+        """Both blobs, because litellm_params wins over model_info wherever the two are merged."""
+        written = update_db_model(
+            db_model=_deployment_with_stored_ptu(),
+            updated_patch=updateDeployment(model_name="gpt-4o-renamed"),
+        )
+        for blob in ("model_info", "litellm_params"):
+            stored = json.loads(written[blob])
+            assert all(stored[field] == 0 for field in _PTU_ZEROED_PRICING_FIELDS), blob
+
+    def test_an_unrelated_patch_of_a_ptu_row_that_carries_a_price_is_not_refused(self):
+        """The pause toggle and the credential-rotation modal send no pricing at all. Refusing
+        them because the stored row is mispriced blocks flows that cannot fix it."""
+        priced_ptu = Deployment(
+            model_name="gpt-4o",
+            litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=5e-07),
+            model_info=ModelInfo(
+                id="dep-0",
+                team_id="t",
+                ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                **self.PTU,
+            ),
+        )
+        written = update_db_model(db_model=priced_ptu, updated_patch=updateDeployment(model_name="renamed"))
+        assert written["model_name"] == "renamed"
+        assert json.loads(written["litellm_params"])["input_cost_per_token"] == 0
+
+    def test_removing_ptu_config_hands_per_token_billing_back(self):
+        """Left behind, the zeros this rule wrote would serve the deployment for free forever."""
+        zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
+        written = update_db_model(
+            db_model=Deployment(
+                model_name="gpt-4o",
+                litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros),
+                model_info=ModelInfo(
+                    id="dep-0",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                    **zeros,
+                ),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None)
+            ),
+        )
+        for blob in ("model_info", "litellm_params"):
+            stored = json.loads(written[blob])
+            assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS), blob
+
+    def test_the_dashboard_clear_releases_the_zeros_it_echoes_back(self):
+        """The edit form re-sends the whole stored model_info on every save, so the clearing
+        patch carries the zeros this rule wrote. Treating those as a rate the operator chose
+        left the deployment serving free and reading as a free model to the budget checks."""
+        zeros = dict.fromkeys(_PTU_ZEROED_PRICING_FIELDS, 0.0)
+        written = update_db_model(
+            db_model=Deployment(
+                model_name="gpt-4o",
+                litellm_params=LiteLLM_Params(model="openai/gpt-4o", **zeros),
+                model_info=ModelInfo(
+                    id="dep-0",
+                    team_id="t",
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                    **zeros,
+                ),
+            ),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(id="dep-0", ptu_count=None, cost_per_ptu_per_hour=None, **zeros)
+            ),
+        )
+        stored = json.loads(written["model_info"])
+        assert not any(field in stored for field in _PTU_ZEROED_PRICING_FIELDS)
+
+    def test_a_deployment_that_never_had_ptu_keeps_a_price_its_operator_set_to_zero(self):
+        """The dashboard sends both PTU keys as null on every save while the feature is on, so a
+        release keyed on the patch alone would strip a deliberate zero rate from any model."""
+        free = Deployment(
+            model_name="free-model",
+            litellm_params=LiteLLM_Params(model="openai/gpt-4o", input_cost_per_token=0.0),
+            model_info=ModelInfo(id="dep-free", team_id="t", input_cost_per_token=0.0),
+        )
+        written = update_db_model(
+            db_model=free,
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(id="dep-free", ptu_count=None, cost_per_ptu_per_hour=None)
+            ),
+        )
+        for blob in ("model_info", "litellm_params"):
+            assert json.loads(written[blob])["input_cost_per_token"] == 0, blob
+
+    def test_a_patch_pricing_a_ptu_deployment_is_refused(self):
+        with pytest.raises(HTTPException) as exc:
+            update_db_model(
+                db_model=_deployment_with_stored_ptu(),
+                updated_patch=updateDeployment(
+                    litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07)
+                ),
+            )
+        assert exc.value.status_code == 400
+
+    def test_a_price_the_client_only_echoes_back_is_not_read_as_an_attempt_to_charge(self):
+        """/model/info fills missing rates from the public cost map and the edit form re-sends the
+        whole blob, so a model_info price is one the server wrote. Reading it as the operator's
+        refused every attempt to put an existing deployment on PTU from the dashboard."""
+        written = update_db_model(
+            db_model=_deployment_without_ptu(),
+            updated_patch=updateDeployment(
+                model_info=ModelInfo(
+                    id="dep-0",
+                    team_id="t",
+                    input_cost_per_token=3e-07,
+                    output_cost_per_token=2.5e-06,
+                    ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                    **self.PTU,
+                )
+            ),
+        )
+        stored = json.loads(written["model_info"])
+        assert stored["ptu_count"] == 15
+        assert stored["input_cost_per_token"] == 0
+        assert stored["output_cost_per_token"] == 0
+
+    def test_adding_ptu_config_to_an_already_priced_deployment_is_refused(self):
+        priced = Deployment(
+            model_name="gpt-4o",
+            litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
+            model_info=ModelInfo(id="dep-0", team_id="t"),
+        )
+        with pytest.raises(HTTPException) as exc:
+            update_db_model(
+                db_model=priced,
+                updated_patch=updateDeployment(
+                    litellm_params=updateLiteLLMParams(model="openai/gpt-4o", input_cost_per_token=5e-07),
+                    model_info=ModelInfo(
+                        id="dep-0",
+                        team_id="t",
+                        ptu_effective_from=datetime.datetime(2020, 1, 1, tzinfo=datetime.timezone.utc),
+                        **self.PTU,
+                    ),
+                ),
+            )
+        assert exc.value.status_code == 400
+
+    def test_a_deployment_without_ptu_config_keeps_its_pricing_through_a_patch(self):
+        priced = Deployment(
+            model_name="gpt-4o",
+            litellm_params=LiteLLM_Params(model="openai/gpt-4o"),
+            model_info=ModelInfo(id="dep-0", team_id="t", input_cost_per_token=5e-07),
+        )
+        stored = json.loads(
+            update_db_model(db_model=priced, updated_patch=updateDeployment(model_name="renamed"))["model_info"]
+        )
+        assert stored["input_cost_per_token"] == 5e-07
+
+    @pytest.mark.asyncio
+    async def test_model_new_stores_zero_pricing_on_both_blobs(self):
+        (_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model")
+        admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+
+        with ExitStack() as stack:
+            for active_patch in patches:
+                stack.enter_context(active_patch)
+            await add_new_model(
+                model_params=TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model"),
+                user_api_key_dict=admin,
+            )
+
+        written = add_team_model_to_db.call_args.kwargs["model_params"]
+        assert all(getattr(written.model_info, field, None) == 0 for field in SPECIAL_MODEL_INFO_PARAMS)
+        assert all(written.litellm_params.get(field) == 0 for field in _PTU_ZEROED_PRICING_FIELDS)
+
+    @pytest.mark.asyncio
+    async def test_model_new_refuses_a_priced_ptu_deployment(self):
+        (_, add_team_model_to_db), patches = TestAddNewModelPtuGate._patched_proxy("ptu-priced-model")
+        admin = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
+        base = TestAddNewModelPtuGate._ptu_deployment("ptu-priced-model")
+        deployment = base.model_copy(
+            update={"litellm_params": base.litellm_params.model_copy(update={"input_cost_per_token": 5e-07})}
+        )
+
+        with ExitStack() as stack:
+            for active_patch in patches:
+                stack.enter_context(active_patch)
+            with pytest.raises(Exception) as exc:
+                await add_new_model(model_params=deployment, user_api_key_dict=admin)
+
+        assert "input_cost_per_token" in str(exc.value)
+        add_team_model_to_db.assert_not_called()
diff --git a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
index d17f6293cc3..736fc13d137 100644
--- a/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
+++ b/tests/test_litellm/proxy/spend_tracking/test_ptu_flat_cost_rollup.py
@@ -657,6 +657,87 @@ async def test_scheduled_rollup_stays_quiet_when_every_charge_landed():
     alert.assert_not_awaited()
 
 
+@pytest.mark.asyncio
+async def test_scheduled_rollup_alerts_once_a_ptu_window_has_closed():
+    """Reserved capacity is billed until the deployment is deleted, so a closed window stops
+    the attribution without stopping the charge. Nobody notices unless it is escalated."""
+    ptu = {
+        "ptu_count": 5,
+        "cost_per_ptu_per_hour": 2.0,
+        "team_id": "t",
+        "ptu_effective_from": "2020-01-01T00:00:00Z",
+        "ptu_effective_to": "2020-02-01T00:00:00Z",
+    }
+    prisma, _ = _prisma_with_models([_model_row(model_id="dep-lapsed", model_info=ptu)])
+    alert = AsyncMock()
+
+    result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
+
+    assert result.lapsed == ("gpt-4o-mini-ptu",)
+    alert.assert_awaited_once()
+    message = alert.await_args.args[0]
+    assert "window has closed" in message
+    assert "gpt-4o-mini-ptu" in message
+
+
+@pytest.mark.asyncio
+async def test_a_model_name_cannot_smuggle_slack_markup_into_the_alert():
+    """The alert lands in an operator channel and a model name is operator-supplied, so an
+    unescaped name could post a channel-wide mention."""
+    ptu = {
+        "ptu_count": 5,
+        "cost_per_ptu_per_hour": 2.0,
+        "team_id": "t",
+        "ptu_effective_from": "2020-01-01T00:00:00Z",
+        "ptu_effective_to": "2020-02-01T00:00:00Z",
+    }
+    row = _model_row(model_id="dep-x", model_name="<!channel> & <https://evil.example|click>", model_info=ptu)
+    prisma, _ = _prisma_with_models([row])
+    alert = AsyncMock()
+
+    await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
+
+    message = alert.await_args.args[0]
+    assert "<!channel>" not in message
+    assert "&lt;!channel&gt;" in message
+
+
+@pytest.mark.asyncio
+async def test_an_open_ptu_window_raises_no_lapsed_alert():
+    ptu = {
+        "ptu_count": 5,
+        "cost_per_ptu_per_hour": 2.0,
+        "team_id": "t",
+        "ptu_effective_from": "2020-01-01T00:00:00Z",
+        "ptu_effective_to": "2999-01-01T00:00:00Z",
+    }
+    prisma, _ = _prisma_with_models([_model_row(model_id="dep-open", model_info=ptu)])
+    alert = AsyncMock()
+
+    result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
+
+    assert result.lapsed == ()
+    alert.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_an_open_ended_ptu_window_raises_no_lapsed_alert():
+    """No end bound means the operator never asked the attribution to stop."""
+    ptu = {
+        "ptu_count": 5,
+        "cost_per_ptu_per_hour": 2.0,
+        "team_id": "t",
+        "ptu_effective_from": "2020-01-01T00:00:00Z",
+    }
+    prisma, _ = _prisma_with_models([_model_row(model_id="dep-forever", model_info=ptu)])
+    alert = AsyncMock()
+
+    result = await run_scheduled_ptu_rollup(prisma, target_date=DAY, alert=alert)
+
+    assert result.lapsed == ()
+    alert.assert_not_awaited()
+
+
 @pytest.mark.asyncio
 async def test_a_broken_alert_channel_does_not_fail_the_rollup():
     rows = [_model_row(model_info={"ptu_count": 5, "cost_per_ptu_per_hour": 2.0, "team_id": "t"})]
diff --git a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
index 5d7196a84fe..df20924d321 100644
--- a/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
+++ b/ui/litellm-dashboard/src/components/add_model/advanced_settings.tsx
@@ -14,6 +14,7 @@ import {
   PTU_RATE_FIELD,
   PTU_START_FIELD,
   ptuCountRules,
+  ptuNoUsageCostRule,
   ptuPairRule,
   ptuRateRules,
   ptuStartRequiredRule,
@@ -261,7 +262,8 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
                     <Form.Item
                       label="Input Cost (per 1M tokens)"
                       name="input_cost_per_token"
-                      rules={[{ validator: validateNumber }]}
+                      dependencies={[PTU_COUNT_FIELD]}
+                      rules={[{ validator: validateNumber }, ptuNoUsageCostRule(PTU_COUNT_FIELD)]}
                       className="mb-4"
                     >
                       <TextInput />
@@ -269,7 +271,8 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
                     <Form.Item
                       label="Output Cost (per 1M tokens)"
                       name="output_cost_per_token"
-                      rules={[{ validator: validateNumber }]}
+                      dependencies={[PTU_COUNT_FIELD]}
+                      rules={[{ validator: validateNumber }, ptuNoUsageCostRule(PTU_COUNT_FIELD)]}
                       className="mb-4"
                     >
                       <TextInput />
@@ -277,7 +280,8 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
                     <Form.Item
                       label="Cache Read Cost (per 1M tokens)"
                       name="cache_read_input_token_cost"
-                      rules={[{ validator: validateNumber }]}
+                      dependencies={[PTU_COUNT_FIELD]}
+                      rules={[{ validator: validateNumber }, ptuNoUsageCostRule(PTU_COUNT_FIELD)]}
                       tooltip="If left blank, defaults to Input Cost."
                       className="mb-4"
                     >
@@ -286,7 +290,8 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
                     <Form.Item
                       label="Cache Write Cost (per 1M tokens)"
                       name="cache_creation_input_token_cost"
-                      rules={[{ validator: validateNumber }]}
+                      dependencies={[PTU_COUNT_FIELD]}
+                      rules={[{ validator: validateNumber }, ptuNoUsageCostRule(PTU_COUNT_FIELD)]}
                       tooltip="If left blank, defaults to Input Cost (the backend falls back to input_cost_per_token when no cache-write rate is set)."
                       className="mb-4"
                     >
@@ -297,7 +302,8 @@ const AdvancedSettings: React.FC<AdvancedSettingsProps> = ({
                   <Form.Item
                     label="Cost Per Second"
                     name="input_cost_per_second"
-                    rules={[{ validator: validateNumber }]}
+                    dependencies={[PTU_COUNT_FIELD]}
+                    rules={[{ validator: validateNumber }, ptuNoUsageCostRule(PTU_COUNT_FIELD)]}
                     className="mb-4"
                   >
                     <TextInput />
diff --git a/ui/litellm-dashboard/src/components/model_info_view.test.tsx b/ui/litellm-dashboard/src/components/model_info_view.test.tsx
index a3ce40494eb..5eb17fee2c4 100644
--- a/ui/litellm-dashboard/src/components/model_info_view.test.tsx
+++ b/ui/litellm-dashboard/src/components/model_info_view.test.tsx
@@ -617,9 +617,14 @@ describe("ModelInfoView", () => {
   describe("PTU cost attribution gate", () => {
     const ptuModelData = {
       ...defaultModelData,
+      // Zero per-token pricing is what the backend stores for a PTU deployment, since the flat
+      // cost of its reserved capacity already covers the traffic that capacity serves.
+      litellm_params: { ...defaultModelData.litellm_params, input_cost_per_token: 0, output_cost_per_token: 0 },
       model_info: {
         ...defaultModelData.model_info,
         team_id: "team-1",
+        input_cost_per_token: 0,
+        output_cost_per_token: 0,
         ptu_count: 15,
         cost_per_ptu_per_hour: 2,
         ptu_effective_from: "2026-07-01T00:00:00+00:00",
@@ -683,6 +688,70 @@ describe("ModelInfoView", () => {
       expect(modelInfo).not.toHaveProperty("ptu_effective_to");
     });
 
+    it("shows a zeroed PTU price as 0.0000 rather than Not Set", async () => {
+      mockUsePtuCostAttributionEnabled.mockReturnValue(true);
+      renderWithPtuModel();
+
+      await waitFor(() => {
+        expect(screen.getByText("Input Cost (per 1M tokens)")).toBeInTheDocument();
+      });
+      for (const label of ["Input Cost (per 1M tokens)", "Output Cost (per 1M tokens)"]) {
+        expect(screen.getByText(label).parentElement).toHaveTextContent("0.0000");
+      }
+    });
+
+    it("blocks the save once the operator types a non-zero per-token cost alongside PTU config", async () => {
+      mockUsePtuCostAttributionEnabled.mockReturnValue(true);
+      const user = userEvent.setup();
+      renderWithPtuModel();
+
+      await waitFor(() => {
+        expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
+      });
+      await user.click(screen.getByRole("button", { name: /edit settings/i }));
+
+      await waitFor(() => {
+        expect(screen.getByPlaceholderText("Enter input cost")).toBeInTheDocument();
+      });
+      await user.clear(screen.getByPlaceholderText("Enter input cost"));
+      await user.type(screen.getByPlaceholderText("Enter input cost"), "2.5");
+      await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+      await waitFor(() => {
+        expect(screen.getByText(/bills by reserved capacity/i)).toBeInTheDocument();
+      });
+      expect(mockModelPatchUpdateCall).not.toHaveBeenCalled();
+    });
+
+    it("lets the operator put a cost-map-priced deployment on PTU without clearing the seeded rate", async () => {
+      // A rate the form seeded from /model/info is the server's own, so refusing it blocked
+      // every attempt to enable PTU from the dashboard.
+      mockUsePtuCostAttributionEnabled.mockReturnValue(true);
+      const seededModel = {
+        ...defaultModelData,
+        model_info: { ...defaultModelData.model_info, team_id: "team-1", input_cost_per_token: 0.0000003 },
+      };
+      mockUseModelsInfo.mockReturnValue({ data: { data: [seededModel] }, isLoading: false, error: null });
+      mockModelInfoV1Call.mockResolvedValue({ data: [seededModel] });
+      const user = userEvent.setup();
+      render(<ModelInfoView {...DEFAULT_ADMIN_PROPS} />, { wrapper });
+
+      await waitFor(() => {
+        expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument();
+      });
+      await user.click(screen.getByRole("button", { name: /edit settings/i }));
+
+      await waitFor(() => {
+        expect(screen.getByPlaceholderText("e.g. 15")).toBeInTheDocument();
+      });
+      await user.type(screen.getByPlaceholderText("e.g. 15"), "15");
+      await user.click(screen.getByRole("button", { name: /save changes/i }));
+
+      await waitFor(() => {
+        expect(screen.queryByText(/bills by reserved capacity/i)).not.toBeInTheDocument();
+      });
+    });
+
     it("sends the PTU fields on save when enabled", async () => {
       mockUsePtuCostAttributionEnabled.mockReturnValue(true);
       const user = userEvent.setup();
diff --git a/ui/litellm-dashboard/src/components/model_info_view.tsx b/ui/litellm-dashboard/src/components/model_info_view.tsx
index e1a4d4311ab..a0c545e17aa 100644
--- a/ui/litellm-dashboard/src/components/model_info_view.tsx
+++ b/ui/litellm-dashboard/src/components/model_info_view.tsx
@@ -25,6 +25,7 @@ import {
   PTU_COUNT_FIELD,
   PTU_RATE_FIELD,
   ptuCountRules,
+  ptuNoUsageCostRule,
   ptuPairRule,
   ptuRateRules,
   ptuStartRequiredRule,
@@ -128,6 +129,12 @@ const PTU_EDIT_FIELDS: PtuEditField[] = [
   },
 ];
 
+/** Per-1M-token rate for the first rate that is set, so a deliberate 0 seeds the form as 0. */
+const perMillionTokens = (...rates: (number | null | undefined)[]): number | null => {
+  const rate = rates.find((candidate) => candidate != null);
+  return rate == null ? null : rate * 1_000_000;
+};
+
 const ptuFieldDependencies = ({ isStart, pairedWith, windowPeer }: PtuEditField): string[] | undefined => {
   const deps = [
     ...(isStart ? [PTU_COUNT_FIELD] : []),
@@ -227,6 +234,8 @@ export default function ModelInfoView({
   const { data: modelHubData } = useModelHub();
   const { data: teams } = useTeams();
   const ptuCostAttributionEnabled = usePtuCostAttributionEnabled();
+  const ptuCostRule = (field: string) =>
+    ptuCostAttributionEnabled ? [ptuNoUsageCostRule(PTU_COUNT_FIELD, field)] : [];
 
   // Transform the model data
   const getProviderFromModel = (model: string) => {
@@ -835,12 +844,14 @@ export default function ModelInfoView({
                     max_retries: localModelData.litellm_params.max_retries,
                     timeout: localModelData.litellm_params.timeout,
                     stream_timeout: localModelData.litellm_params.stream_timeout,
-                    input_cost: localModelData.litellm_params.input_cost_per_token
-                      ? localModelData.litellm_params.input_cost_per_token * 1_000_000
-                      : localModelData.model_info?.input_cost_per_token * 1_000_000 || null,
-                    output_cost: localModelData.litellm_params?.output_cost_per_token
-                      ? localModelData.litellm_params.output_cost_per_token * 1_000_000
-                      : localModelData.model_info?.output_cost_per_token * 1_000_000 || null,
+                    input_cost: perMillionTokens(
+                      localModelData.litellm_params.input_cost_per_token,
+                      localModelData.model_info?.input_cost_per_token,
+                    ),
+                    output_cost: perMillionTokens(
+                      localModelData.litellm_params?.output_cost_per_token,
+                      localModelData.model_info?.output_cost_per_token,
+                    ),
                     ptu_count: localModelData.model_info?.ptu_count ?? null,
                     cost_per_ptu_per_hour: localModelData.model_info?.cost_per_ptu_per_hour ?? null,
                     ptu_effective_from: utcIsoToPickerValue(localModelData.model_info?.ptu_effective_from),
@@ -917,14 +928,19 @@ export default function ModelInfoView({
                       <div>
                         <Text className="font-medium">Input Cost (per 1M tokens)</Text>
                         {isEditing ? (
-                          <Form.Item name="input_cost" className="mb-0">
+                          <Form.Item
+                            name="input_cost"
+                            className="mb-0"
+                            dependencies={[PTU_COUNT_FIELD]}
+                            rules={ptuCostRule("input_cost")}
+                          >
                             <NumericalInput placeholder="Enter input cost" />
                           </Form.Item>
                         ) : (
                           <div className="mt-1 p-2 bg-gray-50 rounded-sm">
-                            {localModelData?.litellm_params?.input_cost_per_token
-                              ? (localModelData.litellm_params?.input_cost_per_token * 1_000_000).toFixed(4)
-                              : localModelData?.model_info?.input_cost_per_token
+                            {localModelData?.litellm_params?.input_cost_per_token != null
+                              ? (localModelData.litellm_params.input_cost_per_token * 1_000_000).toFixed(4)
+                              : localModelData?.model_info?.input_cost_per_token != null
                                 ? (localModelData.model_info.input_cost_per_token * 1_000_000).toFixed(4)
                                 : "Not Set"}
                           </div>
@@ -934,14 +950,19 @@ export default function ModelInfoView({
                       <div>
                         <Text className="font-medium">Output Cost (per 1M tokens)</Text>
                         {isEditing ? (
-                          <Form.Item name="output_cost" className="mb-0">
+                          <Form.Item
+                            name="output_cost"
+                            className="mb-0"
+                            dependencies={[PTU_COUNT_FIELD]}
+                            rules={ptuCostRule("output_cost")}
+                          >
                             <NumericalInput placeholder="Enter output cost" />
                           </Form.Item>
                         ) : (
                           <div className="mt-1 p-2 bg-gray-50 rounded-sm">
-                            {localModelData?.litellm_params?.output_cost_per_token
+                            {localModelData?.litellm_params?.output_cost_per_token != null
                               ? (localModelData.litellm_params.output_cost_per_token * 1_000_000).toFixed(4)
-                              : localModelData?.model_info?.output_cost_per_token
+                              : localModelData?.model_info?.output_cost_per_token != null
                                 ? (localModelData.model_info.output_cost_per_token * 1_000_000).toFixed(4)
                                 : "Not Set"}
                           </div>
@@ -995,6 +1016,8 @@ export default function ModelInfoView({
                           <Form.Item
                             name="cache_read_cost"
                             className="mb-0"
+                            dependencies={[PTU_COUNT_FIELD]}
+                            rules={ptuCostRule("cache_read_cost")}
                             tooltip="If left blank on save, defaults to Input Cost."
                           >
                             <NumericalInput placeholder="Defaults to Input Cost if blank" />
@@ -1018,6 +1041,8 @@ export default function ModelInfoView({
                           <Form.Item
                             name="cache_write_cost"
                             className="mb-0"
+                            dependencies={[PTU_COUNT_FIELD]}
+                            rules={ptuCostRule("cache_write_cost")}
                             tooltip="If left blank on save, defaults to Input Cost (backend falls back to input_cost_per_token)."
                           >
                             <NumericalInput placeholder="Defaults to Input Cost if blank" />
diff --git a/ui/litellm-dashboard/src/utils/ptuValidation.test.ts b/ui/litellm-dashboard/src/utils/ptuValidation.test.ts
index 1338c2a39ae..ed30e6a7b20 100644
--- a/ui/litellm-dashboard/src/utils/ptuValidation.test.ts
+++ b/ui/litellm-dashboard/src/utils/ptuValidation.test.ts
@@ -3,6 +3,7 @@ import {
   PTU_COUNT_FIELD,
   PTU_RATE_FIELD,
   ptuCountRules,
+  ptuNoUsageCostRule,
   ptuPairRule,
   ptuRateRules,
   ptuStartRequiredRule,
@@ -40,6 +41,39 @@ describe("ptuCountRules", () => {
   });
 });
 
+describe("ptuNoUsageCostRule", () => {
+  const check = (value: unknown, count: unknown) =>
+    ptuNoUsageCostRule(PTU_COUNT_FIELD)({ getFieldValue: () => count }).validator(null, value);
+
+  it("leaves a deployment without PTU config free to carry any price", async () => {
+    await expect(check("2.5", "")).resolves.toBeUndefined();
+    await expect(check(0.5, null)).resolves.toBeUndefined();
+  });
+
+  it("accepts a blank or zero price alongside PTU config, which is what the backend stores", async () => {
+    await expect(check("", 15)).resolves.toBeUndefined();
+    await expect(check(null, 15)).resolves.toBeUndefined();
+    await expect(check(0, 15)).resolves.toBeUndefined();
+    await expect(check("0", 15)).resolves.toBeUndefined();
+  });
+
+  it("rejects a non-zero price alongside PTU config, which the backend answers with a 400", async () => {
+    await expect(check("2.5", 15)).rejects.toThrow("bills by reserved capacity");
+    await expect(check(0.000001, 15)).rejects.toThrow("bills by reserved capacity");
+  });
+
+  it("reads the count by the field name it was given", () => {
+    const seen: string[] = [];
+    ptuNoUsageCostRule(PTU_COUNT_FIELD)({
+      getFieldValue: (name: string) => {
+        seen.push(name);
+        return 15;
+      },
+    }).validator(null, 0);
+    expect(seen).toEqual([PTU_COUNT_FIELD]);
+  });
+});
+
 describe("ptuPairRule", () => {
   const rule = (sibling: unknown) => ptuPairRule(PTU_RATE_FIELD)({ getFieldValue: () => sibling });
   const check = (value: unknown, sibling: unknown) => rule(sibling).validator(null, value);
diff --git a/ui/litellm-dashboard/src/utils/ptuValidation.ts b/ui/litellm-dashboard/src/utils/ptuValidation.ts
index 9fabf39ae23..da367f233b6 100644
--- a/ui/litellm-dashboard/src/utils/ptuValidation.ts
+++ b/ui/litellm-dashboard/src/utils/ptuValidation.ts
@@ -4,6 +4,7 @@ interface ValidatorRule {
 
 interface FormInstance {
   getFieldValue: (name: string) => unknown;
+  isFieldTouched?: (name: string) => boolean;
 }
 
 export const PTU_COUNT_FIELD = "ptu_count";
@@ -71,6 +72,25 @@ export const ptuPairRule =
         : Promise.reject(new Error("PTU Count and Cost per PTU / Hour must be set together")),
   });
 
+/**
+ * A PTU deployment is billed by the flat cost of its reserved capacity, so the backend refuses
+ * a non-zero per-token price alongside PTU config and stores 0 when none is given. Pair this
+ * with `dependencies` on the count so the error clears once the price or the PTU config goes.
+ */
+export const ptuNoUsageCostRule =
+  (countField: string, thisField?: string) =>
+  ({ getFieldValue, isFieldTouched }: FormInstance): ValidatorRule => ({
+    validator: (_, value) => {
+      // A cost the operator never typed was seeded from the rate /model/info resolved, which
+      // for an unpriced deployment is the public cost map. Refusing it would block every
+      // attempt to put an existing deployment on PTU, and the save omits it anyway.
+      const echoed = thisField !== undefined && isFieldTouched !== undefined && !isFieldTouched(thisField);
+      return echoed || !isFilled(getFieldValue(countField)) || !isFilled(value) || Number(value) === 0
+        ? Promise.resolve()
+        : Promise.reject(new Error("A PTU deployment bills by reserved capacity, so this cost must be 0 or blank"));
+    },
+  });
+
 /**
  * The backend requires an effective start whenever PTU is configured, since flat cost
  * accrues from that instant and an inferred start would bill days a deployment did not

From 7a519e26ec8707f45762eb7bbb9e0b0e7173b269 Mon Sep 17 00:00:00 2001
From: lostmartian <62000079+lostmartian@users.noreply.github.com>
Date: Fri, 14 Aug 2026 09:18:16 +0530
Subject: [PATCH 30/31] fix(proxy): track spend for OpenAI passthrough
 /v1/embeddings (#36660)

* fix(proxy): track spend for OpenAI passthrough /v1/embeddings

OpenAI passthrough embeddings returned 200 but wrote no spend because the
route was unsupported and Cohere's /v1/embed prefix stole the match.

* fix(proxy): clear embeddings lint and Greptile comment nits

Inline embeddings cost tracking to avoid new LIT001/002 hits, trim
redundant doc comments, and cover the Cohere /v1/embeddings collision.

* fix(proxy): drop unreachable embeddings TypeError guard

convert_to_model_response_object with response_type=embedding already
returns EmbeddingResponse; the isinstance check was dead patch coverage.
---
 .../cohere_passthrough_logging_handler.py     |   2 +-
 .../openai_passthrough_logging_handler.py     |  40 +-
 .../pass_through_endpoints/success_handler.py |  11 +-
 ...test_cohere_passthrough_logging_handler.py |  48 +-
 ...test_openai_passthrough_logging_handler.py | 499 +++++++++---------
 5 files changed, 338 insertions(+), 262 deletions(-)

diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py
index 132af097a55..597c2e742b3 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/cohere_passthrough_logging_handler.py
@@ -82,7 +82,7 @@ class CoherePassthroughLoggingHandler(BasePassthroughLoggingHandler):
         Handle Cohere passthrough logging with route detection and cost tracking.
         """
         # Check if this is an embed endpoint
-        if "/v1/embed" in url_route:
+        if "/v1/embed" in url_route and "/v1/embeddings" not in url_route:
             model: Final = request_body.get("model", response_body.get("model", ""))
             try:
                 cohere_embed_config: Final = CohereEmbeddingConfig()
diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py
index 65ebc2728c6..1c8bce28454 100644
--- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py
+++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py
@@ -31,8 +31,8 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
     EndpointType,
     PassthroughStandardLoggingPayload,
 )
-from litellm.types.utils import ImageResponse, LlmProviders, PassthroughCallTypes
-from litellm.utils import ModelResponse, TextCompletionResponse
+from litellm.types.utils import EmbeddingResponse, ImageResponse, LlmProviders, PassthroughCallTypes
+from litellm.utils import ModelResponse, TextCompletionResponse, convert_to_model_response_object
 
 # Hostnames that route to OpenAI-compatible APIs.
 #
@@ -143,6 +143,14 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
             "/v1/responses" in parsed_url.path or "/responses" in parsed_url.path
         )
 
+    @staticmethod
+    def is_openai_embeddings_route(url_route: str) -> bool:
+        """Check if the URL route is an OpenAI embeddings endpoint."""
+        if not url_route:
+            return False
+        parsed_url: Final = urlparse(url_route)
+        return _is_openai_compatible_host(parsed_url.hostname) and "/v1/embeddings" in parsed_url.path
+
     def _get_user_from_metadata(
         self,
         passthrough_logging_payload: PassthroughStandardLoggingPayload,
@@ -271,22 +279,21 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
         **kwargs,
     ) -> PassThroughEndpointLoggingTypedDict:
         """
-        Handle OpenAI passthrough logging with cost tracking for chat completions, image generation, image editing, and responses API.
+        Handle OpenAI passthrough logging with cost tracking for chat completions,
+        embeddings, image generation, image editing, and responses API.
         """
-        # Check if this is a supported endpoint for cost tracking
         is_chat_completions: Final = OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
+        is_embeddings: Final = OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route)
         is_image_generation: Final = OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route)
         is_image_editing: Final = OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
         is_responses: Final = OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route)
 
-        if not (is_chat_completions or is_image_generation or is_image_editing or is_responses):
-            # For unsupported endpoints, return None to let the system fall back to generic behavior
+        if not (is_chat_completions or is_embeddings or is_image_generation or is_image_editing or is_responses):
             return {
                 "result": None,
                 "kwargs": kwargs,
             }
 
-        # Extract model from request or response
         model: Final = request_body.get("model", response_body.get("model", ""))
         if not model:
             verbose_proxy_logger.warning("No model found in request or response for OpenAI passthrough cost tracking")
@@ -307,7 +314,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
         try:
             response_cost = 0.0
             litellm_model_response: (
-                ModelResponse | TextCompletionResponse | ImageResponse | ResponsesAPIResponse | None
+                ModelResponse | TextCompletionResponse | EmbeddingResponse | ImageResponse | ResponsesAPIResponse | None
             ) = None
             handler_instance: Final = OpenAIPassthroughLoggingHandler()
 
@@ -338,6 +345,19 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
                     model=model,
                     custom_llm_provider=custom_llm_provider,
                 )
+            elif is_embeddings:
+                litellm_model_response = convert_to_model_response_object(
+                    response_object=response_body,
+                    model_response_object=EmbeddingResponse(),
+                    response_type="embedding",
+                )
+                response_cost = litellm.completion_cost(
+                    completion_response=litellm_model_response,
+                    model=model,
+                    custom_llm_provider=custom_llm_provider,
+                    call_type="aembedding",
+                )
+                litellm_model_response._hidden_params["response_cost"] = response_cost
             elif is_image_generation:
                 # Handle image generation cost calculation
                 response_cost = OpenAIPassthroughLoggingHandler._calculate_image_generation_cost(
@@ -432,9 +452,13 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler):
             endpoint_type: Final = (
                 "chat_completions"
                 if is_chat_completions
+                else "embeddings"
+                if is_embeddings
                 else "image_generation"
                 if is_image_generation
                 else "image_editing"
+                if is_image_editing
+                else "responses"
             )
             verbose_proxy_logger.debug(
                 f"OpenAI passthrough cost tracking - Endpoint: {endpoint_type}, Model: {model}, Cost: ${response_cost:.6f}"
diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py
index 3dcc8257b82..34286b203c7 100644
--- a/litellm/proxy/pass_through_endpoints/success_handler.py
+++ b/litellm/proxy/pass_through_endpoints/success_handler.py
@@ -349,10 +349,14 @@ class PassThroughEndpointLogging:
                 return True
         return False
 
-    def is_cohere_route(self, url_route: str):
+    def is_cohere_route(self, url_route: str) -> bool:
         for route in self.TRACKED_COHERE_ROUTES:
-            if route in url_route:
-                return True
+            if route not in url_route:
+                continue
+            if route == "/v1/embed" and "/v1/embeddings" in url_route:
+                continue
+            return True
+        return False
 
     def is_assemblyai_route(self, url_route: str):
         parsed_url: Final = urlparse(url_route)
@@ -429,6 +433,7 @@ class PassThroughEndpointLogging:
 
         return (
             OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(url_route)
+            or OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(url_route)
             or OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(url_route)
             or OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(url_route)
             or OpenAIPassthroughLoggingHandler.is_openai_responses_route(url_route)
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py
index 887aedaf0aa..6d7011fe10c 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_cohere_passthrough_logging_handler.py
@@ -7,9 +7,7 @@ from unittest.mock import MagicMock, patch
 import httpx
 import pytest
 
-sys.path.insert(
-    0, os.path.abspath("../../..")
-)  # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../../.."))  # Adds the parent directory to the system path
 
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.cohere_passthrough_logging_handler import (
@@ -69,12 +67,8 @@ class TestCoherePassthroughLoggingHandler:
         )
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    @patch(
-        "litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response"
-    )
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    @patch("litellm.llms.cohere.embed.v1_transformation.CohereEmbeddingConfig._transform_response")
     def test_cohere_embed_passthrough_cost_tracking(
         self, mock_transform_response, mock_get_standard_logging, mock_completion_cost
     ):
@@ -92,9 +86,7 @@ class TestCoherePassthroughLoggingHandler:
         mock_embedding_response.object = "list"
         from litellm.types.utils import Usage
 
-        mock_embedding_response.usage = Usage(
-            prompt_tokens=3, completion_tokens=0, total_tokens=3
-        )
+        mock_embedding_response.usage = Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3)
 
         mock_transform_response.return_value = mock_embedding_response
         mock_completion_cost.return_value = 3.6e-07  # Expected cost for embed-v4.0
@@ -151,6 +143,38 @@ class TestCoherePassthroughLoggingHandler:
         assert hasattr(result["result"], "model")
         assert result["result"].model == "embed-english-v3.0"
 
+    @patch(
+        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler.BasePassthroughLoggingHandler.passthrough_chat_handler"
+    )
+    @patch("litellm.completion_cost")
+    def test_openai_embeddings_route_does_not_use_cohere_embed_path(self, mock_completion_cost, mock_chat_handler):
+        mock_chat_handler.return_value = {"result": None, "kwargs": {}}
+        response_body = {
+            "object": "list",
+            "model": "text-embedding-3-small",
+            "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
+            "usage": {"prompt_tokens": 6, "total_tokens": 6},
+        }
+        result = self.handler.cohere_passthrough_handler(
+            httpx_response=self._create_mock_httpx_response(response_body),
+            response_body=response_body,
+            logging_obj=self._create_mock_logging_obj(),
+            url_route="https://api.openai.com/v1/embeddings",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"},
+            passthrough_logging_payload=PassthroughStandardLoggingPayload(
+                url="https://api.openai.com/v1/embeddings",
+                request_body={"model": "text-embedding-3-small", "input": "PROOF_SENTINEL_TEXT"},
+                request_method="POST",
+            ),
+        )
+        mock_completion_cost.assert_not_called()
+        mock_chat_handler.assert_called_once()
+        assert result == {"result": None, "kwargs": {}}
+
 
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py
index 05051ab3745..664015003e4 100644
--- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py
+++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py
@@ -8,9 +8,7 @@ from unittest.mock import AsyncMock, MagicMock, patch
 import httpx
 import pytest
 
-sys.path.insert(
-    0, os.path.abspath("../../..")
-)  # Adds the parent directory to the system path
+sys.path.insert(0, os.path.abspath("../../.."))  # Adds the parent directory to the system path
 
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
@@ -70,9 +68,7 @@ class TestOpenAIPassthroughLoggingHandler:
         mock_response.headers = {"content-type": "application/json"}
         return mock_response
 
-    def _create_passthrough_logging_payload(
-        self, user: str = "test_user"
-    ) -> PassthroughStandardLoggingPayload:
+    def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
         """Create a mock passthrough logging payload"""
         return PassthroughStandardLoggingPayload(
             url="https://api.openai.com/v1/chat/completions",
@@ -113,9 +109,7 @@ class TestOpenAIPassthroughLoggingHandler:
 
         # Negative cases
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
-                "https://api.openai.com/v1/models"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.openai.com/v1/models")
             == False
         )
         assert (
@@ -125,15 +119,10 @@ class TestOpenAIPassthroughLoggingHandler:
             == False
         )
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
-                "https://api.anthropic.com/v1/messages"
-            )
-            == False
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("")
+            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("https://api.anthropic.com/v1/messages")
             == False
         )
+        assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route("") == False
 
     def test_is_openai_image_generation_route(self):
         """Test OpenAI image generation route detection"""
@@ -159,9 +148,7 @@ class TestOpenAIPassthroughLoggingHandler:
             == False
         )
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(
-                "https://api.openai.com/v1/images/edits"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("https://api.openai.com/v1/images/edits")
             == False
         )
         assert (
@@ -170,32 +157,23 @@ class TestOpenAIPassthroughLoggingHandler:
             )
             == False
         )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("")
-            == False
-        )
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route("") == False
 
     def test_is_openai_image_editing_route(self):
         """Test OpenAI image editing route detection"""
         # Positive cases
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
-                "https://api.openai.com/v1/images/edits"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/images/edits")
             == True
         )
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
-                "https://openai.azure.com/v1/images/edits"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://openai.azure.com/v1/images/edits")
             == True
         )
 
         # Negative cases
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
-                "https://api.openai.com/v1/chat/completions"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("https://api.openai.com/v1/chat/completions")
             == False
         )
         assert (
@@ -210,118 +188,91 @@ class TestOpenAIPassthroughLoggingHandler:
             )
             == False
         )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False
-        )
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route("") == False
 
     def test_is_openai_responses_route(self):
         """Test OpenAI responses API route detection"""
         # Positive cases
+        assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/responses") == True
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "https://api.openai.com/v1/responses"
-            )
-            == True
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "https://openai.azure.com/v1/responses"
-            )
-            == True
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "https://api.openai.com/responses"
-            )
-            == True
+            OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://openai.azure.com/v1/responses") == True
         )
+        assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/responses") == True
 
         # Negative cases
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "https://api.openai.com/v1/chat/completions"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/chat/completions")
             == False
         )
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "https://api.openai.com/v1/images/generations"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_responses_route("https://api.openai.com/v1/images/generations")
             == False
         )
         assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                "http://localhost:4000/openai/v1/responses"
-            )
+            OpenAIPassthroughLoggingHandler.is_openai_responses_route("http://localhost:4000/openai/v1/responses")
             == False
         )
         assert OpenAIPassthroughLoggingHandler.is_openai_responses_route("") == False
 
+    def test_is_openai_embeddings_route(self):
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/embeddings") is True
+        )
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://openai.azure.com/v1/embeddings") is True
+        )
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
+                "https://my-resource.cognitiveservices.azure.com/v1/embeddings"
+            )
+            is True
+        )
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
+                "https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings"
+            )
+            is False
+        )
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("https://api.openai.com/v1/chat/completions")
+            is False
+        )
+        assert (
+            OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(
+                "http://localhost:4000/openai_passthrough/v1/embeddings"
+            )
+            is False
+        )
+        assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route("") is False
+
     def test_is_openai_route_recognizes_cognitiveservices_azure_com(self):
         """Azure OpenAI resources created via the newer "Azure AI Foundry" /
         Cognitive Services pathway live on `*.cognitiveservices.azure.com`
-        subdomains rather than the older `openai.azure.com`. All four
+        subdomains rather than the older `openai.azure.com`. The
         is_openai_*_route methods must recognize both Azure subdomains so
         cost tracking applies regardless of which Azure naming the user's
         resource happens to be on.
         """
-        cognitive_chat = (
-            "https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
-        )
-        cognitive_images_gen = (
-            "https://my-resource.cognitiveservices.azure.com/v1/images/generations"
-        )
-        cognitive_images_edit = (
-            "https://my-resource.cognitiveservices.azure.com/v1/images/edits"
-        )
-        cognitive_responses = (
-            "https://my-resource.cognitiveservices.azure.com/v1/responses"
-        )
+        cognitive_chat = "https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
+        cognitive_images_gen = "https://my-resource.cognitiveservices.azure.com/v1/images/generations"
+        cognitive_images_edit = "https://my-resource.cognitiveservices.azure.com/v1/images/edits"
+        cognitive_responses = "https://my-resource.cognitiveservices.azure.com/v1/responses"
+        cognitive_embeddings = "https://my-resource.cognitiveservices.azure.com/v1/embeddings"
 
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
-                cognitive_chat
-            )
-            is True
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(
-                cognitive_images_gen
-            )
-            is True
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(
-                cognitive_images_edit
-            )
-            is True
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(
-                cognitive_responses
-            )
-            is True
-        )
+        assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_chat) is True
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_generation_route(cognitive_images_gen) is True
+        assert OpenAIPassthroughLoggingHandler.is_openai_image_editing_route(cognitive_images_edit) is True
+        assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_responses) is True
+        assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_embeddings) is True
 
         # Cross-route negatives still hold for cognitiveservices hosts.
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(
-                cognitive_responses
-            )
-            is False
-        )
-        assert (
-            OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat)
-            is False
-        )
+        assert OpenAIPassthroughLoggingHandler.is_openai_chat_completions_route(cognitive_responses) is False
+        assert OpenAIPassthroughLoggingHandler.is_openai_responses_route(cognitive_chat) is False
+        assert OpenAIPassthroughLoggingHandler.is_openai_embeddings_route(cognitive_chat) is False
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    def test_openai_passthrough_handler_success(
-        self, mock_get_standard_logging, mock_completion_cost
-    ):
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    def test_openai_passthrough_handler_success(self, mock_get_standard_logging, mock_completion_cost):
         """Test successful cost tracking for OpenAI chat completions"""
         # Arrange
         mock_completion_cost.return_value = 0.000045
@@ -370,9 +321,7 @@ class TestOpenAIPassthroughLoggingHandler:
         assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
 
     @patch("litellm.completion_cost")
-    def test_openai_passthrough_handler_non_chat_completions(
-        self, mock_completion_cost
-    ):
+    def test_openai_passthrough_handler_non_chat_completions(self, mock_completion_cost):
         """Test that non-chat-completions routes fall back to base handler"""
         # Arrange
         mock_httpx_response = self._create_mock_httpx_response()
@@ -406,12 +355,8 @@ class TestOpenAIPassthroughLoggingHandler:
         # The important thing is that our specific OpenAI handler logic didn't run
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    def test_openai_passthrough_handler_with_user_tracking(
-        self, mock_get_standard_logging, mock_completion_cost
-    ):
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    def test_openai_passthrough_handler_with_user_tracking(self, mock_get_standard_logging, mock_completion_cost):
         """Test cost tracking with user information"""
         # Arrange
         mock_completion_cost.return_value = 0.000123
@@ -464,15 +409,10 @@ class TestOpenAIPassthroughLoggingHandler:
         assert "litellm_params" in result["kwargs"]
         assert "proxy_server_request" in result["kwargs"]["litellm_params"]
         assert "body" in result["kwargs"]["litellm_params"]["proxy_server_request"]
-        assert (
-            result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"]
-            == "test_user_123"
-        )
+        assert result["kwargs"]["litellm_params"]["proxy_server_request"]["body"]["user"] == "test_user_123"
 
     @patch("litellm.completion_cost")
-    def test_openai_passthrough_handler_cost_calculation_error(
-        self, mock_completion_cost
-    ):
+    def test_openai_passthrough_handler_cost_calculation_error(self, mock_completion_cost):
         """Test error handling in cost calculation"""
         # Arrange
         mock_completion_cost.side_effect = Exception("Cost calculation failed")
@@ -521,9 +461,7 @@ class TestOpenAIPassthroughLoggingHandler:
 
     @patch(f"{OpenAIPassthroughLoggingHandler.__module__}.get_standard_logging_object_payload")
     @patch("litellm.completion_cost", return_value=3.3e-06)
-    def test_streaming_responses_cost_uses_completed_response(
-        self, mock_completion_cost, mock_get_standard_logging
-    ):
+    def test_streaming_responses_cost_uses_completed_response(self, mock_completion_cost, mock_get_standard_logging):
         response_id = "resp_PROOFSENTINEL0123456789abcdef"
         completed_event = {
             "type": "response.completed",
@@ -796,12 +734,8 @@ class TestOpenAIPassthroughLoggingHandler:
         mock_completion_cost.assert_not_called()
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    def test_different_models_cost_tracking(
-        self, mock_get_standard_logging, mock_completion_cost
-    ):
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    def test_different_models_cost_tracking(self, mock_get_standard_logging, mock_completion_cost):
         """Test cost tracking for different OpenAI models"""
         # Arrange
         mock_get_standard_logging.return_value = {"test": "logging_payload"}
@@ -868,12 +802,8 @@ class TestOpenAIPassthroughLoggingHandler:
         assert handler.get_provider_config("gpt-4o") is not None
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    def test_azure_passthrough_tags_metadata_model_provider(
-        self, mock_get_standard_logging, mock_completion_cost
-    ):
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    def test_azure_passthrough_tags_metadata_model_provider(self, mock_get_standard_logging, mock_completion_cost):
         """Test that tags, metadata, model, and custom_llm_provider are preserved for Azure passthrough in UI"""
         # Arrange
         mock_completion_cost.return_value = 0.000045
@@ -929,9 +859,7 @@ class TestOpenAIPassthroughLoggingHandler:
 
         # Verify model and custom_llm_provider are set correctly
         assert result["kwargs"]["model"] == "gpt-4o"
-        assert (
-            result["kwargs"]["custom_llm_provider"] == "azure"
-        )  # Should preserve Azure, not default to "openai"
+        assert result["kwargs"]["custom_llm_provider"] == "azure"  # Should preserve Azure, not default to "openai"
         assert result["kwargs"]["response_cost"] == 0.000045
 
         # Verify metadata tags are preserved in litellm_params
@@ -955,12 +883,8 @@ class TestOpenAIPassthroughLoggingHandler:
         assert call_args[1]["custom_llm_provider"] == "azure"
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
-    @patch(
-        "litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response"
-    )
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    @patch("litellm.llms.openai.responses.transformation.OpenAIResponsesAPIConfig.transform_response_api_response")
     def test_responses_api_cost_tracking(
         self,
         mock_transform_responses,
@@ -1052,9 +976,7 @@ class TestOpenAIPassthroughLoggingHandler:
         assert mock_logging_obj.model_call_details["custom_llm_provider"] == "openai"
 
     @patch("litellm.completion_cost")
-    @patch(
-        "litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload"
-    )
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
     def test_responses_api_uses_responses_transformer_not_chat_completions(
         self, mock_get_standard_logging, mock_completion_cost
     ):
@@ -1185,9 +1107,7 @@ class TestOpenAIPassthroughIntegration:
         mock_response.headers = {"content-type": "application/json"}
         return mock_response
 
-    def _create_passthrough_logging_payload(
-        self, user: str = "test_user"
-    ) -> PassthroughStandardLoggingPayload:
+    def _create_passthrough_logging_payload(self, user: str = "test_user") -> PassthroughStandardLoggingPayload:
         """Create a mock passthrough logging payload"""
         return PassthroughStandardLoggingPayload(
             url="https://api.openai.com/v1/chat/completions",
@@ -1201,59 +1121,32 @@ class TestOpenAIPassthroughIntegration:
     def test_is_openai_route_detection(self):
         """Test OpenAI route detection in the main success handler"""
         # Positive cases
-        assert (
-            self.handler.is_openai_route("https://api.openai.com/v1/chat/completions")
-            == True
-        )
-        assert (
-            self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions")
-            == True
-        )
+        assert self.handler.is_openai_route("https://api.openai.com/v1/chat/completions") == True
+        assert self.handler.is_openai_route("https://openai.azure.com/v1/chat/completions") == True
         assert self.handler.is_openai_route("https://api.openai.com/v1/models") == True
         # Azure OpenAI on the shared Cognitive Services domain, identified by an
         # OpenAI-style path segment.
         assert (
-            self.handler.is_openai_route(
-                "https://my-resource.cognitiveservices.azure.com/v1/chat/completions"
-            )
-            == True
+            self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/v1/chat/completions") == True
         )
 
         # Negative cases
-        assert (
-            self.handler.is_openai_route(
-                "http://localhost:4000/openai/v1/chat/completions"
-            )
-            == False
-        )
-        assert (
-            self.handler.is_openai_route("https://api.anthropic.com/v1/messages")
-            == False
-        )
-        assert (
-            self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript")
-            == False
-        )
+        assert self.handler.is_openai_route("http://localhost:4000/openai/v1/chat/completions") == False
+        assert self.handler.is_openai_route("https://api.anthropic.com/v1/messages") == False
+        assert self.handler.is_openai_route("https://api.assemblyai.com/v2/transcript") == False
         # Non-OpenAI Azure Cognitive Services share the `cognitiveservices.azure.com`
         # domain but must NOT be classified as OpenAI routes (no OpenAI path segment).
         assert (
-            self.handler.is_openai_route(
-                "https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize"
-            )
+            self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/speechtotext/v3.1/recognize")
             == False
         )
         assert (
-            self.handler.is_openai_route(
-                "https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze"
-            )
-            == False
+            self.handler.is_openai_route("https://my-resource.cognitiveservices.azure.com/vision/v3.2/analyze") == False
         )
         # A look-alike domain that merely contains an OpenAI host as a substring
         # must be rejected by the suffix-based hostname match.
         assert (
-            self.handler.is_openai_route(
-                "https://cognitiveservices.azure.com.attacker.example/v1/chat/completions"
-            )
+            self.handler.is_openai_route("https://cognitiveservices.azure.com.attacker.example/v1/chat/completions")
             == False
         )
         assert self.handler.is_openai_route("") == False
@@ -1274,52 +1167,188 @@ class TestOpenAIPassthroughIntegration:
         remove Responses from the OR-chain without a test failure.
         """
         # Responses must be supported on api.openai.com and openai.azure.com.
-        assert (
-            self.handler._is_supported_openai_endpoint(
-                "https://api.openai.com/v1/responses"
-            )
-            is True
-        )
-        assert (
-            self.handler._is_supported_openai_endpoint(
-                "https://openai.azure.com/v1/responses"
-            )
-            is True
-        )
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/responses") is True
+        assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/responses") is True
         # The other supported endpoints stay supported (no regression).
-        assert (
-            self.handler._is_supported_openai_endpoint(
-                "https://api.openai.com/v1/chat/completions"
-            )
-            is True
-        )
-        assert (
-            self.handler._is_supported_openai_endpoint(
-                "https://api.openai.com/v1/images/generations"
-            )
-            is True
-        )
-        assert (
-            self.handler._is_supported_openai_endpoint(
-                "https://api.openai.com/v1/images/edits"
-            )
-            is True
-        )
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/chat/completions") is True
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/generations") is True
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/images/edits") is True
         # Unsupported OpenAI endpoints (e.g. /v1/models) still return False.
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/models") is False
         assert (
             self.handler._is_supported_openai_endpoint(
-                "https://api.openai.com/v1/models"
+                "https://my-resource.openai.azure.com/openai/deployments/text-embedding-3-small/embeddings"
             )
             is False
         )
 
+    def test_is_supported_openai_endpoint_includes_embeddings(self):
+        assert self.handler._is_supported_openai_endpoint("https://api.openai.com/v1/embeddings") is True
+        assert self.handler._is_supported_openai_endpoint("https://openai.azure.com/v1/embeddings") is True
+
+    def test_is_cohere_route_does_not_match_openai_embeddings(self):
+        assert self.handler.is_cohere_route("https://api.cohere.com/v1/embed") is True
+        assert self.handler.is_cohere_route("https://api.cohere.com/v2/chat") is True
+        assert self.handler.is_cohere_route("https://api.openai.com/v1/embeddings") is False
+        assert self.handler.is_cohere_route("https://api.cohere.com/v1/rerank") is False
+        assert self.handler.is_cohere_route("http://localhost:4000/openai_passthrough/v1/embeddings") is False
+
+    @patch("litellm.completion_cost")
+    @patch("litellm.litellm_core_utils.litellm_logging.get_standard_logging_object_payload")
+    def test_openai_passthrough_handler_embeddings_sets_response_cost(
+        self, mock_get_standard_logging, mock_completion_cost
+    ):
+        mock_completion_cost.return_value = 2.8e-07
+        mock_get_standard_logging.return_value = {"test": "logging_payload"}
+
+        response_body = {
+            "object": "list",
+            "model": "text-embedding-3-small",
+            "data": [
+                {
+                    "object": "embedding",
+                    "index": 0,
+                    "embedding": [0.1, 0.2],
+                }
+            ],
+            "usage": {"prompt_tokens": 14, "total_tokens": 14},
+        }
+        mock_httpx_response = self._create_mock_httpx_response(response_body)
+        mock_logging_obj = self._create_mock_logging_obj()
+        passthrough_payload = PassthroughStandardLoggingPayload(
+            url="https://api.openai.com/v1/embeddings",
+            request_body={
+                "model": "text-embedding-3-small",
+                "input": "PROOF_SENTINEL_TEXT",
+            },
+            request_method="POST",
+        )
+        kwargs = {
+            "passthrough_logging_payload": passthrough_payload,
+            "litellm_params": {},
+        }
+
+        result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
+            httpx_response=mock_httpx_response,
+            response_body=response_body,
+            logging_obj=mock_logging_obj,
+            url_route="https://api.openai.com/v1/embeddings",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body={
+                "model": "text-embedding-3-small",
+                "input": "PROOF_SENTINEL_TEXT",
+            },
+            **kwargs,
+        )
+
+        assert result["result"] is not None
+        assert result["kwargs"]["response_cost"] == 2.8e-07
+        assert result["kwargs"]["model"] == "text-embedding-3-small"
+        assert result["kwargs"]["custom_llm_provider"] == "openai"
+        assert result["result"]._hidden_params["response_cost"] == 2.8e-07
+        mock_completion_cost.assert_called_once()
+        assert mock_completion_cost.call_args.kwargs["call_type"] == "aembedding"
+        assert mock_logging_obj.model_call_details["response_cost"] == 2.8e-07
+
+    @patch(
+        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.passthrough_chat_handler"
+    )
+    @patch("litellm.completion_cost")
+    def test_openai_passthrough_handler_embeddings_without_model_falls_back(
+        self, mock_completion_cost, mock_chat_handler
+    ):
+        mock_chat_handler.return_value = {"result": None, "kwargs": {}}
+        response_body = {
+            "object": "list",
+            "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
+            "usage": {"prompt_tokens": 1, "total_tokens": 1},
+        }
+        result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler(
+            httpx_response=self._create_mock_httpx_response(response_body),
+            response_body=response_body,
+            logging_obj=self._create_mock_logging_obj(),
+            url_route="https://api.openai.com/v1/embeddings",
+            result="",
+            start_time=self.start_time,
+            end_time=self.end_time,
+            cache_hit=False,
+            request_body={"input": "PROOF_SENTINEL_TEXT"},
+            passthrough_logging_payload=PassthroughStandardLoggingPayload(
+                url="https://api.openai.com/v1/embeddings",
+                request_body={"input": "PROOF_SENTINEL_TEXT"},
+                request_method="POST",
+            ),
+        )
+        mock_completion_cost.assert_not_called()
+        mock_chat_handler.assert_called_once()
+        assert result == {"result": None, "kwargs": {}}
+
     @patch(
         "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler"
     )
     @pytest.mark.asyncio
-    async def test_success_handler_dispatches_responses_api_to_openai_handler(
-        self, mock_openai_handler
-    ):
+    async def test_success_handler_dispatches_embeddings_to_openai_handler(self, mock_openai_handler):
+        mock_openai_handler.return_value = {
+            "result": {"object": "list"},
+            "kwargs": {
+                "response_cost": 2.8e-07,
+                "model": "text-embedding-3-small",
+                "custom_llm_provider": "openai",
+            },
+        }
+
+        mock_httpx_response = MagicMock(spec=httpx.Response)
+        mock_httpx_response.text = (
+            '{"object":"list","model":"text-embedding-3-small",'
+            '"data":[{"object":"embedding","index":0,"embedding":[0.1]}],'
+            '"usage":{"prompt_tokens":14,"total_tokens":14}}'
+        )
+
+        mock_logging_obj = AsyncMock()
+        mock_logging_obj.model_call_details = {}
+        mock_logging_obj.async_success_handler = AsyncMock()
+
+        passthrough_payload = PassthroughStandardLoggingPayload(
+            url="https://api.openai.com/v1/embeddings",
+            request_body={
+                "model": "text-embedding-3-small",
+                "input": "PROOF_SENTINEL_TEXT",
+            },
+            request_method="POST",
+        )
+
+        await self.handler.pass_through_async_success_handler(
+            httpx_response=mock_httpx_response,
+            response_body={
+                "object": "list",
+                "model": "text-embedding-3-small",
+                "data": [{"object": "embedding", "index": 0, "embedding": [0.1]}],
+                "usage": {"prompt_tokens": 14, "total_tokens": 14},
+            },
+            logging_obj=mock_logging_obj,
+            url_route="https://api.openai.com/v1/embeddings",
+            result="",
+            start_time=datetime.now(),
+            end_time=datetime.now(),
+            cache_hit=False,
+            request_body={
+                "model": "text-embedding-3-small",
+                "input": "PROOF_SENTINEL_TEXT",
+            },
+            passthrough_logging_payload=passthrough_payload,
+        )
+
+        mock_openai_handler.assert_called_once()
+        assert mock_openai_handler.call_args.kwargs["url_route"] == "https://api.openai.com/v1/embeddings"
+
+    @patch(
+        "litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler.OpenAIPassthroughLoggingHandler.openai_passthrough_handler"
+    )
+    @pytest.mark.asyncio
+    async def test_success_handler_dispatches_responses_api_to_openai_handler(self, mock_openai_handler):
         """End-to-end dispatch test for the Responses API path.
 
         Pre-fix: `_is_supported_openai_endpoint` returned False for
@@ -1395,9 +1424,7 @@ class TestOpenAIPassthroughIntegration:
         }
 
         mock_httpx_response = MagicMock(spec=httpx.Response)
-        mock_httpx_response.text = (
-            '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}'
-        )
+        mock_httpx_response.text = '{"id": "chatcmpl-123", "choices": [{"message": {"content": "Hello"}}]}'
 
         mock_logging_obj = AsyncMock()
         mock_logging_obj.model_call_details = {}
@@ -1590,14 +1617,10 @@ class TestOpenAIPassthroughIntegration:
         # Test the _response_cost_calculator method
         calculated_cost = logging_obj._response_cost_calculator(result=image_response)
 
-        assert (
-            calculated_cost == test_cost
-        ), f"Expected {test_cost}, got {calculated_cost}"
+        assert calculated_cost == test_cost, f"Expected {test_cost}, got {calculated_cost}"
 
     @patch("litellm.cost_calculator.default_image_cost_calculator")
-    def test_openai_passthrough_handler_image_generation(
-        self, mock_image_cost_calculator
-    ):
+    def test_openai_passthrough_handler_image_generation(self, mock_image_cost_calculator):
         """Test successful cost tracking for OpenAI image generation"""
         # Arrange
         mock_image_cost_calculator.return_value = 0.040

From 6704a105eea862b9ebfc7f43c6da6d37dc838081 Mon Sep 17 00:00:00 2001
From: Yassin Kortam <yassin@berri.ai>
Date: Thu, 13 Aug 2026 21:45:36 -0700
Subject: [PATCH 31/31] fix(access groups): sync assigned_team_ids from the
 team write paths (#36825)

---
 .../access_group_endpoints.py                 |  20 +-
 .../management_endpoints/team_endpoints.py    |  31 +-
 .../access_group_team_sync.py                 | 155 ++++++++
 .../test_access_group_team_sync.py            | 230 +++++++++++
 .../test_team_endpoints.py                    | 374 ++++++++++++++++++
 .../proxy/management_endpoints/test_ui_sso.py |  18 +
 .../test_access_group_team_sync.py            |  39 ++
 7 files changed, 845 insertions(+), 22 deletions(-)
 create mode 100644 litellm/proxy/management_helpers/access_group_team_sync.py
 create mode 100644 tests/proxy_admin_ui_tests/test_access_group_team_sync.py
 create mode 100644 tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py

diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py
index c00c2a5ba4c..2271501d480 100644
--- a/litellm/proxy/management_endpoints/access_group_endpoints.py
+++ b/litellm/proxy/management_endpoints/access_group_endpoints.py
@@ -14,11 +14,11 @@ from litellm.proxy.auth.auth_checks import (
     _cache_access_object,
     _cache_key_object,
     _cache_team_object,
-    _delete_cache_access_object,
     _get_team_object_from_cache,
 )
 from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
 from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
+from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_cache
 from litellm.proxy.utils import get_prisma_client_or_throw
 from litellm.repositories.table_repositories import AccessGroupRepository
 from litellm.types.access_group import (
@@ -146,22 +146,6 @@ async def _cache_access_group_record(record: _AccessGroupRecord) -> None:
     )
 
 
-async def _invalidate_cache_access_group(access_group_id: str) -> None:
-    """
-    Invalidate (delete) an access group entry from both in-memory and Redis caches.
-
-    Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
-    to avoid circular imports, following the same pattern as key_management_endpoints.
-    """
-    from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
-
-    await _delete_cache_access_object(
-        access_group_id=access_group_id,
-        user_api_key_cache=user_api_key_cache,
-        proxy_logging_obj=proxy_logging_obj,
-    )
-
-
 # ---------------------------------------------------------------------------
 # DB sync helpers (called inside a Prisma transaction)
 # ---------------------------------------------------------------------------
@@ -595,7 +579,7 @@ async def delete_access_group(
 
         from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
 
-        await _invalidate_cache_access_group(access_group_id)
+        await invalidate_access_group_cache(access_group_id)
         await _patch_team_caches_remove_access_group(
             affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj
         )
diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py
index b72b5d218b2..3d7f0808fb9 100644
--- a/litellm/proxy/management_endpoints/team_endpoints.py
+++ b/litellm/proxy/management_endpoints/team_endpoints.py
@@ -15,6 +15,7 @@ import math
 import traceback
 from collections.abc import Mapping, Sequence
 from datetime import datetime, timezone
+from types import MappingProxyType
 from typing import Annotated, Final, Protocol, TypedDict, TypeVar, cast
 
 import fastapi
@@ -106,6 +107,12 @@ from litellm.proxy.management_endpoints.organization_endpoints import (
 from litellm.proxy.management_endpoints.tag_management_endpoints import (
     get_daily_activity,
 )
+from litellm.proxy.management_helpers.access_group_team_sync import (
+    AccessGroupSyncTx,
+    invalidate_access_group_caches,
+    reconcile_team_access_group_membership,
+    sync_team_access_group_membership,
+)
 from litellm.proxy.management_helpers.object_permission_utils import (
     _set_object_permission,
     enforce_all_proxy_mcp_servers_grant_is_admin_only,
@@ -315,10 +322,17 @@ class _TeamIdInFilter(TypedDict, total=False):
     team_id: Mapping[str, Sequence[str]]
 
 
+class _TeamCreateTx(AccessGroupSyncTx, Protocol):
+    @property
+    def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
+
+
 _STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
 UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(teams)
 """
 
+_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
+
 
 def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
     return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
@@ -1511,10 +1525,15 @@ async def new_team(
         complete_team_data_dict = prisma_client.jsonify_team_object(db_data=complete_team_data_dict)
         team_creation_data: Final[Mapping[str, object]] = complete_team_data_dict
 
-        team_row: Final[LiteLLM_TeamTable] = await _team_db(prisma_client).create(
-            data=team_creation_data,
-            include={"litellm_model_table": True},
-        )
+        tx: _TeamCreateTx
+        async with prisma_client.db.tx() as tx:
+            team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
+                data=team_creation_data,
+                include=_INCLUDE_MODEL_TABLE,
+            )
+            affected_access_groups: Final = await reconcile_team_access_group_membership(tx, team_row.team_id)
+
+        await invalidate_access_group_caches(affected_access_groups)
 
         ## ADD TEAM ID TO USER TABLE ##
         team_member_add_request: Final = TeamMemberAddRequest(
@@ -2217,6 +2236,7 @@ async def update_team(
             )
 
         verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id)
+        await sync_team_access_group_membership(prisma_client=prisma_client, team_id=team_row.team_id)
         await _refresh_cached_team(
             team_row=team_row,
             user_api_key_cache=user_api_key_cache,
@@ -3850,6 +3870,9 @@ async def delete_team(
     # keeping the first one means a failure here still leaves a team the admin can retry deleting.
     await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
 
+    for deleted_team in team_rows:
+        await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
+
     return deleted_teams
 
 
diff --git a/litellm/proxy/management_helpers/access_group_team_sync.py b/litellm/proxy/management_helpers/access_group_team_sync.py
new file mode 100644
index 00000000000..55c0346e375
--- /dev/null
+++ b/litellm/proxy/management_helpers/access_group_team_sync.py
@@ -0,0 +1,155 @@
+"""
+Reverse sync for the team side of the team <-> access group relationship.
+
+`litellm_accessgrouptable.assigned_team_ids` and `litellm_teamtable.access_group_ids`
+are two copies of the same relationship, and both are read: the access group's
+attached-teams view reads the former, and so does the key-side grant check in
+`auth_checks.get_authorized_resources_from_key_access_groups`. The access-group
+endpoints maintain both copies already; this module is what the team write paths
+call so an edit from that side is mirrored back.
+
+It deliberately lives outside `access_group_endpoints`, which is a lazily
+registered feature router (see `_lazy_features.LAZY_FEATURES`). Importing that
+module eagerly from `team_endpoints` would put it in `sys.modules` without its
+router ever being included, which drops its routes from the OpenAPI schema.
+"""
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from typing import Final, Protocol
+
+from pydantic import BaseModel, TypeAdapter
+
+from litellm.proxy.auth.auth_checks import _delete_cache_access_object
+
+# hashtext collisions only cost two unrelated teams a little serialization, and the
+# lock is never taken by the access-group endpoints, so it cannot join their
+# access-group-then-team lock order to form a cycle.
+_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
+
+_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
+
+# The groups the team is on either side of the reconcile, so the cache step is driven by
+# desired state rather than by which rows this attempt happened to change. A retry after a
+# failed invalidation finds the same set even though its statements are already no-ops.
+_AFFECTED_SQL: Final = """
+SELECT access_group_id FROM "LiteLLM_AccessGroupTable"
+WHERE access_group_id = ANY($2::TEXT[])
+   OR $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
+"""
+
+_ATTACH_SQL: Final = """
+UPDATE "LiteLLM_AccessGroupTable"
+SET assigned_team_ids = array_append(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]), $1)
+WHERE access_group_id = ANY($2::TEXT[])
+  AND NOT ($1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[])))
+RETURNING access_group_id
+"""
+
+_DETACH_SQL: Final = """
+UPDATE "LiteLLM_AccessGroupTable"
+SET assigned_team_ids = array_remove(assigned_team_ids, $1)
+WHERE $1 = ANY(COALESCE(assigned_team_ids, ARRAY[]::TEXT[]))
+  AND NOT (access_group_id = ANY($2::TEXT[]))
+RETURNING access_group_id
+"""
+
+
+class _AffectedGroup(BaseModel):
+    access_group_id: str
+
+
+class _TeamGroups(BaseModel):
+    access_group_ids: tuple[str, ...] | None = None
+
+
+_AffectedGroups: Final = TypeAdapter(tuple[_AffectedGroup, ...])
+_TeamRows: Final = TypeAdapter(tuple[_TeamGroups, ...])
+
+
+class AccessGroupSyncTx(Protocol):
+    async def query_raw(self, query: str, *args: object) -> Sequence[Mapping[str, object]]: ...
+
+
+class _Transaction(Protocol):
+    async def __aenter__(self) -> AccessGroupSyncTx: ...
+
+    async def __aexit__(self, *exc_info: object) -> None: ...
+
+
+class _PrismaDb(Protocol):
+    def tx(self) -> _Transaction: ...
+
+
+class _PrismaClient(Protocol):
+    @property
+    def db(self) -> _PrismaDb: ...
+
+
+async def invalidate_access_group_cache(access_group_id: str) -> None:
+    """
+    Drop an access group entry from both the in-memory and Redis caches.
+
+    Uses a lazy import of user_api_key_cache and proxy_logging_obj from proxy_server
+    to avoid circular imports, following the same pattern as key_management_endpoints.
+    """
+    from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache
+
+    await _delete_cache_access_object(
+        access_group_id=access_group_id,
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=proxy_logging_obj,
+    )
+
+
+async def invalidate_access_group_caches(access_group_ids: Sequence[str]) -> None:
+    """
+    Drop every given access group from the caches, then raise if any drop failed.
+
+    Every entry is attempted even when one raises, so a single unreachable cache cannot
+    leave the rest of the reconciled groups serving a grant the admin revoked.
+    """
+    outcomes: Final = await asyncio.gather(
+        *(invalidate_access_group_cache(access_group_id) for access_group_id in access_group_ids),
+        return_exceptions=True,
+    )
+    for outcome in outcomes:
+        if isinstance(outcome, BaseException):
+            raise outcome
+
+
+async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id: str) -> tuple[str, ...]:
+    """
+    Reconcile every access group's `assigned_team_ids` against the team's own
+    `access_group_ids`, and return the groups whose cache the caller has to drop once the
+    transaction commits.
+
+    Call this inside the transaction that writes the team row, or after that row is
+    written or deleted: a team with no row reconciles to an empty set, which detaches it
+    from every group.
+
+    The team row is read here rather than passed in, under an advisory lock held for the
+    rest of the transaction. That is what makes concurrent writes to the same team
+    converge, since each mirror reconciles against the row as the transaction sees it
+    instead of against the snapshot its own caller happened to see. It also means a retry
+    heals a sync that failed partway, where a before/after delta would compute nothing.
+
+    Both mirror statements are set-based and mutate the array inside the statement, so a
+    concurrent write for a different team cannot be lost the way a read-modify-write of
+    the whole array can, and the pair commits together or not at all.
+    """
+    await tx.query_raw(_LOCK_TEAM_SQL, team_id)
+    team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
+    desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
+    affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))
+    await tx.query_raw(_ATTACH_SQL, team_id, desired)
+    await tx.query_raw(_DETACH_SQL, team_id, desired)
+    return tuple(group.access_group_id for group in affected)
+
+
+async def sync_team_access_group_membership(prisma_client: _PrismaClient, team_id: str) -> None:
+    """Reconcile the mirror for an already committed team write, in its own transaction."""
+    async with prisma_client.db.tx() as tx:
+        affected: Final = await reconcile_team_access_group_membership(tx, team_id)
+
+    await invalidate_access_group_caches(affected)
diff --git a/tests/proxy_admin_ui_tests/test_access_group_team_sync.py b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py
new file mode 100644
index 00000000000..629d77f20fc
--- /dev/null
+++ b/tests/proxy_admin_ui_tests/test_access_group_team_sync.py
@@ -0,0 +1,230 @@
+"""
+Real-Postgres coverage for the team -> access group mirror.
+
+`sync_team_access_group_membership` reconciles `assigned_team_ids` with two raw
+statements, and a mocked prisma cannot tell whether that SQL is right: a fake has to
+reimplement the array semantics in Python, so it passes no matter what the SQL says.
+These tests run the statements against the same Postgres CI seeds for the admin UI
+suite, which is the only place a `NOT (... = ANY(...))` guard going missing shows up.
+"""
+
+import asyncio
+import os
+import sys
+from contextlib import asynccontextmanager
+from datetime import timedelta
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+from litellm.proxy.management_helpers.access_group_team_sync import (
+    reconcile_team_access_group_membership,
+    sync_team_access_group_membership,
+)
+
+TEAM = "ags-team-a"
+OTHER_TEAM = "ags-team-b"
+GROUPS = ("ags-group-1", "ags-group-2", "ags-group-3")
+_DELETE_SEEDED = 'DELETE FROM "LiteLLM_AccessGroupTable" WHERE access_group_id = ANY($1::TEXT[])'
+_DELETE_TEAMS = 'DELETE FROM "LiteLLM_TeamTable" WHERE team_id = ANY($1::TEXT[])'
+
+
+@asynccontextmanager
+async def _clean_db():
+    """Connects inside the running test's loop. An async fixture would be torn up on a
+    different loop than the test body, which prisma's engine lock refuses outright."""
+    from prisma import Prisma
+
+    if not os.getenv("DATABASE_URL"):
+        pytest.fail("DATABASE_URL is required; these tests must not silently skip")
+
+    db = Prisma()
+    await db.connect()
+    try:
+        await db.execute_raw(_DELETE_SEEDED, list(GROUPS))
+        await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM])
+        yield db
+    finally:
+        await db.execute_raw(_DELETE_SEEDED, list(GROUPS))
+        await db.execute_raw(_DELETE_TEAMS, [TEAM, OTHER_TEAM])
+        await db.disconnect()
+
+
+async def _seed(db, assignments):
+    for group_id, team_ids in assignments.items():
+        await db.litellm_accessgrouptable.create(
+            data={
+                "access_group_id": group_id,
+                "access_group_name": group_id,
+                "assigned_team_ids": team_ids,
+            }
+        )
+
+
+async def _read(db):
+    rows = await db.query_raw(
+        'SELECT access_group_id, assigned_team_ids FROM "LiteLLM_AccessGroupTable" '
+        "WHERE access_group_id = ANY($1::TEXT[])",
+        list(GROUPS),
+    )
+    return {row["access_group_id"]: sorted(row["assigned_team_ids"] or []) for row in rows}
+
+
+async def _set_team_groups(db, team_id, access_group_ids):
+    """The mirror reads the committed team row, so the desired state is written there."""
+    if access_group_ids is None:
+        await db.execute_raw(_DELETE_TEAMS, [team_id])
+        return
+    await db.litellm_teamtable.upsert(
+        where={"team_id": team_id},
+        data={
+            "create": {"team_id": team_id, "access_group_ids": list(access_group_ids)},
+            "update": {"access_group_ids": list(access_group_ids)},
+        },
+    )
+
+
+async def _sync(db, team_id, access_group_ids):
+    await _set_team_groups(db, team_id, access_group_ids)
+    with patch(
+        "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+        new_callable=AsyncMock,
+    ) as invalidate:
+        await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=team_id)
+    return {call.args[0] for call in invalidate.call_args_list}
+
+
+@pytest.mark.asyncio
+async def test_reconcile_attaches_and_detaches_without_touching_other_teams():
+    """The detach must be scoped to groups the team dropped. Losing that scope would
+    strip the team from the very groups it just kept, silently revoking live grants."""
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]})
+
+        invalidated = await _sync(db, TEAM, [GROUPS[1], GROUPS[2]])
+
+        assert await _read(db) == {
+            GROUPS[0]: [OTHER_TEAM],
+            GROUPS[1]: [TEAM],
+            GROUPS[2]: sorted([TEAM, OTHER_TEAM]),
+        }
+        assert invalidated == {GROUPS[0], GROUPS[1], GROUPS[2]}
+
+
+@pytest.mark.asyncio
+async def test_reconcile_is_idempotent_so_a_retry_heals_rather_than_duplicates():
+    """Reconciling to the same desired state twice must leave the rows alone and still name
+    the team's groups for the cache step, so a retry after a failed cache drop reaches them.
+    A delta-based mirror would instead go quiet once the rows match, leaving the caches
+    serving a grant the admin already revoked."""
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [], GROUPS[1]: [TEAM], GROUPS[2]: []})
+
+        first = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]])
+        after_first = await _read(db)
+        second = await _sync(db, TEAM, [GROUPS[0], GROUPS[1]])
+
+        assert after_first == {GROUPS[0]: [TEAM], GROUPS[1]: [TEAM], GROUPS[2]: []}
+        assert await _read(db) == after_first
+        assert first == {GROUPS[0], GROUPS[1]}
+        assert second == first
+
+
+@pytest.mark.asyncio
+async def test_reconcile_handles_a_null_array_column():
+    """`assigned_team_ids` is nullable in Postgres. Without COALESCE both statements
+    evaluate their guard to NULL, skip the row, and the grant silently never syncs."""
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [], GROUPS[1]: []})
+        await db.execute_raw(
+            'UPDATE "LiteLLM_AccessGroupTable" SET assigned_team_ids = NULL WHERE access_group_id = $1',
+            GROUPS[0],
+        )
+
+        await _sync(db, TEAM, [GROUPS[0]])
+
+        assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []}
+
+
+@pytest.mark.asyncio
+async def test_passing_none_detaches_the_team_from_every_group():
+    """Team deletion. A group the deleted row never listed must still let the team go,
+    otherwise the id dangles under Attached Teams and grants again if it is reused."""
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [TEAM, OTHER_TEAM], GROUPS[1]: [TEAM], GROUPS[2]: [OTHER_TEAM]})
+
+        invalidated = await _sync(db, TEAM, None)
+
+        assert await _read(db) == {GROUPS[0]: [OTHER_TEAM], GROUPS[1]: [], GROUPS[2]: [OTHER_TEAM]}
+        assert invalidated == {GROUPS[0], GROUPS[1]}
+
+
+@pytest.mark.asyncio
+async def test_a_failed_mirror_takes_the_new_team_row_with_it():
+    """`/team/new` inserts the team and mirrors it in one transaction. Mirroring in a
+    transaction of its own instead leaves a committed team whose groups never learned about
+    it, and the retry with that same team id comes back as a duplicate."""
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]})
+
+        with pytest.raises(RuntimeError):
+            async with db.tx() as tx:
+                await tx.litellm_teamtable.create(data={"team_id": TEAM, "access_group_ids": [GROUPS[0]]})
+                await reconcile_team_access_group_membership(tx, TEAM)
+                raise RuntimeError("the cache handoff blew up")
+
+        assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [OTHER_TEAM]}
+        assert await db.litellm_teamtable.find_unique(where={"team_id": TEAM}) is None
+
+
+@pytest.mark.asyncio
+async def test_a_concurrent_writer_cannot_replay_a_stale_team_row_over_a_newer_one():
+    """
+    Two writers edit one team at once. Whichever team row commits last is the admin's
+    final intent and the mirror must match it, so the mirror has to hold the team's
+    advisory lock across its read and its writes.
+
+    A second connection holds that lock and changes the team underneath, which pins the
+    interleaving instead of hoping a sleep lands in the gap. With the lock the sync waits
+    and then reads the new row. Without it the sync reads the old row and writes a group
+    the admin already moved off, which keeps granting to that team.
+    """
+    from prisma import Prisma
+
+    async with _clean_db() as db:
+        await _seed(db, {GROUPS[0]: [], GROUPS[1]: []})
+        await _sync(db, TEAM, [GROUPS[0]])
+        assert await _read(db) == {GROUPS[0]: [TEAM], GROUPS[1]: []}
+
+        blocker = Prisma()
+        await blocker.connect()
+        sync_started = asyncio.Event()
+
+        async def competing_sync():
+            sync_started.set()
+            with patch(
+                "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+                new_callable=AsyncMock,
+            ):
+                await sync_team_access_group_membership(prisma_client=SimpleNamespace(db=db), team_id=TEAM)
+
+        try:
+            async with blocker.tx(timeout=timedelta(seconds=30)) as held:
+                await held.query_raw("SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked", TEAM)
+                task = asyncio.create_task(competing_sync())
+                await sync_started.wait()
+                await asyncio.sleep(0.2)
+                assert not task.done(), "the mirror did not wait on the team's advisory lock"
+                await held.execute_raw(
+                    'UPDATE "LiteLLM_TeamTable" SET access_group_ids = $1 WHERE team_id = $2',
+                    [GROUPS[1]],
+                    TEAM,
+                )
+            await asyncio.wait_for(task, timeout=30)
+        finally:
+            await blocker.disconnect()
+
+        assert await _read(db) == {GROUPS[0]: [], GROUPS[1]: [TEAM]}
diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
index f038af4c7a1..c6960ecda5a 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py
@@ -2,7 +2,9 @@ import asyncio
 import json
 import os
 import sys
+from contextlib import asynccontextmanager
 from datetime import datetime, timezone
+from types import SimpleNamespace
 from typing import Optional, cast
 from unittest.mock import AsyncMock, MagicMock, call, patch
 
@@ -68,6 +70,21 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
 # Setup TestClient
 client = TestClient(app)
 
+
+def _wire_team_create_tx(prisma_client):
+    """`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
+    so a mocked client has to hand its team table back out of `db.tx()`."""
+
+    @asynccontextmanager
+    async def _tx():
+        yield SimpleNamespace(
+            litellm_teamtable=prisma_client.db.litellm_teamtable,
+            query_raw=AsyncMock(return_value=[]),
+        )
+
+    prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
+
+
 # Mock prisma_client
 mock_prisma_client = MagicMock()
 # Set up async mock for db operations
@@ -400,6 +417,7 @@ async def test_new_team_rejects_a_duration_that_never_advances(
     mock_team_create = AsyncMock()
     mock_db_client.db.litellm_teamtable = MagicMock()
     mock_db_client.db.litellm_teamtable.create = mock_team_create
+    _wire_team_create_tx(mock_db_client)
 
     with pytest.raises(ProxyException) as exc_info:
         await new_team(
@@ -481,6 +499,7 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth):
     mock_team_count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable = MagicMock()
     mock_db_client.db.litellm_teamtable.create = mock_team_create
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = mock_team_count
     mock_db_client.db.litellm_teamtable.update = AsyncMock(
         return_value=team_create_result
@@ -570,6 +589,7 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut
     mock_db_client.db.litellm_teamtable.create = AsyncMock(
         return_value=team_create_result
     )
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable.update = AsyncMock(
         return_value=team_create_result
@@ -663,6 +683,7 @@ async def test_new_team_disable_auto_add_proxy_admin_flag(
     mock_db_client.db.litellm_teamtable.create = AsyncMock(
         return_value=team_create_result
     )
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_usertable = MagicMock()
     mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
@@ -4430,6 +4451,7 @@ async def test_new_team_max_budget_within_user_limit():
         mock_prisma.db.litellm_teamtable.create = AsyncMock(
             return_value=mock_created_team
         )
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_teamtable.update = AsyncMock(
             return_value=mock_created_team
         )
@@ -4573,6 +4595,7 @@ async def test_new_team_org_scoped_budget_bypasses_user_limit():
         mock_prisma.db.litellm_teamtable.create = AsyncMock(
             return_value=mock_created_team
         )
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_teamtable.update = AsyncMock(
             return_value=mock_created_team
         )
@@ -4721,6 +4744,7 @@ async def test_new_team_org_scoped_models_bypasses_user_limit():
         mock_prisma.db.litellm_teamtable.create = AsyncMock(
             return_value=mock_created_team
         )
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_teamtable.update = AsyncMock(
             return_value=mock_created_team
         )
@@ -6567,6 +6591,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit():
         mock_created_team.rpm_limit = 1000
         mock_created_team.metadata = None
         mock_created_team.members_with_roles = []
+        mock_created_team.access_group_ids = None
         mock_created_team.model_dump.return_value = {
             "team_id": "new-bypass-team-id",
             "team_alias": "org-bypass-test-team",
@@ -6578,6 +6603,7 @@ async def test_new_team_org_scoped_tpm_rpm_bypasses_user_limit():
         mock_prisma.db.litellm_teamtable.create = AsyncMock(
             return_value=mock_created_team
         )
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_teamtable.update = AsyncMock(
             return_value=mock_created_team
         )
@@ -6856,6 +6882,7 @@ async def test_update_team_org_scoped_tpm_rpm_bypasses_user_limit():
         mock_updated_team.team_id = "org-team-update-bypass-123"
         mock_updated_team.tpm_limit = 10000
         mock_updated_team.rpm_limit = 1000
+        mock_updated_team.access_group_ids = None
         mock_updated_team.model_dump.return_value = {
             "team_id": "org-team-update-bypass-123",
             "tpm_limit": 10000,
@@ -7009,6 +7036,7 @@ async def test_update_team_guardrails_with_org_id():
             "guardrails": ["aporia-pre-call", "aporia-post-call"]
         }
         mock_updated_team.litellm_model_table = None
+        mock_updated_team.access_group_ids = None
         mock_updated_team.model_dump.return_value = {
             "team_id": "team-guardrails-123",
             "organization_id": "test-org-guardrails",
@@ -7937,6 +7965,7 @@ async def test_new_team_soft_budget_validation(
         mock_prisma.db.litellm_teamtable.create = AsyncMock(
             return_value=mock_created_team
         )
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_teamtable.update = AsyncMock(
             return_value=mock_created_team
         )
@@ -8236,6 +8265,7 @@ async def test_new_team_with_router_settings(mock_db_client, mock_admin_auth):
     mock_team_count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable = MagicMock()
     mock_db_client.db.litellm_teamtable.create = mock_team_create
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = mock_team_count
     mock_db_client.db.litellm_teamtable.update = AsyncMock(
         return_value=team_create_result
@@ -9626,6 +9656,7 @@ async def test_new_team_encrypts_callback_vars(
     team_create_result.model_dump.return_value = {"team_id": "team-456"}
     mock_team_create = AsyncMock(return_value=team_create_result)
     mock_db_client.db.litellm_teamtable.create = mock_team_create
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable.update = AsyncMock(
         return_value=team_create_result
@@ -10786,6 +10817,7 @@ async def test_new_team_validator_runs_without_metadata_and_rejection_blocks_cre
     ):
         mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
         mock_prisma.db.litellm_teamtable.create = AsyncMock()
+        _wire_team_create_tx(mock_prisma)
         mock_license.is_team_count_over_limit.return_value = False
 
         with pytest.raises(ProxyException) as exc_info:
@@ -10820,6 +10852,7 @@ async def test_new_team_validator_accept_proceeds_to_create(mock_db_client, mock
     team_create_result.model_dump.return_value = {"team_id": "team-accept-1"}
     mock_db_client.db.litellm_teamtable = MagicMock()
     mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=team_create_result)
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result)
     mock_db_client.db.litellm_usertable = MagicMock()
@@ -10859,6 +10892,7 @@ async def test_new_team_rejection_precedes_model_alias_write():
     ):
         mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
         mock_prisma.db.litellm_teamtable.create = AsyncMock()
+        _wire_team_create_tx(mock_prisma)
         mock_prisma.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model-1"))
         mock_license.is_team_count_over_limit.return_value = False
 
@@ -11626,6 +11660,7 @@ def _wire_new_team_prisma(mock_db_client):
     mock_db_client.db.litellm_teamtable = MagicMock()
     mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_db_client.db.litellm_teamtable.create = AsyncMock(return_value=created_team)
+    _wire_team_create_tx(mock_db_client)
     mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=created_team)
     mock_db_client.db.litellm_usertable = MagicMock()
     mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock())
@@ -11714,3 +11749,342 @@ async def test_new_team_explicit_null_max_budget_still_takes_configured_default(
 
     team_data = mock_team_create.call_args.kwargs["data"]
     assert team_data.get("max_budget") == 100.0
+
+
+class _FakeMirrorDb:
+    """Stands in for prisma inside the access-group mirror.
+
+    Dispatches on the statement so a change to the SQL's shape is visible here, but it
+    cannot validate the SQL itself: it reimplements the array semantics in Python, so it
+    passes whatever the statement says. Correctness of the SQL is pinned against a real
+    Postgres in tests/proxy_admin_ui_tests/test_access_group_team_sync.py.
+    """
+
+    def __init__(self, access_groups, teams, plain_lists=False):
+        self._access_groups = access_groups
+        self._teams = teams
+        self._plain_lists = plain_lists
+        self.transactions = []
+
+    def _team_ids(self, group_id):
+        stored = self._access_groups[group_id]
+        return stored if self._plain_lists else stored["assigned_team_ids"]
+
+    async def _query_raw(self, sql, *args):
+        assert self._open, "mirror statement ran outside a transaction"
+        if "pg_advisory_xact_lock" in sql:
+            self.transactions[-1].append("lock")
+            return [{"locked": False}]
+        if "LiteLLM_TeamTable" in sql:
+            self.transactions[-1].append("read")
+            team_id = args[0]
+            if team_id not in self._teams:
+                return []
+            return [{"access_group_ids": list(self._teams[team_id])}]
+
+        team_id, desired = args
+        if sql.lstrip().startswith("SELECT"):
+            self.transactions[-1].append("affected")
+            affected = [g for g in self._access_groups if g in desired or team_id in self._team_ids(g)]
+            return [{"access_group_id": group_id} for group_id in affected]
+
+        if "array_append" in sql:
+            self.transactions[-1].append("attach")
+            changed = [
+                g for g in desired if g in self._access_groups and team_id not in self._team_ids(g)
+            ]
+            for group_id in changed:
+                self._team_ids(group_id).append(team_id)
+        else:
+            self.transactions[-1].append("detach")
+            changed = [
+                g for g in self._access_groups if team_id in self._team_ids(g) and g not in desired
+            ]
+            for group_id in changed:
+                self._team_ids(group_id).remove(team_id)
+        return [{"access_group_id": group_id} for group_id in changed]
+
+    async def _create_team(self, data, include=None):
+        self.transactions[-1].append("create")
+        team_id = data["team_id"]
+        self._teams[team_id] = list(data.get("access_group_ids") or ())
+        return SimpleNamespace(
+            team_id=team_id,
+            access_group_ids=list(self._teams[team_id]),
+            model_dump=lambda: {"team_id": team_id},
+        )
+
+    def tx(self, *_args, **_kwargs):
+        outer = self
+
+        class _Tx:
+            async def __aenter__(self):
+                outer.transactions.append([])
+                outer._open = True
+                return SimpleNamespace(
+                    query_raw=outer._query_raw,
+                    litellm_teamtable=SimpleNamespace(create=outer._create_team),
+                )
+
+            async def __aexit__(self, *_exc_info):
+                outer._open = False
+                return None
+
+        return _Tx()
+
+    _open = False
+
+
+@pytest.mark.asyncio
+async def test_update_team_syncs_access_group_assigned_team_ids_in_both_directions():
+    """
+    A team-side edit of `access_group_ids` must be mirrored onto every affected access
+    group's `assigned_team_ids`, in one transaction, in both directions.
+
+    `assigned_team_ids` is not display-only. `get_authorized_resources_from_key_access_groups`
+    reads it as an authorization input, so a group the team dropped must stop granting its
+    resources to keys on that team, and a group the team added must start granting them.
+    A single-direction assertion would pass against a fix that only ever removes (or only
+    ever adds), so this covers add, remove, untouched, and the authorization consequence.
+    """
+    from unittest.mock import Mock
+
+    from fastapi import Request
+
+    from litellm.proxy._types import LiteLLM_AccessGroupTable
+    from litellm.proxy.auth.auth_checks import (
+        get_authorized_resources_from_key_access_groups,
+    )
+
+    access_groups = {
+        "ag-drop": {"assigned_team_ids": ["team-a"], "access_model_names": ["dropped-model"]},
+        "ag-keep": {"assigned_team_ids": ["team-a"], "access_model_names": ["kept-model"]},
+        "ag-add": {"assigned_team_ids": [], "access_model_names": ["added-model"]},
+        "ag-other-team": {"assigned_team_ids": ["team-b"], "access_model_names": ["other-model"]},
+    }
+    committed_team_groups = ["ag-keep", "ag-add"]
+    fake_db = _FakeMirrorDb(access_groups, {"team-a": committed_team_groups})
+
+    existing_team = MagicMock()
+    existing_team.access_group_ids = ["ag-drop", "ag-keep"]
+    existing_team.metadata = {}
+    existing_team.max_budget = None
+    existing_team.organization_id = None
+    existing_team.team_alias = "team-a"
+    existing_team.model_dump.return_value = {"team_id": "team-a", "team_alias": "team-a"}
+
+    updated_team = MagicMock()
+    updated_team.team_id = "team-a"
+    updated_team.access_group_ids = committed_team_groups
+    updated_team.model_dump.return_value = {"team_id": "team-a"}
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client") as prisma,
+        patch("litellm.proxy.proxy_server.llm_router"),
+        patch("litellm.proxy.proxy_server.user_api_key_cache"),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj"),
+        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
+        patch("litellm.proxy.management_endpoints.team_endpoints._refresh_cached_team"),
+        patch(
+            "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ) as invalidate_cache,
+    ):
+        prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team)
+        prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team)
+        prisma.db.tx = fake_db.tx
+        prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
+
+        await update_team(
+            data=UpdateTeamRequest(team_id="team-a", access_group_ids=committed_team_groups),
+            http_request=Mock(spec=Request),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
+        )
+
+    assert access_groups["ag-drop"]["assigned_team_ids"] == []
+    assert access_groups["ag-add"]["assigned_team_ids"] == ["team-a"]
+    assert access_groups["ag-keep"]["assigned_team_ids"] == ["team-a"]
+    assert access_groups["ag-other-team"]["assigned_team_ids"] == ["team-b"]
+
+    assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]]
+    assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-drop", "ag-keep", "ag-add"}
+
+    async def _get_access_object(*, access_group_id, **_kwargs):
+        stored = access_groups[access_group_id]
+        return LiteLLM_AccessGroupTable(
+            access_group_id=access_group_id,
+            access_group_name=access_group_id,
+            access_model_names=list(stored["access_model_names"]),
+            assigned_team_ids=list(stored["assigned_team_ids"]),
+            assigned_key_ids=[],
+        )
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+        patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()),
+        patch(
+            "litellm.proxy.auth.auth_checks.get_access_object",
+            new_callable=AsyncMock,
+            side_effect=_get_access_object,
+        ),
+    ):
+        authorized_models = await get_authorized_resources_from_key_access_groups(
+            valid_token=UserAPIKeyAuth(
+                token="sk-hash",
+                models=[],
+                team_id="team-a",
+                access_group_ids=["ag-drop", "ag-keep", "ag-add"],
+            ),
+            team_object=LiteLLM_TeamTable(team_id="team-a", models=[]),
+            resource_field="access_model_names",
+        )
+
+    assert sorted(authorized_models) == ["added-model", "kept-model"]
+
+
+@pytest.mark.asyncio
+async def test_sync_reads_the_committed_team_row_rather_than_the_callers_snapshot():
+    """
+    The mirror takes no desired-state argument on purpose. It locks the team and reads
+    the row as committed, so two concurrent writers for one team converge on the row the
+    last one committed instead of each replaying its own stale snapshot. Reconciling also
+    means a retry heals a half-applied sync, where a before/after delta computes nothing.
+
+    The same holds for the cache step: the groups to drop come from the reconciled set,
+    not from the rows this attempt happened to change, so a retry after an unreachable
+    cache still drops the entries even though its statements are now no-ops.
+
+    A team with no row at all is deletion, and must detach from every group.
+    """
+    from litellm.proxy.management_helpers.access_group_team_sync import (
+        sync_team_access_group_membership,
+    )
+
+    access_groups = {"ag-1": ["team-a", "team-b"], "ag-2": ["team-a"], "ag-3": []}
+    teams = {"team-a": ["ag-2", "ag-3"]}
+    fake_db = _FakeMirrorDb(access_groups, teams, plain_lists=True)
+    prisma_client = SimpleNamespace(db=SimpleNamespace(tx=fake_db.tx))
+
+    with patch(
+        "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+        new_callable=AsyncMock,
+        side_effect=[ConnectionError("redis unreachable"), None, None],
+    ) as invalidate_cache:
+        with pytest.raises(ConnectionError):
+            await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
+        assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]}
+        assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1", "ag-2", "ag-3"}
+
+        invalidate_cache.reset_mock()
+        invalidate_cache.side_effect = None
+        await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
+        assert access_groups == {"ag-1": ["team-b"], "ag-2": ["team-a"], "ag-3": ["team-a"]}
+        assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"}
+
+        invalidate_cache.reset_mock()
+        del teams["team-a"]
+        await sync_team_access_group_membership(prisma_client=prisma_client, team_id="team-a")
+        assert access_groups == {"ag-1": ["team-b"], "ag-2": [], "ag-3": []}
+        assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-2", "ag-3"}
+
+    assert fake_db.transactions == [["lock", "read", "affected", "attach", "detach"]] * 3
+
+
+@pytest.mark.asyncio
+async def test_new_team_and_delete_team_both_drive_the_mirror():
+    """Every writer of `team.access_group_ids` has to reach the mirror, not just update.
+    These pin the wiring on the other two paths; the mirror's own behavior is covered above.
+
+    Creation has to insert the team row and mirror it in one transaction. With the mirror
+    in a transaction of its own, a sync that fails leaves a committed team whose groups
+    never learned about it, and the retry is rejected as a duplicate team id."""
+    from unittest.mock import Mock
+
+    from fastapi import Request
+
+    from litellm.proxy._types import DeleteTeamRequest, NewTeamRequest
+    from litellm.proxy.management_endpoints.team_endpoints import delete_team, new_team
+
+    access_groups = {"ag-1": [], "ag-2": []}
+    fake_db = _FakeMirrorDb(access_groups, {}, plain_lists=True)
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client") as prisma,
+        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
+        patch("litellm.proxy.proxy_server.user_api_key_cache"),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj"),
+        patch("litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", new_callable=AsyncMock),
+        patch(
+            "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+            new_callable=AsyncMock,
+        ) as invalidate_cache,
+    ):
+        prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
+        prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
+        prisma.db.tx = fake_db.tx
+        prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
+        prisma.get_data = AsyncMock(return_value=None)
+
+        await new_team(
+            data=NewTeamRequest(team_id="team-new", team_alias="new", access_group_ids=["ag-1"]),
+            http_request=Mock(spec=Request),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
+        )
+
+    assert access_groups == {"ag-1": ["team-new"], "ag-2": []}
+    assert fake_db.transactions == [["create", "lock", "read", "affected", "attach", "detach"]]
+    assert {call.args[0] for call in invalidate_cache.call_args_list} == {"ag-1"}
+
+    team_row = LiteLLM_TeamTable(team_id="team-gone", models=[], access_group_ids=["ag-1"])
+
+    with (
+        patch("litellm.proxy.proxy_server.prisma_client") as prisma,
+        patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin"),
+        patch("litellm.proxy.proxy_server.llm_router", None),
+        patch("litellm.proxy.management_endpoints.team_endpoints._persist_deleted_team_records", new_callable=AsyncMock),
+        patch("litellm.proxy.management_endpoints.team_endpoints._verify_team_access", new_callable=AsyncMock),
+        patch(
+            "litellm.proxy.management_endpoints.team_endpoints.sync_team_access_group_membership",
+            new_callable=AsyncMock,
+        ) as sync,
+    ):
+        prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row)
+        prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
+        prisma.delete_data = AsyncMock(return_value=[team_row])
+        prisma.db.execute_raw = AsyncMock(return_value=0)
+        prisma.db.litellm_teammembership.delete_many = AsyncMock(return_value=0)
+
+        await delete_team(
+            data=DeleteTeamRequest(team_ids=["team-gone"]),
+            http_request=Mock(spec=Request),
+            user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"),
+        )
+
+    assert sync.await_args_list[0].kwargs["team_id"] == "team-gone"
+
+
+@pytest.mark.asyncio
+async def test_invalidate_access_group_cache_deletes_the_cached_object():
+    """The mirror's cache step is what stops a revoked group granting from cache until TTL,
+    so pin that it actually reaches the delete rather than only being called."""
+    from litellm.proxy.management_helpers.access_group_team_sync import (
+        invalidate_access_group_cache,
+    )
+
+    cache, logging_obj = MagicMock(), MagicMock()
+    with (
+        patch("litellm.proxy.proxy_server.user_api_key_cache", cache),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj", logging_obj),
+        patch(
+            "litellm.proxy.management_helpers.access_group_team_sync._delete_cache_access_object",
+            new_callable=AsyncMock,
+        ) as delete_cached,
+    ):
+        await invalidate_access_group_cache("ag-1")
+
+    assert delete_cached.await_args.kwargs == {
+        "access_group_id": "ag-1",
+        "user_api_key_cache": cache,
+        "proxy_logging_obj": logging_obj,
+    }
diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
index da805b864fc..b83b862d6b8 100644
--- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
+++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py
@@ -2,6 +2,7 @@ import asyncio
 import json
 import os
 import sys
+from contextlib import asynccontextmanager
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -37,6 +38,20 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
 )
 
 
+def _wire_team_create_tx(prisma_client):
+    """`/team/new` inserts the team and mirrors it onto the access groups in one transaction,
+    so a mocked client has to hand its team table back out of `db.tx()`."""
+
+    @asynccontextmanager
+    async def _tx():
+        yield SimpleNamespace(
+            litellm_teamtable=prisma_client.db.litellm_teamtable,
+            query_raw=AsyncMock(return_value=[]),
+        )
+
+    prisma_client.db.tx = lambda *_args, **_kwargs: _tx()
+
+
 def test_microsoft_sso_handler_openid_from_response_user_principal_name():
     # Arrange
     # Create a mock response similar to what Microsoft SSO would return
@@ -577,6 +592,7 @@ async def test_default_team_params(team_params):
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
     mock_prisma.db.litellm_teamtable.create = AsyncMock()
+    _wire_team_create_tx(mock_prisma)
     mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_prisma.get_data = AsyncMock(return_value=None)
     mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object)
@@ -624,6 +640,7 @@ async def test_default_team_params_organization_id_reaches_sso_created_team(team
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
     mock_prisma.db.litellm_teamtable.create = AsyncMock()
+    _wire_team_create_tx(mock_prisma)
     mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_prisma.get_data = AsyncMock(return_value=None)
     mock_prisma.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data)
@@ -671,6 +688,7 @@ async def test_create_team_without_default_params():
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_teamtable.find_first = AsyncMock(return_value=None)
     mock_prisma.db.litellm_teamtable.create = AsyncMock()
+    _wire_team_create_tx(mock_prisma)
     mock_prisma.db.litellm_teamtable.count = AsyncMock(return_value=0)
     mock_prisma.get_data = AsyncMock(return_value=None)
     mock_prisma.jsonify_team_object = MagicMock(side_effect=mock_jsonify_team_object)
diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py
new file mode 100644
index 00000000000..eb11292cf42
--- /dev/null
+++ b/tests/test_litellm/proxy/management_helpers/test_access_group_team_sync.py
@@ -0,0 +1,39 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from litellm.proxy.management_helpers.access_group_team_sync import (
+    invalidate_access_group_caches,
+)
+
+
+@pytest.mark.asyncio
+async def test_one_unreachable_cache_does_not_skip_the_other_groups(monkeypatch):
+    """
+    `assigned_team_ids` is an authorization input, so a group whose cache still holds the
+    revoked grant keeps serving it until the entry is dropped.
+
+    A sequential loop would stop at the first failing group and leave the groups behind it
+    serving stale grants, and swallowing the failure would report success to the admin for
+    a revoke that never took effect. Every group has to be attempted, and the endpoint has
+    to fail so the caller can retry.
+    """
+    attempted: list[str] = []
+
+    async def _invalidate(access_group_id: str) -> None:
+        attempted.append(access_group_id)
+        if access_group_id == "ag-redis-down":
+            raise ConnectionError("redis unreachable")
+
+    monkeypatch.setattr(
+        "litellm.proxy.management_helpers.access_group_team_sync.invalidate_access_group_cache",
+        _invalidate,
+    )
+
+    with pytest.raises(ConnectionError):
+        await invalidate_access_group_caches(("ag-redis-down", "ag-2", "ag-3"))
+
+    assert attempted == ["ag-redis-down", "ag-2", "ag-3"]