From 1eb072fd4366b61d2f1abe9f9c9c819d7408f9e6 Mon Sep 17 00:00:00 2001
From: JonSnow <zzzhaichao@gmail.com>
Date: Fri, 28 Feb 2025 19:33:47 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20the=20edges=20between=20the=20nodes=20in?=
 =?UTF-8?q?side=20the=20copied=20iteration=20node=20are=E2=80=A6=20(#12692?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../workflow/hooks/use-nodes-interactions.ts  | 32 ++++++++++++++++++-
 .../nodes/iteration/use-interactions.ts       | 11 +++++--
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/web/app/components/workflow/hooks/use-nodes-interactions.ts b/web/app/components/workflow/hooks/use-nodes-interactions.ts
index 8962333311..b445cdfdc0 100644
--- a/web/app/components/workflow/hooks/use-nodes-interactions.ts
+++ b/web/app/components/workflow/hooks/use-nodes-interactions.ts
@@ -1117,9 +1117,12 @@ export const useNodesInteractions = () => {
     const {
       getNodes,
       setNodes,
+      edges,
+      setEdges,
     } = store.getState()
 
     const nodesToPaste: Node[] = []
+    const edgesToPaste: Edge[] = []
     const nodes = getNodes()
 
     if (clipboardElements.length) {
@@ -1128,6 +1131,7 @@ export const useNodesInteractions = () => {
       const currentPosition = screenToFlowPosition({ x: mousePosition.pageX, y: mousePosition.pageY })
       const offsetX = currentPosition.x - x
       const offsetY = currentPosition.y - y
+      let idMapping: Record<string, string> = {}
       clipboardElements.forEach((nodeToPaste, index) => {
         const nodeType = nodeToPaste.data.type
 
@@ -1159,7 +1163,13 @@ export const useNodesInteractions = () => {
           newIterationStartNode!.parentId = newNode.id;
           (newNode.data as IterationNodeType).start_node_id = newIterationStartNode!.id
 
-          newChildren = handleNodeIterationChildrenCopy(nodeToPaste.id, newNode.id)
+          const oldIterationStartNode = nodes
+            .find(n => n.parentId === nodeToPaste.id && n.type === CUSTOM_ITERATION_START_NODE)
+          idMapping[oldIterationStartNode!.id] = newIterationStartNode!.id
+
+          const { copyChildren, newIdMapping } = handleNodeIterationChildrenCopy(nodeToPaste.id, newNode.id, idMapping)
+          newChildren = copyChildren
+          idMapping = newIdMapping
           newChildren.forEach((child) => {
             newNode.data._children?.push(child.id)
           })
@@ -1172,7 +1182,27 @@ export const useNodesInteractions = () => {
           nodesToPaste.push(...newChildren)
       })
 
+      edges.forEach((edge) => {
+        const sourceId = idMapping[edge.source]
+        const targetId = idMapping[edge.target]
+
+        if (sourceId && targetId) {
+          const newEdge: Edge = {
+            ...edge,
+            id: `${sourceId}-${edge.sourceHandle}-${targetId}-${edge.targetHandle}`,
+            source: sourceId,
+            target: targetId,
+            data: {
+              ...edge.data,
+              _connectedNodeIsSelected: false,
+            },
+          }
+          edgesToPaste.push(newEdge)
+        }
+      })
+
       setNodes([...nodes, ...nodesToPaste])
+      setEdges([...edges, ...edgesToPaste])
       saveStateToHistory(WorkflowHistoryEvent.NodePaste)
       handleSyncWorkflowDraft()
     }
diff --git a/web/app/components/workflow/nodes/iteration/use-interactions.ts b/web/app/components/workflow/nodes/iteration/use-interactions.ts
index f8e3640cc4..c0c005bdf6 100644
--- a/web/app/components/workflow/nodes/iteration/use-interactions.ts
+++ b/web/app/components/workflow/nodes/iteration/use-interactions.ts
@@ -105,12 +105,13 @@ export const useNodeIterationInteractions = () => {
       handleNodeIterationRerender(parentId)
   }, [store, handleNodeIterationRerender])
 
-  const handleNodeIterationChildrenCopy = useCallback((nodeId: string, newNodeId: string) => {
+  const handleNodeIterationChildrenCopy = useCallback((nodeId: string, newNodeId: string, idMapping: Record<string, string>) => {
     const { getNodes } = store.getState()
     const nodes = getNodes()
     const childrenNodes = nodes.filter(n => n.parentId === nodeId && n.type !== CUSTOM_ITERATION_START_NODE)
+    const newIdMapping = { ...idMapping }
 
-    return childrenNodes.map((child, index) => {
+    const copyChildren = childrenNodes.map((child, index) => {
       const childNodeType = child.data.type as BlockEnum
       const nodesWithSameType = nodes.filter(node => node.data.type === childNodeType)
       const { newNode } = generateNewNode({
@@ -131,8 +132,14 @@ export const useNodeIterationInteractions = () => {
         zIndex: child.zIndex,
       })
       newNode.id = `${newNodeId}${newNode.id + index}`
+      newIdMapping[child.id] = newNode.id
       return newNode
     })
+
+    return {
+      copyChildren,
+      newIdMapping,
+    }
   }, [store, t])
 
   return {