diff --git a/src/store/__tests__/undoRedo.test.ts b/src/store/__tests__/undoRedo.test.ts index b5425fe5..5bbda035 100644 --- a/src/store/__tests__/undoRedo.test.ts +++ b/src/store/__tests__/undoRedo.test.ts @@ -215,7 +215,7 @@ describe("Undo/Redo integration", () => { }); describe("delete node via onNodesChange + onEdgesChange restores connections", () => { - it("single undo restores both node and its connected edges", async () => { + it("single undo restores both node and its connected edges", () => { let store = useWorkflowStore.getState(); // Create two nodes and connect them @@ -243,16 +243,17 @@ describe("Undo/Redo integration", () => { 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 + // Simulate pressing Delete: React Flow v12 fires onEdgesChange(remove) + // BEFORE onNodesChange(remove) — both synchronously in the same cycle. + const edgeId = store.edges[0].id; act(() => { + store.onEdgesChange([{ type: "remove", id: edgeId }]); 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(); + // Advance past the setTimeout(0) that clears deleteCheckpointActive + act(() => { + vi.advanceTimersByTime(0); }); store = useWorkflowStore.getState(); @@ -281,7 +282,7 @@ describe("Undo/Redo integration", () => { expect(store.nodes.length).toBe(2); }); - it("single undo restores image-source node and does not create extra entries from clearStaleInputImages", async () => { + it("single undo restores image-source node and does not create extra entries from clearStaleInputImages", () => { let store = useWorkflowStore.getState(); // Create imageInput -> nanoBanana (image connection triggers clearStaleInputImages on delete) @@ -309,20 +310,18 @@ describe("Undo/Redo integration", () => { 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 + // Delete the imageInput node — React Flow v12 fires edges first, then nodes. + // This triggers clearStaleInputImages which calls updateNodeData on the + // nanoBanana target as a side effect. + const edgeId = store.edges[0].id; act(() => { + store.onEdgesChange([{ type: "remove", id: edgeId }]); 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 () => { + // Advance past the setTimeout(0) that clears deleteCheckpointActive + // and any debounced data-change timer (500ms) + act(() => { vi.advanceTimersByTime(600); }); @@ -349,6 +348,54 @@ describe("Undo/Redo integration", () => { expect(store.edges.length).toBe(0); expect(store.nodes.length).toBe(2); }); + + it("standalone edge removal via onEdgesChange still creates undo entry", () => { + let store = useWorkflowStore.getState(); + + 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.edges.length).toBe(1); + const edgeId = store.edges[0].id; + + // Only remove the edge (no node removal) — should still be undoable + act(() => { + store.onEdgesChange([{ type: "remove", id: edgeId }]); + }); + + act(() => { + vi.advanceTimersByTime(0); + }); + + store = useWorkflowStore.getState(); + expect(store.edges.length).toBe(0); + + act(() => { + store.undo(); + }); + + store = useWorkflowStore.getState(); + expect(store.edges.length).toBe(1); + expect(store.edges[0].id).toBe(edgeId); + }); }); describe("new action clears redo stack", () => { diff --git a/src/store/workflowStore.ts b/src/store/workflowStore.ts index b2311411..92df2075 100644 --- a/src/store/workflowStore.ts +++ b/src/store/workflowStore.ts @@ -397,9 +397,15 @@ const undoManager = new UndoManager(); let isDragging = false; let pendingDataSnapshot: UndoSnapshot | null = null; let dataChangeTimer: ReturnType | 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; +// When true, a remove-checkpoint was already pushed in the current React Flow +// deleteElements cycle. React Flow v12 fires onEdgesChange(remove) BEFORE +// onNodesChange(remove) — both happen synchronously inside the same microtask +// after an internal `await`. The flag is set by whichever handler fires first +// and checked by the second so only ONE checkpoint is recorded. It is also +// checked by updateNodeData to suppress debounced snapshots from side-effects +// like clearStaleInputImages. Cleared via setTimeout(0) (macrotask) so it +// survives all microtasks / Promise continuations in the current event-loop turn. +let deleteCheckpointActive = false; // RAF debounce for hover updates — coalesces rapid mouseenter/mouseleave events // into a single store update per animation frame @@ -682,9 +688,9 @@ const workflowStoreImpl: StateCreator = (set, get) => ({ updateNodeData: (nodeId: string, data: Partial) => { const node = get().nodes.find((n) => n.id === nodeId); - // Debounced undo tracking: skip during execution and during node removal + // Debounced undo tracking: skip during execution and during node/edge deletion // (clearStaleInputImages calls updateNodeData as a side effect of deletion) - if (!get().isRunning && !nodeRemoveCheckpointActive) { + if (!get().isRunning && !deleteCheckpointActive) { if (!pendingDataSnapshot) { pendingDataSnapshot = captureUndoSnapshot(get()); } @@ -751,13 +757,13 @@ const workflowStoreImpl: StateCreator = (set, get) => ({ isDragging = false; } - // 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) { + // Undo: capture snapshot before node removal — but skip if onEdgesChange + // already pushed a checkpoint in this same deleteElements cycle + // (React Flow v12 fires edge removals BEFORE node removals). + if (hasRemoveChange && !deleteCheckpointActive) { pushUndoCheckpoint(get, set); - nodeRemoveCheckpointActive = true; - Promise.resolve().then(() => { nodeRemoveCheckpointActive = false; }); + deleteCheckpointActive = true; + setTimeout(() => { deleteCheckpointActive = false; }, 0); } set((state) => ({ @@ -777,10 +783,14 @@ const workflowStoreImpl: StateCreator = (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 — but skip if a node-remove - // checkpoint was already pushed in this same React Flow event cycle - if (hasRemoveChange && !nodeRemoveCheckpointActive) { + // Undo: capture snapshot before edge removal — but skip if a checkpoint + // was already pushed in this same React Flow deleteElements cycle. + // React Flow v12 fires onEdgesChange(remove) BEFORE onNodesChange(remove), + // so this is typically the first handler to set the flag. + if (hasRemoveChange && !deleteCheckpointActive) { pushUndoCheckpoint(get, set); + deleteCheckpointActive = true; + setTimeout(() => { deleteCheckpointActive = false; }, 0); } // Capture removed edges before applyEdgeChanges removes them