fix(ui): address Virtual Keys redesign review nits (#33112)

* fix(ui): address Virtual Keys redesign review nits

Restore sorting by budget on the merged Spend / Budget column. The column now
uses a new DataTableMultiSortHeader whose chevron opens a menu offering Spend
and Budget in both directions plus Reset, so the progress-bar cell stays merged
while the sort field becomes an explicit choice. Sorting is server-side, so the
chosen field id (spend or max_budget, both accepted by /key/list) flows straight
through as sort_by

Fill the DataTable to its container width when column resizing is on. The table
width was pinned to the sum of column widths, so hiding columns left an empty
gutter on the right. It now keeps that width as a minimum and stretches to 100%
on underflow while still scrolling on overflow, which also covers the same gap
in TeamVirtualKeysTable since both share the component

Drop the dark background box behind the page-header icon so the Virtual Keys
header reads like the Teams header, and pull the 4-line inline filter lambda in
SearchSelect out into a named matchesQuery helper

Extends the DataTable and VirtualKeysTable tests to cover the new multi-field
sort menu (field id maps to sort_by, active indicator, reset) and the
fill-to-container width

* fix(ui): emphasize the active field in the Spend / Budget sort header

The merged Spend / Budget header always read "Spend / Budget" regardless of
which field drove the sort, so after picking Budget descending there was no way
to tell what was sorted without reopening the menu. The header now builds its
label from the sort fields and emphasizes whichever one is active (bold,
full-strength text) while muting the other, so the sorted column reads at a
glance alongside the direction chevron. Drops the now-redundant title prop since
the label is derived from the fields

* fix(ui): remove w-full so the keys page content stops overflowing by 32px

The virtual keys content wrapper used "w-full mx-4", which sets the width to
100% of the parent and then adds 16px of horizontal margin on each side, so its
margin-box came to 100% + 32px and overflowed the scrollable main region by
exactly 32px. That surfaced as a horizontal scrollbar along the bottom of the
whole content area, under the pagination. A block div is already full-width, so
dropping w-full lets mx-4 inset it correctly with no overflow

* fix(ui): darken the clickable Key cell on hover so it reads as clickable

The Key cell was the click target that opens the key detail, but hovering only
faded the chevron in with no change to the cell itself, so there was no cue that
the area was clickable. Give the cell a subtle muted background and a pointer
cursor on hover. The button spans the full cell (a negative inline margin plus a
matching width offset so the hover fill reaches both cell edges while the title
stays aligned with the other columns)

ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx

@@ -266,7 +266,8 @@ it("should render the redesigned table headers", () => {
   expect(screen.getByText("Key")).toBeInTheDocument();
   expect(screen.getByText("Team")).toBeInTheDocument();
   expect(screen.getByText("Models")).toBeInTheDocument();
-  expect(screen.getByText("Spend / Budget")).toBeInTheDocument();
+  expect(screen.getByText("Spend", { selector: "[data-sort-field='spend']" })).toBeInTheDocument();
+  expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" })).toBeInTheDocument();
 });
 
 it("sorts by the backend key_alias field (not the column label) when the Key header is clicked", async () => {
@@ -280,6 +281,51 @@ it("sorts by the backend key_alias field (not the column label) when the Key hea
   });
 });
 
+it("sorts by the backend max_budget field when 'Budget descending' is chosen from the Spend / Budget menu", async () => {
+  const user = userEvent.setup();
+  renderWithProviders(<VirtualKeysTable />);
+
+  await user.click(screen.getByTestId("sort-trigger-spend"));
+  await user.click(await screen.findByText("Budget descending"));
+
+  await waitFor(() => {
+    expect(mockUseKeys).toHaveBeenLastCalledWith(
+      1,
+      50,
+      expect.objectContaining({ sortBy: "max_budget", sortOrder: "desc" }),
+    );
+  });
+});
+
+it("emphasizes the active field in the Spend / Budget header so the sorted column reads without opening the menu", async () => {
+  const user = userEvent.setup();
+  renderWithProviders(<VirtualKeysTable />);
+
+  await user.click(screen.getByTestId("sort-trigger-spend"));
+  await user.click(await screen.findByText("Budget descending"));
+
+  await waitFor(() => {
+    expect(screen.getByText("Budget", { selector: "[data-sort-field='max_budget']" }).className).toContain(
+      "font-semibold",
+    );
+  });
+  expect(screen.getByText("Spend", { selector: "[data-sort-field='spend']" }).className).toContain(
+    "text-muted-foreground",
+  );
+});
+
+it("sorts by spend ascending when 'Spend ascending' is chosen from the Spend / Budget menu", async () => {
+  const user = userEvent.setup();
+  renderWithProviders(<VirtualKeysTable />);
+
+  await user.click(screen.getByTestId("sort-trigger-spend"));
+  await user.click(await screen.findByText("Spend ascending"));
+
+  await waitFor(() => {
+    expect(mockUseKeys).toHaveBeenLastCalledWith(1, 50, expect.objectContaining({ sortBy: "spend", sortOrder: "asc" }));
+  });
+});
+
 it("should open KeyInfoView when clicking the key cell", async () => {
   renderWithProviders(<VirtualKeysTable />);
 

ui/litellm-dashboard/src/components/VirtualKeysPage/keyTableColumns.tsx

@@ -4,7 +4,7 @@ import { InfoCircleOutlined } from "@ant-design/icons";
 import { ColumnDef } from "@tanstack/react-table";
 import { Popover, Typography } from "antd";
 
-import { DataTableSortHeader } from "@/components/shared/DataTable";
+import { DataTableMultiSortHeader, DataTableSortHeader, type DataTableSortField } from "@/components/shared/DataTable";
 import { Skeleton } from "@/components/ui/skeleton";
 import {
   DateCell,
@@ -26,6 +26,11 @@ interface KeyStatus {
   tooltip?: string;
 }
 
+const SPEND_BUDGET_SORT_FIELDS: DataTableSortField[] = [
+  { id: "spend", label: "Spend" },
+  { id: "max_budget", label: "Budget" },
+];
+
 const getKeyStatus = (key: KeyResponse): KeyStatus => {
   if (key.blocked === true) {
     const isScimBlocked = (key.metadata as Record<string, unknown> | null | undefined)?.scim_blocked === true;
@@ -291,7 +296,7 @@ export const getKeyTableColumns = ({
     id: "spend",
     accessorKey: "spend",
     meta: { title: "Spend / Budget", skeleton: "meter" },
-    header: ({ column }) => <DataTableSortHeader column={column} title="Spend / Budget" variant="header-cycle" />,
+    header: ({ table }) => <DataTableMultiSortHeader table={table} fields={SPEND_BUDGET_SORT_FIELDS} />,
     size: 180,
     enableSorting: true,
     cell: ({ row }) => {

ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx

@@ -5,7 +5,7 @@ import { useState } from "react";
 import { describe, expect, it, vi } from "vitest";
 
 import { DataTable } from "./DataTable";
-import { DataTableSortHeader } from "./DataTableSortHeader";
+import { DataTableMultiSortHeader, DataTableSortHeader } from "./DataTableSortHeader";
 import { DataTableViewOptions } from "./DataTableViewOptions";
 
 interface Person {
@@ -55,6 +55,23 @@ const dropdownSortColumns: ColumnDef<Person, unknown>[] = [
   },
 ];
 
+const multiSortColumns: ColumnDef<Person, unknown>[] = [
+  {
+    id: "spend",
+    accessorKey: "name",
+    header: ({ table }) => (
+      <DataTableMultiSortHeader
+        table={table}
+        fields={[
+          { id: "spend", label: "Spend" },
+          { id: "max_budget", label: "Budget" },
+        ]}
+      />
+    ),
+    cell: ({ row }) => <span data-testid="name-cell">{row.original.name}</span>,
+  },
+];
+
 const nameEmailColumns: ColumnDef<Person, unknown>[] = [
   {
     accessorKey: "name",
@@ -165,6 +182,60 @@ describe("DataTable sorting", () => {
     await user.click(await screen.findByText("Reset"));
     expect(names()).toEqual(["Charlie", "Alice", "Bob"]);
   });
+
+  it("multi-sort header emits the chosen field id (not the column id) as the sort key", async () => {
+    const user = userEvent.setup();
+    const onSortingChange = vi.fn();
+    render(
+      <DataTable
+        data={CHARLIE_ALICE_BOB}
+        columns={multiSortColumns}
+        sortingMode="server"
+        sorting={[]}
+        onSortingChange={onSortingChange}
+      />,
+    );
+
+    await user.click(screen.getByTestId("sort-trigger-spend"));
+    await user.click(await screen.findByText("Budget descending"));
+    expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "max_budget", desc: true }]);
+
+    await user.click(screen.getByTestId("sort-trigger-spend"));
+    await user.click(await screen.findByText("Spend ascending"));
+    expect(onSortingChange).toHaveBeenLastCalledWith([{ id: "spend", desc: false }]);
+  });
+
+  it("multi-sort header reflects the active field and direction, and Reset clears it", async () => {
+    const user = userEvent.setup();
+    const onSortingChange = vi.fn();
+    render(
+      <DataTable
+        data={CHARLIE_ALICE_BOB}
+        columns={multiSortColumns}
+        sortingMode="server"
+        sorting={[{ id: "max_budget", desc: true }]}
+        onSortingChange={onSortingChange}
+      />,
+    );
+
+    await user.click(screen.getByTestId("sort-trigger-spend"));
+    // The header trigger shows the active (descending) indicator while sorted by a field it owns.
+    expect(screen.getByTestId("sort-trigger-spend").querySelector("[data-sort-indicator='desc']")).not.toBeNull();
+
+    await user.click(await screen.findByText("Reset"));
+    expect(onSortingChange).toHaveBeenLastCalledWith([]);
+  });
+});
+
+describe("DataTable layout", () => {
+  it("stretches the table to fill the container when resizing is on, so hidden columns leave no right-side gap", () => {
+    const { container } = render(<DataTable data={CHARLIE_ALICE_BOB} columns={nameCellColumns} enableColumnResizing />);
+
+    const table = container.querySelector("table");
+    expect(table).not.toBeNull();
+    // width pins the natural column total (horizontal scroll on overflow); minWidth:100% fills the gap on underflow.
+    expect(table?.style.minWidth).toBe("100%");
+  });
 });
 
 describe("DataTable pagination", () => {

ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx

@@ -532,7 +532,7 @@ export function DataTable<TData extends RowData, TValue>(props: DataTableProps<T
   const rows = table.getRowModel().rows;
   const visibleColumnCount = table.getVisibleLeafColumns().length;
   const stickyHeader = maxBodyHeight !== undefined;
-  const tableStyle = enableColumnResizing ? { width: table.getTotalSize() } : undefined;
+  const tableStyle = enableColumnResizing ? { width: table.getTotalSize(), minWidth: "100%" } : undefined;
 
   const renderPagination = (): React.ReactNode => {
     if (paginationSlot !== undefined) {

ui/litellm-dashboard/src/components/shared/DataTable/DataTableSortHeader.tsx

@@ -1,8 +1,8 @@
 "use client";
 
 import { Menu } from "@base-ui/react/menu";
-import type { Column, SortDirection } from "@tanstack/react-table";
-import { ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react";
+import type { Column, SortDirection, Table } from "@tanstack/react-table";
+import { Check, ChevronDown, ChevronsUpDown, ChevronUp, X } from "lucide-react";
 import type * as React from "react";
 
 import { cn } from "@/lib/cva.config";
@@ -94,3 +94,101 @@ export function DataTableSortHeader<TData, TValue>({
     </button>
   );
 }
+
+export interface DataTableSortField {
+  /** Backend sort column, sent verbatim as the sorting state id (e.g. "spend", "max_budget"). */
+  id: string;
+  label: string;
+}
+
+interface DataTableMultiSortHeaderProps<TData> {
+  table: Table<TData>;
+  fields: DataTableSortField[];
+  className?: string;
+}
+
+/**
+ * Sort header for a column that merges several backend-sortable fields into one cell
+ * (e.g. a combined Spend / Budget cell). The header label is the field labels joined by " / ",
+ * with the field currently driving the sort emphasized so the active column reads at a glance
+ * without opening the menu. The chevron opens a menu offering each field in both directions.
+ */
+export function DataTableMultiSortHeader<TData>({ table, fields, className }: DataTableMultiSortHeaderProps<TData>) {
+  const active = table.getState().sorting[0];
+  const activeField = active !== undefined && fields.some((field) => field.id === active.id) ? active : undefined;
+  const activeDirection: SortDirection = activeField?.desc === true ? "desc" : "asc";
+  const sorted: false | SortDirection = activeField === undefined ? false : activeDirection;
+
+  const options = fields.flatMap((field) => [
+    { key: `${field.id}-asc`, id: field.id, desc: false, label: `${field.label} ascending`, Icon: ChevronUp },
+    { key: `${field.id}-desc`, id: field.id, desc: true, label: `${field.label} descending`, Icon: ChevronDown },
+  ]);
+
+  const segmentClass = (isActive: boolean): string => {
+    if (isActive) return "font-semibold text-foreground";
+    if (activeField) return "text-muted-foreground";
+    return "";
+  };
+
+  const labelSegments = fields.flatMap((field, index) => {
+    const isActive = activeField?.id === field.id;
+    const segment = (
+      <span key={field.id} data-sort-field={field.id} className={segmentClass(isActive)}>
+        {field.label}
+      </span>
+    );
+    if (index === 0) return [segment];
+    return [
+      <span key={`sep-${field.id}`} className="text-muted-foreground">
+        {" / "}
+      </span>,
+      segment,
+    ];
+  });
+
+  return (
+    <div className={cn("flex items-center gap-1", className)}>
+      <span className="font-medium">{labelSegments}</span>
+      <Menu.Root>
+        <Menu.Trigger
+          render={
+            <button
+              type="button"
+              data-testid={`sort-trigger-${fields[0]?.id ?? "field"}`}
+              aria-label={`Sort options for ${fields.map((field) => field.label).join(" or ")}`}
+              onClick={(event) => event.stopPropagation()}
+              className={cn(
+                "inline-flex size-6 items-center justify-center rounded-md hover:bg-muted",
+                sorted ? "text-primary" : "text-muted-foreground",
+              )}
+            >
+              <SortIndicator sorted={sorted} />
+            </button>
+          }
+        />
+        <Menu.Portal>
+          <Menu.Positioner side="bottom" align="start" sideOffset={4} className="isolate z-50">
+            <Menu.Popup className="min-w-[9rem] rounded-md bg-popover p-1 text-sm text-popover-foreground shadow-md ring-1 ring-foreground/10 outline-hidden">
+              {options.map((option) => {
+                const isActive = activeField?.id === option.id && activeField.desc === option.desc;
+                return (
+                  <Menu.Item
+                    key={option.key}
+                    className={cn(MENU_ITEM_CLASS, isActive ? "text-primary" : "")}
+                    onClick={() => table.setSorting([{ id: option.id, desc: option.desc }])}
+                  >
+                    <option.Icon className="size-3.5" /> {option.label}
+                    {isActive && <Check className="ml-auto size-3.5" />}
+                  </Menu.Item>
+                );
+              })}
+              <Menu.Item className={MENU_ITEM_CLASS} onClick={() => table.setSorting([])}>
+                <X className="size-3.5" /> Reset
+              </Menu.Item>
+            </Menu.Popup>
+          </Menu.Positioner>
+        </Menu.Portal>
+      </Menu.Root>
+    </div>
+  );
+}

ui/litellm-dashboard/src/components/shared/DataTable/index.ts

@@ -5,7 +5,12 @@ export { DataTableFilterDrawer, DataTableFilterField, type FilterDraft } from ".
 export { DataTablePagination, DEFAULT_PAGE_SIZE_OPTIONS } from "./DataTablePagination";
 export { DataTableToolbar } from "./DataTableToolbar";
 export { DataTableViewOptions } from "./DataTableViewOptions";
-export { DataTableSortHeader, type DataTableSortVariant } from "./DataTableSortHeader";
+export {
+  DataTableSortHeader,
+  DataTableMultiSortHeader,
+  type DataTableSortVariant,
+  type DataTableSortField,
+} from "./DataTableSortHeader";
 export type { DataTablePaginationProps } from "./DataTablePagination";
 export type {
   ColumnPinnedSide,

ui/litellm-dashboard/src/components/shared/PageHeader.tsx

@@ -12,12 +12,8 @@ interface PageHeaderProps {
 export function PageHeader({ title, subtitle, icon, actions }: PageHeaderProps) {
   return (
     <div className="flex flex-wrap items-start justify-between gap-4">
-      <div className="flex items-center gap-3">
-        {icon != null && (
-          <span className="flex size-9 flex-none items-center justify-center rounded-lg bg-primary text-primary-foreground">
-            {icon}
-          </span>
-        )}
+      <div className="flex items-center gap-2.5">
+        {icon != null && <span className="flex flex-none items-center text-foreground">{icon}</span>}
         <div className="min-w-0">
           <h1 className="text-xl font-semibold tracking-tight text-foreground">{title}</h1>
           {subtitle != null && <p className="mt-0.5 text-sm text-muted-foreground">{subtitle}</p>}

ui/litellm-dashboard/src/components/shared/SearchSelect.tsx

@@ -26,6 +26,12 @@ interface SearchSelectProps {
   className?: string;
 }
 
+const matchesQuery = (option: SearchSelectOption, query: string): boolean => {
+  const q = query.trim().toLowerCase();
+  if (!q) return true;
+  return option.label.toLowerCase().includes(q) || (option.sublabel?.toLowerCase().includes(q) ?? false);
+};
+
 export function SearchSelect({
   options,
   value,
@@ -44,11 +50,7 @@ export function SearchSelect({
       onValueChange={(item: SearchSelectOption | null) => onValueChange(item?.value ?? "")}
       isItemEqualToValue={(a: SearchSelectOption, b: SearchSelectOption) => a.value === b.value}
       itemToStringLabel={(item: SearchSelectOption) => item.label}
-      filter={(item: SearchSelectOption, query: string) => {
-        const q = query.trim().toLowerCase();
-        if (!q) return true;
-        return item.label.toLowerCase().includes(q) || (item.sublabel?.toLowerCase().includes(q) ?? false);
-      }}
+      filter={matchesQuery}
       disabled={disabled}
     >
       <ComboboxInput

ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.test.tsx

@@ -26,12 +26,15 @@ describe("IdentityCell", () => {
     expect(screen.queryByRole("button")).not.toBeInTheDocument();
   });
 
-  it("renders a clickable button and fires onClick", async () => {
+  it("renders a clickable button that signals interactivity and fires onClick", async () => {
     const onClick = vi.fn();
     const user = userEvent.setup();
     render(<IdentityCell title="prod-gateway" subtitle="sk-...v0Pw" onClick={onClick} />);
     const button = screen.getByRole("button");
     expect(button.querySelector(".lucide-chevron-right")).not.toBeNull();
+    // The clickable area must read as clickable: a hover background and a pointer cursor.
+    expect(button.className).toContain("hover:bg-muted");
+    expect(button.className).toContain("cursor-pointer");
     await user.click(button);
     expect(onClick).toHaveBeenCalledTimes(1);
   });

ui/litellm-dashboard/src/components/shared/table_cells/identity_cell.tsx

@@ -36,7 +36,10 @@ export function IdentityCell({ title, subtitle, badge, onClick, className, title
       <button
         type="button"
         onClick={onClick}
-        className={cn("group flex w-full items-center gap-2 rounded-md py-1 text-left", className)}
+        className={cn(
+          "group -mx-2 flex w-[calc(100%+1rem)] cursor-pointer items-center gap-2 rounded-md px-2 py-1 text-left transition-colors hover:bg-muted",
+          className,
+        )}
       >
         {body}
         <ChevronRight className="ml-auto size-4 shrink-0 text-muted-foreground opacity-0 transition-opacity group-hover:opacity-100" />

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

@@ -285,7 +285,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
   const canCreateKey = userRole !== "Admin Viewer" && userRole !== "proxy_admin_viewer";
 
   return (
-    <div className="w-full mx-4 h-[75vh]">
+    <div className="mx-4 h-[75vh]">
       <Grid numItems={1} className="gap-2 p-8 w-full mt-2">
         <Col numColSpan={1} className="flex flex-col gap-2">
           <VirtualKeysTable