feat: remember the schema visualizer's manual table layout (#1555)

Dragged node positions are now persisted in a localStorage-backed
TanStack DB collection (uiPersistentStateCollection) scoped per schema,
so a manual arrangement survives leaving the visualizer and full page
reloads. Saved positions are merged over the ELK auto-layout on mount,
so tables without a remembered position (for example newly created
ones) still fall back to auto-layout, and the existing Reset layout
header action now also clears the remembered manual layout.

Closes #1397

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: sorenbs

.changeset/schema-visualizer-remember-layout.md

@@ -0,0 +1,5 @@
+---
+"@prisma/studio-core": minor
+---
+
+Remember the schema visualizer's manual table layout. Dragged node positions are now stored in localStorage-backed UI state scoped per schema, so a manual arrangement survives leaving the visualizer and full page reloads. Tables without a remembered position (for example newly created ones) fall back to ELK auto-layout, and the `Reset layout` action clears the remembered layout.

Architecture/ui-state.md

@@ -40,6 +40,10 @@ Studio context provides the canonical stores in [`ui/studio/context.tsx`](../ui/
   - `filteredRowCount`
 - `uiLocalStateCollection` (`localOnlyCollectionOptions`)
   - General scoped UI state (for example DataGrid selection machine state).
+- `uiPersistentStateCollection` (`localStorageCollectionOptions`)
+  - General scoped UI state that must survive page reloads, accessed through
+    `useUiState` with `{ persistent: true }` (for example the schema
+    visualizer's manually arranged node positions).
 - `sqlEditorStateCollection` (`localStorageCollectionOptions`)
   - Persisted SQL editor draft state:
   - `queryText`
@@ -135,6 +139,9 @@ The following are valid examples of UI state and where they belong:
 - Schema visualizer node positions and layout state: `uiLocalStateCollection` via `useUiState`
   - Scoped by active schema plus the current visualized table set so returning to the same schema graph restores dragged positions without leaking across schemas.
   - Includes the stored ELK baseline positions and reset-layout request token used by the header action.
+- Schema visualizer manually arranged node positions: `uiPersistentStateCollection` via `useUiState({ persistent: true })`
+  - Scoped by active schema name (`schema-visualizer:${schema}:manual-layout:node-positions`) with per-table entries, so manual layouts survive reloads, new tables fall back to the ELK auto-layout, and different schemas do not collide.
+  - The header `Reset layout` action clears this store in addition to re-applying the ELK baseline.
 - Command-palette `x more...` handoff into table browsing: the same navigation table-name search `useUiState` entry, not a second command-palette-specific table-filter store
 
 If new UI state is shared across components, it MUST be assigned to one of these stores (or a new TanStack DB collection added in Studio context).

FEATURES.md

@@ -150,7 +150,7 @@ Cards keep a fixed width inside their own horizontal scroller, the header toggle
 
 Studio includes a schema graph view with table nodes, column metadata, and detected foreign-key relationships labeled as 1:1 or 1:n.
 The visualizer now runs ELK auto-layout with component-aware spacing so disconnected tables do not collapse into the same visual band, and orthogonal step edges leave clearer corridors between nodes.
-Dragged node positions persist when you leave and return to the same schema view, and a header-level `Reset layout` action re-applies the current ELK baseline when you want to discard manual placement.
+Dragged node positions are remembered in localStorage-backed UI state scoped per schema, so a manual arrangement survives leaving the visualizer, switching views, and full page reloads; tables without a remembered position (for example newly created ones) fall back to ELK auto-layout. A header-level `Reset layout` action re-applies the current ELK baseline and forgets the stored manual placement.
 Users can pan/zoom, inspect key and nullable markers, and jump from a node directly to that table’s data view.
 
 ## Query Insights

ui/hooks/use-ui-state.context.test.tsx

@@ -13,6 +13,7 @@ const useOptionalStudioMock = vi.fn<
   () =>
     | {
         uiLocalStateCollection: ReturnType<typeof createUiCollection>;
+        uiPersistentStateCollection?: ReturnType<typeof createUiCollection>;
       }
     | undefined
 >();
@@ -27,10 +28,14 @@ vi.mock("../studio/context", () => {
   globalThis as { IS_REACT_ACT_ENVIRONMENT?: boolean }
 ).IS_REACT_ACT_ENVIRONMENT = true;
 
+let collectionInstanceCounter = 0;
+
 function createUiCollection() {
+  collectionInstanceCounter += 1;
+
   return createCollection(
     localOnlyCollectionOptions<StudioLocalUiState>({
-      id: "use-ui-state-context-test",
+      id: `use-ui-state-context-test-${collectionInstanceCounter}`,
       getKey(item) {
         return item.id;
       },
@@ -87,6 +92,48 @@ describe("useUiState with Studio context collection", () => {
     container.remove();
   });
 
+  it("routes persistent state into the persistent ui collection", () => {
+    const persistentCollection = createUiCollection();
+
+    useOptionalStudioMock.mockReturnValue({
+      uiLocalStateCollection: uiCollection,
+      uiPersistentStateCollection: persistentCollection,
+    });
+
+    const key = "context-persistent-state";
+    const container = document.createElement("div");
+    document.body.appendChild(container);
+    const root = createRoot(container);
+
+    let latestState: ReturnType<typeof useUiState<string>> | undefined;
+
+    function Harness() {
+      latestState = useUiState<string>(key, "alpha", { persistent: true });
+      return null;
+    }
+
+    act(() => {
+      root.render(<Harness />);
+    });
+
+    expect(latestState?.[0]).toBe("alpha");
+    expect(persistentCollection.get(key)?.value).toBe("alpha");
+    expect(uiCollection.has(key)).toBe(false);
+
+    act(() => {
+      latestState?.[1]("beta");
+    });
+
+    expect(latestState?.[0]).toBe("beta");
+    expect(persistentCollection.get(key)?.value).toBe("beta");
+    expect(uiCollection.has(key)).toBe(false);
+
+    act(() => {
+      root.unmount();
+    });
+    container.remove();
+  });
+
   it("does not mutate the shared ui collection for cleanup-on-unmount state", () => {
     const key = "context-cleanup-state";
     const insertSpy = vi.spyOn(uiCollection, "insert");

ui/hooks/use-ui-state.ts

@@ -14,6 +14,11 @@ type Updater<T> = T | ((previous: T) => T);
 
 export interface UseUiStateOptions {
   cleanupOnUnmount?: boolean;
+  /**
+   * When enabled, state is stored in the localStorage-backed persistent UI
+   * state collection so it survives page reloads.
+   */
+  persistent?: boolean;
 }
 
 const fallbackUiStateCollection = instrumentTanStackCollectionMutations(
@@ -98,15 +103,18 @@ export function useUiState<T>(
   initialValue: T,
   options: UseUiStateOptions = {},
 ) {
-  const { cleanupOnUnmount = false } = options;
+  const { cleanupOnUnmount = false, persistent = false } = options;
   const [volatileValue, setVolatileValue] = useState<T>(() =>
     cloneValue(initialValue),
   );
   const previousVolatileKeyRef = useRef<string | undefined>(key);
   const studioContext = useOptionalStudio();
   const uiLocalStateCollection =
-    (studioContext?.uiLocalStateCollection as typeof fallbackUiStateCollection) ??
-    fallbackUiStateCollection;
+    ((persistent
+      ? studioContext?.uiPersistentStateCollection
+      : studioContext?.uiLocalStateCollection) as
+      | typeof fallbackUiStateCollection
+      | undefined) ?? fallbackUiStateCollection;
 
   const { data: stateRow } = useLiveQuery(
     (q) => {

ui/studio/context.tsx

@@ -44,6 +44,7 @@ const STUDIO_UI_STATE_ID = "studio-ui-state";
 const STUDIO_UI_STORAGE_KEY = "prisma-studio-ui-state-v1";
 const SQL_EDITOR_STATE_ID = "studio-sql-editor-state";
 const SQL_EDITOR_STORAGE_KEY = "prisma-studio-sql-editor-state-v1";
+const UI_PERSISTENT_STATE_STORAGE_KEY = "prisma-studio-persistent-ui-state-v1";
 const DEFAULT_TABLE_PAGE_SIZE = 25;
 export const DEFAULT_NAVIGATION_WIDTH = 192;
 export const MIN_NAVIGATION_WIDTH = 192;
@@ -280,6 +281,7 @@ interface StudioContextValue {
   tableUiStateCollection: Collection<TableUiState, string | number>;
   tableQueryMetaCollection: Collection<TableQueryMetaState, string | number>;
   uiLocalStateCollection: Collection<StudioLocalUiState, string | number>;
+  uiPersistentStateCollection: Collection<StudioLocalUiState, string | number>;
   sqlEditorStateCollection: Collection<SqlEditorState, string | number>;
   navigationTableNamesCollection: Collection<
     NavigationTableNameState,
@@ -394,6 +396,20 @@ export function StudioContextProvider(props: StudioContextProviderProps) {
       { collectionName: "studio-local-ui-state" },
     ),
   );
+  const uiPersistentStateCollectionRef = useRef(
+    instrumentTanStackCollectionMutations(
+      createCollection(
+        localStorageCollectionOptions<StudioLocalUiState>({
+          id: "studio-persistent-ui-state",
+          storageKey: UI_PERSISTENT_STATE_STORAGE_KEY,
+          getKey(item) {
+            return item.id;
+          },
+        }),
+      ),
+      { collectionName: "studio-persistent-ui-state" },
+    ),
+  );
   const sqlEditorStateCollectionRef = useRef(
     instrumentTanStackCollectionMutations(
       createCollection(
@@ -429,6 +445,7 @@ export function StudioContextProvider(props: StudioContextProviderProps) {
   const tableUiStateCollection = tableUiStateCollectionRef.current;
   const tableQueryMetaCollection = tableQueryMetaCollectionRef.current;
   const uiLocalStateCollection = uiLocalStateCollectionRef.current;
+  const uiPersistentStateCollection = uiPersistentStateCollectionRef.current;
   const sqlEditorStateCollection = sqlEditorStateCollectionRef.current;
   const navigationTableNamesCollection =
     navigationTableNamesCollectionRef.current;
@@ -865,6 +882,7 @@ export function StudioContextProvider(props: StudioContextProviderProps) {
           tableUiStateCollection,
           tableQueryMetaCollection,
           uiLocalStateCollection,
+          uiPersistentStateCollection,
           sqlEditorStateCollection,
           navigationTableNamesCollection,
           getOrCreateRowsCollection,

ui/studio/views/schema/SchemaView.test.tsx

@@ -10,18 +10,23 @@ function cloneMockValue<T>(value: T): T {
   return structuredClone(value);
 }
 
-const { uiStateStore, useNavigationMock, useSchemaVisualizationMock } =
-  vi.hoisted(() => ({
-    uiStateStore: new Map<string, unknown>(),
-    useNavigationMock: vi.fn<
-      () => {
-        metadata: {
-          activeSchema: { name: string };
-        };
-      }
-    >(),
-    useSchemaVisualizationMock: vi.fn<() => SchemaVisualizationData>(),
-  }));
+const {
+  persistentUiStateStore,
+  uiStateStore,
+  useNavigationMock,
+  useSchemaVisualizationMock,
+} = vi.hoisted(() => ({
+  persistentUiStateStore: new Map<string, unknown>(),
+  uiStateStore: new Map<string, unknown>(),
+  useNavigationMock: vi.fn<
+    () => {
+      metadata: {
+        activeSchema: { name: string };
+      };
+    }
+  >(),
+  useSchemaVisualizationMock: vi.fn<() => SchemaVisualizationData>(),
+}));
 
 vi.mock("@/ui/hooks/use-navigation", () => ({
   useNavigation: useNavigationMock,
@@ -31,15 +36,20 @@ vi.mock("@/ui/hooks/use-ui-state", async () => {
   const React = await vi.importActual<typeof import("react")>("react");
 
   return {
-    useUiState: <T,>(key: string, initialValue: T) => {
+    useUiState: <T,>(
+      key: string,
+      initialValue: T,
+      options?: { persistent?: boolean },
+    ) => {
+      const store = options?.persistent ? persistentUiStateStore : uiStateStore;
+
       const [value, setValue] = React.useState<T>(() => {
-        if (!uiStateStore.has(key)) {
-          uiStateStore.set(key, cloneMockValue(initialValue));
+        if (!store.has(key)) {
+          store.set(key, cloneMockValue(initialValue));
         }
 
         return (
-          (uiStateStore.get(key) as T | undefined) ??
-          cloneMockValue(initialValue)
+          (store.get(key) as T | undefined) ?? cloneMockValue(initialValue)
         );
       });
 
@@ -51,11 +61,11 @@ vi.mock("@/ui/hooks/use-ui-state", async () => {
                 ? (updater as (previous: T) => T)(previous)
                 : updater;
 
-            uiStateStore.set(key, cloneMockValue(nextValue));
+            store.set(key, cloneMockValue(nextValue));
             return cloneMockValue(nextValue);
           });
         },
-        [key],
+        [key, store],
       );
 
       return [value, setSharedValue] as const;
@@ -93,6 +103,7 @@ vi.mock("./Visualiser", () => ({
 describe("SchemaView", () => {
   beforeEach(() => {
     uiStateStore.clear();
+    persistentUiStateStore.clear();
     useNavigationMock.mockReturnValue({
       metadata: {
         activeSchema: { name: "public" },
@@ -157,6 +168,12 @@ describe("SchemaView", () => {
       "schema-visualizer:public:posts|users:reset-layout-version",
       0,
     );
+    persistentUiStateStore.set(
+      "schema-visualizer:public:manual-layout:node-positions",
+      {
+        users: { x: 333, y: 444 },
+      },
+    );
 
     const container = document.createElement("div");
     document.body.appendChild(container);
@@ -190,6 +207,11 @@ describe("SchemaView", () => {
         "schema-visualizer:public:posts|users:reset-layout-version",
       ),
     ).toBe(1);
+    expect(
+      persistentUiStateStore.get(
+        "schema-visualizer:public:manual-layout:node-positions",
+      ),
+    ).toEqual({});
     expect(container.textContent).not.toContain("Reset layout");
 
     act(() => {

ui/studio/views/schema/SchemaView.tsx

@@ -9,6 +9,7 @@ import { useSchemaVisualization } from "../../../hooks/use-schema-visualization"
 import { StudioHeader } from "../../StudioHeader";
 import { ViewProps } from "../View";
 import {
+  createSchemaVisualizerPersistentStateScope,
   createSchemaVisualizerStateScope,
   createSchemaVisualizerUiStateKey,
   doSchemaNodePositionsDiffer,
@@ -32,10 +33,19 @@ export function SchemaView(_props: ViewProps) {
         .sort((left, right) => left.localeCompare(right)),
     [tables],
   );
+  const persistentStateScope = useMemo(
+    () => createSchemaVisualizerPersistentStateScope(activeSchema?.name),
+    [activeSchema?.name],
+  );
   const [nodePositions, setNodePositions] = useUiState<SchemaNodePositions>(
     createSchemaVisualizerUiStateKey(stateScope, "node-positions"),
     {},
   );
+  const [, setManualNodePositions] = useUiState<SchemaNodePositions>(
+    createSchemaVisualizerUiStateKey(persistentStateScope, "node-positions"),
+    {},
+    { persistent: true },
+  );
   const [autoLayoutPositions] = useUiState<SchemaNodePositions>(
     createSchemaVisualizerUiStateKey(stateScope, "auto-layout-node-positions"),
     {},
@@ -54,6 +64,7 @@ export function SchemaView(_props: ViewProps) {
       variant="outline"
       onClick={() => {
         setNodePositions(autoLayoutPositions);
+        setManualNodePositions({});
         setResetLayoutVersion((currentVersion) => currentVersion + 1);
       }}
     >