From d638c911c8564088f664f018aef6f5f0c3782070 Mon Sep 17 00:00:00 2001
From: clioo <jesus_acosta1996@hotmail.com>
Date: Tue, 4 Aug 2026 14:02:14 -0600
Subject: [PATCH] Fix organize undo and viewport restoration

---
 src/App.tsx                 | 57 ++++++++++++++++++++++++++++---------
 src/index.css               | 26 +++++++++++++++--
 src/store.test.ts           |  7 +++++
 src/store.ts                | 22 ++++++++++++++
 tests/e2e/opengraph.spec.ts | 34 ++++++++++++++++++++--
 5 files changed, 129 insertions(+), 17 deletions(-)

diff --git a/src/App.tsx b/src/App.tsx
index febd730..c97a3f2 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -229,6 +229,7 @@ function App() {
     commit,
     setNodesLive,
     setEdgesLive,
+    setViewportLive,
     updateViewport,
     setSelected,
     setSettingsOpen,
@@ -251,7 +252,7 @@ function App() {
     archiveGraph,
     restoreGraph,
   } = useOpenGraphStore();
-  const { screenToFlowPosition, fitView } = useReactFlow();
+  const { screenToFlowPosition, fitView, getViewport, setViewport } = useReactFlow();
   const canvasRef = useRef<HTMLDivElement>(null);
   const companionRef = useRef<CompanionBridge | null>(null);
   const [editingName, setEditingName] = useState(false);
@@ -283,6 +284,35 @@ function App() {
       ? document.edges.find((edge) => edge.id === selected.id)
       : undefined;
 
+  const restoreHistoryViewport = useCallback(() => {
+    skipViewportCommitUntilRef.current = Date.now() + 500;
+    window.requestAnimationFrame(() => {
+      const viewport = useOpenGraphStore.getState().document.viewport;
+      void setViewport(viewport, { duration: 200 });
+    });
+  }, [setViewport]);
+
+  const handleUndo = useCallback(
+    (message = "Change undone") => {
+      if (!useOpenGraphStore.getState().past.length) return;
+      undo();
+      setOrganizationRevision(null);
+      setOrganizationToast(null);
+      setToast(message);
+      restoreHistoryViewport();
+    },
+    [restoreHistoryViewport, setToast, undo],
+  );
+
+  const handleRedo = useCallback(() => {
+    if (!useOpenGraphStore.getState().future.length) return;
+    redo();
+    setOrganizationRevision(null);
+    setOrganizationToast(null);
+    setToast("Change redone");
+    restoreHistoryViewport();
+  }, [redo, restoreHistoryViewport, setToast]);
+
   useEffect(() => {
     (
       window as Window & { __opengraphDocument?: typeof document }
@@ -360,6 +390,7 @@ function App() {
   }, [appearance]);
   useEffect(() => {
     const fitOnNarrow = () => {
+      skipViewportCommitUntilRef.current = Date.now() + 250;
       const frame = window.requestAnimationFrame(() =>
         fitView({
           padding: window.innerWidth <= 800 ? 0.18 : 0.14,
@@ -397,11 +428,11 @@ function App() {
       }
       if (command && event.key.toLowerCase() === "z") {
         event.preventDefault();
-        event.shiftKey ? redo() : undo();
+        event.shiftKey ? handleRedo() : handleUndo();
       }
       if (command && event.key.toLowerCase() === "y") {
         event.preventDefault();
-        redo();
+        handleRedo();
       }
       if (command && event.key.toLowerCase() === "d") {
         event.preventDefault();
@@ -422,10 +453,10 @@ function App() {
     createGraph,
     deleteSelected,
     duplicateSelected,
-    redo,
+    handleRedo,
+    handleUndo,
     setActiveTool,
     setSelected,
-    undo,
   ]);
   const onNodesChange = useCallback(
     (changes: NodeChange<GraphNode>[]) => {
@@ -646,6 +677,7 @@ function App() {
   };
   const fit = () => fitView({ padding: 0.18, duration: 200 });
   const organize = useCallback(() => {
+    setViewportLive(getViewport());
     const current = useOpenGraphStore.getState().document;
     const result = organizeGraphDocument(current);
     if (result.document === current) {
@@ -666,7 +698,7 @@ function App() {
     window.requestAnimationFrame(() =>
       fitView({ padding: 0.18, duration: 200 }),
     );
-  }, [commit, fitView, setSelected, setToast]);
+  }, [commit, fitView, getViewport, setSelected, setToast, setViewportLive]);
   const activeGraphs = useMemo(
     () => graphs.filter((graph) => !graph.archived),
     [graphs],
@@ -819,10 +851,10 @@ function App() {
           <span>Organize</span>
         </button>
         <div className="history-actions">
-          <IconButton label="Undo" onClick={undo} disabled={!past.length}>
+          <IconButton label="Undo" onClick={() => handleUndo()} disabled={!past.length}>
             <Icon name="undo" />
           </IconButton>
-          <IconButton label="Redo" onClick={redo} disabled={!future.length}>
+          <IconButton label="Redo" onClick={handleRedo} disabled={!future.length}>
             <Icon name="redo" />
           </IconButton>
         </div>
@@ -846,9 +878,10 @@ function App() {
           className="primary-button"
           onClick={handleCopy}
           disabled={copying}
+          aria-label={copying ? "Preparing graph" : "Copy graph"}
         >
           <Icon name="copy" size={16} />
-          {copying ? "Preparing…" : "Copy graph"}
+          <span>{copying ? "Preparing…" : "Copy graph"}</span>
         </button>
         <button
           className="model-settings-button"
@@ -1113,6 +1146,7 @@ function App() {
             }}
             onMoveEnd={(_event, viewport) => {
               if (Date.now() < skipViewportCommitUntilRef.current) {
+                setViewportLive(viewport);
                 return;
               }
               updateViewport(viewport);
@@ -1231,10 +1265,7 @@ function App() {
             <button
               className="toast-action"
               onClick={() => {
-                undo();
-                setOrganizationRevision(null);
-                setOrganizationToast(null);
-                setToast("Graph organization undone");
+                handleUndo("Graph organization undone");
               }}
             >
               Undo
diff --git a/src/index.css b/src/index.css
index c7f1572..886b21a 100644
--- a/src/index.css
+++ b/src/index.css
@@ -1194,7 +1194,6 @@ select:focus-visible,
   }
   .brand-divider,
   .graph-name-wrap,
-  .history-actions,
   .appearance-switch {
     display: none;
   }
@@ -1230,9 +1229,15 @@ select:focus-visible,
     display: none;
   }
   .primary-button {
-    padding: 0 12px;
+    width: 44px;
+    min-width: 44px;
+    padding: 0;
+    justify-content: center;
     margin-left: auto;
   }
+  .primary-button span {
+    display: none;
+  }
   .model-settings-button {
     width: 44px;
     min-width: 44px;
@@ -1292,6 +1297,23 @@ select:focus-visible,
     overflow-wrap: anywhere;
   }
 }
+
+@media (max-width: 500px) {
+  .brand-mark,
+  .history-actions .icon-button:last-child,
+  .appearance-switch {
+    display: none;
+  }
+  .brand-lockup {
+    gap: 0;
+  }
+}
+
+@media (max-width: 340px) {
+  .brand-lockup {
+    display: none;
+  }
+}
 @media (prefers-reduced-motion: reduce) {
   *,
   *::before,
diff --git a/src/store.test.ts b/src/store.test.ts
index d29477a..55c2dfc 100644
--- a/src/store.test.ts
+++ b/src/store.test.ts
@@ -133,6 +133,13 @@ describe("OpenGraph store", () => {
     store.setAppearance("dark");
     store.setToast("Saved");
     store.setSelected(null);
+    store.setViewportLive({ x: 25, y: 35, zoom: 1.1 });
+    expect(useOpenGraphStore.getState().document.viewport).toEqual({
+      x: 25,
+      y: 35,
+      zoom: 1.1,
+    });
+    expect(store.canUndo()).toBe(false);
     expect(useOpenGraphStore.getState()).toMatchObject({
       settingsOpen: true,
       activeTool: "connect",
diff --git a/src/store.ts b/src/store.ts
index ff8df63..f90ba05 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -311,6 +311,7 @@ type OpenGraphState = {
   ) => GraphSnapshot;
   setNodesLive: (nodes: GraphNode[]) => void;
   setEdgesLive: (edges: GraphEdge[]) => void;
+  setViewportLive: (viewport: GraphDocument["viewport"]) => void;
   updateViewport: (viewport: GraphDocument["viewport"]) => void;
   setSelected: (selected: OpenGraphState["selected"]) => void;
   setSettingsOpen: (open: boolean) => void;
@@ -409,6 +410,27 @@ export const useOpenGraphStore = create<OpenGraphState>((set, get) => ({
       document: { ...state.document, edges },
       transientBefore: state.transientBefore ?? cloneDocument(state.document),
     })),
+  setViewportLive: (viewport) =>
+    set((state) => {
+      if (
+        state.document.viewport.x === viewport.x &&
+        state.document.viewport.y === viewport.y &&
+        state.document.viewport.zoom === viewport.zoom
+      )
+        return state;
+      const document = { ...state.document, viewport };
+      const graphs = state.graphs.map((graph) =>
+        graph.id === state.activeGraphId ? { ...graph, document } : graph,
+      );
+      const saved =
+        persistDocument(document) &&
+        persistGraphLibrary(graphs, state.activeGraphId);
+      return {
+        document,
+        graphs,
+        saveStatus: saved ? "saved" : "error",
+      };
+    }),
   updateViewport: (viewport) => {
     get().transact((document) => ({ ...document, viewport }), "local");
   },
diff --git a/tests/e2e/opengraph.spec.ts b/tests/e2e/opengraph.spec.ts
index e21cdaf..281f891 100644
--- a/tests/e2e/opengraph.spec.ts
+++ b/tests/e2e/opengraph.spec.ts
@@ -163,6 +163,9 @@ test.describe("OpenGraph", () => {
     await canvas.dblclick({ position: { x: 440, y: 480 } });
     await expect(page.locator(".react-flow__node-workflow")).toHaveCount(3);
     await expect(page.locator(".react-flow__edge")).toHaveCount(0);
+    const beforePositions = await page.locator(".react-flow__node-workflow").evaluateAll(
+      (elements) => elements.map((element) => (element as HTMLElement).style.transform),
+    );
 
     await page.getByRole("button", { name: "Organize" }).click();
 
@@ -176,9 +179,36 @@ test.describe("OpenGraph", () => {
     expect(organizedPositions).toHaveLength(3);
     expect(new Set(organizedPositions).size).toBe(3);
 
-    await page.getByRole("status").getByRole("button", { name: "Undo" }).click();
+    await page.waitForTimeout(700);
+    await page.keyboard.press("Control+z");
     await expect(page.locator(".react-flow__edge")).toHaveCount(0);
-    await expect(page.getByRole("status")).toContainText("Graph organization undone");
+    await expect.poll(() =>
+      page.locator(".react-flow__node-workflow").evaluateAll(
+        (elements) => elements.map((element) => (element as HTMLElement).style.transform),
+      ),
+    ).toEqual(beforePositions);
+    await expect.poll(() =>
+      page.locator(".react-flow__node-workflow").evaluateAll((elements) =>
+        elements.filter((element) => {
+          const box = element.getBoundingClientRect();
+          return (
+            box.right > 0 &&
+            box.bottom > 80 &&
+            box.left < window.innerWidth &&
+            box.top < window.innerHeight
+          );
+        }).length,
+      ),
+    ).toBe(3);
+
+    await page.getByRole("button", { name: "Redo" }).click();
+    await expect(page.locator(".react-flow__edge")).toHaveCount(2);
+    await page.getByRole("button", { name: "Undo", exact: true }).click();
+    await expect(page.locator(".react-flow__edge")).toHaveCount(0);
+
+    await page.setViewportSize({ width: 700, height: 760 });
+    await expect(page.getByRole("button", { name: "Undo", exact: true })).toBeVisible();
+    await expect(page.getByRole("button", { name: "Redo", exact: true })).toBeVisible();
   });
 
   test("starts a blank graph from a canvas double-click", async ({ page }) => {