feat(ui): add controlled row selection to the shared DataTable (#34167)

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

@@ -18,6 +18,7 @@ import {
   type OnChangeFn,
   type Row,
   type RowData,
+  type RowSelectionState,
   type Table,
   type TableOptions,
   useReactTable,
@@ -70,6 +71,8 @@ export function validateDataTableConfig<TData extends RowData, TValue>(
   const bothSortingSources = props.defaultSorting !== undefined && props.sorting !== undefined;
   const bothFilterSources = props.defaultColumnFilters !== undefined && props.columnFilters !== undefined;
 
+  const controlledSelectionIncomplete = props.rowSelection !== undefined && props.onRowSelectionChange === undefined;
+
   return [
     serverSortingIncomplete ? "sortingMode='server' requires both `sorting` and `onSortingChange`." : null,
     serverPaginationIncomplete
@@ -80,6 +83,9 @@ export function validateDataTableConfig<TData extends RowData, TValue>(
     bothFilterSources
       ? "Provide either `defaultColumnFilters` (uncontrolled) or `columnFilters` (controlled), not both."
       : null,
+    controlledSelectionIncomplete
+      ? "Controlled `rowSelection` requires `onRowSelectionChange`; without it selection changes are dropped."
+      : null,
   ].filter((message): message is string => message !== null);
 }
 
@@ -448,6 +454,9 @@ function useDataTableInstance<TData extends RowData, TValue>(props: DataTablePro
     renderSubComponent,
     expanded,
     onExpandedChange,
+    enableRowSelection,
+    rowSelection,
+    onRowSelectionChange,
   } = props;
 
   const sortingState = useControllable(sorting, onSortingChange, defaultSorting ?? []);
@@ -462,6 +471,7 @@ function useDataTableInstance<TData extends RowData, TValue>(props: DataTablePro
   );
   const globalFilterState = useControllable<string>(globalFilter, onGlobalFilterChange, "");
   const expandedState = useControllable<ExpandedState>(expanded, onExpandedChange, {});
+  const rowSelectionState = useControllable<RowSelectionState>(rowSelection, onRowSelectionChange, {});
   const [columnVisibility, setColumnVisibility] = useState<VisibilityState>(defaultColumnVisibility ?? {});
   const [columnSizing, setColumnSizing] = useState<ColumnSizingState>({});
   const columnPinning = React.useMemo(() => derivePinning(columns), [columns]);
@@ -476,6 +486,7 @@ function useDataTableInstance<TData extends RowData, TValue>(props: DataTablePro
       columnFilters: filterState.value,
       globalFilter: globalFilterState.value,
       expanded: expandedState.value,
+      rowSelection: rowSelectionState.value,
       columnVisibility,
       columnSizing,
     },
@@ -491,11 +502,13 @@ function useDataTableInstance<TData extends RowData, TValue>(props: DataTablePro
     onColumnFiltersChange: filterState.onChange,
     onGlobalFilterChange: globalFilterState.onChange,
     onExpandedChange: expandedState.onChange,
+    onRowSelectionChange: rowSelectionState.onChange,
     onColumnVisibilityChange: setColumnVisibility,
     onColumnSizingChange: setColumnSizing,
     getCoreRowModel: getCoreRowModel(),
     ...buildRowModels(sortingMode, paginationMode, filterMode, expansionGuard),
     ...(getRowId !== undefined ? { getRowId } : {}),
+    ...(enableRowSelection !== undefined ? { enableRowSelection } : {}),
     ...(paginationMode === "server" && rowCount !== undefined ? { rowCount } : {}),
   };
 

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

@@ -0,0 +1,137 @@
+import type { ColumnDef, RowSelectionState } from "@tanstack/react-table";
+import { render, screen } from "@testing-library/react";
+import userEvent from "@testing-library/user-event";
+import { useState } from "react";
+import { describe, expect, it } from "vitest";
+
+import { createSelectionColumn, DataTable, validateDataTableConfig } from "./index";
+
+interface Model {
+  id: string;
+  name: string;
+}
+
+const data: Model[] = [
+  { id: "m1", name: "Alpha" },
+  { id: "m2", name: "Beta" },
+  { id: "m3", name: "Gamma" },
+];
+
+const columns: ColumnDef<Model, unknown>[] = [
+  createSelectionColumn<Model>({ rowAriaLabel: (row) => `Select ${row.original.name}` }),
+  { id: "name", accessorKey: "name", header: "Name", enableSorting: false },
+];
+
+const selectAll = () => screen.getByTestId("datatable-select-all");
+const rowBox = (id: string) => screen.getByTestId(`datatable-select-row-${id}`);
+const selectedCount = () => screen.getByTestId("count");
+
+function ControlledHarness() {
+  const [rowSelection, setRowSelection] = useState<RowSelectionState>({});
+
+  return (
+    <>
+      <span data-testid="keys">
+        {Object.keys(rowSelection)
+          .filter((key) => rowSelection[key])
+          .sort()
+          .join(",")}
+      </span>
+      <button type="button" data-testid="clear" onClick={() => setRowSelection({})}>
+        clear
+      </button>
+      <DataTable
+        data={data}
+        columns={columns}
+        getRowId={(row) => row.id}
+        rowSelection={rowSelection}
+        onRowSelectionChange={setRowSelection}
+      />
+    </>
+  );
+}
+
+describe("DataTable row selection", () => {
+  it("supports uncontrolled per-row toggle, select-all, and indeterminate", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <DataTable
+        data={data}
+        columns={columns}
+        getRowId={(row) => row.id}
+        toolbar={(table) => <span data-testid="count">{table.getSelectedRowModel().rows.length}</span>}
+      />,
+    );
+
+    expect(selectedCount()).toHaveTextContent("0");
+
+    await user.click(rowBox("m1"));
+    expect(selectedCount()).toHaveTextContent("1");
+    expect(selectAll()).toHaveAttribute("aria-checked", "mixed");
+
+    await user.click(selectAll());
+    expect(selectedCount()).toHaveTextContent("3");
+    expect(selectAll()).toHaveAttribute("aria-checked", "true");
+
+    await user.click(selectAll());
+    expect(selectedCount()).toHaveTextContent("0");
+  });
+
+  it("keys controlled selection by getRowId so the parent can map back to entities", async () => {
+    const user = userEvent.setup();
+    render(<ControlledHarness />);
+
+    await user.click(rowBox("m2"));
+    expect(screen.getByTestId("keys")).toHaveTextContent("m2");
+
+    await user.click(rowBox("m3"));
+    expect(screen.getByTestId("keys")).toHaveTextContent("m2,m3");
+  });
+
+  it("lets the parent clear the selection, the pattern an external pager needs", async () => {
+    const user = userEvent.setup();
+    render(<ControlledHarness />);
+
+    await user.click(selectAll());
+    expect(screen.getByTestId("keys")).toHaveTextContent("m1,m2,m3");
+
+    await user.click(screen.getByTestId("clear"));
+    expect(screen.getByTestId("keys")).toBeEmptyDOMElement();
+    expect(rowBox("m1")).toHaveAttribute("aria-checked", "false");
+  });
+
+  it("respects an enableRowSelection predicate", async () => {
+    const user = userEvent.setup();
+
+    render(
+      <DataTable
+        data={data}
+        columns={columns}
+        getRowId={(row) => row.id}
+        enableRowSelection={(row) => row.original.id !== "m2"}
+        toolbar={(table) => <span data-testid="count">{table.getSelectedRowModel().rows.length}</span>}
+      />,
+    );
+
+    expect(rowBox("m2")).toHaveAttribute("aria-disabled", "true");
+
+    await user.click(rowBox("m2"));
+    expect(selectedCount()).toHaveTextContent("0");
+
+    await user.click(rowBox("m1"));
+    expect(selectedCount()).toHaveTextContent("1");
+  });
+
+  it("rejects controlled rowSelection without onRowSelectionChange", () => {
+    const errors = validateDataTableConfig<Model, unknown>({ data, columns, rowSelection: { m1: true } });
+
+    expect(errors).toContain(
+      "Controlled `rowSelection` requires `onRowSelectionChange`; without it selection changes are dropped.",
+    );
+  });
+
+  it("does not complain when selection is left uncontrolled", () => {
+    expect(validateDataTableConfig<Model, unknown>({ data, columns })).toHaveLength(0);
+  });
+});

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

@@ -0,0 +1,53 @@
+"use client";
+
+import type { ColumnDef, Row, RowData, Table } from "@tanstack/react-table";
+
+import { Checkbox } from "@/components/ui/checkbox";
+
+interface SelectionColumnOptions<TData> {
+  rowAriaLabel?: (row: Row<TData>) => string;
+}
+
+function SelectAllCheckbox<TData>({ table }: { table: Table<TData> }) {
+  const allSelected = table.getIsAllPageRowsSelected();
+  const someSelected = table.getIsSomePageRowsSelected();
+
+  return (
+    <Checkbox
+      aria-label="Select all rows"
+      data-testid="datatable-select-all"
+      checked={allSelected}
+      indeterminate={someSelected && !allSelected}
+      onCheckedChange={(checked) => table.toggleAllPageRowsSelected(Boolean(checked))}
+    />
+  );
+}
+
+function SelectRowCheckbox<TData>({ row, label }: { row: Row<TData>; label: string }) {
+  return (
+    <Checkbox
+      aria-label={label}
+      data-testid={`datatable-select-row-${row.id}`}
+      checked={row.getIsSelected()}
+      disabled={!row.getCanSelect()}
+      onCheckedChange={(checked) => row.toggleSelected(Boolean(checked))}
+    />
+  );
+}
+
+export function createSelectionColumn<TData extends RowData>(
+  options: SelectionColumnOptions<TData> = {},
+): ColumnDef<TData, unknown> {
+  const { rowAriaLabel } = options;
+
+  return {
+    id: "select",
+    size: 44,
+    enableSorting: false,
+    enableHiding: false,
+    enableResizing: false,
+    meta: { title: "Select", className: "w-11", headerClassName: "w-11" },
+    header: ({ table }) => <SelectAllCheckbox table={table} />,
+    cell: ({ row }) => <SelectRowCheckbox row={row} label={rowAriaLabel?.(row) ?? "Select row"} />,
+  };
+}

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

@@ -3,6 +3,7 @@ import "./columnMeta";
 export { DataTable, DataTableConfigError, validateDataTableConfig } from "./DataTable";
 export { DataTableFilterDrawer, DataTableFilterField, type FilterDraft } from "./DataTableFilterDrawer";
 export { DataTablePagination, DEFAULT_PAGE_SIZE_OPTIONS } from "./DataTablePagination";
+export { createSelectionColumn } from "./DataTableSelectionColumn";
 export { DataTableToolbar } from "./DataTableToolbar";
 export { DataTableViewOptions } from "./DataTableViewOptions";
 export {

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

@@ -6,6 +6,7 @@ import type {
   PaginationState,
   Row,
   RowData,
+  RowSelectionState,
   SortingState,
   Table,
   VisibilityState,
@@ -59,6 +60,10 @@ export interface DataTableProps<TData extends RowData, TValue> {
   expanded?: ExpandedState;
   onExpandedChange?: OnChangeFn<ExpandedState>;
 
+  enableRowSelection?: boolean | ((row: Row<TData>) => boolean);
+  rowSelection?: RowSelectionState;
+  onRowSelectionChange?: OnChangeFn<RowSelectionState>;
+
   onRowClick?: (row: TData) => void;
 
   rowClassName?: (row: Row<TData>) => string;

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

@@ -0,0 +1,28 @@
+"use client";
+
+import { Checkbox as CheckboxPrimitive } from "@base-ui/react/checkbox";
+
+import { cn } from "@/lib/cva.config";
+import { CheckIcon } from "lucide-react";
+
+function Checkbox({ className, ...props }: CheckboxPrimitive.Root.Props) {
+  return (
+    <CheckboxPrimitive.Root
+      data-slot="checkbox"
+      className={cn(
+        "peer relative flex size-4 shrink-0 items-center justify-center rounded-[4px] border border-input shadow-xs transition-shadow outline-none group-has-disabled/field:opacity-50 after:absolute after:-inset-x-3 after:-inset-y-2 focus-visible:border-ring focus-visible:ring-3 focus-visible:ring-ring/50 disabled:cursor-not-allowed disabled:opacity-50 aria-invalid:border-destructive aria-invalid:ring-3 aria-invalid:ring-destructive/20 aria-invalid:aria-checked:border-primary dark:bg-input/30 dark:aria-invalid:border-destructive/50 dark:aria-invalid:ring-destructive/40 data-checked:border-primary data-checked:bg-primary data-checked:text-primary-foreground dark:data-checked:bg-primary",
+        className,
+      )}
+      {...props}
+    >
+      <CheckboxPrimitive.Indicator
+        data-slot="checkbox-indicator"
+        className="grid place-content-center text-current transition-none [&>svg]:size-3.5"
+      >
+        <CheckIcon />
+      </CheckboxPrimitive.Indicator>
+    </CheckboxPrimitive.Root>
+  );
+}
+
+export { Checkbox };