From c02758ea89e337480f502dfd8ea0cf11d99501ed Mon Sep 17 00:00:00 2001
From: shrimbly <falloon.willie@gmail.com>
Date: Tue, 31 Mar 2026 06:46:26 +1300
Subject: [PATCH] fix: prevent clearStaleInputImages from creating extra undo
 entries on node delete

When deleting a connected node, clearStaleInputImages calls updateNodeData
on target nodes to clear stale input images. Each updateNodeData call was
creating a debounced undo snapshot, resulting in 3-4 undos needed instead
of 1. Now updateNodeData skips undo tracking when nodeRemoveCheckpointActive
is true, since the node removal checkpoint already captured full pre-deletion state.

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

diff --git a/src/store/__tests__/undoRedo.test.ts b/src/store/__tests__/undoRedo.test.ts
index ca140a0d..b5425fe5 100644
--- a/src/store/__tests__/undoRedo.test.ts
+++ b/src/store/__tests__/undoRedo.test.ts
@@ -5,7 +5,7 @@
  * for all undoable actions (add/remove nodes, connect edges, etc.).
  */
 
-import { describe, it, expect, vi, beforeEach } from "vitest";
+import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { act } from "@testing-library/react";
 import { useWorkflowStore } from "../workflowStore";
 
@@ -52,9 +52,14 @@ function resetStore() {
 
 describe("Undo/Redo integration", () => {
   beforeEach(() => {
+    vi.useFakeTimers();
     resetStore();
   });
 
+  afterEach(() => {
+    vi.useRealTimers();
+  });
+
   it("starts with canUndo and canRedo both false", () => {
     const { canUndo, canRedo } = useWorkflowStore.getState();
     expect(canUndo).toBe(false);
@@ -263,6 +268,86 @@ describe("Undo/Redo integration", () => {
       expect(store.nodes.length).toBe(2);
       expect(store.edges.length).toBe(1);
       expect(store.nodes.find((n) => n.id === promptId)).toBeDefined();
+
+      // The delete should have been exactly one undo entry — undoing again
+      // should go back to before the edge was connected, not to some
+      // intermediate state caused by clearStaleInputImages side effects
+      act(() => {
+        store.undo();
+      });
+      store = useWorkflowStore.getState();
+      // Previous checkpoint was onConnect, so edges should be gone
+      expect(store.edges.length).toBe(0);
+      expect(store.nodes.length).toBe(2);
+    });
+
+    it("single undo restores image-source node and does not create extra entries from clearStaleInputImages", async () => {
+      let store = useWorkflowStore.getState();
+
+      // Create imageInput -> nanoBanana (image connection triggers clearStaleInputImages on delete)
+      act(() => {
+        store.addNode("imageInput", { x: 0, y: 0 });
+      });
+      store = useWorkflowStore.getState();
+      act(() => {
+        store.addNode("nanoBanana", { x: 300, y: 0 });
+      });
+      store = useWorkflowStore.getState();
+      const imageInputId = store.nodes[0].id;
+      const genId = store.nodes[1].id;
+
+      act(() => {
+        store.onConnect({
+          source: imageInputId,
+          target: genId,
+          sourceHandle: "image",
+          targetHandle: "image",
+        });
+      });
+
+      store = useWorkflowStore.getState();
+      expect(store.nodes.length).toBe(2);
+      expect(store.edges.length).toBe(1);
+
+      // Delete the imageInput node — this triggers clearStaleInputImages
+      // which calls updateNodeData on the nanoBanana target
+      act(() => {
+        store.onNodesChange([{ type: "remove", id: imageInputId }]);
+        store.onEdgesChange([{ type: "remove", id: store.edges[0].id }]);
+      });
+
+      // Wait for microtask to clear nodeRemoveCheckpointActive
+      await act(async () => {
+        await Promise.resolve();
+      });
+
+      // Also advance past the debounced data-change timer (500ms)
+      await act(async () => {
+        vi.advanceTimersByTime(600);
+      });
+
+      store = useWorkflowStore.getState();
+      expect(store.nodes.length).toBe(1);
+      expect(store.edges.length).toBe(0);
+
+      // Single undo should restore the node and 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 === imageInputId)).toBeDefined();
+
+      // Verify no extra undo entries from clearStaleInputImages:
+      // Next undo should go back to before onConnect (2 nodes, 0 edges)
+      act(() => {
+        store.undo();
+      });
+      store = useWorkflowStore.getState();
+      expect(store.edges.length).toBe(0);
+      expect(store.nodes.length).toBe(2);
     });
   });
 
diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts
index 642745df..b2311411 100644
--- a/src/store/workflowStore.ts
+++ b/src/store/workflowStore.ts
@@ -682,8 +682,9 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
   updateNodeData: (nodeId: string, data: Partial<WorkflowNodeData>) => {
     const node = get().nodes.find((n) => n.id === nodeId);
 
-    // Debounced undo tracking: skip during execution
-    if (!get().isRunning) {
+    // Debounced undo tracking: skip during execution and during node removal
+    // (clearStaleInputImages calls updateNodeData as a side effect of deletion)
+    if (!get().isRunning && !nodeRemoveCheckpointActive) {
       if (!pendingDataSnapshot) {
         pendingDataSnapshot = captureUndoSnapshot(get());
       }