Add undo/redo functionality to flowchart graph

src/feature/flow_chart_panel/FlowChartKeyboardShortcuts.tsx

@@ -2,9 +2,11 @@ import useKeyboardShortcut from "@src/hooks/useKeyboardShortcut";
 import { useSave } from "@src/hooks/useSave";
 import { useLoadApp } from "@src/hooks/useLoadApp";
 import { useReactFlow } from "reactflow";
+import { useFlowChartGraph } from "@src/hooks/useFlowChartGraph";
 
 const FlowChartKeyboardShortcuts = () => {
   const { zoomIn, zoomOut } = useReactFlow();
+  const { undo, redo } = useFlowChartGraph();
   const save = useSave();
   const openFileSelector = useLoadApp();
 
@@ -20,6 +22,12 @@ const FlowChartKeyboardShortcuts = () => {
   useKeyboardShortcut("ctrl", "o", openFileSelector);
   useKeyboardShortcut("meta", "o", openFileSelector);
 
+  useKeyboardShortcut("ctrl", "z", undo);
+  useKeyboardShortcut("meta", "z", undo);
+
+  useKeyboardShortcut("ctrl", "y", redo);
+  useKeyboardShortcut("meta", "y", redo);
+
   return <div></div>;
 };
 

src/feature/flow_chart_panel/FlowChartTabView.tsx

@@ -33,7 +33,7 @@ import { getEdgeTypes, isCompatibleType } from "@src/utils/TypeCheck";
 import { CenterObserver } from "./components/CenterObserver";
 import useNodeTypes from "./hooks/useNodeTypes";
 import { Separator } from "@src/components/ui/separator";
-import { Pencil, Text, Workflow, X } from "lucide-react";
+import { Pencil, Text, Workflow, X, Undo, Redo } from "lucide-react";
 import { GalleryModal } from "@src/components/gallery/GalleryModal";
 import { toast } from "sonner";
 import { useTheme } from "@src/providers/themeProvider";
@@ -86,6 +86,11 @@ const FlowChartTab = () => {
     setEdges,
     selectedNode,
     unSelectedNodes,
+    canRedo,
+    canUndo,
+    recordState,
+    redo,
+    undo,
   } = useFlowChartGraph();
   const nodesMetadataMap = useNodesMetadata();
   const manifest = useManifest();
@@ -108,8 +113,10 @@ const FlowChartTab = () => {
     setNodes,
     getTakenNodeLabels,
     nodesMetadataMap,
+    recordState,
   );
-  const addTextNode = useAddTextNode();
+
+  const addTextNode = useAddTextNode(recordState);
 
   const duplicateNode = (node: Node<ElementsData>) => {
     const funcName = node.data.func;
@@ -197,6 +204,7 @@ const FlowChartTab = () => {
       nodes[nodeIndex] = node;
       setHasUnsavedChanges(true);
     });
+    recordState();
   };
 
   const onNodesChange: OnNodesChange = useCallback(
@@ -223,7 +231,9 @@ const FlowChartTab = () => {
       setEdges((eds) => {
         if (manifest) {
           const [sourceType, targetType] = getEdgeTypes(manifest, connection);
+
           if (isCompatibleType(sourceType, targetType)) {
+            recordState();
             return addEdge(connection, eds);
           }
 
@@ -245,6 +255,7 @@ const FlowChartTab = () => {
         prev.filter((node) => !selectedNodeIds.includes(node.id)),
       );
       setHasUnsavedChanges(true);
+      recordState();
     },
     [setNodes, setHasUnsavedChanges],
   );
@@ -254,7 +265,7 @@ const FlowChartTab = () => {
     setEdges([]);
     setHasUnsavedChanges(true);
     setProgramResults([]);
-
+    recordState();
     sendEventToMix("Canvas cleared", "");
   }, [setNodes, setEdges, setHasUnsavedChanges, setProgramResults]);
 
@@ -317,6 +328,29 @@ const FlowChartTab = () => {
               setIsGalleryOpen={setIsGalleryOpen}
             />
             <div className="grow" />
+            {canUndo && (
+              <Button
+                variant="ghost"
+                className="gap-2"
+                onClick={undo}
+                data-testid="undo-button"
+              >
+                <Undo size={18} className="stroke-muted-foreground" />
+                Undo
+              </Button>
+            )}
+
+            {canRedo && (
+              <Button
+                variant="ghost"
+                className="gap-2"
+                onClick={redo}
+                data-testid="redo-button"
+              >
+                <Redo size={18} className="stroke-muted-foreground" />
+                Redo
+              </Button>
+            )}
             {selectedNode && (
               <>
                 {!isEditMode ? (

src/feature/flow_chart_panel/hooks/useAddNewNode.tsx

@@ -21,6 +21,7 @@ export const useAddNewNode = (
   ) => void,
   getTakenNodeLabels: (func: string) => string[][],
   nodesMetadataMap: BlocksMetadataMap | null,
+  recordState: () => void,
 ) => {
   const center = useAtomValue(centerPositionAtom);
   const setHasUnsavedChanges = useSetAtom(unsavedChangesAtom);
@@ -89,6 +90,7 @@ export const useAddNewNode = (
         position: nodePosition,
       };
       setNodes((els) => els.concat(newNode));
+      recordState();
       setHasUnsavedChanges(true);
       sendEventToMix("Node Added", newNode.data.label);
     },

src/feature/flow_chart_panel/hooks/useAddTextNode.tsx

@@ -6,7 +6,7 @@ import { useCallback } from "react";
 import { v4 as uuidv4 } from "uuid";
 import { sendEventToMix } from "@src/services/MixpanelServices";
 
-export const useAddTextNode = () => {
+export const useAddTextNode = (recordState: () => void) => {
   const { setTextNodes } = useFlowChartGraph();
   const center = useAtomValue(centerPositionAtom);
 
@@ -23,5 +23,6 @@ export const useAddTextNode = () => {
       }),
     );
     sendEventToMix("Text Node Added", "");
+    recordState();
   }, [setTextNodes, center]);
 };

src/hooks/useFlowChartGraph.ts

@@ -1,7 +1,7 @@
 import { ElementsData } from "@/types";
-import { useAtom } from "jotai";
+import { atom, useAtom } from "jotai";
 import { atomWithImmer } from "jotai-immer";
-import { useCallback, useEffect, useMemo } from "react";
+import { useCallback, useEffect, useMemo, useRef } from "react";
 import { Edge, Node, ReactFlowJsonObject } from "reactflow";
 import * as RECIPES from "../data/RECIPES";
 import * as galleryItems from "../data/apps";
@@ -11,18 +11,99 @@ import { TextData } from "@src/types/node";
 import { sendEventToMix } from "@src/services/MixpanelServices";
 
 const project = resolveDefaultProjectReference();
-const projectData = resolveProjectReference(project) || RECIPES.NOISY_SINE;
+const projectData = resolveProjectReference(project!) || RECIPES.NOISY_SINE;
 const initialNodes: Node<ElementsData>[] = projectData.nodes;
 const initialEdges: Edge[] = projectData.edges;
 
 const nodesAtom = atomWithImmer<Node<ElementsData>[]>(initialNodes);
 export const textNodesAtom = atomWithImmer<Node<TextData>[]>([]);
 const edgesAtom = atomWithImmer<Edge[]>(initialEdges);
 
+type UndoRedoStackItem = {
+  edges: Edge[];
+  nodes: Node<ElementsData>[];
+  textNodes: Node<TextData>[];
+};
+
+const undoStackAtom = atom<UndoRedoStackItem[]>([]);
+const redoStackAtom = atom<UndoRedoStackItem[]>([]);
+
 export const useFlowChartGraph = () => {
+  const initialNodesRef = useRef(initialNodes);
+  const initialEdgesRef = useRef(initialEdges);
+
   const [nodes, setNodes] = useAtom(nodesAtom);
   const [textNodes, setTextNodes] = useAtom(textNodesAtom);
   const [edges, setEdges] = useAtom(edgesAtom);
+
+  const [undoStack, setUndoStack] = useAtom(undoStackAtom);
+  const [redoStack, setRedoStack] = useAtom(redoStackAtom);
+
+  const recordState = useCallback(() => {
+    setUndoStack((prev) => [
+      ...prev,
+      {
+        edges,
+        nodes,
+        textNodes,
+      },
+    ]);
+    setRedoStack([]);
+  }, [edges, nodes, textNodes]);
+
+  const canRedo = useMemo(() => redoStack.length > 0, [redoStack]);
+  const canUndo = useMemo(() => undoStack.length > 0, [undoStack]);
+
+  const undo = useCallback(() => {
+    if (!canUndo) return;
+
+    setUndoStack((prev) => {
+      const prevState = prev[prev.length - 2] || {
+        edges: initialEdgesRef.current,
+        nodes: initialNodesRef.current,
+        textNodes: [],
+      };
+      setNodes(prevState.nodes);
+      setEdges(prevState.edges);
+      setTextNodes(prevState.textNodes);
+      return prev.slice(0, -1);
+    });
+
+    setRedoStack((prev) => [
+      ...prev,
+      {
+        edges,
+        nodes,
+        textNodes,
+      },
+    ]);
+  }, [edges, nodes, textNodes]);
+
+  const redo = useCallback(() => {
+    if (!canRedo) return;
+
+    setRedoStack((prev) => {
+      const nextState = prev[prev.length - 1];
+
+      if (nextState) {
+        setNodes(nextState.nodes);
+        setEdges(nextState.edges);
+        setTextNodes(nextState.textNodes);
+      }
+
+      return prev.slice(0, -1);
+    });
+
+    setUndoStack((prev) => [
+      ...prev,
+      {
+        edges,
+        nodes,
+        textNodes,
+      },
+    ]);
+  }, [edges, nodes, textNodes]);
+
   const { selectedNodes, unSelectedNodes } = useMemo(() => {
     const selectedNodes: Node<ElementsData>[] = [];
     const unSelectedNodes: Node<ElementsData>[] = [];
@@ -33,6 +114,7 @@ export const useFlowChartGraph = () => {
     }
     return { selectedNodes, unSelectedNodes };
   }, [nodes]);
+
   const selectedNode = selectedNodes.length > 0 ? selectedNodes[0] : null;
 
   const loadFlowExportObject = useCallback(
@@ -143,6 +225,11 @@ export const useFlowChartGraph = () => {
     updateInitCtrlInputDataForNode,
     loadFlowExportObject,
     handleTitleChange,
+    canRedo,
+    canUndo,
+    recordState,
+    redo,
+    undo,
   };
 };
 
@@ -157,7 +244,7 @@ function resolveProjectReference(project: string) {
 }
 
 function resolveDefaultProjectReference() {
-  let project;
+  let project: string | undefined;
   if (typeof window !== "undefined") {
     const query = new URLSearchParams(window.location.search);
     project = query.get("project") || process.env.DEFAULT_PROJECT;