feat: Add post-update callback to setNode and improve node internals handling (#4696)

* ✨ (use-handle-new-value.tsx): Introduce useUpdateNodeInternals hook from reactflow to update node internals when handling new values
♻️ (flowStore.ts): Refactor set method in useFlowStore to accept a callback function to be executed after updating nodes
📝 (flowStore.ts): Update FlowStoreType to include a callback function in the update method signature

* 📝 (flowStore.ts): add comment to explain deferring callback execution to ensure it runs after state updates are fully applied

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) => {
+    setNode(
-      const newData = cloneDeep(oldNode.data);
+      nodeId,
-      newData.node = newNode;
+      (oldNode) => {
-      return {
+        const newData = cloneDeep(oldNode.data);
-        ...oldNode,
+        newData.node = newNode;
-        data: newData,
+        return {
-      };
+          ...oldNode,
-    });
+          data: newData,
+        };
+      },
+      true,
+      () => {
+        updateNodeInternals(nodeId);
+      },
+    );
   };
 
   return { handleOnNewValue };

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);

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;