From 2f7b47a1b444d05a3b81baf1d10039faaee727c2 Mon Sep 17 00:00:00 2001 From: AliFozooni <72443248+AliFozooni@users.noreply.github.com> Date: Sat, 18 Jul 2026 05:25:41 -0700 Subject: [PATCH] fix(ui): restore committed states in workflow undo/redo (#550) --- .../workflow/[workflowId]/RenderWorkflow.tsx | 2 + .../[workflowId]/hooks/useWorkflowState.ts | 6 + .../[workflowId]/stores/workflowStore.test.ts | 151 ++++++++++++++ .../[workflowId]/stores/workflowStore.ts | 184 +++++++++++------- 4 files changed, 271 insertions(+), 72 deletions(-) create mode 100644 ui/src/app/workflow/[workflowId]/stores/workflowStore.test.ts diff --git a/ui/src/app/workflow/[workflowId]/RenderWorkflow.tsx b/ui/src/app/workflow/[workflowId]/RenderWorkflow.tsx index 55638dc2..7c0d624d 100644 --- a/ui/src/app/workflow/[workflowId]/RenderWorkflow.tsx +++ b/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} diff --git a/ui/src/app/workflow/[workflowId]/hooks/useWorkflowState.ts b/ui/src/app/workflow/[workflowId]/hooks/useWorkflowState.ts index c96b29cc..77e8f44c 100644 --- a/ui/src/app/workflow/[workflowId]/hooks/useWorkflowState.ts +++ b/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, diff --git a/ui/src/app/workflow/[workflowId]/stores/workflowStore.test.ts b/ui/src/app/workflow/[workflowId]/stores/workflowStore.test.ts new file mode 100644 index 00000000..3fa5f351 --- /dev/null +++ b/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); + }); +}); diff --git a/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts b/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts index fa8cb278..7af8f160 100644 --- a/ui/src/app/workflow/[workflowId]/stores/workflowStore.ts +++ b/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, + 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((set, get) => ({ // Initial state @@ -134,27 +158,14 @@ export const useWorkflowStore = create((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((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((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((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((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) => {