feat(ui): let Lens runs pick their columns from the table header

Adds an icon-only variant to the shared DataTableViewOptions picker so any
InspectorTable can host it in a header cell, and wires it into the runs list
with visibility saved per table. Time and Agent stay pinned visible.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.test.tsx

@@ -1,5 +1,6 @@
 import { fireEvent, render, screen, within } from "@testing-library/react";
-import { describe, expect, it, vi } from "vitest";
+import userEvent from "@testing-library/user-event";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 
 import { renderWithProviders } from "../../../../../tests/test-utils";
 import { Inspector } from "@/components/shared/Inspector";
@@ -108,3 +109,46 @@ describe("AgentTracesTable virtualization", () => {
     expect(screen.queryByText("question 0")).not.toBeInTheDocument();
   });
 });
+
+describe("AgentTracesTable column picker", () => {
+  const runs = (traceList as TracePage).data as TraceSummary[];
+  const renderRuns = () =>
+    renderWithProviders(
+      inList(
+        <AgentTracesTable
+          traces={runs}
+          isLoading={false}
+          error={null}
+          hasMore={false}
+          onLoadMore={vi.fn()}
+          onSetUpTracing={vi.fn()}
+        />,
+      ),
+    );
+  const headers = () => screen.getAllByRole("columnheader").map((header) => header.textContent);
+
+  beforeEach(() => localStorage.clear());
+
+  it("hides a column from the header picker and keeps it hidden after a remount", async () => {
+    const user = userEvent.setup();
+    const { unmount } = renderRuns();
+    const before = headers().length;
+    expect(headers()).toContain("Cost");
+
+    await user.click(screen.getByRole("button", { name: "Columns" }));
+    expect(screen.queryByTestId("view-option-time")).not.toBeInTheDocument();
+    expect(screen.queryByTestId("view-option-agent")).not.toBeInTheDocument();
+    expect(screen.queryByTestId("view-option-open")).not.toBeInTheDocument();
+    await user.click(await screen.findByTestId("view-option-cost"));
+
+    expect(headers()).not.toContain("Cost");
+    expect(headers()).toHaveLength(before - 1);
+    const [firstRow] = screen.getAllByTestId("agent-trace-row");
+    expect(within(firstRow).getAllByRole("cell")).toHaveLength(before - 1);
+
+    unmount();
+    renderRuns();
+    expect(headers()).not.toContain("Cost");
+    expect(headers()).toHaveLength(before - 1);
+  });
+});

ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.tsx

@@ -5,6 +5,8 @@ import { ArrowDown, ChevronRight } from "lucide-react";
 import { useEffect } from "react";
 import { useInView } from "react-intersection-observer";
 
+import { DataTableViewOptions } from "@/components/shared/DataTable/DataTableViewOptions";
+import { usePersistedColumnVisibility } from "@/components/shared/DataTable/usePersistedColumnVisibility";
 import { InspectorTable, useInspectorTable } from "@/components/shared/InspectorTable";
 import { Button } from "@/components/ui/button";
 import { Skeleton } from "@/components/ui/skeleton";
@@ -81,6 +83,7 @@ const RUN_COLUMNS: ColumnDef<TraceSummary>[] = [
   {
     id: "time",
     size: 170,
+    enableHiding: false,
     header: () => (
       <span className="inline-flex items-center gap-1 whitespace-nowrap">
         Time <ArrowDown className="size-2.5" />
@@ -95,6 +98,7 @@ const RUN_COLUMNS: ColumnDef<TraceSummary>[] = [
       </span>
     ),
     meta: {
+      title: "Time",
       className: "font-mono tabular-nums text-muted-foreground",
       renderSkeleton: () => <Skeleton className="h-3 w-24" />,
     },
@@ -102,6 +106,7 @@ const RUN_COLUMNS: ColumnDef<TraceSummary>[] = [
   {
     id: "agent",
     size: 160,
+    enableHiding: false,
     header: "Agent",
     cell: ({ row }) => <AgentCell run={row.original} />,
     meta: {
@@ -157,9 +162,10 @@ const RUN_COLUMNS: ColumnDef<TraceSummary>[] = [
   {
     id: "open",
     size: 32,
-    header: "",
+    enableHiding: false,
+    header: ({ table }) => <DataTableViewOptions table={table} label="Columns" iconOnly />,
     cell: () => <ChevronRight className="size-3 text-muted-foreground/60" />,
-    meta: { className: "px-0", renderSkeleton: () => null },
+    meta: { className: "px-0", headerClassName: "px-1", renderSkeleton: () => null },
   },
 ];
 
@@ -214,11 +220,14 @@ export function AgentTracesTable({
   const settled = !isLoading && !error;
   const isEmpty = settled && !hasMore && traces.length === 0;
   const autoContinue = settled && hasMore && traces.length > 0;
+  const { columnVisibility, onColumnVisibilityChange } = usePersistedColumnVisibility("lens-traces");
   const tableOptions: TableOptions<TraceSummary> = {
     data: traces,
     columns: RUN_COLUMNS,
     getRowId: runKey,
     autoResetAll: false,
+    state: { columnVisibility },
+    onColumnVisibilityChange,
     getCoreRowModel: getCoreRowModel(),
   };
   const table = useReactTable(tableOptions);

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

@@ -5,14 +5,21 @@ import type { Table } from "@tanstack/react-table";
 import { Check, Columns3 } from "lucide-react";
 
 import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/cva.config";
 
 interface DataTableViewOptionsProps<TData> {
   table: Table<TData>;
   label?: string;
+  iconOnly?: boolean;
   className?: string;
 }
 
-export function DataTableViewOptions<TData>({ table, label = "View", className }: DataTableViewOptionsProps<TData>) {
+export function DataTableViewOptions<TData>({
+  table,
+  label = "View",
+  iconOnly = false,
+  className,
+}: DataTableViewOptionsProps<TData>) {
   const hideableColumns = table.getAllLeafColumns().filter((column) => column.getCanHide());
 
   if (hideableColumns.length === 0) {
@@ -23,10 +30,26 @@ export function DataTableViewOptions<TData>({ table, label = "View", className }
     <Menu.Root>
       <Menu.Trigger
         render={
-          <Button variant="outline" size="sm" className={className} data-testid="view-options-trigger">
-            <Columns3 />
-            {label}
-          </Button>
+          iconOnly ? (
+            <Button
+              variant="ghost"
+              size="icon-xs"
+              aria-label={label}
+              title={label}
+              className={cn(
+                "text-muted-foreground/60 hover:text-foreground aria-expanded:bg-muted aria-expanded:text-foreground",
+                className,
+              )}
+              data-testid="view-options-trigger"
+            >
+              <Columns3 />
+            </Button>
+          ) : (
+            <Button variant="outline" size="sm" className={className} data-testid="view-options-trigger">
+              <Columns3 />
+              {label}
+            </Button>
+          )
         }
       />
       <Menu.Portal>