fix(ui): read persisted column visibility from storage instead of a mounted copy

usePersistedColumnVisibility kept a useState copy seeded once at mount, so a later tableId or defaults change showed the old table's columns and saved them under the new key. It now reads localStorage through useSyncExternalStore, keeping only writes that storage refused in memory, so the hook has no copy to go stale. Stored choices are layered over the defaults on every read, and changes saved in another tab show up.

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

@@ -1,3 +1,4 @@
+import type { VisibilityState } from "@tanstack/react-table";
 import { act, renderHook } from "@testing-library/react";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
@@ -10,6 +11,9 @@ const stored = (tableId: string): unknown => {
   return raw === null ? null : JSON.parse(raw);
 };
 
+const showEveryColumn = (previous: VisibilityState): VisibilityState =>
+  Object.fromEntries(Object.keys(previous).map((column) => [column, true]));
+
 describe("usePersistedColumnVisibility", () => {
   beforeEach(() => {
     localStorage.clear();
@@ -55,6 +59,15 @@ describe("usePersistedColumnVisibility", () => {
     expect(stored("keys")).toEqual({ email: false, name: false });
   });
 
+  it("hands a function updater the default-hidden columns, so showing every column sticks", () => {
+    const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false }));
+
+    act(() => result.current.onColumnVisibilityChange(showEveryColumn));
+
+    expect(result.current.columnVisibility).toEqual({ spend: true });
+    expect(stored("keys")).toEqual({ spend: true });
+  });
+
   it.each([
     ["truncated JSON", '{"email":fal'],
     ["a JSON scalar", "42"],
@@ -80,19 +93,72 @@ describe("usePersistedColumnVisibility", () => {
     expect(stored("teams")).toBeNull();
   });
 
+  it("reads and writes the new table's columns after the tableId changes", () => {
+    localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false }));
+    localStorage.setItem(keyFor("teams"), JSON.stringify({ spend: false }));
+    const { result, rerender } = renderHook(({ tableId }) => usePersistedColumnVisibility(tableId), {
+      initialProps: { tableId: "keys" },
+    });
+
+    rerender({ tableId: "teams" });
+    expect(result.current.columnVisibility).toEqual({ spend: false });
+
+    act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, name: false })));
+    expect(stored("teams")).toEqual({ spend: false, name: false });
+    expect(stored("keys")).toEqual({ email: false });
+  });
+
+  it("applies new defaults passed after mount", () => {
+    const { result, rerender } = renderHook(({ defaults }) => usePersistedColumnVisibility("keys", defaults), {
+      initialProps: { defaults: { spend: false } },
+    });
+
+    rerender({ defaults: { name: false } });
+
+    expect(result.current.columnVisibility).toEqual({ name: false });
+  });
+
+  it("shows a change another tab saved for the same table", () => {
+    const { result } = renderHook(() => usePersistedColumnVisibility("keys"));
+
+    act(() => {
+      localStorage.setItem(keyFor("keys"), JSON.stringify({ email: false }));
+      window.dispatchEvent(new StorageEvent("storage", { key: keyFor("keys") }));
+    });
+
+    expect(result.current.columnVisibility).toEqual({ email: false });
+  });
+
+  it("keeps a toggle that storage refused, and saves the next one once storage accepts it", () => {
+    localStorage.setItem(keyFor("full"), JSON.stringify({ spend: false }));
+    vi.spyOn(console, "warn").mockImplementation(() => {});
+    vi.spyOn(Storage.prototype, "setItem").mockImplementationOnce(() => {
+      throw new Error("QuotaExceededError");
+    });
+    const { result } = renderHook(() => usePersistedColumnVisibility("full"));
+
+    act(() => result.current.onColumnVisibilityChange({ email: false }));
+    expect(result.current.columnVisibility).toEqual({ email: false });
+    expect(stored("full")).toEqual({ spend: false });
+
+    act(() => result.current.onColumnVisibilityChange({ name: false }));
+    expect(result.current.columnVisibility).toEqual({ name: false });
+    expect(stored("full")).toEqual({ name: false });
+  });
+
   it("returns the defaults without throwing when storage is unavailable", () => {
     vi.spyOn(console, "warn").mockImplementation(() => {});
     vi.spyOn(Storage.prototype, "getItem").mockImplementation(() => {
       throw new Error("SecurityError");
     });
     vi.spyOn(Storage.prototype, "setItem").mockImplementation(() => {
-      throw new Error("QuotaExceededError");
+      throw new Error("SecurityError");
     });
 
-    const { result } = renderHook(() => usePersistedColumnVisibility("keys", { spend: false }));
+    const { result } = renderHook(() => usePersistedColumnVisibility("blocked", { spend: false }));
     expect(result.current.columnVisibility).toEqual({ spend: false });
 
-    act(() => result.current.onColumnVisibilityChange({ email: false }));
-    expect(result.current.columnVisibility).toEqual({ email: false });
+    act(() => result.current.onColumnVisibilityChange((previous) => ({ ...previous, email: false })));
+    expect(result.current.columnVisibility).toEqual({ spend: false, email: false });
   });
 });

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

@@ -1,16 +1,46 @@
 import type { OnChangeFn, VisibilityState } from "@tanstack/react-table";
-import { useCallback, useState } from "react";
+import { useCallback, useMemo, useSyncExternalStore } from "react";
 
-import { getLocalStorageItem, setLocalStorageItem } from "@/utils/localStorageUtils";
+import {
+  LOCAL_STORAGE_EVENT,
+  emitLocalStorageChange,
+  getLocalStorageItem,
+  setLocalStorageItem,
+} from "@/utils/localStorageUtils";
 
 const STORAGE_KEY_PREFIX = "litellm_table_columns_";
 
 const EMPTY_VISIBILITY: VisibilityState = {};
 
+const unsavedWrites = new Map<string, string>();
+
 function storageKey(tableId: string): string {
   return `${STORAGE_KEY_PREFIX}${tableId}`;
 }
 
+function subscribe(onChange: () => void): () => void {
+  window.addEventListener("storage", onChange);
+  window.addEventListener(LOCAL_STORAGE_EVENT, onChange);
+  return () => {
+    window.removeEventListener("storage", onChange);
+    window.removeEventListener(LOCAL_STORAGE_EVENT, onChange);
+  };
+}
+
+function readRaw(key: string): string | null {
+  return unsavedWrites.get(key) ?? getLocalStorageItem(key);
+}
+
+function writeRaw(key: string, raw: string): void {
+  setLocalStorageItem(key, raw);
+  if (getLocalStorageItem(key) === raw) {
+    unsavedWrites.delete(key);
+  } else {
+    unsavedWrites.set(key, raw);
+  }
+  emitLocalStorageChange(key);
+}
+
 function isVisibilityState(value: unknown): value is VisibilityState {
   if (typeof value !== "object" || value === null || Array.isArray(value)) {
     return false;
@@ -18,8 +48,7 @@ function isVisibilityState(value: unknown): value is VisibilityState {
   return Object.values(value).every((visible) => typeof visible === "boolean");
 }
 
-function readStoredVisibility(tableId: string, defaults: VisibilityState): VisibilityState {
-  const raw = getLocalStorageItem(storageKey(tableId));
+function parseVisibility(raw: string | null, defaults: VisibilityState): VisibilityState {
   if (raw === null) {
     return defaults;
   }
@@ -35,19 +64,20 @@ export function usePersistedColumnVisibility(
   tableId: string,
   defaults: VisibilityState = EMPTY_VISIBILITY,
 ): { columnVisibility: VisibilityState; onColumnVisibilityChange: OnChangeFn<VisibilityState> } {
-  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>(() =>
-    readStoredVisibility(tableId, defaults),
+  const key = storageKey(tableId);
+  const raw = useSyncExternalStore(
+    subscribe,
+    () => readRaw(key),
+    () => null,
   );
+  const columnVisibility = useMemo(() => parseVisibility(raw, defaults), [raw, defaults]);
 
   const onColumnVisibilityChange = useCallback<OnChangeFn<VisibilityState>>(
     (updater) => {
-      setColumnVisibility((previous) => {
-        const next = typeof updater === "function" ? updater(previous) : updater;
-        setLocalStorageItem(storageKey(tableId), JSON.stringify(next));
-        return next;
-      });
+      const next = typeof updater === "function" ? updater(parseVisibility(readRaw(key), defaults)) : updater;
+      writeRaw(key, JSON.stringify(next));
     },
-    [tableId],
+    [key, defaults],
   );
 
   return { columnVisibility, onColumnVisibilityChange };