fix(ui): restore committed states in workflow undo/redo (#550)

ui/src/app/workflow/[workflowId]/RenderWorkflow.tsx

@@ -117,6 +117,7 @@ function RenderWorkflow({
         onConnect,
         onEdgesChange,
         onNodesChange,
+        onDelete,
     } = useWorkflowState({
         initialWorkflowName,
         workflowId,
@@ -514,6 +515,7 @@ function RenderWorkflow({
                                 edges={edges}
                                 onNodesChange={onNodesChange}
                                 onEdgesChange={onEdgesChange}
+                                onDelete={onDelete}
                                 nodeTypes={nodeTypes}
                                 edgeTypes={edgeTypes}
                                 onConnect={isViewingHistoricalVersion ? undefined : onConnect}

ui/src/app/workflow/[workflowId]/hooks/useWorkflowState.ts

@@ -136,6 +136,7 @@ export const useWorkflowState = ({
         templateContextVariables,
         workflowConfigurations,
         initializeWorkflow,
+        commitDeletion,
         setNodes,
         setEdges,
         setWorkflowName,
@@ -493,6 +494,10 @@ export const useWorkflowState = ({
         [setNodes],
     );
 
+    const onDelete = useCallback(() => {
+        commitDeletion();
+    }, [commitDeletion]);
+
     const onRun = async (mode: string) => {
         if (!user?.id) return;
         const workflowRunName = `WR-${getRandomId()}`;
@@ -638,6 +643,7 @@ export const useWorkflowState = ({
         onConnect,
         onEdgesChange,
         onNodesChange,
+        onDelete,
         onRun,
         saveTemplateContextVariables,
         saveWorkflowConfigurations,

ui/src/app/workflow/[workflowId]/stores/workflowStore.test.ts

@@ -0,0 +1,151 @@
+import { beforeEach, describe, expect, it } from 'vitest';
+
+import { FlowEdge, FlowNode } from '@/components/flow/types';
+
+import { useWorkflowStore } from './workflowStore';
+
+const createNode = (id: string, x = 0): FlowNode => ({
+  id,
+  type: 'agentNode',
+  position: { x, y: 0 },
+  data: { name: id },
+});
+
+const createEdge = (id: string, source: string, target: string): FlowEdge => ({
+  id,
+  source,
+  target,
+  data: { condition: 'always', label: id },
+});
+
+const nodeIds = () => useWorkflowStore.getState().nodes.map((node) => node.id);
+
+describe('workflow history', () => {
+  beforeEach(() => {
+    useWorkflowStore.getState().clearStore();
+    useWorkflowStore.getState().initializeWorkflow(1, 'Initial', [], []);
+  });
+
+  it('keeps every committed node state available to undo and redo', () => {
+    const store = useWorkflowStore.getState();
+    store.addNode(createNode('a'));
+    store.addNode(createNode('b'));
+
+    store.undo();
+    expect(nodeIds()).toEqual(['a']);
+
+    store.undo();
+    expect(nodeIds()).toEqual([]);
+
+    store.redo();
+    store.redo();
+    expect(nodeIds()).toEqual(['a', 'b']);
+  });
+
+  it('undoes a node and connected-edge deletion in one step', () => {
+    const firstNode = createNode('a');
+    const secondNode = createNode('b');
+    const edge = createEdge('a-b', 'a', 'b');
+    useWorkflowStore.getState().initializeWorkflow(
+      1,
+      'Initial',
+      [firstNode, secondNode],
+      [edge]
+    );
+
+    // React Flow emits connected-edge removals before node removals, then onDelete.
+    useWorkflowStore.getState().setEdges(
+      [],
+      [{ id: edge.id, type: 'remove' }]
+    );
+    useWorkflowStore.getState().setNodes(
+      [secondNode],
+      [{ id: firstNode.id, type: 'remove' }]
+    );
+    useWorkflowStore.getState().commitDeletion();
+    expect(nodeIds()).toEqual(['b']);
+    expect(useWorkflowStore.getState().edges).toEqual([]);
+    expect(useWorkflowStore.getState().history).toHaveLength(2);
+
+    useWorkflowStore.getState().undo();
+    expect(nodeIds()).toEqual(['a', 'b']);
+    expect(useWorkflowStore.getState().edges).toEqual([edge]);
+
+    useWorkflowStore.getState().redo();
+    expect(nodeIds()).toEqual(['b']);
+    expect(useWorkflowStore.getState().edges).toEqual([]);
+  });
+
+  it('coalesces active dragging into one committed final position', () => {
+    const initialNode = createNode('a');
+    useWorkflowStore.getState().initializeWorkflow(1, 'Initial', [initialNode], []);
+
+    for (const x of [10, 20]) {
+      useWorkflowStore.getState().setNodes(
+        [createNode('a', x)],
+        [{ id: 'a', type: 'position', position: { x, y: 0 }, dragging: true }]
+      );
+    }
+    useWorkflowStore.getState().setNodes(
+      [createNode('a', 30)],
+      [{ id: 'a', type: 'position', position: { x: 30, y: 0 }, dragging: false }]
+    );
+
+    expect(useWorkflowStore.getState().history).toHaveLength(2);
+    useWorkflowStore.getState().undo();
+    expect(useWorkflowStore.getState().nodes[0].position.x).toBe(0);
+    useWorkflowStore.getState().redo();
+    expect(useWorkflowStore.getState().nodes[0].position.x).toBe(30);
+  });
+
+  it('truncates redo history after a new edit', () => {
+    useWorkflowStore.getState().addNode(createNode('a'));
+    useWorkflowStore.getState().addNode(createNode('b'));
+    useWorkflowStore.getState().undo();
+    useWorkflowStore.getState().addNode(createNode('c'));
+
+    expect(nodeIds()).toEqual(['a', 'c']);
+    expect(useWorkflowStore.getState().canRedo()).toBe(false);
+
+    useWorkflowStore.getState().undo();
+    expect(nodeIds()).toEqual(['a']);
+    useWorkflowStore.getState().redo();
+    expect(nodeIds()).toEqual(['a', 'c']);
+  });
+
+  it('keeps the current state at the end of the bounded history', () => {
+    for (let index = 1; index <= 55; index += 1) {
+      useWorkflowStore.getState().setWorkflowName(`Edit ${index}`);
+    }
+
+    expect(useWorkflowStore.getState().history).toHaveLength(50);
+    expect(useWorkflowStore.getState().historyIndex).toBe(49);
+
+    let undoCount = 0;
+    while (useWorkflowStore.getState().canUndo()) {
+      useWorkflowStore.getState().undo();
+      undoCount += 1;
+    }
+    expect(undoCount).toBe(49);
+    expect(useWorkflowStore.getState().workflowName).toBe('Edit 6');
+
+    while (useWorkflowStore.getState().canRedo()) {
+      useWorkflowStore.getState().redo();
+    }
+    expect(useWorkflowStore.getState().workflowName).toBe('Edit 55');
+  });
+
+  it('does not track selection-only changes', () => {
+    const initialNode = createNode('a');
+    useWorkflowStore.getState().initializeWorkflow(1, 'Initial', [initialNode], []);
+
+    useWorkflowStore.getState().setNodes(
+      [{ ...initialNode, selected: true }],
+      [{ id: 'a', type: 'select', selected: true }]
+    );
+
+    expect(useWorkflowStore.getState().nodes[0].selected).toBe(true);
+    expect(useWorkflowStore.getState().history).toHaveLength(1);
+    expect(useWorkflowStore.getState().isDirty).toBe(false);
+  });
+});

ui/src/app/workflow/[workflowId]/stores/workflowStore.ts

@@ -54,7 +54,7 @@ interface WorkflowActions {
   ) => void;
 
   // History management
-  pushToHistory: () => void;
+  commitDeletion: () => void;
   undo: () => void;
   redo: () => void;
   canUndo: () => boolean;
@@ -99,6 +99,30 @@ type WorkflowStore = WorkflowState & WorkflowActions;
 
 const MAX_HISTORY_SIZE = 50;
 
+const commitHistory = (
+  state: Pick<WorkflowState, 'history' | 'historyIndex'>,
+  snapshot: HistoryState
+): Pick<
+  WorkflowState,
+  'nodes' | 'edges' | 'workflowName' | 'history' | 'historyIndex' | 'isDirty'
+> => {
+  const history = [
+    ...state.history.slice(0, state.historyIndex + 1),
+    snapshot,
+  ];
+
+  if (history.length > MAX_HISTORY_SIZE) {
+    history.shift();
+  }
+
+  return {
+    ...snapshot,
+    history,
+    historyIndex: history.length - 1,
+    isDirty: true,
+  };
+};
+
 // Create the store
 export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
   // Initial state
@@ -134,27 +158,14 @@ export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
     });
   },
 
-  pushToHistory: () => {
-    const state = get();
-    const currentState: HistoryState = {
-      nodes: state.nodes,
-      edges: state.edges,
-      workflowName: state.workflowName,
-    };
-
-    // Remove any forward history if we're not at the end
-    const newHistory = state.history.slice(0, state.historyIndex + 1);
-    newHistory.push(currentState);
-
-    // Limit history size
-    if (newHistory.length > MAX_HISTORY_SIZE) {
-      newHistory.shift();
-    }
-
-    set({
-      history: newHistory,
-      historyIndex: newHistory.length - 1,
-    });
+  commitDeletion: () => {
+    set((state) =>
+      commitHistory(state, {
+        nodes: state.nodes,
+        edges: state.edges,
+        workflowName: state.workflowName,
+      })
+    );
   },
 
   undo: () => {
@@ -200,10 +211,10 @@ export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
   setNodes: (nodes, changes) => {
     // Determine whether to push to history and set isDirty based on change types
     if (changes && changes.length > 0) {
-      // Check for add/remove changes (always push to history)
-      const hasAddRemoveChanges = changes.some(change =>
-        change.type === 'add' || change.type === 'remove'
-      );
+      const hasAddChanges = changes.some(change => change.type === 'add');
+      // React Flow emits edge and node removals separately. They are committed
+      // together from onDelete after both live-state updates have completed.
+      const hasRemoveChanges = changes.some(change => change.type === 'remove');
 
       // Check for position changes - only push to history when drag ENDS (dragging: false)
       // but still mark as dirty during dragging
@@ -214,11 +225,16 @@ export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
         change.type === 'position' && change.dragging === true
       );
 
-      if (hasAddRemoveChanges || hasDragEndChanges) {
-        get().pushToHistory();
-        set({ nodes, isDirty: true });
-      } else if (isActiveDragging) {
-        // During active dragging, update nodes but don't push to history
+      if (hasAddChanges || hasDragEndChanges) {
+        set((state) =>
+          commitHistory(state, {
+            nodes,
+            edges: state.edges,
+            workflowName: state.workflowName,
+          })
+        );
+      } else if (hasRemoveChanges || isActiveDragging) {
+        // During active dragging or deletion, update nodes before committing.
         set({ nodes, isDirty: true });
       } else {
         // For selection changes or dimension updates, don't push to history or set dirty
@@ -231,49 +247,62 @@ export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
   },
 
   addNode: (node) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      nodes: [...state.nodes, node],
-      isDirty: true
+    set((state) => {
+      const nodes = [...state.nodes, node];
+      return commitHistory(state, {
+        nodes,
+        edges: state.edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   updateNode: (nodeId, updates) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      nodes: state.nodes.map((node) =>
+    set((state) => {
+      const nodes = state.nodes.map((node) =>
         node.id === nodeId ? { ...node, ...updates } : node
-      ),
-      isDirty: true,
+      );
+      return commitHistory(state, {
+        nodes,
+        edges: state.edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   deleteNode: (nodeId) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      nodes: state.nodes.filter((node) => node.id !== nodeId),
-      edges: state.edges.filter(
+    set((state) => {
+      const nodes = state.nodes.filter((node) => node.id !== nodeId);
+      const edges = state.edges.filter(
         (edge) => edge.source !== nodeId && edge.target !== nodeId
-      ),
-      isDirty: true,
+      );
+      return commitHistory(state, {
+        nodes,
+        edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   setEdges: (edges, changes) => {
     // Determine whether to push to history and set isDirty based on change types
     if (changes && changes.length > 0) {
-      // Check if any changes are user-initiated (not just selections)
-      const hasDirtyChanges = changes.some(change =>
+      const hasImmediateHistoryChanges = changes.some(change =>
         change.type === 'add' ||
-        change.type === 'remove' ||
         change.type === 'replace'
       );
+      const hasRemoveChanges = changes.some(change => change.type === 'remove');
 
-      if (hasDirtyChanges) {
-        get().pushToHistory();
+      if (hasImmediateHistoryChanges) {
+        set((state) =>
+          commitHistory(state, {
+            nodes: state.nodes,
+            edges,
+            workflowName: state.workflowName,
+          })
+        );
+      } else if (hasRemoveChanges) {
+        // React Flow calls onDelete after all edge and node removals.
         set({ edges, isDirty: true });
       } else {
         // For selection changes, don't push to history
@@ -286,37 +315,48 @@ export const useWorkflowStore = create<WorkflowStore>((set, get) => ({
   },
 
   addEdge: (edge) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      edges: [...state.edges, edge],
-      isDirty: true
+    set((state) => {
+      const edges = [...state.edges, edge];
+      return commitHistory(state, {
+        nodes: state.nodes,
+        edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   updateEdge: (edgeId, updates) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      edges: state.edges.map((edge) =>
+    set((state) => {
+      const edges = state.edges.map((edge) =>
         edge.id === edgeId ? { ...edge, ...updates } : edge
-      ),
-      isDirty: true,
+      );
+      return commitHistory(state, {
+        nodes: state.nodes,
+        edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   deleteEdge: (edgeId) => {
-    const state = get();
-    get().pushToHistory();
-    set({
-      edges: state.edges.filter((edge) => edge.id !== edgeId),
-      isDirty: true,
+    set((state) => {
+      const edges = state.edges.filter((edge) => edge.id !== edgeId);
+      return commitHistory(state, {
+        nodes: state.nodes,
+        edges,
+        workflowName: state.workflowName,
+      });
     });
   },
 
   setWorkflowName: (workflowName) => {
-    get().pushToHistory();
-    set({ workflowName, isDirty: true });
+    set((state) =>
+      commitHistory(state, {
+        nodes: state.nodes,
+        edges: state.edges,
+        workflowName,
+      })
+    );
   },
 
   setTemplateContextVariables: (templateContextVariables) => {