diff --git a/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx b/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
index a8faffa49..de34135e7 100644
--- a/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-handle-new-value.tsx
@@ -6,6 +6,7 @@ import useFlowsManagerStore from "@/stores/flowsManagerStore";
 import { APIClassType, InputFieldType } from "@/types/api";
 import { NodeType } from "@/types/flow";
 import { cloneDeep } from "lodash";
+import { useUpdateNodeInternals } from "reactflow";
 import { mutateTemplate } from "../helpers/mutate-template";
 
 export type handleOnNewValueType = (
@@ -33,6 +34,7 @@ const useHandleOnNewValue = ({
   const takeSnapshot = useFlowsManagerStore((state) => state.takeSnapshot);
 
   const setNode = setNodeExternal ?? useFlowStore((state) => state.setNode);
+  const updateNodeInternals = useUpdateNodeInternals();
 
   const setErrorData = useAlertStore((state) => state.setErrorData);
 
@@ -90,14 +92,21 @@ const useHandleOnNewValue = ({
       );
     }
 
-    setNode(nodeId, (oldNode) => {
-      const newData = cloneDeep(oldNode.data);
-      newData.node = newNode;
-      return {
-        ...oldNode,
-        data: newData,
-      };
-    });
+    setNode(
+      nodeId,
+      (oldNode) => {
+        const newData = cloneDeep(oldNode.data);
+        newData.node = newNode;
+        return {
+          ...oldNode,
+          data: newData,
+        };
+      },
+      true,
+      () => {
+        updateNodeInternals(nodeId);
+      },
+    );
   };
 
   return { handleOnNewValue };
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 9aa1a1f54..1a18fb156 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -257,13 +257,15 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     id: string,
     change: Node | ((oldState: Node) => Node),
     isUserChange: boolean = true,
+    callback?: () => void,
   ) => {
     let newChange =
       typeof change === "function"
         ? change(get().nodes.find((node) => node.id === id)!)
         : change;
-    get().setNodes((oldNodes) =>
-      oldNodes.map((node) => {
+
+    set((state) => {
+      const newNodes = state.nodes.map((node) => {
         if (node.id === id) {
           if (isUserChange) {
             if ((node.data as NodeDataType).node?.frozen) {
@@ -273,8 +275,18 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
           return newChange;
         }
         return node;
-      }),
-    );
+      });
+
+      if (callback) {
+        // Defer the callback execution to ensure it runs after state updates are fully applied.
+        queueMicrotask(callback);
+      }
+
+      return {
+        ...state,
+        nodes: newNodes,
+      };
+    });
   },
   getNode: (id: string) => {
     return get().nodes.find((node) => node.id === id);
diff --git a/src/frontend/src/types/zustand/flow/index.ts b/src/frontend/src/types/zustand/flow/index.ts
index 9e72648e1..b8aa9415e 100644
--- a/src/frontend/src/types/zustand/flow/index.ts
+++ b/src/frontend/src/types/zustand/flow/index.ts
@@ -106,6 +106,7 @@ export type FlowStoreType = {
     id: string,
     update: Node | ((oldState: Node) => Node),
     isUserChange?: boolean,
+    callback?: () => void,
   ) => void;
   getNode: (id: string) => Node | undefined;
   deleteNode: (nodeId: string | Array<string>) => void;