feat: keep connections when change node (#31982)

web/app/components/workflow/hooks/use-nodes-interactions.ts

@@ -1422,21 +1422,136 @@ export const useNodesInteractions = () => {
         extent: currentNode.extent,
         zIndex: currentNode.zIndex,
       })
-      const nodesConnectedSourceOrTargetHandleIdsMap
+      const parentNode = nodes.find(node => node.id === currentNode.parentId)
-        = getNodesConnectedSourceOrTargetHandleIdsMap(
+      const newNodeIsInIteration
-          connectedEdges.map(edge => ({ type: 'remove', edge })),
+        = !!parentNode && parentNode.data.type === BlockEnum.Iteration
-          nodes,
+      const newNodeIsInLoop
-        )
+        = !!parentNode && parentNode.data.type === BlockEnum.Loop
-      const newNodes = produce(nodes, (draft) => {
+      const outgoingEdges = connectedEdges.filter(
+        edge => edge.source === currentNodeId,
+      )
+      const normalizedSourceHandle = sourceHandle || 'source'
+      const outgoingHandles = new Set(
+        outgoingEdges.map(edge => edge.sourceHandle || 'source'),
+      )
+      const branchSourceHandle = currentNode.data._targetBranches?.[0]?.id
+      let outgoingHandleToPreserve = normalizedSourceHandle
+      if (!outgoingHandles.has(outgoingHandleToPreserve)) {
+        if (branchSourceHandle && outgoingHandles.has(branchSourceHandle))
+          outgoingHandleToPreserve = branchSourceHandle
+        else if (outgoingHandles.has('source'))
+          outgoingHandleToPreserve = 'source'
+        else
+          outgoingHandleToPreserve = outgoingEdges[0]?.sourceHandle || 'source'
+      }
+      const outgoingEdgesToPreserve = outgoingEdges.filter(
+        edge => (edge.sourceHandle || 'source') === outgoingHandleToPreserve,
+      )
+      const outgoingEdgeIds = new Set(
+        outgoingEdgesToPreserve.map(edge => edge.id),
+      )
+      const newNodeSourceHandle = newCurrentNode.data._targetBranches?.[0]?.id || 'source'
+      const reconnectedEdges = connectedEdges.reduce<Edge[]>(
+        (acc, edge) => {
+          if (outgoingEdgeIds.has(edge.id)) {
+            const originalTargetNode = nodes.find(
+              node => node.id === edge.target,
+            )
+            const targetNodeForEdge
+              = originalTargetNode && originalTargetNode.id !== currentNodeId
+                ? originalTargetNode
+                : newCurrentNode
+            if (!targetNodeForEdge)
+              return acc
+
+            const targetHandle = edge.targetHandle || 'target'
+            const targetParentNode
+              = targetNodeForEdge.id === newCurrentNode.id
+                ? parentNode || null
+                : nodes.find(node => node.id === targetNodeForEdge.parentId)
+                  || null
+            const isInIteration
+              = !!targetParentNode
+                && targetParentNode.data.type === BlockEnum.Iteration
+            const isInLoop
+              = !!targetParentNode
+                && targetParentNode.data.type === BlockEnum.Loop
+
+            acc.push({
+              ...edge,
+              id: `${newCurrentNode.id}-${newNodeSourceHandle}-${targetNodeForEdge.id}-${targetHandle}`,
+              source: newCurrentNode.id,
+              sourceHandle: newNodeSourceHandle,
+              target: targetNodeForEdge.id,
+              targetHandle,
+              type: CUSTOM_EDGE,
+              data: {
+                ...(edge.data || {}),
+                sourceType: newCurrentNode.data.type,
+                targetType: targetNodeForEdge.data.type,
+                isInIteration,
+                iteration_id: isInIteration
+                  ? targetNodeForEdge.parentId
+                  : undefined,
+                isInLoop,
+                loop_id: isInLoop ? targetNodeForEdge.parentId : undefined,
+                _connectedNodeIsSelected: false,
+              },
+              zIndex: targetNodeForEdge.parentId
+                ? isInIteration
+                  ? ITERATION_CHILDREN_Z_INDEX
+                  : LOOP_CHILDREN_Z_INDEX
+                : 0,
+            })
+          }
+
+          if (
+            edge.target === currentNodeId
+            && edge.source !== currentNodeId
+            && !outgoingEdgeIds.has(edge.id)
+          ) {
+            const sourceNode = nodes.find(node => node.id === edge.source)
+            if (!sourceNode)
+              return acc
+
+            const targetHandle = edge.targetHandle || 'target'
+            const sourceHandle = edge.sourceHandle || 'source'
+
+            acc.push({
+              ...edge,
+              id: `${sourceNode.id}-${sourceHandle}-${newCurrentNode.id}-${targetHandle}`,
+              source: sourceNode.id,
+              sourceHandle,
+              target: newCurrentNode.id,
+              targetHandle,
+              type: CUSTOM_EDGE,
+              data: {
+                ...(edge.data || {}),
+                sourceType: sourceNode.data.type,
+                targetType: newCurrentNode.data.type,
+                isInIteration: newNodeIsInIteration,
+                iteration_id: newNodeIsInIteration
+                  ? newCurrentNode.parentId
+                  : undefined,
+                isInLoop: newNodeIsInLoop,
+                loop_id: newNodeIsInLoop ? newCurrentNode.parentId : undefined,
+                _connectedNodeIsSelected: false,
+              },
+              zIndex: newCurrentNode.parentId
+                ? newNodeIsInIteration
+                  ? ITERATION_CHILDREN_Z_INDEX
+                  : LOOP_CHILDREN_Z_INDEX
+                : 0,
+            })
+          }
+
+          return acc
+        },
+        [],
+      )
+      const nodesWithNewNode = produce(nodes, (draft) => {
         draft.forEach((node) => {
           node.data.selected = false
-
-          if (nodesConnectedSourceOrTargetHandleIdsMap[node.id]) {
-            node.data = {
-              ...node.data,
-              ...nodesConnectedSourceOrTargetHandleIdsMap[node.id],
-            }
-          }
         })
         const index = draft.findIndex(node => node.id === currentNodeId)
 
@@ -1446,18 +1561,32 @@ export const useNodesInteractions = () => {
         if (newLoopStartNode)
           draft.push(newLoopStartNode)
       })
-      setNodes(newNodes)
+      const nodesConnectedSourceOrTargetHandleIdsMap
-      const newEdges = produce(edges, (draft) => {
+        = getNodesConnectedSourceOrTargetHandleIdsMap(
-        const filtered = draft.filter(
+          [
-          edge =>
+            ...connectedEdges.map(edge => ({ type: 'remove', edge })),
-            !connectedEdges.find(
+            ...reconnectedEdges.map(edge => ({ type: 'add', edge })),
-              connectedEdge => connectedEdge.id === edge.id,
+          ],
-            ),
+          nodesWithNewNode,
         )
-
+      const newNodes = produce(nodesWithNewNode, (draft) => {
-        return filtered
+        draft.forEach((node) => {
+          if (nodesConnectedSourceOrTargetHandleIdsMap[node.id]) {
+            node.data = {
+              ...node.data,
+              ...nodesConnectedSourceOrTargetHandleIdsMap[node.id],
+            }
+          }
+        })
       })
-      setEdges(newEdges)
+      setNodes(newNodes)
+      const remainingEdges = edges.filter(
+        edge =>
+          !connectedEdges.find(
+            connectedEdge => connectedEdge.id === edge.id,
+          ),
+      )
+      setEdges([...remainingEdges, ...reconnectedEdges])
       if (nodeType === BlockEnum.TriggerWebhook) {
         handleSyncWorkflowDraft(true, true, {
           onSuccess: () => autoGenerateWebhookUrl(newCurrentNode.id),