diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index 27c550db8..eb6d27137 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -1,5 +1,5 @@
 import { useContext, useEffect, useState } from "react";
-import { NodeToolbar, useUpdateNodeInternals } from "reactflow";
+import { NodeToolbar } from "reactflow";
 import ShadTooltip from "../../components/ShadTooltipComponent";
 import Tooltip from "../../components/TooltipComponent";
 import IconComponent from "../../components/genericIconComponent";
@@ -29,7 +29,6 @@ export default function GenericNode({
   xPos: number;
   yPos: number;
 }): JSX.Element {
-  const updateNodeInternals = useUpdateNodeInternals();
   const { types } = useContext(typesContext);
   const { deleteNode, setNode } = useContext(FlowsContext);
   const name = nodeIconsLucide[data.type] ? data.type : types[data.type];
@@ -419,7 +418,6 @@ export default function GenericNode({
                           },
                         },
                       }));
-                      updateNodeInternals(data.id);
                     }
                   }}
                 />
diff --git a/src/frontend/src/modals/EditNodeModal/index.tsx b/src/frontend/src/modals/EditNodeModal/index.tsx
index fea51b765..0b6ab5a39 100644
--- a/src/frontend/src/modals/EditNodeModal/index.tsx
+++ b/src/frontend/src/modals/EditNodeModal/index.tsx
@@ -1,6 +1,5 @@
 import { cloneDeep } from "lodash";
 import { forwardRef, useContext, useEffect, useState } from "react";
-import { useUpdateNodeInternals } from "reactflow";
 import ShadTooltip from "../../components/ShadTooltipComponent";
 import CodeAreaComponent from "../../components/codeAreaComponent";
 import DictComponent from "../../components/dictComponent";
@@ -58,11 +57,9 @@ const EditNodeModal = forwardRef(
     },
     ref
   ) => {
-    const updateNodeInternals = useUpdateNodeInternals();
-
     const [myData, setMyData] = useState(data);
 
-    const { setPending, tabId, edges } = useContext(FlowsContext);
+    const { setPending, edges, setNode } = useContext(FlowsContext);
     const { setModalContextOpen } = useContext(alertContext);
 
     function changeAdvanced(n) {
@@ -80,7 +77,6 @@ const EditNodeModal = forwardRef(
         newData.node!.template[name].value = newValue;
         return newData;
       });
-      updateNodeInternals(data.id);
     };
 
     useEffect(() => {
@@ -534,7 +530,13 @@ const EditNodeModal = forwardRef(
             id={"saveChangesBtn"}
             className="mt-3"
             onClick={() => {
-              data.node = myData.node;
+              setNode(data.id, (old) => ({
+                ...old,
+                data: {
+                  ...old.data,
+                  node: myData.node,
+                },
+              }));
               setPending(true);
               setOpen(false);
             }}
diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 5b4d5ed38..7e04211b8 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -1,6 +1,6 @@
 import { cloneDeep } from "lodash";
 import { useContext, useEffect, useState } from "react";
-import { useReactFlow, useUpdateNodeInternals } from "reactflow";
+import { useReactFlow } from "reactflow";
 import ShadTooltip from "../../../../components/ShadTooltipComponent";
 import IconComponent from "../../../../components/genericIconComponent";
 import {
@@ -50,7 +50,6 @@ export default function NodeToolbarComponent({
           data.node.template[templateField].type === "NestedDict")
     ).length
   );
-  const updateNodeInternals = useUpdateNodeInternals();
   const { getNodeId } = useContext(FlowsContext);
   const { hasApiKey, validApiKey, hasStore } = useContext(StoreContext);
 
@@ -98,7 +97,6 @@ export default function NodeToolbarComponent({
       case "show":
         takeSnapshot();
         setShowNode(data.showNode ?? true ? false : true);
-        updateNodeInternals(data.id);
         break;
       case "Download":
         downloadNode(createFlowComponent(cloneDeep(data), version));