diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 4c81e29a5..ed8930c16 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -71,7 +71,7 @@ export default function ParameterComponent({
   const { setErrorData, modalContextOpen } = useContext(alertContext);
   const updateNodeInternals = useUpdateNodeInternals();
   const [position, setPosition] = useState(0);
-  const { setTabsState, tabId, flows, tabsState, updateFlow } =
+  const { setTabsState, tabId, flows, tabsState, updateFlow, nodes, edges, setEdges } =
     useContext(FlowsContext);
 
   const [dataRef, setDataRef] = useState(data);
@@ -96,10 +96,9 @@ export default function ParameterComponent({
     setDataRef(data);
   }, [modalContextOpen]);
 
-  const { reactFlowInstance, setFilterEdge } = useContext(typesContext);
+  const { setFilterEdge } = useContext(typesContext);
   let disabled =
-    reactFlowInstance
-      ?.getEdges()
+    edges
       .some(
         (edge) =>
           edge.targetHandle ===
@@ -181,14 +180,14 @@ export default function ParameterComponent({
     let flow = flows.find((flow) => flow.id === tabId);
     setTimeout(() => {
       //timeout necessary because ReactFlow updates are not async
-      if (reactFlowInstance && flow && flow.data) {
+      if (flow && flow.data) {
         cleanEdges({
           flow: {
             edges: flow.data!.edges,
             nodes: flow.data!.nodes,
           },
           updateEdge: (edge) => {
-            reactFlowInstance.setEdges(edge);
+            setEdges(edge);
             updateNodeInternals(data.id);
           },
         });
@@ -323,7 +322,7 @@ export default function ParameterComponent({
                   : scapedJSONStringfy(id)
               }
               isValidConnection={(connection) =>
-                isValidConnection(connection, reactFlowInstance!)
+                isValidConnection(connection, nodes, edges)
               }
               className={classNames(
                 left ? "my-12 -ml-0.5 " : " my-12 -mr-0.5 ",
@@ -396,7 +395,7 @@ export default function ParameterComponent({
                       : scapedJSONStringfy(id)
                   }
                   isValidConnection={(connection) =>
-                    isValidConnection(connection, reactFlowInstance!)
+                    isValidConnection(connection, nodes, edges)
                   }
                   className={classNames(
                     left ? "-ml-0.5 " : "-mr-0.5 ",
diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index ca92a74cf..3e7554023 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -33,7 +33,7 @@ export default function GenericNode({
   const { updateFlow, flows, tabId, saveCurrentFlow } =
     useContext(FlowsContext);
   const updateNodeInternals = useUpdateNodeInternals();
-  const { types, deleteNode, reactFlowInstance, setFilterEdge, getFilterEdge } =
+  const { types, deleteNode } =
     useContext(typesContext);
   const name = nodeIconsLucide[data.type] ? data.type : types[data.type];
   const [inputName, setInputName] = useState(false);
diff --git a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
index f1e374a03..4a64e04db 100644
--- a/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
+++ b/src/frontend/src/components/chatComponent/buildTrigger/index.tsx
@@ -25,8 +25,7 @@ export default function BuildTrigger({
   isBuilt: boolean;
 }): JSX.Element {
   const { updateSSEData, isBuilding, setIsBuilding, sseData } = useSSE();
-  const { reactFlowInstance } = useContext(typesContext);
-  const { setTabsState, saveFlow } = useContext(FlowsContext);
+  const { setTabsState, saveFlow, nodes, edges } = useContext(FlowsContext);
   const { setErrorData, setSuccessData } = useContext(alertContext);
   const [isIconTouched, setIsIconTouched] = useState(false);
   const eventClick = isBuilding ? "pointer-events-none" : "";
@@ -38,8 +37,8 @@ export default function BuildTrigger({
         return;
       }
       const errors = validateNodes(
-        reactFlowInstance!.getNodes(),
-        reactFlowInstance!.getEdges()
+        nodes,
+        edges
       );
       if (errors.length > 0) {
         setErrorData({
diff --git a/src/frontend/src/components/codeTabsComponent/index.tsx b/src/frontend/src/components/codeTabsComponent/index.tsx
index bf9f4ff74..09b90860e 100644
--- a/src/frontend/src/components/codeTabsComponent/index.tsx
+++ b/src/frontend/src/components/codeTabsComponent/index.tsx
@@ -41,6 +41,7 @@ import { classNames } from "../../utils/utils";
 import DictComponent from "../dictComponent";
 import IconComponent from "../genericIconComponent";
 import KeypairListComponent from "../keypairListComponent";
+import { FlowsContext } from "../../contexts/flowsContext";
 
 export default function CodeTabsComponent({
   flow,
@@ -54,7 +55,7 @@ export default function CodeTabsComponent({
   const [data, setData] = useState(flow ? flow["data"]!["nodes"] : null);
   const [openAccordion, setOpenAccordion] = useState<string[]>([]);
   const { dark } = useContext(darkContext);
-  const { reactFlowInstance } = useContext(typesContext);
+  const {setNodes} = useContext(FlowsContext);
   const [errorDuplicateKey, setErrorDuplicateKey] = useState(false);
 
   useEffect(() => {
@@ -68,7 +69,7 @@ export default function CodeTabsComponent({
       unselectAllNodes({
         data,
         updateNodes: (nodes) => {
-          reactFlowInstance?.setNodes(nodes);
+          setNodes(nodes);
         },
       });
     }
diff --git a/src/frontend/src/contexts/flowsContext.tsx b/src/frontend/src/contexts/flowsContext.tsx
index 3978aa8bc..974ce6595 100644
--- a/src/frontend/src/contexts/flowsContext.tsx
+++ b/src/frontend/src/contexts/flowsContext.tsx
@@ -3,6 +3,7 @@ import _, { cloneDeep } from "lodash";
 import {
   ReactNode,
   createContext,
+  useCallback,
   useContext,
   useEffect,
   useRef,
@@ -10,10 +11,14 @@ import {
 } from "react";
 import {
   Edge,
+  EdgeChange,
   Node,
+  NodeChange,
   ReactFlowJsonObject,
   XYPosition,
   addEdge,
+  useEdgesState,
+  useNodesState,
 } from "reactflow";
 import ShortUniqueId from "short-unique-id";
 import {
@@ -93,8 +98,12 @@ const FlowsContextInitialValue: FlowsContextType = {
   saveComponent: async (component: NodeDataType, override: boolean) => "",
   deleteComponent: (key: string) => {},
   version: "",
-  nodesOnFlow: "",
-  setNodesOnFlow: (nodes: string) => "",
+  nodes: [],
+  setNodes: () => {},
+  onNodesChange: () => {},
+  edges: [],
+  setEdges: () => {},
+  onEdgesChange: () => {},
 };
 
 export const FlowsContext = createContext<FlowsContextType>(
@@ -102,18 +111,15 @@ export const FlowsContext = createContext<FlowsContextType>(
 );
 
 export function FlowsProvider({ children }: { children: ReactNode }) {
-  const { setErrorData, setNoticeData, setSuccessData } =
+  const { setErrorData, setSuccessData } =
     useContext(alertContext);
   const { getAuthentication, isAuthenticated } = useContext(AuthContext);
 
   const [tabId, setTabId] = useState("");
-
   const [isLoading, setIsLoading] = useState(false);
-  const [nodesOnFlow, setNodesOnFlow] = useState("");
-
   const [flows, setFlows] = useState<Array<FlowType>>([]);
   const [id, setId] = useState(uid());
-  const { reactFlowInstance, setData, data } = useContext(typesContext);
+  const { reactFlowInstance, setData } = useContext(typesContext);
   const [lastCopiedSelection, setLastCopiedSelection] = useState<{
     nodes: any;
     edges: any;
@@ -121,6 +127,50 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
   const [tabsState, setTabsState] = useState<FlowsState>({});
   const [getTweak, setTweak] = useState<tweakType>([]);
 
+  const [nodes, setNodesInternal, onNodesChangeInternal] = useNodesState([]);
+
+  const [edges, setEdgesInternal, onEdgesChangeInternal] = useEdgesState([]);
+
+  const onNodesChange = useCallback((nodes: NodeChange[]) => {
+    onNodesChangeInternal(nodes);
+    //@ts-ignore
+    setTabsState((prev: FlowsState) => {
+      return {
+        ...prev,
+        [tabId]: {
+          ...prev[tabId],
+          isPending: true,
+        },
+      };
+    });
+    //SAVE FLOW IN FLOWS
+  }, [onNodesChangeInternal, setTabsState, tabId]);
+
+  const onEdgesChange = useCallback(
+    (edges: EdgeChange[]) => {
+      onEdgesChangeInternal(edges);
+      //@ts-ignore
+      setTabsState((prev: FlowsState) => {
+        return {
+          ...prev,
+          [tabId]: {
+            ...prev[tabId],
+            isPending: true,
+          },
+        };
+      });
+    },
+    [onEdgesChangeInternal, setTabsState, tabId]
+  );
+
+  const setNodes = (nodes: Node[] | ((oldState: Node[]) => Node[])) => {
+    setNodesInternal(nodes)
+  }
+
+  const setEdges = (edges: Edge[] | ((oldState: Edge[]) => Edge[])) => {
+    setEdgesInternal(edges);
+  };
+
   useEffect(() => {
     if (!isAuthenticated) {
       hardReset();
@@ -128,6 +178,7 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
   }, [isAuthenticated]);
 
   const newNodeId = useRef(uid());
+
   function incrementNodeId() {
     newNodeId.current = uid();
     return newNodeId.current;
@@ -425,8 +476,8 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
     let minimumX = Infinity;
     let minimumY = Infinity;
     let idsMap = {};
-    let nodes: Node<NodeDataType>[] = reactFlowInstance!.getNodes();
-    let edges = reactFlowInstance!.getEdges();
+    let newNodes: Node<NodeDataType>[] = nodes;
+    let newEdges = edges;
     selectionInstance.nodes.forEach((node: Node) => {
       if (node.position.y < minimumY) {
         minimumY = node.position.y;
@@ -463,11 +514,11 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
       };
 
       // Add the new node to the list of nodes in state
-      nodes = nodes
+      newNodes = nodes
         .map((node) => ({ ...node, selected: false }))
         .concat({ ...newNode, selected: false });
     });
-    reactFlowInstance!.setNodes(nodes);
+    setNodes(newNodes);
 
     selectionInstance.edges.forEach((edge: Edge) => {
       let source = idsMap[edge.source];
@@ -498,7 +549,7 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
         "-" +
         target +
         targetHandle;
-      edges = addEdge(
+      newEdges = addEdge(
         {
           source,
           target,
@@ -514,10 +565,10 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
           animated: targetHandleObject.type === "Text",
           selected: false,
         },
-        edges.map((edge) => ({ ...edge, selected: false }))
+        newEdges.map((edge) => ({ ...edge, selected: false }))
       );
     });
-    reactFlowInstance!.setEdges(edges);
+    setEdges(newEdges);
   }
 
   const addFlow = async (
@@ -744,8 +795,12 @@ export function FlowsProvider({ children }: { children: ReactNode }) {
         isLoading,
         saveComponent,
         deleteComponent,
-        nodesOnFlow,
-        setNodesOnFlow,
+        nodes,
+        setNodes,
+        onNodesChange,
+        edges,
+        setEdges,
+        onEdgesChange,
       }}
     >
       {children}
diff --git a/src/frontend/src/contexts/undoRedoContext.tsx b/src/frontend/src/contexts/undoRedoContext.tsx
index f02be6257..ce925206f 100644
--- a/src/frontend/src/contexts/undoRedoContext.tsx
+++ b/src/frontend/src/contexts/undoRedoContext.tsx
@@ -29,7 +29,7 @@ const defaultOptions: UseUndoRedoOptions = {
 export const undoRedoContext = createContext<undoRedoContextType>(initialValue);
 
 export function UndoRedoProvider({ children }) {
-  const { tabId, flows } = useContext(FlowsContext);
+  const { tabId, flows, setNodes, setEdges, nodes, edges } = useContext(FlowsContext);
 
   const [past, setPast] = useState<HistoryItem[][]>(flows.map(() => []));
   const [future, setFuture] = useState<HistoryItem[][]>(flows.map(() => []));
@@ -48,14 +48,12 @@ export function UndoRedoProvider({ children }) {
     setTabIndex(flows.findIndex((flow) => flow.id === tabId));
   }, [flows, tabId]);
 
-  const { setNodes, setEdges, getNodes, getEdges } = useReactFlow();
-
   const takeSnapshot = useCallback(() => {
     // push the current graph to the past state
     let newPast = cloneDeep(past);
     let newState = {
-      nodes: cloneDeep(getNodes()),
-      edges: cloneDeep(getEdges()),
+      nodes: cloneDeep(nodes),
+      edges: cloneDeep(edges),
     };
     if (
       past[tabIndex] &&
@@ -76,7 +74,7 @@ export function UndoRedoProvider({ children }) {
       newFuture[tabIndex] = [];
       return newFuture;
     });
-  }, [getNodes, getEdges, past, future, flows, tabId, setPast, setFuture]);
+  }, [nodes, edges, past, future, flows, tabId, setPast, setFuture]);
 
   const undo = useCallback(() => {
     // get the last state that we want to go back to
@@ -93,7 +91,7 @@ export function UndoRedoProvider({ children }) {
       setFuture((old) => {
         let newFuture = cloneDeep(old);
         newFuture[tabIndex] = old[tabIndex];
-        newFuture[tabIndex].push({ nodes: getNodes(), edges: getEdges() });
+        newFuture[tabIndex].push({ nodes: nodes, edges: edges });
         return newFuture;
       });
       // now we can set the graph to the past state
@@ -103,8 +101,8 @@ export function UndoRedoProvider({ children }) {
   }, [
     setNodes,
     setEdges,
-    getNodes,
-    getEdges,
+    nodes,
+    edges,
     future,
     past,
     setFuture,
@@ -124,7 +122,7 @@ export function UndoRedoProvider({ children }) {
       setPast((old) => {
         let newPast = cloneDeep(old);
         newPast[tabIndex] = old[tabIndex];
-        newPast[tabIndex].push({ nodes: getNodes(), edges: getEdges() });
+        newPast[tabIndex].push({ nodes: nodes, edges: edges });
         return newPast;
       });
       setNodes(futureState.nodes);
@@ -137,8 +135,8 @@ export function UndoRedoProvider({ children }) {
     setPast,
     setNodes,
     setEdges,
-    getNodes,
-    getEdges,
+    nodes,
+    edges,
     future,
     tabIndex,
   ]);
diff --git a/src/frontend/src/modals/EditNodeModal/index.tsx b/src/frontend/src/modals/EditNodeModal/index.tsx
index 26dc15fd4..b9576594b 100644
--- a/src/frontend/src/modals/EditNodeModal/index.tsx
+++ b/src/frontend/src/modals/EditNodeModal/index.tsx
@@ -62,8 +62,7 @@ const EditNodeModal = forwardRef(
 
     const [myData, setMyData] = useState(data);
 
-    const { setTabsState, tabId } = useContext(FlowsContext);
-    const { reactFlowInstance } = useContext(typesContext);
+    const { setTabsState, tabId, edges } = useContext(FlowsContext);
     const { setModalContextOpen } = useContext(alertContext);
 
     function changeAdvanced(n) {
@@ -159,7 +158,7 @@ const EditNodeModal = forwardRef(
                             fieldName: templateParam,
                           };
                           let disabled =
-                            reactFlowInstance?.getEdges().some(
+                            edges.some(
                               (edge) =>
                                 edge.targetHandle ===
                                 scapedJSONStringfy(
diff --git a/src/frontend/src/modals/codeAreaModal/index.tsx b/src/frontend/src/modals/codeAreaModal/index.tsx
index f488704c8..c700b701b 100644
--- a/src/frontend/src/modals/codeAreaModal/index.tsx
+++ b/src/frontend/src/modals/codeAreaModal/index.tsx
@@ -28,7 +28,6 @@ export default function CodeAreaModal({
 }: codeAreaModalPropsType): JSX.Element {
   const [code, setCode] = useState(value);
   const { dark } = useContext(darkContext);
-  const { reactFlowInstance } = useContext(typesContext);
   const [height, setHeight] = useState<string | null>(null);
   const { setErrorData, setSuccessData } = useContext(alertContext);
   const [error, setError] = useState<{
diff --git a/src/frontend/src/modals/exportModal/index.tsx b/src/frontend/src/modals/exportModal/index.tsx
index 5ea23b644..f2d2424a4 100644
--- a/src/frontend/src/modals/exportModal/index.tsx
+++ b/src/frontend/src/modals/exportModal/index.tsx
@@ -13,7 +13,6 @@ import BaseModal from "../baseModal";
 const ExportModal = forwardRef(
   (props: { children: ReactNode }, ref): JSX.Element => {
     const { flows, tabId, downloadFlow, version } = useContext(FlowsContext);
-    const { reactFlowInstance } = useContext(typesContext);
     const { setNoticeData } = useContext(alertContext);
     const [checked, setChecked] = useState(true);
     const flow = flows.find((f) => f.id === tabId);
diff --git a/src/frontend/src/modals/formModal/index.tsx b/src/frontend/src/modals/formModal/index.tsx
index 571cbca06..891a1dfa6 100644
--- a/src/frontend/src/modals/formModal/index.tsx
+++ b/src/frontend/src/modals/formModal/index.tsx
@@ -38,7 +38,7 @@ export default function FormModal({
   setOpen: (open: boolean) => void;
   flow: FlowType;
 }): JSX.Element {
-  const { tabsState, setTabsState } = useContext(FlowsContext);
+  const { tabsState, setTabsState, nodes, edges } = useContext(FlowsContext);
   const [chatValue, setChatValue] = useState(() => {
     try {
       const { formKeysData } = tabsState[flow.id];
@@ -62,7 +62,6 @@ export default function FormModal({
 
   const [chatHistory, setChatHistory] = useState<ChatMessageType[]>([]);
   const template = useRef(tabsState[flow.id].formKeysData.template);
-  const { reactFlowInstance } = useContext(typesContext);
   const { accessToken } = useContext(AuthContext);
   const { setErrorData } = useContext(alertContext);
   const ws = useRef<WebSocket | null>(null);
@@ -385,8 +384,8 @@ export default function FormModal({
 
   function sendMessage(): void {
     let nodeValidationErrors = validateNodes(
-      reactFlowInstance!.getNodes(),
-      reactFlowInstance!.getEdges()
+      nodes,
+      edges
     );
     if (nodeValidationErrors.length === 0) {
       setLockChat(true);
diff --git a/src/frontend/src/modals/shareModal/index.tsx b/src/frontend/src/modals/shareModal/index.tsx
index 0edbe80be..3363777d6 100644
--- a/src/frontend/src/modals/shareModal/index.tsx
+++ b/src/frontend/src/modals/shareModal/index.tsx
@@ -40,10 +40,9 @@ export default function ShareModal({
   setOpen?: (open: boolean) => void;
   disabled?: boolean;
 }): JSX.Element {
-  const { version, addFlow, downloadFlow } = useContext(FlowsContext);
+  const { version } = useContext(FlowsContext);
   const { hasApiKey, hasStore } = useContext(StoreContext);
   const { setSuccessData, setErrorData } = useContext(alertContext);
-  const { reactFlowInstance } = useContext(typesContext);
   const [internalOpen, internalSetOpen] = useState(children ? false : true);
   const [openConfirmationModal, setOpenConfirmationModal] = useState(false);
   const nameComponent = is_component ? "component" : "flow";
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index d035812f0..ed1700be4 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -82,7 +82,7 @@ export default function Page({
   const reactFlowWrapper = useRef<HTMLDivElement>(null);
 
   const { takeSnapshot } = useContext(undoRedoContext);
-  const { nodesOnFlow, setNodesOnFlow } = useContext(FlowsContext);
+  const { nodes, edges, setNodes, setEdges, onNodesChange, onEdgesChange } = useContext(FlowsContext);
 
   const position = useRef({ x: 0, y: 0 });
   const [lastSelection, setLastSelection] =
@@ -164,13 +164,7 @@ export default function Page({
 
   const { setExtraComponent, setExtraNavigation } = useContext(locationContext);
   const { setErrorData } = useContext(alertContext);
-  const [nodes, setNodes, onNodesChange] = useNodesState(
-    flow.data?.nodes ?? []
-  );
 
-  const [edges, setEdges, onEdgesChange] = useEdgesState(
-    flow.data?.edges ?? []
-  );
   const { setViewport } = useReactFlow();
   const edgeUpdateSuccessful = useRef(true);
 
@@ -210,50 +204,6 @@ export default function Page({
     };
   }, [flow, flow.data]);
 
-  const onEdgesChangeMod = useCallback(
-    (change: EdgeChange[]) => {
-      onEdgesChange(change);
-      //@ts-ignore
-      setTabsState((prev: FlowsState) => {
-        return {
-          ...prev,
-          [tabId]: {
-            ...prev[tabId],
-            isPending: true,
-          },
-        };
-      });
-      saveCurrentFlowTimeout();
-    },
-    [onEdgesChange, setNodes, setTabsState, saveCurrentFlowTimeout, tabId]
-  );
-
-  const onNodesChangeMod = useCallback(
-    (change: NodeChange[]) => {
-      const changeString = JSON.stringify(change);
-      if (changeString !== nodesOnFlow) {
-        onNodesChange(change);
-        updateNodeFlow(changeString);
-        //@ts-ignore
-        setTabsState((prev: FlowsState) => {
-          return {
-            ...prev,
-            [tabId]: {
-              ...prev[tabId],
-              isPending: true,
-            },
-          };
-        });
-        saveCurrentFlowTimeout();
-      }
-    },
-    [onNodesChange, setTabsState, tabId, updateNodeFlow, saveCurrentFlowTimeout]
-  );
-
-  function updateNodeFlow(changeString: string) {
-    setNodesOnFlow(changeString);
-  }
-
   const onConnect = useCallback(
     (params: Connection) => {
       takeSnapshot();
@@ -416,7 +366,7 @@ export default function Page({
 
   const onEdgeUpdate = useCallback(
     (oldEdge: Edge, newConnection: Connection) => {
-      if (isValidConnection(newConnection, reactFlowInstance!)) {
+      if (isValidConnection(newConnection, nodes, edges)) {
         edgeUpdateSuccessful.current = true;
         setEdges((els) => updateEdge(oldEdge, newConnection, els));
       }
@@ -498,8 +448,8 @@ export default function Page({
                   nodes={nodes}
                   onMove={onMove}
                   edges={edges}
-                  onNodesChange={onNodesChangeMod}
-                  onEdgesChange={onEdgesChangeMod}
+                  onNodesChange={onNodesChange}
+                  onEdgesChange={onEdgesChange}
                   onConnect={onConnect}
                   disableKeyboardA11y={true}
                   onInit={setReactFlowInstance}
@@ -579,7 +529,7 @@ export default function Page({
                   />
                 </ReactFlow>
                 {!view && (
-                  <Chat flow={flow} reactFlowInstance={reactFlowInstance!} />
+                  <Chat flow={flow} />
                 )}
               </div>
             ) : (
diff --git a/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
index a70132794..84e34a07b 100644
--- a/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/extraSidebarComponent/index.tsx
@@ -26,7 +26,7 @@ import DisclosureComponent from "../DisclosureComponent";
 import SidebarDraggableComponent from "./sideBarDraggableComponent";
 
 export default function ExtraSidebar(): JSX.Element {
-  const { data, templates, getFilterEdge, setFilterEdge, reactFlowInstance } =
+  const { data, templates, getFilterEdge, setFilterEdge } =
     useContext(typesContext);
   const { flows, tabId, uploadFlow, tabsState, saveFlow, isBuilt, version } =
     useContext(FlowsContext);
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 00b5e46c8..5b4d5ed38 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -65,9 +65,8 @@ export default function NodeToolbarComponent({
   const isMinimal = canMinimize();
   const isGroup = data.node?.flow ? true : false;
 
-  const { paste, saveComponent, version, flows } = useContext(FlowsContext);
+  const { paste, saveComponent, version, flows, nodes, edges, setNodes, setEdges } = useContext(FlowsContext);
   const { takeSnapshot } = useContext(undoRedoContext);
-  const reactFlowInstance = useReactFlow();
   const [showModalAdvanced, setShowModalAdvanced] = useState(false);
   const [showconfirmShare, setShowconfirmShare] = useState(false);
   const [selectedValue, setSelectedValue] = useState("");
@@ -115,7 +114,7 @@ export default function NodeToolbarComponent({
       case "ungroup":
         takeSnapshot();
         updateFlowPosition(position, data.node?.flow!);
-        expandGroupNode(data, reactFlowInstance, getNodeId);
+        expandGroupNode(data, getNodeId, nodes, edges, setNodes, setEdges);
         break;
       case "override":
         setShowOverrideModal(true);
@@ -152,14 +151,14 @@ export default function NodeToolbarComponent({
                 event.preventDefault();
                 paste(
                   {
-                    nodes: [reactFlowInstance.getNode(data.id)],
+                    nodes: [nodes.find((node) => node.id === data.id)!],
                     edges: [],
                   },
                   {
                     x: 50,
                     y: 10,
-                    paneX: reactFlowInstance.getNode(data.id)?.position.x,
-                    paneY: reactFlowInstance.getNode(data.id)?.position.y,
+                    paneX: nodes.find((node) => node.id === data.id)?.position.x,
+                    paneY: nodes.find((node) => node.id === data.id)?.position.y,
                   }
                 );
               }}
diff --git a/src/frontend/src/types/chat/index.ts b/src/frontend/src/types/chat/index.ts
index 71a081be7..66422a968 100644
--- a/src/frontend/src/types/chat/index.ts
+++ b/src/frontend/src/types/chat/index.ts
@@ -1,7 +1,7 @@
 import { ReactFlowInstance } from "reactflow";
 import { FlowType } from "../flow";
 
-export type ChatType = { flow: FlowType; reactFlowInstance: ReactFlowInstance };
+export type ChatType = { flow: FlowType; };
 export type ChatMessageType = {
   message: string | Object;
   template?: string;
diff --git a/src/frontend/src/types/tabs/index.ts b/src/frontend/src/types/tabs/index.ts
index d4e120c1c..732ddbb80 100644
--- a/src/frontend/src/types/tabs/index.ts
+++ b/src/frontend/src/types/tabs/index.ts
@@ -1,6 +1,9 @@
-import { XYPosition } from "reactflow";
+import { XYPosition, Node, NodeChange, Edge, EdgeChange } from "reactflow";
 import { tweakType } from "../components";
 import { FlowType, NodeDataType } from "../flow";
+import { Dispatch, SetStateAction } from "react";
+
+type OnChange<ChangesType> = (changes: ChangesType[]) => void;
 
 export type FlowsContextType = {
   saveFlow: (flow: FlowType, silent?: boolean) => Promise<void>;
@@ -56,8 +59,12 @@ export type FlowsContextType = {
   ) => Promise<String | undefined>;
   deleteComponent: (key: string) => void;
   version: string;
-  nodesOnFlow: string;
-  setNodesOnFlow: (nodes: string) => void;
+  nodes: Array<Node>;
+  setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void;
+  onNodesChange: OnChange<NodeChange>;
+  edges: Array<Edge>;
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void;
+  onEdgesChange: OnChange<EdgeChange>;
 };
 
 export type FlowsState = {
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index c8b7120a8..957708dcb 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -86,8 +86,10 @@ export function unselectAllNodes({ updateNodes, data }: unselectAllNodesType) {
 
 export function isValidConnection(
   { source, target, sourceHandle, targetHandle }: Connection,
-  reactFlowInstance: ReactFlowInstance
+  nodes: Node[],
+  edges: Edge[],
 ) {
+  debugger
   const targetHandleObject: targetHandleType = scapeJSONParse(targetHandle!);
   const sourceHandleObject: sourceHandleType = scapeJSONParse(sourceHandle!);
   if (
@@ -100,19 +102,17 @@ export function isValidConnection(
         t === targetHandleObject.type
     )
   ) {
-    let targetNode = reactFlowInstance?.getNode(target!)?.data?.node;
+    let targetNode = nodes.find((node) => node.id === target!)?.data?.node;
     if (!targetNode) {
       if (
-        !reactFlowInstance
-          .getEdges()
+        edges
           .find((e) => e.targetHandle === targetHandle)
       ) {
         return true;
       }
     } else if (
       (!targetNode.template[targetHandleObject.fieldName].list &&
-        !reactFlowInstance
-          .getEdges()
+        edges
           .find((e) => e.targetHandle === targetHandle)) ||
       targetNode.template[targetHandleObject.fieldName].list
     ) {
@@ -543,12 +543,13 @@ export function generateFlow(
 
 export function filterFlow(
   selection: OnSelectionChangeParams,
-  reactFlowInstance: ReactFlowInstance
+  setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void,
 ) {
-  reactFlowInstance.setNodes((nodes) =>
+  setNodes((nodes) =>
     nodes.filter((node) => !selection.nodes.includes(node))
   );
-  reactFlowInstance.setEdges((edges) =>
+  setEdges((edges) =>
     edges.filter((edge) => !selection.edges.includes(edge))
   );
 }
@@ -575,11 +576,12 @@ export function updateFlowPosition(NewPosition: XYPosition, flow: FlowType) {
 
 export function concatFlows(
   flow: FlowType,
-  ReactFlowInstance: ReactFlowInstance
+  setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void,
 ) {
   const { nodes, edges } = flow.data!;
-  ReactFlowInstance.addNodes(nodes);
-  ReactFlowInstance.addEdges(edges);
+  setNodes((old) => [...old, ...nodes]);
+  setEdges((old) => [...old, ...edges]);
 }
 
 export function validateSelection(
@@ -921,13 +923,16 @@ function updateEdgesIds(edges: Edge[], idsMap: { [key: string]: string }) {
 
 export function expandGroupNode(
   groupNode: NodeDataType,
-  ReactFlowInstance: ReactFlowInstance,
-  getNodeId: (type: string) => string
+  getNodeId: (type: string) => string,
+  nodes: Node[],
+  edges: Edge[],
+  setNodes: (update: Node[] | ((oldState: Node[]) => Node[])) => void,
+  setEdges: (update: Edge[] | ((oldState: Edge[]) => Edge[])) => void
 ) {
   const { template, flow } = _.cloneDeep(groupNode.node!);
   const idsMap = updateIds(flow!.data!, getNodeId);
   updateProxyIdsOnTemplate(template, idsMap);
-  let flowEdges = ReactFlowInstance.getEdges();
+  let flowEdges = edges;
   updateEdgesIds(flowEdges, idsMap);
   const gNodes: NodeType[] = flow?.data?.nodes!;
   const gEdges = flow!.data!.edges;
@@ -1008,19 +1013,19 @@ export function expandGroupNode(
     }
   });
 
-  const nodes = [
-    ...ReactFlowInstance.getNodes().filter((n) => n.id !== groupNode.id),
+  const filteredNodes = [
+    ...nodes.filter((n) => n.id !== groupNode.id),
     ...gNodes,
   ];
-  const edges = [
-    ...ReactFlowInstance.getEdges().filter(
+  const filteredEdges = [
+    ...edges.filter(
       (e) => e.target !== groupNode.id && e.source !== groupNode.id
     ),
     ...gEdges,
     ...updatedEdges,
   ];
-  ReactFlowInstance.setNodes(nodes);
-  ReactFlowInstance.setEdges(edges);
+  setNodes(filteredNodes);
+  setEdges(filteredEdges);
 }
 
 export function processFlow(FlowObject: ReactFlowJsonObject) {