From 93e6ccff8d760585e708af9b2eedb3cd517864af Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 4 Oct 2026 16:44:15 -0700 Subject: [PATCH] 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 --- .../traces/list/AgentTracesTable.test.tsx | 46 ++++++++++++++++++- .../lens/traces/list/AgentTracesTable.tsx | 13 +++++- .../shared/DataTable/DataTableViewOptions.tsx | 33 +++++++++++-- 3 files changed, 84 insertions(+), 8 deletions(-) diff --git a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.test.tsx b/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.test.tsx index 52f8544b351..ef9bd287016 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.test.tsx +++ b/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( + , + ), + ); + 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); + }); +}); diff --git a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.tsx b/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.tsx index d1d645d0446..541d92550e9 100644 --- a/ui/litellm-dashboard/src/components/lens/traces/list/AgentTracesTable.tsx +++ b/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[] = [ { id: "time", size: 170, + enableHiding: false, header: () => ( Time @@ -95,6 +98,7 @@ const RUN_COLUMNS: ColumnDef[] = [ ), meta: { + title: "Time", className: "font-mono tabular-nums text-muted-foreground", renderSkeleton: () => , }, @@ -102,6 +106,7 @@ const RUN_COLUMNS: ColumnDef[] = [ { id: "agent", size: 160, + enableHiding: false, header: "Agent", cell: ({ row }) => , meta: { @@ -157,9 +162,10 @@ const RUN_COLUMNS: ColumnDef[] = [ { id: "open", size: 32, - header: "", + enableHiding: false, + header: ({ table }) => , cell: () => , - 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 = { data: traces, columns: RUN_COLUMNS, getRowId: runKey, autoResetAll: false, + state: { columnVisibility }, + onColumnVisibilityChange, getCoreRowModel: getCoreRowModel(), }; const table = useReactTable(tableOptions); diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableViewOptions.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTableViewOptions.tsx index ca60e42d9aa..86d99fd72a6 100644 --- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTableViewOptions.tsx +++ b/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 { table: Table; label?: string; + iconOnly?: boolean; className?: string; } -export function DataTableViewOptions({ table, label = "View", className }: DataTableViewOptionsProps) { +export function DataTableViewOptions({ + table, + label = "View", + iconOnly = false, + className, +}: DataTableViewOptionsProps) { const hideableColumns = table.getAllLeafColumns().filter((column) => column.getCanHide()); if (hideableColumns.length === 0) { @@ -23,10 +30,26 @@ export function DataTableViewOptions({ table, label = "View", className } - - {label} - + iconOnly ? ( + + ) : ( + + ) } />