[Bug] Fix broken output link id during reroute migration (#3244)

src/utils/migration/migrateReroute.ts

@@ -1,3 +1,5 @@
+import _ from 'lodash'
+
 import type {
   ComfyLink,
   ComfyNode,
@@ -67,10 +69,15 @@ export function createNewLinks(
   workflow: WorkflowJSON04,
   rerouteMap: Map<NodeId, RerouteEntry>
 ): {
+  nodes: ComfyNode[]
   reroutes: Reroute[]
   links: ComfyLink[]
   linkExtensions: LinkExtension[]
 } {
+  const nodeById = _.keyBy(
+    workflow.nodes.filter((node) => node.type !== 'Reroute').map(_.cloneDeep),
+    'id'
+  )
   const links: ComfyLink[] = []
   const linkExtensions: LinkExtension[] = []
 
@@ -153,6 +160,23 @@ export function createNewLinks(
 
           validLinkExtensions.push(linkExtension)
           chainedReroutes.forEach((reroute) => validReroutes.add(reroute))
+
+          // Update source node's output slot's link ids to point to the new link.
+          const sourceNode = nodeById[sourceNodeId]
+          if (!sourceNode) {
+            throw new Error(
+              `Corrupted workflow: Source node ${sourceNodeId} not found`
+            )
+          }
+          const outputSlot = sourceNode.outputs?.[sourceSlot]
+          if (!outputSlot) {
+            throw new Error(
+              `Corrupted workflow: Output slot ${sourceSlot} not found`
+            )
+          }
+          outputSlot.links = outputSlot.links?.map((l) =>
+            l === rerouteInputLink[0] ? linkId : l
+          )
         }
         entry = undefined
       }
@@ -160,6 +184,7 @@ export function createNewLinks(
   }
 
   return {
+    nodes: Object.values(nodeById),
     links,
     linkExtensions: validLinkExtensions,
     reroutes: Array.from(validReroutes)
@@ -192,16 +217,14 @@ export const migrateLegacyRerouteNodes = (
   const rerouteMap = createReroutePoints(legacyRerouteNodes)
 
   // Create new links and link extensions
-  const { links, linkExtensions, reroutes } = createNewLinks(
+  const { nodes, links, linkExtensions, reroutes } = createNewLinks(
     workflow,
     rerouteMap
   )
 
   // Update the workflow
   newWorkflow.links = links
-  newWorkflow.nodes = newWorkflow.nodes.filter(
-    (node) => node.type !== 'Reroute'
-  )
+  newWorkflow.nodes = nodes
   newWorkflow.extra.reroutes = reroutes
   newWorkflow.extra.linkExtensions = linkExtensions
 

tests-ui/tests/utils/migration/workflows/reroute/native/branching.json

@@ -2,6 +2,51 @@
   "last_node_id": 27,
   "last_link_id": 34,
   "nodes": [
+    {
+      "id": 4,
+      "type": "CheckpointLoaderSimple",
+      "pos": [
+        47.948699951171875,
+        239.2628173828125
+      ],
+      "size": [
+        315,
+        98
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "slot_index": 0,
+          "links": []
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "slot_index": 1,
+          "links": []
+        },
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "slot_index": 2,
+          "links": [
+            13,
+            21
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CheckpointLoaderSimple"
+      },
+      "widgets_values": [
+        "v1-5-pruned-emaonly.safetensors"
+      ]
+    },
     {
       "id": 12,
       "type": "VAEDecode",
@@ -40,51 +85,6 @@
       },
       "widgets_values": []
     },
-    {
-      "id": 4,
-      "type": "CheckpointLoaderSimple",
-      "pos": [
-        47.948699951171875,
-        239.2628173828125
-      ],
-      "size": [
-        315,
-        98
-      ],
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "MODEL",
-          "type": "MODEL",
-          "slot_index": 0,
-          "links": []
-        },
-        {
-          "name": "CLIP",
-          "type": "CLIP",
-          "slot_index": 1,
-          "links": []
-        },
-        {
-          "name": "VAE",
-          "type": "VAE",
-          "slot_index": 2,
-          "links": [
-            13,
-            31
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "CheckpointLoaderSimple"
-      },
-      "widgets_values": [
-        "v1-5-pruned-emaonly.safetensors"
-      ]
-    },
     {
       "id": 26,
       "type": "VAEDecode",

tests-ui/tests/utils/migration/workflows/reroute/native/single_connected.json

@@ -2,6 +2,50 @@
   "last_node_id": 24,
   "last_link_id": 30,
   "nodes": [
+    {
+      "id": 4,
+      "type": "CheckpointLoaderSimple",
+      "pos": [
+        160,
+        240
+      ],
+      "size": [
+        315,
+        98
+      ],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "MODEL",
+          "type": "MODEL",
+          "slot_index": 0,
+          "links": []
+        },
+        {
+          "name": "CLIP",
+          "type": "CLIP",
+          "slot_index": 1,
+          "links": []
+        },
+        {
+          "name": "VAE",
+          "type": "VAE",
+          "slot_index": 2,
+          "links": [
+            21
+          ]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CheckpointLoaderSimple"
+      },
+      "widgets_values": [
+        "v1-5-pruned-emaonly.safetensors"
+      ]
+    },
     {
       "id": 12,
       "type": "VAEDecode",
@@ -39,50 +83,6 @@
         "Node name for S&R": "VAEDecode"
       },
       "widgets_values": []
-    },
-    {
-      "id": 4,
-      "type": "CheckpointLoaderSimple",
-      "pos": [
-        160,
-        240
-      ],
-      "size": [
-        315,
-        98
-      ],
-      "flags": {},
-      "order": 0,
-      "mode": 0,
-      "inputs": [],
-      "outputs": [
-        {
-          "name": "MODEL",
-          "type": "MODEL",
-          "slot_index": 0,
-          "links": []
-        },
-        {
-          "name": "CLIP",
-          "type": "CLIP",
-          "slot_index": 1,
-          "links": []
-        },
-        {
-          "name": "VAE",
-          "type": "VAE",
-          "slot_index": 2,
-          "links": [
-            13
-          ]
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "CheckpointLoaderSimple"
-      },
-      "widgets_values": [
-        "v1-5-pruned-emaonly.safetensors"
-      ]
     }
   ],
   "links": [