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 ? (
+
+ ) : (
+
+ )
}
/>