From dce5523f5ec8a24b0b1d05403f2d6073a83730e2 Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Mon, 7 Oct 2024 18:04:07 -0300
Subject: [PATCH] fix: Add useUpdateNodeInternals hook to handle node class
 changes (#4058)

* fix: Add useUpdateNodeInternals hook to handle node class changes

* [autofix.ci] apply automated fixes

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
Co-authored-by: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
---
 src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
index fd86f51d1..baf78a781 100644
--- a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
@@ -1,6 +1,7 @@
 import useFlowStore from "@/stores/flowStore";
 import { NodeType } from "@/types/flow";
 import { cloneDeep } from "lodash";
+import { useUpdateNodeInternals } from "reactflow";
 
 const useHandleNodeClass = (
   nodeId: string,
@@ -10,6 +11,7 @@ const useHandleNodeClass = (
   ) => void,
 ) => {
   const setNode = setMyNode ?? useFlowStore((state) => state.setNode);
+  const updateNodeInternals = useUpdateNodeInternals();
 
   const handleNodeClass = (newNodeClass, type?: string) => {
     setNode(nodeId, (oldNode) => {
@@ -23,6 +25,8 @@ const useHandleNodeClass = (
         newNode.data.type = type;
       }
 
+      updateNodeInternals(nodeId);
+
       return newNode;
     });
   };