From 866076cf06ea7ac4ea61b2234854faf7575ae4bf Mon Sep 17 00:00:00 2001
From: Lucas Oliveira <62335616+lucaseduoli@users.noreply.github.com>
Date: Thu, 26 Jun 2025 10:18:32 -0300
Subject: [PATCH] fix: make not connected edges be cleared when entering flow
 (#8678)

* Fixed output not cleaning edges

* Added useEffect to select first output if none is selected

* updated useeffect condition

* Fixed edges connected being cleared

* added handleSelectOutput to dependencies

* Fixed loop disconnecting
---
 .../src/CustomNodes/GenericNode/index.tsx         | 12 ++++++++++++
 src/frontend/src/stores/flowStore.ts              |  2 +-
 src/frontend/src/utils/reactflowUtils.ts          | 15 ++++++++++++---
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/frontend/src/CustomNodes/GenericNode/index.tsx b/src/frontend/src/CustomNodes/GenericNode/index.tsx
index dae1f3fde..b48f6a3b2 100644
--- a/src/frontend/src/CustomNodes/GenericNode/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/index.tsx
@@ -325,6 +325,18 @@ function GenericNode({
     [data.id, setNode, setEdges, updateNodeInternals],
   );
 
+  useEffect(() => {
+    if (
+      data?.selected_output ||
+      (data?.node?.outputs?.filter((output) => !output.group_outputs)?.length ??
+        0) <= 1
+    )
+      return;
+    handleSelectOutput(
+      data.node?.outputs?.find((output) => output.selected) || null,
+    );
+  }, [data.node?.outputs, data?.selected_output, handleSelectOutput]);
+
   const [hasChangedNodeDescription, setHasChangedNodeDescription] =
     useState(false);
 
diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index f52a437a6..839e39bb9 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -226,7 +226,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         localStorage.getItem(`dismiss_${flow?.id}`) ?? "[]",
       ) as string[],
     });
-    unselectAllNodesEdges(nodes, edges);
+    unselectAllNodesEdges(nodes, newEdges);
     if (flow?.id) {
       useTweaksStore.getState().initialSetup(nodes, flow?.id);
     }
diff --git a/src/frontend/src/utils/reactflowUtils.ts b/src/frontend/src/utils/reactflowUtils.ts
index 63a1e6922..6fc3f3b4a 100644
--- a/src/frontend/src/utils/reactflowUtils.ts
+++ b/src/frontend/src/utils/reactflowUtils.ts
@@ -138,9 +138,18 @@ export function cleanEdges(nodes: AllNodeType[], edges: EdgeType[]) {
       const name = parsedSourceHandle.name;
 
       if (sourceNode.type == "genericNode") {
-        const output = sourceNode.data
-          .node!.outputs?.filter((output) => output.selected)
-          .find((output) => output.name === name);
+        const output =
+          sourceNode.data.node!.outputs?.find(
+            (output) => output.name === sourceNode.data.selected_output,
+          ) ??
+          sourceNode.data.node!.outputs?.find(
+            (output) =>
+              (output.selected ||
+                (sourceNode.data.node!.outputs?.filter(
+                  (output) => !output.group_outputs,
+                )?.length ?? 0) <= 1) &&
+              output.name === name,
+          );
 
         if (output) {
           const outputTypes =