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,
           };