node connect

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

@@ -103,15 +103,22 @@ export const useEdgesInteractions = () => {
     if (currentEdgeIndex < 0)
       return
     const currentEdge = edges[currentEdgeIndex]
-    const newNodes = produce(getNodes(), (draft: Node[]) => {
-      const sourceNode = draft.find(node => node.id === currentEdge?.source)
-      const targetNode = draft.find(node => node.id === currentEdge?.target)
-
-      if (sourceNode)
-        sourceNode.data._connectedSourceHandleIds = sourceNode.data._connectedSourceHandleIds?.filter(handleId => handleId !== currentEdge.sourceHandle)
-
-      if (targetNode)
-        targetNode.data._connectedTargetHandleIds = targetNode.data._connectedTargetHandleIds?.filter(handleId => handleId !== currentEdge.targetHandle)
+    const nodes = getNodes()
+    const nodesConnectedSourceOrTargetHandleIdsMap = getNodesConnectedSourceOrTargetHandleIdsMap(
+      [
+        { type: 'remove', edge: currentEdge },
+      ],
+      nodes,
+    )
+    const newNodes = produce(nodes, (draft: Node[]) => {
+      draft.forEach((node) => {
+        if (nodesConnectedSourceOrTargetHandleIdsMap[node.id]) {
+          node.data = {
+            ...node.data,
+            ...nodesConnectedSourceOrTargetHandleIdsMap[node.id],
+          }
+        }
+      })
     })
     setNodes(newNodes)
     const newEdges = produce(edges, (draft) => {

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

@@ -311,7 +311,21 @@ export const useNodesInteractions = () => {
       setEdges,
     } = store.getState()
     const nodes = getNodes()
-    const needDeleteEdges = edges.filter(edge => (edge.source === source && edge.sourceHandle === sourceHandle) || (edge.target === target && edge.targetHandle === targetHandle))
+    const needDeleteEdges = edges.filter((edge) => {
+      if (edge.source === source) {
+        if (edge.sourceHandle)
+          return edge.sourceHandle === sourceHandle
+        else
+          return true
+      }
+      if (edge.target === target) {
+        if (edge.targetHandle)
+          return edge.targetHandle === targetHandle
+        else
+          return true
+      }
+      return false
+    })
     const needDeleteEdgesIds = needDeleteEdges.map(edge => edge.id)
     const newEdge = {
       id: `${source}-${target}`,

web/app/components/workflow/utils.ts

@@ -69,6 +69,12 @@ export const initialEdges = (edges: Edge[], nodes: Node[]) => {
   return edges.map((edge) => {
     edge.type = 'custom'
 
+    if (!edge.sourceHandle)
+      edge.sourceHandle = 'source'
+
+    if (!edge.targetHandle)
+      edge.targetHandle = 'target'
+
     if (!edge.data?.sourceType) {
       edge.data = {
         ...edge.data,
@@ -146,7 +152,7 @@ export const getNodesConnectedSourceOrTargetHandleIdsMap = (changes: ConnectedSo
     }
     const targetNode = nodes.find(node => node.id === edge.target)!
     nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id] = nodesConnectedSourceOrTargetHandleIdsMap[targetNode.id] || {
-      _connectedSourceHandleIds: sourceNode?.data._connectedSourceHandleIds || [],
+      _connectedSourceHandleIds: targetNode?.data._connectedSourceHandleIds || [],
       _connectedTargetHandleIds: targetNode?.data._connectedTargetHandleIds || [],
     }