diff --git a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
index 97c3e5ec2..e6a7af312 100644
--- a/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/nodeToolbarComponent/index.tsx
@@ -26,6 +26,7 @@ import {
   updateFlowPosition,
 } from "../../../../utils/reactflowUtils";
 import { classNames, cn } from "../../../../utils/utils";
+import { useUpdateNodeInternals } from "reactflow";
 
 export default function NodeToolbarComponent({
   data,
@@ -86,6 +87,8 @@ export default function NodeToolbarComponent({
       onCloseAdvancedModal!(false);
     }
   }, [showModalAdvanced]);
+  const updateNodeInternals = useUpdateNodeInternals();
+
 
   useEffect(() => {
     setFlowComponent(createFlowComponent(cloneDeep(data), version));
@@ -191,6 +194,7 @@ export default function NodeToolbarComponent({
 
       return newNode;
     });
+    updateNodeInternals(data.id);
   };
 
   const [openModal, setOpenModal] = useState(false);