From 11e089a6f7babbbdbe1a6868a8f71edcfb70efd6 Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Tue, 31 Mar 2026 06:36:22 +1300
Subject: [PATCH] fix: coalesce node+edge removal into single undo entry

When deleting a node via the canvas (Delete key), React Flow fires
onNodesChange(remove) then onEdgesChange(remove) synchronously.
Previously this created two undo entries, requiring two undos and
leaving an intermediate state with orphaned edges.

Now onNodesChange sets a microtask-scoped flag that tells onEdgesChange
to skip its checkpoint, collapsing both into a single undo entry.
Added integration test verifying single-undo restores node + edges.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 src/store/__tests__/undoRedo.test.ts | 57 ++++++++++++++++++++++++++++
 src/store/workflowStore.ts           | 14 +++++--
 2 files changed, 68 insertions(+), 3 deletions(-)

diff --git a/src/store/__tests__/undoRedo.test.ts b/src/store/__tests__/undoRedo.test.ts
index 3847874b..ca140a0d 100644
--- a/src/store/__tests__/undoRedo.test.ts
+++ b/src/store/__tests__/undoRedo.test.ts
@@ -209,6 +209,63 @@ describe("Undo/Redo integration", () => {
     });
   });
 
+  describe("delete node via onNodesChange + onEdgesChange restores connections", () => {
+    it("single undo restores both node and its connected edges", async () => {
+      let store = useWorkflowStore.getState();
+
+      // Create two nodes and connect them
+      act(() => {
+        store.addNode("prompt", { x: 0, y: 0 });
+      });
+      store = useWorkflowStore.getState();
+      act(() => {
+        store.addNode("nanoBanana", { x: 300, y: 0 });
+      });
+      store = useWorkflowStore.getState();
+      const promptId = store.nodes[0].id;
+      const genId = store.nodes[1].id;
+
+      act(() => {
+        store.onConnect({
+          source: promptId,
+          target: genId,
+          sourceHandle: "text",
+          targetHandle: "text",
+        });
+      });
+
+      store = useWorkflowStore.getState();
+      expect(store.nodes.length).toBe(2);
+      expect(store.edges.length).toBe(1);
+
+      // Simulate pressing Delete: React Flow fires onNodesChange(remove) then
+      // onEdgesChange(remove) synchronously in the same cycle
+      act(() => {
+        store.onNodesChange([{ type: "remove", id: promptId }]);
+        store.onEdgesChange([{ type: "remove", id: store.edges[0].id }]);
+      });
+
+      // Wait for microtask to clear nodeRemoveCheckpointActive
+      await act(async () => {
+        await Promise.resolve();
+      });
+
+      store = useWorkflowStore.getState();
+      expect(store.nodes.length).toBe(1);
+      expect(store.edges.length).toBe(0);
+
+      // Single undo should restore both the node AND the edge
+      act(() => {
+        store.undo();
+      });
+
+      store = useWorkflowStore.getState();
+      expect(store.nodes.length).toBe(2);
+      expect(store.edges.length).toBe(1);
+      expect(store.nodes.find((n) => n.id === promptId)).toBeDefined();
+    });
+  });
+
   describe("new action clears redo stack", () => {
     it("clears redo stack when a new undoable action is performed", () => {
       let store = useWorkflowStore.getState();
diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts
index 04eadaf0..642745df 100644
--- a/src/store/workflowStore.ts
+++ b/src/store/workflowStore.ts
@@ -397,6 +397,9 @@ const undoManager = new UndoManager();
 let isDragging = false;
 let pendingDataSnapshot: UndoSnapshot | null = null;
 let dataChangeTimer: ReturnType<typeof setTimeout> | null = null;
+// When true, onEdgesChange(remove) skips its checkpoint because onNodesChange(remove)
+// already captured one in the same React Flow event cycle.
+let nodeRemoveCheckpointActive = false;
 
 // RAF debounce for hover updates — coalesces rapid mouseenter/mouseleave events
 // into a single store update per animation frame
@@ -747,9 +750,13 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
       isDragging = false;
     }
 
-    // Undo: capture snapshot before node removal
+    // Undo: capture snapshot before node removal.
+    // Also set flag so the consequent onEdgesChange(remove) from React Flow
+    // doesn't push a second checkpoint for the same user action.
     if (hasRemoveChange) {
       pushUndoCheckpoint(get, set);
+      nodeRemoveCheckpointActive = true;
+      Promise.resolve().then(() => { nodeRemoveCheckpointActive = false; });
     }
 
     set((state) => ({
@@ -769,8 +776,9 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
     const hasRemoveChange = changes.some((c) => c.type === "remove");
     const hasAddOrRemove = changes.some((c) => c.type === "add" || c.type === "remove");
 
-    // Undo: capture snapshot before edge removal
-    if (hasRemoveChange) {
+    // Undo: capture snapshot before edge removal — but skip if a node-remove
+    // checkpoint was already pushed in this same React Flow event cycle
+    if (hasRemoveChange && !nodeRemoveCheckpointActive) {
       pushUndoCheckpoint(get, set);
     }