[Bug] Fix converted widget compression on export (#3354)

src/utils/executionUtil.ts

@@ -6,6 +6,8 @@ import type {
   ComfyWorkflowJSON
 } from '@/schemas/comfyWorkflowSchema'
 
+import { compressWidgetInputSlots } from './litegraphUtil'
+
 /**
  * Converts the current graph workflow for sending to the API.
  * Note: Node widgets are updated before serialization to prepare queueing.
@@ -38,12 +40,7 @@ export const graphToPrompt = async (
     }
   }
 
-  // Remove all unconnected widget input slots
-  for (const node of workflow.nodes) {
-    node.inputs = node.inputs?.filter(
-      (input) => !(input.widget && input.link === null)
-    )
-  }
+  compressWidgetInputSlots(workflow)
 
   const output: ComfyApiWorkflow = {}
   // Process nodes in order of execution

src/utils/litegraphUtil.ts

@@ -1,5 +1,6 @@
 import type { ColorOption, LGraph } from '@comfyorg/litegraph'
 import { LGraphGroup, LGraphNode, isColorable } from '@comfyorg/litegraph'
+import type { ISerialisedGraph } from '@comfyorg/litegraph/dist/types/serialisation'
 import type {
   IComboWidget,
   IWidget
@@ -144,3 +145,26 @@ export function fixLinkInputSlots(graph: LGraph) {
     }
   }
 }
+
+/**
+ * Compress widget input slots by removing all unconnected widget input slots.
+ * This should match the serialization format of legacy widget conversion.
+ *
+ * @param graph - The graph to compress widget input slots for.
+ */
+export function compressWidgetInputSlots(graph: ISerialisedGraph) {
+  for (const node of graph.nodes) {
+    node.inputs = node.inputs?.filter(
+      (input) => !(input.widget && input.link === null)
+    )
+
+    for (const [inputIndex, input] of node.inputs?.entries() ?? []) {
+      if (input.link) {
+        const link = graph.links.find((link) => link[0] === input.link)
+        if (link) {
+          link[4] = inputIndex
+        }
+      }
+    }
+  }
+}

tests-ui/tests/utils/litegraphUtil.test.ts

@@ -1,8 +1,12 @@
+import { ISerialisedGraph } from '@comfyorg/litegraph/dist/types/serialisation'
 import type { IWidget } from '@comfyorg/litegraph/dist/types/widgets'
 import { describe, expect, it } from 'vitest'
 
 import type { InputSpec } from '@/schemas/nodeDef/nodeDefSchemaV2'
-import { migrateWidgetsValues } from '@/utils/litegraphUtil'
+import {
+  compressWidgetInputSlots,
+  migrateWidgetsValues
+} from '@/utils/litegraphUtil'
 
 describe('migrateWidgetsValues', () => {
   it('should remove widget values for forceInput inputs', () => {
@@ -86,3 +90,84 @@ describe('migrateWidgetsValues', () => {
     expect(result).toEqual(['first value', 'last value'])
   })
 })
+
+describe('compressWidgetInputSlots', () => {
+  it('should remove unconnected widget input slots', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [
+        {
+          id: 1,
+          type: 'foo',
+          pos: [0, 0],
+          size: [100, 100],
+          flags: {},
+          order: 0,
+          mode: 0,
+          inputs: [
+            { widget: { name: 'foo' }, link: null, type: 'INT', name: 'foo' },
+            { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+            { widget: { name: 'baz' }, link: null, type: 'INT', name: 'baz' }
+          ],
+          outputs: []
+        }
+      ],
+      links: [[2, 1, 0, 1, 0, 'INT']]
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes[0].inputs).toEqual([
+      { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' }
+    ])
+  })
+
+  it('should update link target slots correctly', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [
+        {
+          id: 1,
+          type: 'foo',
+          pos: [0, 0],
+          size: [100, 100],
+          flags: {},
+          order: 0,
+          mode: 0,
+          inputs: [
+            { widget: { name: 'foo' }, link: null, type: 'INT', name: 'foo' },
+            { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+            { widget: { name: 'baz' }, link: 3, type: 'INT', name: 'baz' }
+          ],
+          outputs: []
+        }
+      ],
+      links: [
+        [2, 1, 0, 1, 1, 'INT'],
+        [3, 1, 0, 1, 2, 'INT']
+      ]
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes[0].inputs).toEqual([
+      { widget: { name: 'bar' }, link: 2, type: 'INT', name: 'bar' },
+      { widget: { name: 'baz' }, link: 3, type: 'INT', name: 'baz' }
+    ])
+
+    expect(graph.links).toEqual([
+      [2, 1, 0, 1, 0, 'INT'],
+      [3, 1, 0, 1, 1, 'INT']
+    ])
+  })
+
+  it('should handle graphs with no nodes gracefully', () => {
+    const graph: ISerialisedGraph = {
+      nodes: [],
+      links: []
+    } as unknown as ISerialisedGraph
+
+    compressWidgetInputSlots(graph)
+
+    expect(graph.nodes).toEqual([])
+    expect(graph.links).toEqual([])
+  })
+})