Merge pull request #248 from ShipSecAI/@krishna9358/undo-redo

fix: undo/redo not restoring connections properly

Author: LuD1161

frontend/src/components/workflow/Canvas.tsx

@@ -904,27 +904,41 @@ export function Canvas({
           dedupedEdges.set(edge.id, { ...edge, selected: false });
         });
 
+        // Calculate next state for snapshot before applying changes
+        let nextNodes = nodes;
+        let nextEdges = edges;
+
         if (selectedNodes.length > 0) {
-          setNodes((nds) => nds.filter((node) => !nodeIds.has(node.id)));
-          setEdges((eds) =>
-            eds.filter((edge) => !nodeIds.has(edge.source) && !nodeIds.has(edge.target)),
+          nextNodes = nodes.filter((node) => !nodeIds.has(node.id));
+          nextEdges = edges.filter(
+            (edge) => !nodeIds.has(edge.source) && !nodeIds.has(edge.target),
           );
-          setSelectedNode(null);
         }
 
         if (selectedEdges.length > 0) {
-          setEdges((eds) => eds.filter((edge) => !selectedEdgeIds.has(edge.id)));
+          nextEdges = nextEdges.filter((edge) => !selectedEdgeIds.has(edge.id));
+        }
+
+        // Apply the changes
+        if (selectedNodes.length > 0) {
+          setNodes(nextNodes);
+          setEdges(nextEdges);
+          setSelectedNode(null);
+        } else if (selectedEdges.length > 0) {
+          setEdges(nextEdges);
         }
 
+        // Capture snapshot for undo/redo
         if (selectedNodes.length > 0 || selectedEdges.length > 0) {
+          onSnapshot?.(nextNodes, nextEdges);
           markDirty();
         }
       }
     };
 
     document.addEventListener('keydown', handleKeyPress);
     return () => document.removeEventListener('keydown', handleKeyPress);
-  }, [nodes, edges, setNodes, setEdges, markDirty, mode]);
+  }, [nodes, edges, setNodes, setEdges, markDirty, mode, onSnapshot, toast]);
 
   // Panel width changes are handled by CSS transitions, no manual viewport translation needed
 

frontend/src/features/workflow-builder/WorkflowBuilder.tsx

@@ -291,16 +291,19 @@ function WorkflowBuilderContent() {
 
   const onEdgesChange = useCallback(
     (changes: any[]) => {
-      // Capture snapshot for edge changes
+      // Capture snapshot for edge changes (add/remove)
+      // Note: Edge removals due to node deletion are handled by onNodesChange,
+      // so we only need to capture explicit edge changes here
       if (mode === 'design' && changes.length > 0) {
         const hasStructuralChange = changes.some(
           (c: any) => c.type === 'add' || c.type === 'remove',
         );
         if (hasStructuralChange) {
+          const currentNodes = designNodesRef.current;
           const currentEdges = designEdgesRef.current;
           const nextEdges = applyEdgeChanges(changes, currentEdges);
-          // Pass undefined for nodes to allow merging with pending node changes (e.g. from node deletion)
-          captureSnapshot(undefined, nextEdges);
+          // Pass both nodes and edges to ensure consistent snapshot
+          captureSnapshot(currentNodes, nextEdges);
         }
       }
 

frontend/src/features/workflow-builder/hooks/useWorkflowHistory.ts

@@ -114,7 +114,11 @@ export const useWorkflowHistory = ({
   );
 
   /**
-   * Capture the current graph state as an undoable snapshot
+   * Capture the current graph state as an undoable snapshot.
+   *
+   * IMPORTANT: Always pass BOTH nodes and edges for a consistent snapshot.
+   * The debouncing ensures that rapid changes (like node deletion which triggers
+   * both onNodesChange and onEdgesChange) are captured as a single history entry.
    */
   const captureSnapshot = useCallback(
     (nodesOverride?: ReactFlowNode<FrontendNodeData>[], edgesOverride?: ReactFlowEdge[]) => {
@@ -123,16 +127,17 @@ export const useWorkflowHistory = ({
       }
 
       // Update pending state with new overrides if provided
-      if (nodesOverride) pendingStateRef.current.nodes = nodesOverride;
-      if (edgesOverride) pendingStateRef.current.edges = edgesOverride;
+      // Always prefer the explicitly passed state over refs for consistency
+      if (nodesOverride !== undefined) pendingStateRef.current.nodes = nodesOverride;
+      if (edgesOverride !== undefined) pendingStateRef.current.edges = edgesOverride;
 
       if (debounceRef.current) {
         clearTimeout(debounceRef.current);
       }
 
       debounceRef.current = setTimeout(() => {
-        // Resolve final state: Override (via pending) -> Ref
-        // Note: We use the accumulated pending state as the source of truth for "next" state
+        // Resolve final state from pending or fallback to refs
+        // The pending state should contain the NEXT state after the change
         const nodes = pendingStateRef.current.nodes ?? designGraph.nodesRef.current;
         const edges = pendingStateRef.current.edges ?? designGraph.edgesRef.current;