From f6a9e8fad4afbb3903c432ff1807d10c617750ed Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com>
Date: Wed, 9 Jul 2025 18:32:57 -0300
Subject: [PATCH] fix: clean edges of same type but different name when
 switching outputs, update color when deleting edges (#8959)

* fix edges not being cleared if same type but different name

* removed console log

* removed console log

* made edges change trigger re-render of handle color

* Updated position of edges
---
 .../components/handleRenderComponent/index.tsx        | 11 ++++++-----
 src/frontend/src/utils/reactflowUtils.ts              |  2 +-
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx
index 8a1c88d7a..23463a500 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx
@@ -184,6 +184,8 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
     useShallow((state) => state.currentFlow?.locked),
   );
 
+  const edges = useFlowStore((state) => state.edges);
+
   const {
     setHandleDragging,
     setFilterType,
@@ -266,11 +268,9 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
     const openHandle = filterOpenHandle || draggingOpenHandle;
     const filterPresent = handleDragging || filterType;
 
-    const connectedEdge = useFlowStore
-      .getState()
-      .edges.find(
-        (edge) => edge.target === nodeId && edge.targetHandle === myId,
-      );
+    const connectedEdge = edges.find(
+      (edge) => edge.target === nodeId && edge.targetHandle === myId,
+    );
     const outputType = connectedEdge?.data?.sourceHandle?.output_types?.[0];
     const connectedColor = outputType ? nodeColorsName[outputType] : "gray";
 
@@ -344,6 +344,7 @@ const HandleRenderComponent = memo(function HandleRenderComponent({
     colors,
     colorName,
     tooltipTitle,
+    edges,
   ]);
 
   const handleMouseDown = useCallback(
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index a3d991f74..839c7cdb8 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -157,7 +157,7 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) {
 
           const id: sourceHandleType = {
             id: sourceNode.data.id,
-            name: name,
+            name: output?.name ?? name,
             output_types: outputTypes,
             dataType: sourceNode.data.type,
           };