From 4d30bbce453a06e850cbf6007a42dd7e00803380 Mon Sep 17 00:00:00 2001
From: ryan-crabbe-berri <ryan@berri.ai>
Date: Wed, 16 Sep 2026 09:34:05 -0700
Subject: [PATCH] fix(ui): keep controlled client-side table pages across data
 reloads

Controlled client-mode DataTables no longer let TanStack reset the page index when rows change, since the owner of the pagination state decides the page. Once rows settle, a page past the end snaps back to the last page, matching server mode
---
 .../shared/DataTable/DataTable.test.tsx       | 73 +++++++++++++++++++
 .../components/shared/DataTable/DataTable.tsx | 34 ++++++++-
 2 files changed, 104 insertions(+), 3 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
index 336fc43d695..a7e4befa6ac 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.test.tsx
@@ -357,6 +357,79 @@ describe("DataTable pagination", () => {
     expect(onChange).not.toHaveBeenCalled();
   });
 
+  type ClientPageHarnessProps = {
+    data: Person[];
+    isLoading?: boolean;
+    initialPageIndex: number;
+    onChange: (next: PaginationState) => void;
+  };
+
+  function ClientPageHarness({ data, isLoading = false, initialPageIndex, onChange }: ClientPageHarnessProps) {
+    const [pagination, setPagination] = useState<PaginationState>({ pageIndex: initialPageIndex, pageSize: 2 });
+    const handleChange: OnChangeFn<PaginationState> = (updater) => {
+      const next = typeof updater === "function" ? updater(pagination) : updater;
+      onChange(next);
+      setPagination(next);
+    };
+    return (
+      <DataTable
+        data={data}
+        columns={nameCellColumns}
+        paginationMode="client"
+        pageSizeOptions={[2]}
+        pagination={pagination}
+        onPaginationChange={handleChange}
+        isLoading={isLoading}
+      />
+    );
+  }
+
+  it("client mode keeps a controlled page when rows arrive after loading and when they are refetched", async () => {
+    const onChange = vi.fn();
+    const { rerender } = render(<ClientPageHarness data={[]} isLoading initialPageIndex={1} onChange={onChange} />);
+
+    rerender(<ClientPageHarness data={fivePeople} initialPageIndex={1} onChange={onChange} />);
+    await new Promise((resolve) => setTimeout(resolve, 20));
+    expect(names()).toEqual(["P2", "P3"]);
+
+    rerender(<ClientPageHarness data={[...fivePeople]} initialPageIndex={1} onChange={onChange} />);
+    await new Promise((resolve) => setTimeout(resolve, 20));
+    expect(names()).toEqual(["P2", "P3"]);
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
+  it("client mode snaps a controlled page past the end back to the last page", async () => {
+    const onChange = vi.fn();
+    render(<ClientPageHarness data={fivePeople} initialPageIndex={5} onChange={onChange} />);
+
+    await waitFor(() => expect(onChange).toHaveBeenCalledWith({ pageIndex: 2, pageSize: 2 }));
+    expect(onChange).toHaveBeenCalledTimes(1);
+    expect(names()).toEqual(["P4"]);
+  });
+
+  it("client mode leaves a controlled page alone while there are no rows to page through", async () => {
+    const onChange = vi.fn();
+    render(<ClientPageHarness data={[]} initialPageIndex={3} onChange={onChange} />);
+
+    await new Promise((resolve) => setTimeout(resolve, 20));
+    expect(onChange).not.toHaveBeenCalled();
+  });
+
+  it("client mode without a controlled page still returns to the first page when the rows change", async () => {
+    const user = userEvent.setup();
+    const { rerender } = render(
+      <DataTable data={fivePeople} columns={nameCellColumns} paginationMode="client" pageSizeOptions={[2]} />,
+    );
+
+    await user.click(screen.getByTestId("pagination-next"));
+    expect(names()).toEqual(["P2", "P3"]);
+
+    rerender(
+      <DataTable data={[...fivePeople]} columns={nameCellColumns} paginationMode="client" pageSizeOptions={[2]} />,
+    );
+    await waitFor(() => expect(names()).toEqual(["P0", "P1"]));
+  });
+
   it("server mode resumes clamping once the error clears and a real rowCount arrives", async () => {
     const onChange = vi.fn();
     const { rerender } = render(<ServerPageHarness rowCount={0} isError initialPageIndex={2} onChange={onChange} />);
diff --git a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
index e0f57ae1052..340f8d4f44f 100644
--- a/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
+++ b/ui/litellm-dashboard/src/components/shared/DataTable/DataTable.tsx
@@ -425,7 +425,7 @@ function useControllable<T>(
   return { value: internal, onChange: setInternal };
 }
 
-function useServerPageClamp(
+function usePageClamp(
   active: boolean,
   rowCount: number | undefined,
   pagination: { value: PaginationState; onChange: OnChangeFn<PaginationState> },
@@ -484,7 +484,6 @@ function useDataTableInstance<TData extends RowData, TValue>(
     pageIndex: 0,
     pageSize: pageSizeOptions[0] ?? 25,
   });
-  useServerPageClamp(paginationMode === "server" && !isLoading && !isError, rowCount, paginationState);
   const filterState = useControllable<ColumnFiltersState>(
     columnFilters,
     onColumnFiltersChange,
@@ -536,9 +535,38 @@ function useDataTableInstance<TData extends RowData, TValue>(
     ...(getRowId !== undefined ? { getRowId } : {}),
     ...(enableRowSelection !== undefined ? { enableRowSelection } : {}),
     ...(paginationMode === "server" && rowCount !== undefined ? { rowCount } : {}),
+    autoResetPageIndex: pagination === undefined && paginationMode !== "server",
   };
 
-  return useReactTable(tableOptions);
+  const table = useReactTable(tableOptions);
+  const clampOptions: SettledPageClampOptions = {
+    paginationMode,
+    controlled: pagination !== undefined,
+    settled: !isLoading && !isError,
+    rowCount,
+    pagination: paginationState,
+  };
+  useSettledPageClamp(table, clampOptions);
+  return table;
+}
+
+type SettledPageClampOptions = {
+  paginationMode: PaginationMode;
+  controlled: boolean;
+  settled: boolean;
+  rowCount: number | undefined;
+  pagination: { value: PaginationState; onChange: OnChangeFn<PaginationState> };
+};
+
+function useSettledPageClamp<TData extends RowData>(table: Table<TData>, options: SettledPageClampOptions): void {
+  const { paginationMode, controlled, settled, rowCount, pagination } = options;
+  const clientRowCount = paginationMode === "client" ? table.getPrePaginationRowModel().rows.length : 0;
+  const clientPageIsClampable = paginationMode === "client" && controlled && clientRowCount > 0;
+  usePageClamp(
+    settled && (paginationMode === "server" || clientPageIsClampable),
+    paginationMode === "server" ? rowCount : clientRowCount,
+    pagination,
+  );
 }
 
 export function DataTable<TData extends RowData, TValue>(props: DataTableProps<TData, TValue>) {