From b4f7285b335e77ce5519cff1b7911e00799df7d4 Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <lucas.edu.oli@hotmail.com>
Date: Sat, 6 Jan 2024 11:34:36 -0300
Subject: [PATCH] Implemented undoRedo on Zustand

---
 .../components/parameterComponent/index.tsx   |   4 +-
 .../src/CustomNodes/GenericNode/index.tsx     |   4 +-
 .../components/menuBar/index.tsx              |   4 +-
 .../src/components/inputComponent/index.tsx   |   3 -
 .../components/inputFileComponent/index.tsx   |   2 +-
 .../components/inputListComponent/index.tsx   |   6 -
 .../src/components/intComponent/index.tsx     |   2 +-
 .../components/keypairListComponent/index.tsx |   6 -
 .../src/components/promptComponent/index.tsx  |   2 +-
 src/frontend/src/contexts/index.tsx           |   4 +-
 src/frontend/src/contexts/undoRedoContext.tsx | 192 ------------------
 .../components/PageComponent/index.tsx        |  25 ++-
 .../components/nodeToolbarComponent/index.tsx |   8 +-
 src/frontend/src/stores/flowsManagerStore.ts  |  69 ++++++-
 .../src/types/zustand/flowsManager/index.ts   |   3 +
 15 files changed, 100 insertions(+), 234 deletions(-)
 delete mode 100644 src/frontend/src/contexts/undoRedoContext.tsx

diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 959865766..c0a462571 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -26,7 +26,6 @@ import {
   LANGFLOW_SUPPORTED_TYPES,
   TOOLTIP_EMPTY,
 } from "../../../../constants/constants";
-import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 import { postCustomComponentUpdate } from "../../../../controllers/API";
 import useAlertStore from "../../../../stores/alertStore";
 import useFlowStore from "../../../../stores/flowStore";
@@ -88,7 +87,7 @@ export default function ParameterComponent({
 
   const myData = useTypesStore((state) => state.data);
 
-  const { takeSnapshot } = useContext(undoRedoContext);
+  const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
 
   const handleUpdateValues = async (name: string, data: NodeDataType) => {
     const code = data.node?.template["code"]?.value;
@@ -528,7 +527,6 @@ export default function ParameterComponent({
               duplicateKey={errorDuplicateKey}
               onChange={(newValue) => {
                 const valueToNumbers = convertValuesToNumbers(newValue);
-                data.node!.template[name].value = valueToNumbers;
                 setErrorDuplicateKey(hasDuplicateKeys(valueToNumbers));
                 handleOnNewValue(valueToNumbers);
               }}
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 3a90c9ce4..339eb73ac 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -7,7 +7,6 @@ import InputComponent from "../../components/inputComponent";
 import { Textarea } from "../../components/ui/textarea";
 import { priorityFields } from "../../constants/constants";
 import { useSSE } from "../../contexts/SSEContext";
-import { undoRedoContext } from "../../contexts/undoRedoContext";
 import NodeToolbarComponent from "../../pages/FlowPage/components/nodeToolbarComponent";
 import useFlowStore from "../../stores/flowStore";
 import { validationStatusType } from "../../types/components";
@@ -17,6 +16,7 @@ import { nodeColors, nodeIconsLucide } from "../../utils/styleUtils";
 import { classNames, cn, getFieldTitle } from "../../utils/utils";
 import ParameterComponent from "./components/parameterComponent";
 import { useTypesStore } from "../../stores/typesStore";
+import useFlowsManagerStore from "../../stores/flowsManagerStore";
 
 export default function GenericNode({
   data,
@@ -44,7 +44,7 @@ export default function GenericNode({
   const [handles, setHandles] = useState<boolean[] | []>([]);
   let numberOfInputs: boolean[] = [];
 
-  const { takeSnapshot } = useContext(undoRedoContext);
+  const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
 
   function countHandles(): void {
     numberOfInputs = Object.keys(data.node!.template)
diff --git a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
index e2fd345c5..20ac20f04 100644
--- a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
+++ b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx
@@ -8,7 +8,6 @@ import {
 } from "../../../ui/dropdown-menu";
 
 import { useNavigate } from "react-router-dom";
-import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 import FlowSettingsModal from "../../../../modals/flowSettingsModal";
 import useAlertStore from "../../../../stores/alertStore";
 import useFlowsManagerStore from "../../../../stores/flowsManagerStore";
@@ -19,7 +18,8 @@ export const MenuBar = (): JSX.Element => {
   const addFlow = useFlowsManagerStore((state) => state.addFlow);
   const currentFlow = useFlowsManagerStore((state) => state.currentFlow);
   const setErrorData = useAlertStore((state) => state.setErrorData);
-  const { undo, redo } = useContext(undoRedoContext);
+  const undo = useFlowsManagerStore((state) => state.undo);
+  const redo = useFlowsManagerStore((state) => state.redo);
   const [openSettings, setOpenSettings] = useState(false);
 
   const navigate = useNavigate();
diff --git a/src/frontend/src/components/inputComponent/index.tsx b/src/frontend/src/components/inputComponent/index.tsx
index 0efd29c0d..8017c5458 100644
--- a/src/frontend/src/components/inputComponent/index.tsx
+++ b/src/frontend/src/components/inputComponent/index.tsx
@@ -59,9 +59,6 @@ export default function InputComponent({
               e.preventDefault();
             }}
             onKeyDown={(e) => {
-              if (e.ctrlKey && e.key === "c") {
-                // Perform any actions you need when Ctrl+C is detected
-              }
               handleKeyDown(e, value, "");
               if (blurOnEnter && e.key === "Enter") refInput.current?.blur();
             }}
diff --git a/src/frontend/src/components/inputFileComponent/index.tsx b/src/frontend/src/components/inputFileComponent/index.tsx
index a365b188c..0d7b180d2 100644
--- a/src/frontend/src/components/inputFileComponent/index.tsx
+++ b/src/frontend/src/components/inputFileComponent/index.tsx
@@ -20,7 +20,7 @@ export default function InputFileComponent({
 
   // Clear component state
   useEffect(() => {
-    if (disabled) {
+    if (disabled && value !== "") {
       setMyValue("");
       onChange("");
       onFileChange("");
diff --git a/src/frontend/src/components/inputListComponent/index.tsx b/src/frontend/src/components/inputListComponent/index.tsx
index ae154484d..e98bba304 100644
--- a/src/frontend/src/components/inputListComponent/index.tsx
+++ b/src/frontend/src/components/inputListComponent/index.tsx
@@ -44,12 +44,6 @@ export default function InputListComponent({
                 newInputList[idx] = event.target.value;
                 onChange(newInputList);
               }}
-              onKeyDown={(e) => {
-                if (e.ctrlKey && e.key === "Backspace") {
-                  e.preventDefault();
-                  e.stopPropagation();
-                }
-              }}
             />
             {idx === value.length - 1 ? (
               <button
diff --git a/src/frontend/src/components/intComponent/index.tsx b/src/frontend/src/components/intComponent/index.tsx
index 633d7bd8d..fd6d91ffd 100644
--- a/src/frontend/src/components/intComponent/index.tsx
+++ b/src/frontend/src/components/intComponent/index.tsx
@@ -14,7 +14,7 @@ export default function IntComponent({
 
   // Clear component state
   useEffect(() => {
-    if (disabled) {
+    if (disabled && value !== "") {
       onChange("");
     }
   }, [disabled, onChange]);
diff --git a/src/frontend/src/components/keypairListComponent/index.tsx b/src/frontend/src/components/keypairListComponent/index.tsx
index 82e46d2c7..034628643 100644
--- a/src/frontend/src/components/keypairListComponent/index.tsx
+++ b/src/frontend/src/components/keypairListComponent/index.tsx
@@ -65,12 +65,6 @@ export default function KeypairListComponent({
                 )}
                 placeholder="Type key..."
                 onChange={(event) => handleChangeKey(event, index)}
-                onKeyDown={(e) => {
-                  if (e.ctrlKey && e.key === "Backspace") {
-                    e.preventDefault();
-                    e.stopPropagation();
-                  }
-                }}
               />
 
               <Input
diff --git a/src/frontend/src/components/promptComponent/index.tsx b/src/frontend/src/components/promptComponent/index.tsx
index 5a2ced971..784e84fdb 100644
--- a/src/frontend/src/components/promptComponent/index.tsx
+++ b/src/frontend/src/components/promptComponent/index.tsx
@@ -17,7 +17,7 @@ export default function PromptAreaComponent({
   readonly = false,
 }: PromptAreaComponentType): JSX.Element {
   useEffect(() => {
-    if (disabled) {
+    if (disabled && value !== "") {
       onChange("");
     }
   }, [disabled]);
diff --git a/src/frontend/src/contexts/index.tsx b/src/frontend/src/contexts/index.tsx
index 49b8106e6..7090b7c40 100644
--- a/src/frontend/src/contexts/index.tsx
+++ b/src/frontend/src/contexts/index.tsx
@@ -7,8 +7,6 @@ import { SSEProvider } from "./SSEContext";
 import { AuthProvider } from "./authContext";
 import { LocationProvider } from "./locationContext";
 
-import { UndoRedoProvider } from "./undoRedoContext";
-
 export default function ContextWrapper({ children }: { children: ReactNode }) {
   //element to wrap all context
   return (
@@ -20,7 +18,7 @@ export default function ContextWrapper({ children }: { children: ReactNode }) {
               <LocationProvider>
                 <ApiInterceptor />
                 <SSEProvider>
-                  <UndoRedoProvider>{children}</UndoRedoProvider>
+                  {children}
                 </SSEProvider>
               </LocationProvider>
             </ReactFlowProvider>
diff --git a/src/frontend/src/contexts/undoRedoContext.tsx b/src/frontend/src/contexts/undoRedoContext.tsx
deleted file mode 100644
index 095878b8c..000000000
--- a/src/frontend/src/contexts/undoRedoContext.tsx
+++ /dev/null
@@ -1,192 +0,0 @@
-import { cloneDeep } from "lodash";
-import {
-  createContext,
-  useCallback,
-  useContext,
-  useEffect,
-  useState,
-} from "react";
-import useFlowStore from "../stores/flowStore";
-import {
-  HistoryItem,
-  UseUndoRedoOptions,
-  undoRedoContextType,
-} from "../types/typesContext";
-import { isWrappedWithClass } from "../utils/utils";
-import useFlowsManagerStore from "../stores/flowsManagerStore";
-
-const initialValue = {
-  undo: () => {},
-  redo: () => {},
-  takeSnapshot: () => {},
-};
-
-const defaultOptions: UseUndoRedoOptions = {
-  maxHistorySize: 100,
-  enableShortcuts: true,
-};
-
-export const undoRedoContext = createContext<undoRedoContextType>(initialValue);
-
-export function UndoRedoProvider({ children }) {
-  const flows = useFlowsManagerStore((state) => state.flows);
-  const currentFlowId = useFlowsManagerStore((state) => state.currentFlowId);
-
-  const setNodes = useFlowStore((state) => state.setNodes);
-  const setEdges = useFlowStore((state) => state.setEdges);
-  const nodes = useFlowStore((state) => state.nodes);
-  const edges = useFlowStore((state) => state.edges);
-
-  const [past, setPast] = useState<HistoryItem[][]>(flows.map(() => []));
-  const [future, setFuture] = useState<HistoryItem[][]>(flows.map(() => []));
-  const [tabIndex, setTabIndex] = useState(
-    flows.findIndex((flow) => flow.id === currentFlowId)
-  );
-
-  useEffect(() => {
-    // whenever the flows variable changes, we need to add one array to the past and future states
-    setPast((old) =>
-      flows.map((flow, index) => (old[index] ? old[index] : []))
-    );
-    setFuture((old) =>
-      flows.map((flow, index) => (old[index] ? old[index] : []))
-    );
-    setTabIndex(flows.findIndex((flow) => flow.id === currentFlowId));
-  }, [flows, currentFlowId]);
-
-  const takeSnapshot = useCallback(() => {
-    // push the current graph to the past state
-    let newPast = cloneDeep(past);
-    let newState = {
-      nodes: cloneDeep(nodes),
-      edges: cloneDeep(edges),
-    };
-    if (
-      past[tabIndex] &&
-      JSON.stringify(past[tabIndex][past[tabIndex].length - 1]) !==
-        JSON.stringify(newState)
-    ) {
-      newPast[tabIndex] = past[tabIndex].slice(
-        past[tabIndex].length - defaultOptions.maxHistorySize + 1,
-        past[tabIndex].length
-      );
-      newPast[tabIndex].push(newState);
-    }
-    setPast(newPast);
-
-    // whenever we take a new snapshot, the redo operations need to be cleared to avoid state mismatches
-    setFuture((old) => {
-      let newFuture = cloneDeep(old);
-      newFuture[tabIndex] = [];
-      return newFuture;
-    });
-  }, [nodes, edges, past, future, flows, currentFlowId, setPast, setFuture]);
-
-  const undo = useCallback(() => {
-    // get the last state that we want to go back to
-    const pastState = past[tabIndex][past[tabIndex].length - 1];
-
-    if (pastState) {
-      // first we remove the state from the history
-      setPast((old) => {
-        let newPast = cloneDeep(old);
-        newPast[tabIndex] = old[tabIndex].slice(0, old[tabIndex].length - 1);
-        return newPast;
-      });
-      // we store the current graph for the redo operation
-      setFuture((old) => {
-        let newFuture = cloneDeep(old);
-        newFuture[tabIndex] = old[tabIndex];
-        newFuture[tabIndex].push({ nodes: nodes, edges: edges });
-        return newFuture;
-      });
-      // now we can set the graph to the past state
-      setNodes(pastState.nodes);
-      setEdges(pastState.edges);
-    }
-  }, [
-    setNodes,
-    setEdges,
-    nodes,
-    edges,
-    future,
-    past,
-    setFuture,
-    setPast,
-    tabIndex,
-  ]);
-
-  const redo = useCallback(() => {
-    const futureState = future[tabIndex][future[tabIndex].length - 1];
-
-    if (futureState) {
-      setFuture((old) => {
-        let newFuture = cloneDeep(old);
-        newFuture[tabIndex] = old[tabIndex].slice(0, old[tabIndex].length - 1);
-        return newFuture;
-      });
-      setPast((old) => {
-        let newPast = cloneDeep(old);
-        newPast[tabIndex] = old[tabIndex];
-        newPast[tabIndex].push({ nodes: nodes, edges: edges });
-        return newPast;
-      });
-      setNodes(futureState.nodes);
-      setEdges(futureState.edges);
-    }
-  }, [
-    future,
-    past,
-    setFuture,
-    setPast,
-    setNodes,
-    setEdges,
-    nodes,
-    edges,
-    future,
-    tabIndex,
-  ]);
-
-  useEffect(() => {
-    // this effect is used to attach the global event handlers
-    if (!defaultOptions.enableShortcuts) {
-      return;
-    }
-
-    const keyDownHandler = (event: KeyboardEvent) => {
-      if (!isWrappedWithClass(event, "noundo")) {
-        if (
-          event.key === "z" &&
-          (event.ctrlKey || event.metaKey) &&
-          event.shiftKey
-        ) {
-          event.preventDefault();
-          redo();
-        } else if (event.key === "y" && (event.ctrlKey || event.metaKey)) {
-          event.preventDefault(); // prevent the default action
-          redo();
-        } else if (event.key === "z" && (event.ctrlKey || event.metaKey)) {
-          event.preventDefault();
-          undo();
-        }
-      }
-    };
-
-    document.addEventListener("keydown", keyDownHandler);
-
-    return () => {
-      document.removeEventListener("keydown", keyDownHandler);
-    };
-  }, [undo, redo]);
-  return (
-    <undoRedoContext.Provider
-      value={{
-        undo,
-        redo,
-        takeSnapshot,
-      }}
-    >
-      {children}
-    </undoRedoContext.Provider>
-  );
-}
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index f3dd8f6bd..d53b49846 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -22,7 +22,6 @@ import GenericNode from "../../../../CustomNodes/GenericNode";
 import Chat from "../../../../components/chatComponent";
 import Loading from "../../../../components/ui/loading";
 import { locationContext } from "../../../../contexts/locationContext";
-import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 import useAlertStore from "../../../../stores/alertStore";
 import useFlowStore from "../../../../stores/flowStore";
 import useFlowsManagerStore from "../../../../stores/flowsManagerStore";
@@ -66,8 +65,6 @@ export default function Page({
     edges: any;
   } | null>(null);
 
-  const { takeSnapshot } = useContext(undoRedoContext);
-
   const reactFlowInstance = useFlowStore((state) => state.reactFlowInstance);
   const setReactFlowInstance = useFlowStore(
     (state) => state.setReactFlowInstance
@@ -83,6 +80,9 @@ export default function Page({
   const deleteNode = useFlowStore((state) => state.deleteNode);
   const deleteEdge = useFlowStore((state) => state.deleteEdge);
   const setPending = useFlowStore((state) => state.setPending);
+  const undo = useFlowsManagerStore((state) => state.undo);
+  const redo = useFlowsManagerStore((state) => state.redo);
+  const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
   const isPending = useFlowStore((state) => state.isPending);
   const paste = useFlowStore((state) => state.paste);
 
@@ -92,6 +92,19 @@ export default function Page({
 
   useEffect(() => {
     const onKeyDown = (event: KeyboardEvent) => {
+      if (!isWrappedWithClass(event, "noundo")) {
+        if (
+          (event.key === "y" || (event.key === "z" && event.shiftKey)) &&
+          (event.ctrlKey || event.metaKey)
+        ) {
+          event.preventDefault(); // prevent the default action
+          redo();
+        } else if (event.key === "z" && (event.ctrlKey || event.metaKey)) {
+          console.log("nao era pra tar aqui");
+          event.preventDefault();
+          undo();
+        }
+      }
       if (
         !isWrappedWithClass(event, "nocopy") &&
         window.getSelection()?.toString().length === 0
@@ -103,8 +116,7 @@ export default function Page({
         ) {
           event.preventDefault();
           setLastCopiedSelection(_.cloneDeep(lastSelection));
-        }
-        if (
+        } else if (
           (event.ctrlKey || event.metaKey) &&
           event.key === "v" &&
           lastCopiedSelection
@@ -115,8 +127,7 @@ export default function Page({
             x: position.current.x,
             y: position.current.y,
           });
-        }
-        if (
+        } else if (
           (event.ctrlKey || event.metaKey) &&
           event.key === "g" &&
           lastSelection
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 84b00208a..15b602808 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -8,7 +8,6 @@ import {
   SelectItem,
   SelectTrigger,
 } from "../../../../components/ui/select-custom";
-import { undoRedoContext } from "../../../../contexts/undoRedoContext";
 import ConfirmationModal from "../../../../modals/ConfirmationModal";
 import EditNodeModal from "../../../../modals/EditNodeModal";
 import ShareModal from "../../../../modals/shareModal";
@@ -34,7 +33,7 @@ export default function NodeToolbarComponent({
   numberOfHandles,
   showNode,
 }: nodeToolbarPropsType): JSX.Element {
-  const [nodeLength, setNodeLength] = useState(
+  const nodeLength = 
     Object.keys(data.node!.template).filter(
       (templateField) =>
         templateField.charAt(0) !== "_" &&
@@ -49,8 +48,7 @@ export default function NodeToolbarComponent({
           data.node.template[templateField].type === "int" ||
           data.node.template[templateField].type === "dict" ||
           data.node.template[templateField].type === "NestedDict")
-    ).length
-  );
+    ).length;
 
   const hasStore = useStoreStore((state) => state.hasStore);
   const hasApiKey = useStoreStore((state) => state.hasApiKey);
@@ -76,7 +74,7 @@ export default function NodeToolbarComponent({
   const saveComponent = useFlowsManagerStore((state) => state.saveComponent);
   const flows = useFlowsManagerStore((state) => state.flows);
   const version = useDarkStore((state) => state.version);
-  const { takeSnapshot } = useContext(undoRedoContext);
+  const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
   const [showModalAdvanced, setShowModalAdvanced] = useState(false);
   const [showconfirmShare, setShowconfirmShare] = useState(false);
   const [selectedValue, setSelectedValue] = useState("");
diff --git a/src/frontend/src/stores/flowsManagerStore.ts b/src/frontend/src/stores/flowsManagerStore.ts
index 6b1ffa974..210987909 100644
--- a/src/frontend/src/stores/flowsManagerStore.ts
+++ b/src/frontend/src/stores/flowsManagerStore.ts
@@ -10,6 +10,7 @@ import {
 } from "../controllers/API";
 import { FlowType, NodeDataType } from "../types/flow";
 import { FlowState } from "../types/tabs";
+import { UseUndoRedoOptions } from "../types/typesContext";
 import { FlowsManagerStoreType } from "../types/zustand/flowsManager";
 import {
   addVersionToDuplicates,
@@ -19,12 +20,20 @@ import {
   processFlows,
 } from "../utils/reactflowUtils";
 import useAlertStore from "./alertStore";
+import { useDarkStore } from "./darkStore";
 import useFlowStore from "./flowStore";
 import { useTypesStore } from "./typesStore";
-import { useDarkStore } from "./darkStore";
 
 let saveTimeoutId: NodeJS.Timeout | null = null;
 
+const defaultOptions: UseUndoRedoOptions = {
+  maxHistorySize: 100,
+  enableShortcuts: true,
+};
+
+const past = {};
+const future = {};
+
 const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
   currentFlowId: "",
   setCurrentFlowId: (currentFlowId: string) => {
@@ -328,7 +337,63 @@ const useFlowsManagerStore = create<FlowsManagerStoreType>((set, get) => ({
   },
   saveComponent: (component: NodeDataType, override: boolean) => {
     component.node!.official = false;
-    return get().addFlow(true, createFlowComponent(component, useDarkStore.getState().version), override);
+    return get().addFlow(
+      true,
+      createFlowComponent(component, useDarkStore.getState().version),
+      override
+    );
+  },
+  takeSnapshot: () => {
+    const currentFlowId = get().currentFlowId;
+    // push the current graph to the past state
+    const newState = useFlowStore.getState();
+    const pastLength = past[currentFlowId]?.length ?? 0;
+    if (
+      pastLength > 0 &&
+      JSON.stringify(past[currentFlowId][pastLength - 1]) !==
+        JSON.stringify(newState)
+    ) {
+      past[currentFlowId] = past[currentFlowId]
+        .slice(pastLength - defaultOptions.maxHistorySize + 1, pastLength)
+        
+      past[currentFlowId].push(newState);
+    } else {
+      past[currentFlowId] = [newState];
+    }
+
+    future[currentFlowId] = [];
+  },
+  undo: () => {
+    const newState = useFlowStore.getState();
+    const currentFlowId = get().currentFlowId;
+    const pastLength = past[currentFlowId]?.length ?? 0;
+    const pastState = past[currentFlowId][pastLength - 1] ?? null;
+
+    if (pastState) {
+      past[currentFlowId] = past[currentFlowId].slice(0, pastLength - 1);
+
+      if(!future[currentFlowId]) future[currentFlowId] = [];
+      future[currentFlowId].push({ nodes: newState.nodes, edges: newState.edges });
+
+      newState.setNodes(pastState.nodes);
+      newState.setEdges(pastState.edges);
+    }
+  },
+  redo: () => {
+    const newState = useFlowStore.getState();
+    const currentFlowId = get().currentFlowId;
+    const futureLength = future[currentFlowId]?.length ?? 0;
+    const futureState = future[currentFlowId][futureLength - 1] ?? null;
+
+    if (futureState) {
+      future[currentFlowId] = future[currentFlowId].slice(0, futureLength - 1);
+
+      if(!past[currentFlowId]) past[currentFlowId] = [];
+      past[currentFlowId].push({ nodes: newState.nodes, edges: newState.edges });
+
+      newState.setNodes(futureState.nodes);
+      newState.setEdges(futureState.edges);
+    }
   },
 }));
 
diff --git a/src/frontend/src/types/zustand/flowsManager/index.ts b/src/frontend/src/types/zustand/flowsManager/index.ts
index 5a3c01350..d130c00e6 100644
--- a/src/frontend/src/types/zustand/flowsManager/index.ts
+++ b/src/frontend/src/types/zustand/flowsManager/index.ts
@@ -22,4 +22,7 @@ export type FlowsManagerStoreType = {
   deleteComponent: (key: string) => Promise<void>;
   removeFlow: (id: string) => Promise<void>;
   saveComponent: (component: any, override: boolean) => Promise<string | undefined>;
+  undo: () => void;
+  redo: () => void;
+  takeSnapshot: () => void;
 };