[Bug] Fix input link slots (#3349)

Co-authored-by: github-actions <github-actions@github.com>

browser_tests/assets/input_order_swap.json

@@ -0,0 +1,74 @@
+{
+  "id": "51b9b184-770d-40ac-a478-8cc31667ff23",
+  "revision": 0,
+  "last_node_id": 2,
+  "last_link_id": 1,
+  "nodes": [
+    {
+      "id": 1,
+      "type": "CLIPTextEncode",
+      "pos": [904, 466],
+      "size": [400, 200],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "text",
+          "type": "STRING",
+          "widget": {
+            "name": "text"
+          },
+          "link": 1
+        },
+        {
+          "name": "clip",
+          "type": "CLIP",
+          "link": null
+        }
+      ],
+      "outputs": [
+        {
+          "name": "CONDITIONING",
+          "type": "CONDITIONING",
+          "links": null
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "CLIPTextEncode"
+      },
+      "widgets_values": [""]
+    },
+    {
+      "id": 2,
+      "type": "PrimitiveString",
+      "pos": [556.8589477539062, 472.94342041015625],
+      "size": [315, 58],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "STRING",
+          "type": "STRING",
+          "links": [1]
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PrimitiveString"
+      },
+      "widgets_values": ["foo"]
+    }
+  ],
+  "links": [[1, 2, 0, 1, 0, "STRING"]],
+  "groups": [],
+  "config": {},
+  "extra": {
+    "ds": {
+      "scale": 1.7715610000000013,
+      "offset": [-388.521484375, -162.31336975097656]
+    }
+  },
+  "version": 0.4
+}

browser_tests/tests/graph.spec.ts

@@ -0,0 +1,16 @@
+import { expect } from '@playwright/test'
+
+import { comfyPageFixture as test } from '../fixtures/ComfyPage'
+
+test.describe('Graph', () => {
+  // Should be able to fix link input slot index after swap the input order
+  // Ref: https://github.com/Comfy-Org/ComfyUI_frontend/issues/3348
+  test('Fix link input slots', async ({ comfyPage }) => {
+    await comfyPage.loadWorkflow('input_order_swap')
+    expect(
+      await comfyPage.page.evaluate(() => {
+        return window['app'].graph.links.get(1)?.target_slot
+      })
+    ).toBe(1)
+  })
+})

src/scripts/app.ts

@@ -47,7 +47,11 @@ import type { ComfyExtension, MissingNodeType } from '@/types/comfy'
 import { ExtensionManager } from '@/types/extensionTypes'
 import { ColorAdjustOptions, adjustColor } from '@/utils/colorUtil'
 import { graphToPrompt } from '@/utils/executionUtil'
-import { executeWidgetsCallback, isImageNode } from '@/utils/litegraphUtil'
+import {
+  executeWidgetsCallback,
+  fixLinkInputSlots,
+  isImageNode
+} from '@/utils/litegraphUtil'
 import {
   findLegacyRerouteNodes,
   noNativeReroutes
@@ -705,7 +709,9 @@ export class ComfyApp {
   #addAfterConfigureHandler() {
     const app = this
     const onConfigure = app.graph.onConfigure
-    app.graph.onConfigure = function (...args) {
+    app.graph.onConfigure = function (this: LGraph, ...args) {
+      fixLinkInputSlots(this)
+
       // Fire callbacks before the onConfigure, this is used by widget inputs to setup the config
       for (const node of app.graph.nodes) {
         node.onGraphConfigured?.()

src/utils/litegraphUtil.ts

@@ -1,4 +1,4 @@
-import type { ColorOption } from '@comfyorg/litegraph'
+import type { ColorOption, LGraph } from '@comfyorg/litegraph'
 import { LGraphGroup, LGraphNode, isColorable } from '@comfyorg/litegraph'
 import type {
   IComboWidget,
@@ -104,3 +104,43 @@ export function migrateWidgetsValues<TWidgetValue>(
   }
   return widgetsValues
 }
+
+/**
+ * Fix link input slots after loading a graph. Because the node inputs follows
+ * the node definition after 1.16, the node inputs array from previous versions,
+ * might get added items in the middle, which can cause shift to link's slot index.
+ * For example, the node inputs definition is:
+ * "required": {
+ *   "input1": ["INT", { forceInput: true }],
+ *   "input2": ["MODEL", { forceInput: false }],
+ *   "input3": ["MODEL", { forceInput: false }]
+ * }
+ *
+ * previously node inputs array was:
+ * [{name: 'input2'}, {name: 'input3'}, {name: 'input1'}]
+ * because input1 is created as widget first, then convert to input socket after
+ * input 2 and 3.
+ *
+ * Now, the node inputs array just follows the definition order:
+ * [{name: 'input1'}, {name: 'input2'}, {name: 'input3'}]
+ *
+ * We need to update the slot index of corresponding links to match the new
+ * node inputs array order.
+ *
+ * Ref: https://github.com/Comfy-Org/ComfyUI_frontend/issues/3348
+ *
+ * @param graph - The graph to fix links for.
+ */
+export function fixLinkInputSlots(graph: LGraph) {
+  for (const node of graph.nodes) {
+    for (const [inputIndex, input] of node.inputs.entries()) {
+      const linkId = input.link
+      if (!linkId) continue
+
+      const link = graph.links.get(linkId)
+      if (!link) continue
+
+      link.target_slot = inputIndex
+    }
+  }
+}