Browse Source

feat: integrate undo/redo into workflow store with Cmd+Z / Cmd+Shift+Z

Captures snapshots before all undoable mutations (add/remove nodes,
connect/remove edges, paste, group operations, edge style, drag start).
Data edits are debounced (500ms). Execution outputs are not undoable.
History depth: 50 entries. Cleared on load/clear workflow.

Keyboard shortcuts added to WorkflowCanvas and shortcuts dialog.
15 integration tests covering all undo/redo scenarios through the store.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
handoff-20260429-1057
shrimbly 4 months ago
parent
commit
b263dfc4ee
  1. 2
      src/components/KeyboardShortcutsDialog.tsx
  2. 18
      src/components/WorkflowCanvas.tsx
  3. 466
      src/store/__tests__/undoRedo.test.ts
  4. 167
      src/store/workflowStore.ts

2
src/components/KeyboardShortcutsDialog.tsx

@ -22,6 +22,8 @@ const shortcutGroups: ShortcutGroup[] = [
{ keys: [`${modKey}`, "Enter"], description: "Run workflow" }, { keys: [`${modKey}`, "Enter"], description: "Run workflow" },
{ keys: [`${modKey}`, "C"], description: "Copy selected nodes" }, { keys: [`${modKey}`, "C"], description: "Copy selected nodes" },
{ keys: [`${modKey}`, "V"], description: "Paste nodes / image / text" }, { keys: [`${modKey}`, "V"], description: "Paste nodes / image / text" },
{ keys: [`${modKey}`, "Z"], description: "Undo" },
{ keys: [`${modKey}`, "Shift", "Z"], description: "Redo" },
{ keys: ["?"], description: "Show keyboard shortcuts" }, { keys: ["?"], description: "Show keyboard shortcuts" },
], ],
}, },

18
src/components/WorkflowCanvas.tsx

@ -1311,6 +1311,8 @@ export function WorkflowCanvas() {
const pasteNodes = useWorkflowStore((state) => state.pasteNodes); const pasteNodes = useWorkflowStore((state) => state.pasteNodes);
const clearClipboard = useWorkflowStore((state) => state.clearClipboard); const clearClipboard = useWorkflowStore((state) => state.clearClipboard);
const clipboard = useWorkflowStore((state) => state.clipboard); const clipboard = useWorkflowStore((state) => state.clipboard);
const undo = useWorkflowStore((state) => state.undo);
const redo = useWorkflowStore((state) => state.redo);
// Add non-passive wheel listener to handle zoom/pan and prevent browser navigation // Add non-passive wheel listener to handle zoom/pan and prevent browser navigation
// This replaces the onWheel prop which is passive by default and can't preventDefault // This replaces the onWheel prop which is passive by default and can't preventDefault
@ -1418,6 +1420,20 @@ export function WorkflowCanvas() {
return; return;
} }
// Handle undo (Ctrl/Cmd + Z)
if ((event.ctrlKey || event.metaKey) && event.key === "z" && !event.shiftKey) {
event.preventDefault();
undo();
return;
}
// Handle redo (Ctrl/Cmd + Shift + Z)
if ((event.ctrlKey || event.metaKey) && event.key === "z" && event.shiftKey) {
event.preventDefault();
redo();
return;
}
// Helper to get viewport center position in flow coordinates // Helper to get viewport center position in flow coordinates
const getViewportCenter = () => { const getViewportCenter = () => {
const viewport = getViewport(); const viewport = getViewport();
@ -1667,7 +1683,7 @@ export function WorkflowCanvas() {
]); ]);
}); });
} }
}, [nodes, onNodesChange, copySelectedNodes, pasteNodes, clearClipboard, clipboard, getViewport, addNode, updateNodeData, executeWorkflow, setShortcutsDialogOpen]); }, [nodes, onNodesChange, copySelectedNodes, pasteNodes, clearClipboard, clipboard, getViewport, addNode, updateNodeData, executeWorkflow, setShortcutsDialogOpen, undo, redo]);
useEffect(() => { useEffect(() => {
window.addEventListener("keydown", handleKeyDown); window.addEventListener("keydown", handleKeyDown);

466
src/store/__tests__/undoRedo.test.ts

@ -0,0 +1,466 @@
/**
* Integration tests for undo/redo through the Zustand store.
*
* Tests that undo/redo correctly captures and restores workflow state
* for all undoable actions (add/remove nodes, connect edges, etc.).
*/
import { describe, it, expect, vi, beforeEach } from "vitest";
import { act } from "@testing-library/react";
import { useWorkflowStore } from "../workflowStore";
// Mock the Toast hook
vi.mock("@/components/Toast", () => ({
useToast: {
getState: () => ({
show: vi.fn(),
}),
},
}));
// Mock the logger
vi.mock("@/utils/logger", () => ({
logger: {
info: vi.fn(),
warn: vi.fn(),
error: vi.fn(),
startSession: vi.fn().mockResolvedValue(undefined),
endSession: vi.fn().mockResolvedValue(undefined),
getCurrentSession: vi.fn().mockReturnValue(null),
},
}));
// Mock localStorage
const mockLocalStorage: Record<string, string> = {};
vi.stubGlobal("localStorage", {
getItem: vi.fn((key: string) => mockLocalStorage[key] || null),
setItem: vi.fn((key: string, value: string) => {
mockLocalStorage[key] = value;
}),
removeItem: vi.fn((key: string) => {
delete mockLocalStorage[key];
}),
clear: vi.fn(() => {
Object.keys(mockLocalStorage).forEach((key) => delete mockLocalStorage[key]);
}),
});
function resetStore() {
const store = useWorkflowStore.getState();
store.clearWorkflow();
}
describe("Undo/Redo integration", () => {
beforeEach(() => {
resetStore();
});
it("starts with canUndo and canRedo both false", () => {
const { canUndo, canRedo } = useWorkflowStore.getState();
expect(canUndo).toBe(false);
expect(canRedo).toBe(false);
});
describe("addNode + undo", () => {
it("undoes adding a node", () => {
let store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(0);
act(() => {
store.addNode("prompt", { x: 100, y: 100 });
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(1);
expect(store.canUndo).toBe(true);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(0);
expect(store.canUndo).toBe(false);
expect(store.canRedo).toBe(true);
});
it("redoes after undoing addNode", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 100, y: 100 });
});
store = useWorkflowStore.getState();
const nodeId = store.nodes[0].id;
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(0);
act(() => {
store.redo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(1);
expect(store.nodes[0].id).toBe(nodeId);
expect(store.canRedo).toBe(false);
});
});
describe("removeNode + undo", () => {
it("undoes removing a node", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
store = useWorkflowStore.getState();
const nodeId = store.nodes[0].id;
act(() => {
store.removeNode(nodeId);
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(0);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(1);
expect(store.nodes[0].id).toBe(nodeId);
});
});
describe("onConnect + undo", () => {
it("undoes connecting an edge", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
store.addNode("nanoBanana", { x: 300, y: 0 });
});
store = useWorkflowStore.getState();
const promptId = store.nodes[0].id;
const genId = store.nodes[1].id;
const edgesBefore = store.edges.length;
act(() => {
store.onConnect({
source: promptId,
target: genId,
sourceHandle: "text",
targetHandle: "text",
});
});
store = useWorkflowStore.getState();
expect(store.edges.length).toBe(edgesBefore + 1);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.edges.length).toBe(edgesBefore);
});
});
describe("removeEdge + undo", () => {
it("undoes removing an edge", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
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();
const edgeId = store.edges[store.edges.length - 1].id;
act(() => {
store.removeEdge(edgeId);
});
store = useWorkflowStore.getState();
expect(store.edges.find((e) => e.id === edgeId)).toBeUndefined();
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.edges.find((e) => e.id === edgeId)).toBeDefined();
});
});
describe("new action clears redo stack", () => {
it("clears redo stack when a new undoable action is performed", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
act(() => {
store = useWorkflowStore.getState();
store.undo();
});
store = useWorkflowStore.getState();
expect(store.canRedo).toBe(true);
// Perform a new action — should clear redo stack
act(() => {
store.addNode("imageInput", { x: 100, y: 100 });
});
store = useWorkflowStore.getState();
expect(store.canRedo).toBe(false);
});
});
describe("multiple undo/redo", () => {
it("supports multiple undo/redo in sequence", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
act(() => {
store = useWorkflowStore.getState();
store.addNode("imageInput", { x: 100, y: 0 });
});
act(() => {
store = useWorkflowStore.getState();
store.addNode("nanoBanana", { x: 200, y: 0 });
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(3);
// Undo 3 times
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(2);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(1);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(0);
expect(store.canUndo).toBe(false);
// Redo 3 times
act(() => {
store.redo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(1);
act(() => {
store.redo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(2);
act(() => {
store.redo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(3);
expect(store.canRedo).toBe(false);
});
});
describe("undo on empty stack is no-op", () => {
it("does nothing when undo stack is empty", () => {
let store = useWorkflowStore.getState();
const before = store.nodes.length;
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(before);
expect(store.canUndo).toBe(false);
});
});
describe("redo on empty stack is no-op", () => {
it("does nothing when redo stack is empty", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
store = useWorkflowStore.getState();
const nodeCount = store.nodes.length;
act(() => {
store.redo();
});
store = useWorkflowStore.getState();
expect(store.nodes.length).toBe(nodeCount);
});
});
describe("clearWorkflow resets undo history", () => {
it("clears undo and redo stacks on clearWorkflow", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
store.addNode("imageInput", { x: 100, y: 0 });
});
store = useWorkflowStore.getState();
expect(store.canUndo).toBe(true);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.canRedo).toBe(true);
act(() => {
store.clearWorkflow();
});
store = useWorkflowStore.getState();
expect(store.canUndo).toBe(false);
expect(store.canRedo).toBe(false);
});
});
describe("setEdgeStyle + undo", () => {
it("undoes edge style change", () => {
let store = useWorkflowStore.getState();
const originalStyle = store.edgeStyle;
act(() => {
store.setEdgeStyle("angular");
});
store = useWorkflowStore.getState();
expect(store.edgeStyle).toBe("angular");
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.edgeStyle).toBe(originalStyle);
});
});
describe("group operations + undo", () => {
it("undoes creating a group", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
store.addNode("imageInput", { x: 100, y: 0 });
});
store = useWorkflowStore.getState();
const nodeIds = store.nodes.map((n) => n.id);
act(() => {
store.createGroup(nodeIds);
});
store = useWorkflowStore.getState();
expect(Object.keys(store.groups).length).toBe(1);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(Object.keys(store.groups).length).toBe(0);
});
it("undoes deleting a group", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
store = useWorkflowStore.getState();
const nodeIds = store.nodes.map((n) => n.id);
let groupId: string;
act(() => {
groupId = store.createGroup(nodeIds);
});
store = useWorkflowStore.getState();
expect(Object.keys(store.groups).length).toBe(1);
act(() => {
store.deleteGroup(groupId!);
});
store = useWorkflowStore.getState();
expect(Object.keys(store.groups).length).toBe(0);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(Object.keys(store.groups).length).toBe(1);
});
});
describe("hasUnsavedChanges is set on undo/redo", () => {
it("marks as unsaved after undo", () => {
let store = useWorkflowStore.getState();
act(() => {
store.addNode("prompt", { x: 0, y: 0 });
});
// Simulate that the workflow was "saved"
act(() => {
useWorkflowStore.setState({ hasUnsavedChanges: false });
});
store = useWorkflowStore.getState();
expect(store.hasUnsavedChanges).toBe(false);
act(() => {
store.undo();
});
store = useWorkflowStore.getState();
expect(store.hasUnsavedChanges).toBe(true);
});
});
});

167
src/store/workflowStore.ts

@ -26,6 +26,7 @@ import {
MatchMode, MatchMode,
MODEL_DISPLAY_NAMES, MODEL_DISPLAY_NAMES,
} from "@/types"; } from "@/types";
import { UndoManager, UndoSnapshot } from "./undoHistory";
import { useToast } from "@/components/Toast"; import { useToast } from "@/components/Toast";
import { logger } from "@/utils/logger"; import { logger } from "@/utils/logger";
import { externalizeWorkflowMedia, hydrateWorkflowMedia } from "@/utils/mediaStorage"; import { externalizeWorkflowMedia, hydrateWorkflowMedia } from "@/utils/mediaStorage";
@ -208,6 +209,12 @@ interface WorkflowStore {
clipboard: ClipboardData | null; clipboard: ClipboardData | null;
groups: Record<string, NodeGroup>; groups: Record<string, NodeGroup>;
// Undo/Redo
canUndo: boolean;
canRedo: boolean;
undo: () => void;
redo: () => void;
// Settings // Settings
setEdgeStyle: (style: EdgeStyle) => void; setEdgeStyle: (style: EdgeStyle) => void;
@ -385,6 +392,12 @@ let nodeIdCounter = 0;
let groupIdCounter = 0; let groupIdCounter = 0;
let autoSaveIntervalId: ReturnType<typeof setInterval> | null = null; let autoSaveIntervalId: ReturnType<typeof setInterval> | null = null;
// Undo/redo state (module-level, not in Zustand to avoid serialization)
const undoManager = new UndoManager();
let isDragging = false;
let pendingDataSnapshot: UndoSnapshot | null = null;
let dataChangeTimer: ReturnType<typeof setTimeout> | null = null;
// RAF debounce for hover updates — coalesces rapid mouseenter/mouseleave events // RAF debounce for hover updates — coalesces rapid mouseenter/mouseleave events
// into a single store update per animation frame // into a single store update per animation frame
let hoverRafId: number | null = null; let hoverRafId: number | null = null;
@ -450,6 +463,45 @@ function clearStaleInputImages(
} }
} }
/** Capture current undoable state as a deep-cloned snapshot */
function captureUndoSnapshot(state: WorkflowStore): UndoSnapshot {
const cloned = JSON.parse(JSON.stringify({
nodes: state.nodes,
edges: state.edges,
groups: state.groups,
edgeStyle: state.edgeStyle,
})) as UndoSnapshot;
// Strip transient selection state from cloned nodes
for (const node of cloned.nodes) {
delete (node as Record<string, unknown>).selected;
}
return cloned;
}
/** Flush pending debounced data snapshot, capture current state, push to undoManager */
function pushUndoCheckpoint(
get: () => WorkflowStore,
set: (partial: Partial<WorkflowStore>) => void,
): void {
// Flush any pending debounced data snapshot first
if (pendingDataSnapshot) {
undoManager.push(pendingDataSnapshot);
pendingDataSnapshot = null;
if (dataChangeTimer) {
clearTimeout(dataChangeTimer);
dataChangeTimer = null;
}
}
const snapshot = captureUndoSnapshot(get());
undoManager.push(snapshot);
set({ canUndo: undoManager.canUndo, canRedo: undoManager.canRedo });
}
/** Update reactive canUndo/canRedo flags */
function syncUndoFlags(set: (partial: Partial<WorkflowStore>) => void): void {
set({ canUndo: undoManager.canUndo, canRedo: undoManager.canRedo });
}
const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
nodes: [], nodes: [],
edges: [], edges: [],
@ -513,7 +565,58 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
// Skip propagation initial state // Skip propagation initial state
skippedNodeIds: new Set<string>(), skippedNodeIds: new Set<string>(),
// Undo/Redo initial state
canUndo: false,
canRedo: false,
undo: () => {
// Discard any pending debounced data snapshot — it's intermediate state
if (pendingDataSnapshot) {
pendingDataSnapshot = null;
if (dataChangeTimer) {
clearTimeout(dataChangeTimer);
dataChangeTimer = null;
}
}
const current = captureUndoSnapshot(get());
const previous = undoManager.undo(current);
if (previous) {
set({
nodes: previous.nodes,
edges: previous.edges,
groups: previous.groups,
edgeStyle: previous.edgeStyle,
hasUnsavedChanges: true,
});
}
syncUndoFlags(set);
},
redo: () => {
// Discard any pending debounced data snapshot
if (pendingDataSnapshot) {
pendingDataSnapshot = null;
if (dataChangeTimer) {
clearTimeout(dataChangeTimer);
dataChangeTimer = null;
}
}
const current = captureUndoSnapshot(get());
const next = undoManager.redo(current);
if (next) {
set({
nodes: next.nodes,
edges: next.edges,
groups: next.groups,
edgeStyle: next.edgeStyle,
hasUnsavedChanges: true,
});
}
syncUndoFlags(set);
},
setEdgeStyle: (style: EdgeStyle) => { setEdgeStyle: (style: EdgeStyle) => {
pushUndoCheckpoint(get, set);
set({ edgeStyle: style }); set({ edgeStyle: style });
}, },
@ -562,6 +665,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
style: { width, height }, style: { width, height },
}; };
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
nodes: [...state.nodes, newNode], nodes: [...state.nodes, newNode],
hasUnsavedChanges: true, hasUnsavedChanges: true,
@ -574,6 +678,23 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
updateNodeData: (nodeId: string, data: Partial<WorkflowNodeData>) => { updateNodeData: (nodeId: string, data: Partial<WorkflowNodeData>) => {
const node = get().nodes.find((n) => n.id === nodeId); const node = get().nodes.find((n) => n.id === nodeId);
// Debounced undo tracking: skip during execution
if (!get().isRunning) {
if (!pendingDataSnapshot) {
pendingDataSnapshot = captureUndoSnapshot(get());
}
if (dataChangeTimer) clearTimeout(dataChangeTimer);
dataChangeTimer = setTimeout(() => {
if (pendingDataSnapshot) {
undoManager.push(pendingDataSnapshot);
pendingDataSnapshot = null;
syncUndoFlags(set);
}
dataChangeTimer = null;
}, 500);
}
set((state) => ({ set((state) => ({
nodes: state.nodes.map((node) => nodes: state.nodes.map((node) =>
node.id === nodeId node.id === nodeId
@ -592,6 +713,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
removeNode: (nodeId: string) => { removeNode: (nodeId: string) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
nodes: state.nodes.filter((node) => node.id !== nodeId), nodes: state.nodes.filter((node) => node.id !== nodeId),
edges: state.edges.filter( edges: state.edges.filter(
@ -610,6 +732,26 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
// Track manual changes only for remove operations (not position/selection/dimensions) // Track manual changes only for remove operations (not position/selection/dimensions)
const hasRemoveChange = changes.some((c) => c.type === "remove"); const hasRemoveChange = changes.some((c) => c.type === "remove");
// Undo: capture snapshot on drag start
const hasDragStart = changes.some(
(c) => c.type === "position" && (c as { dragging?: boolean }).dragging === true
);
const hasDragEnd = changes.some(
(c) => c.type === "position" && (c as { dragging?: boolean }).dragging === false
);
if (hasDragStart && !isDragging) {
isDragging = true;
pushUndoCheckpoint(get, set);
}
if (hasDragEnd && isDragging) {
isDragging = false;
}
// Undo: capture snapshot before node removal
if (hasRemoveChange) {
pushUndoCheckpoint(get, set);
}
set((state) => ({ set((state) => ({
nodes: applyNodeChanges(changes, state.nodes), nodes: applyNodeChanges(changes, state.nodes),
...(hasMeaningfulChange ? { hasUnsavedChanges: true } : {}), ...(hasMeaningfulChange ? { hasUnsavedChanges: true } : {}),
@ -627,6 +769,11 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
const hasRemoveChange = changes.some((c) => c.type === "remove"); const hasRemoveChange = changes.some((c) => c.type === "remove");
const hasAddOrRemove = changes.some((c) => c.type === "add" || c.type === "remove"); const hasAddOrRemove = changes.some((c) => c.type === "add" || c.type === "remove");
// Undo: capture snapshot before edge removal
if (hasRemoveChange) {
pushUndoCheckpoint(get, set);
}
// Capture removed edges before applyEdgeChanges removes them // Capture removed edges before applyEdgeChanges removes them
let removedEdges: WorkflowEdge[] = []; let removedEdges: WorkflowEdge[] = [];
if (hasRemoveChange) { if (hasRemoveChange) {
@ -653,6 +800,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
onConnect: (connection: Connection, edgeDataOverrides?: Record<string, unknown>) => { onConnect: (connection: Connection, edgeDataOverrides?: Record<string, unknown>) => {
pushUndoCheckpoint(get, set);
set((state) => { set((state) => {
const baseData = buildConnectionEdgeData(connection, state.nodes, state.edges); const baseData = buildConnectionEdgeData(connection, state.nodes, state.edges);
const newEdge = { const newEdge = {
@ -671,6 +819,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
addEdgeWithType: (connection: Connection, edgeType: string, edgeDataOverrides?: Record<string, unknown>) => { addEdgeWithType: (connection: Connection, edgeType: string, edgeDataOverrides?: Record<string, unknown>) => {
pushUndoCheckpoint(get, set);
set((state) => { set((state) => {
const baseData = buildConnectionEdgeData(connection, state.nodes, state.edges); const baseData = buildConnectionEdgeData(connection, state.nodes, state.edges);
const newEdge = { const newEdge = {
@ -687,6 +836,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
removeEdge: (edgeId: string) => { removeEdge: (edgeId: string) => {
pushUndoCheckpoint(get, set);
const removedEdge = get().edges.find((e) => e.id === edgeId); const removedEdge = get().edges.find((e) => e.id === edgeId);
set((state) => ({ set((state) => ({
edges: state.edges.filter((edge) => edge.id !== edgeId), edges: state.edges.filter((edge) => edge.id !== edgeId),
@ -697,6 +847,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
toggleEdgePause: (edgeId: string) => { toggleEdgePause: (edgeId: string) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
edges: state.edges.map((edge) => edges: state.edges.map((edge) =>
edge.id === edgeId edge.id === edgeId
@ -732,6 +883,8 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
if (!clipboard || clipboard.nodes.length === 0) return; if (!clipboard || clipboard.nodes.length === 0) return;
pushUndoCheckpoint(get, set);
// Create a mapping from old node IDs to new node IDs // Create a mapping from old node IDs to new node IDs
const idMapping = new Map<string, string>(); const idMapping = new Map<string, string>();
@ -807,6 +960,8 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
if (nodeIds.length === 0) return ""; if (nodeIds.length === 0) return "";
pushUndoCheckpoint(get, set);
// Get the nodes to group // Get the nodes to group
const nodesToGroup = nodes.filter((n) => nodeIds.includes(n.id)); const nodesToGroup = nodes.filter((n) => nodeIds.includes(n.id));
if (nodesToGroup.length === 0) return ""; if (nodesToGroup.length === 0) return "";
@ -870,6 +1025,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
deleteGroup: (groupId: string) => { deleteGroup: (groupId: string) => {
pushUndoCheckpoint(get, set);
set((state) => { set((state) => {
const { [groupId]: _, ...remainingGroups } = state.groups; const { [groupId]: _, ...remainingGroups } = state.groups;
return { return {
@ -883,6 +1039,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
addNodesToGroup: (nodeIds: string[], groupId: string) => { addNodesToGroup: (nodeIds: string[], groupId: string) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
nodes: state.nodes.map((node) => nodes: state.nodes.map((node) =>
nodeIds.includes(node.id) ? { ...node, groupId } : node nodeIds.includes(node.id) ? { ...node, groupId } : node
@ -892,6 +1049,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
removeNodesFromGroup: (nodeIds: string[]) => { removeNodesFromGroup: (nodeIds: string[]) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
nodes: state.nodes.map((node) => nodes: state.nodes.map((node) =>
nodeIds.includes(node.id) ? { ...node, groupId: undefined } : node nodeIds.includes(node.id) ? { ...node, groupId: undefined } : node
@ -901,6 +1059,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
updateGroup: (groupId: string, updates: Partial<NodeGroup>) => { updateGroup: (groupId: string, updates: Partial<NodeGroup>) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
groups: { groups: {
...state.groups, ...state.groups,
@ -911,6 +1070,7 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
}, },
toggleGroupLock: (groupId: string) => { toggleGroupLock: (groupId: string) => {
pushUndoCheckpoint(get, set);
set((state) => ({ set((state) => ({
groups: { groups: {
...state.groups, ...state.groups,
@ -1812,6 +1972,10 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
get().clearSnapshot(); get().clearSnapshot();
} }
// Clear undo history — loading a workflow is a fresh start
undoManager.clear();
syncUndoFlags(set);
// Recompute dimming after loading workflow // Recompute dimming after loading workflow
get().recomputeDimmedNodes(); get().recomputeDimmedNodes();
}, },
@ -1842,6 +2006,9 @@ const workflowStoreImpl: StateCreator<WorkflowStore> = (set, get) => ({
skippedNodeIds: new Set<string>(), skippedNodeIds: new Set<string>(),
}); });
get().clearSnapshot(); get().clearSnapshot();
// Clear undo history
undoManager.clear();
syncUndoFlags(set);
}, },
addToGlobalHistory: (item: Omit<ImageHistoryItem, "id">) => { addToGlobalHistory: (item: Omit<ImageHistoryItem, "id">) => {

Loading…
Cancel
Save