Merge branch 'main' into glary/fix-primitive-refresh-combo-options

Per repo conventions (no function expressions when a declaration works).

browser_tests/assets/primitive/primitive_combo_sampler_name.json

@@ -0,0 +1,58 @@
+{
+  "last_node_id": 2,
+  "last_link_id": 1,
+  "nodes": [
+    {
+      "id": 2,
+      "type": "KSampler",
+      "pos": [521, 41],
+      "size": [315, 262],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [
+        { "name": "model", "type": "MODEL", "link": null },
+        { "name": "positive", "type": "CONDITIONING", "link": null },
+        { "name": "negative", "type": "CONDITIONING", "link": null },
+        { "name": "latent_image", "type": "LATENT", "link": null },
+        {
+          "name": "sampler_name",
+          "type": "COMBO",
+          "link": 1,
+          "widget": { "name": "sampler_name" }
+        }
+      ],
+      "outputs": [
+        { "name": "LATENT", "type": "LATENT", "links": null, "shape": 3 }
+      ],
+      "properties": { "Node name for S&R": "KSampler" },
+      "widgets_values": [0, "randomize", 20, 8, "euler", "normal", 1]
+    },
+    {
+      "id": 1,
+      "type": "PrimitiveNode",
+      "pos": [15, 46],
+      "size": [400, 110],
+      "flags": {},
+      "order": 0,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "COMBO",
+          "type": "COMBO",
+          "links": [1],
+          "slot_index": 0,
+          "widget": { "name": "sampler_name" }
+        }
+      ],
+      "properties": { "Run widget replace on values": false },
+      "widgets_values": ["euler", "fixed"]
+    }
+  ],
+  "links": [[1, 1, 0, 2, 4, "COMBO"]],
+  "groups": [],
+  "config": {},
+  "extra": { "ds": { "scale": 1, "offset": [0, 0] } },
+  "version": 0.4
+}

browser_tests/tests/primitiveNode.spec.ts

@@ -59,6 +59,51 @@ test.describe('Primitive Node', { tag: ['@screenshot', '@node'] }, () => {
     )
   })
 
+  test('Preserves combo options on the primitive after pressing R to refresh', async ({
+    comfyPage
+  }) => {
+    async function getPrimitiveComboState() {
+      return comfyPage.page.evaluate(() => {
+        const primitive = window.app!.graph!.nodes.find(
+          (node) => node.type === 'PrimitiveNode'
+        )
+        const widget = primitive?.widgets?.[0]
+        const values = widget?.options?.values
+        return {
+          isArray: Array.isArray(values),
+          length: Array.isArray(values) ? values.length : 0,
+          includesEuler: Array.isArray(values)
+            ? values.includes('euler')
+            : false,
+          value: widget?.value
+        }
+      })
+    }
+
+    await comfyPage.workflow.loadWorkflow(
+      'primitive/primitive_combo_sampler_name'
+    )
+
+    await expect.poll(getPrimitiveComboState).toMatchObject({
+      isArray: true,
+      includesEuler: true,
+      value: 'euler'
+    })
+    const before = await getPrimitiveComboState()
+    expect(before.length).toBeGreaterThan(0)
+
+    await comfyPage.canvas.click()
+    await comfyPage.page.keyboard.press('r')
+
+    await expect.poll(getPrimitiveComboState).toMatchObject({
+      isArray: true,
+      includesEuler: true,
+      value: 'euler'
+    })
+    const after = await getPrimitiveComboState()
+    expect(after.length).toBeGreaterThan(0)
+  })
+
   test('Report missing nodes when connect to missing node', async ({
     comfyPage
   }) => {

src/extensions/core/widgetInputs.test.ts

@@ -0,0 +1,271 @@
+import { fromPartial } from '@total-typescript/shoehorn'
+import { describe, expect, it, vi } from 'vitest'
+
+import type { IWidgetLocator } from '@/lib/litegraph/src/interfaces'
+import type {
+  INodeInputSlot,
+  INodeOutputSlot,
+  LGraph,
+  LGraphNode,
+  LLink
+} from '@/lib/litegraph/src/litegraph'
+import type { IBaseWidget } from '@/lib/litegraph/src/types/widgets'
+import type { ComfyNodeDef, InputSpec } from '@/schemas/nodeDefSchema'
+import { GET_CONFIG } from '@/services/litegraphService'
+
+type SlotWidget = IWidgetLocator & {
+  [GET_CONFIG]: () => InputSpec | undefined
+}
+
+vi.mock('@/scripts/app', () => ({
+  app: {
+    canvas: { graph_mouse: [0, 0] },
+    registerExtension: vi.fn()
+  }
+}))
+
+import { PrimitiveNode } from './widgetInputs'
+
+const TARGET_NODE_TYPE = 'KSampler'
+const TARGET_INPUT_NAME = 'sampler_name'
+const ORIGINAL_OPTIONS = ['euler', 'euler_ancestral', 'heun']
+const FRESH_OPTIONS = ['euler', 'euler_ancestral', 'heun', 'lcm']
+
+function createComboWidget(
+  value: IBaseWidget['value'],
+  values: (string | number)[]
+): IBaseWidget {
+  return fromPartial<IBaseWidget>({
+    type: 'combo',
+    name: 'value',
+    value,
+    options: { values },
+    callback: vi.fn()
+  })
+}
+
+function createSlotWidget(config?: InputSpec): SlotWidget | undefined {
+  if (!config) return undefined
+  return fromPartial<SlotWidget>({
+    name: TARGET_INPUT_NAME,
+    [GET_CONFIG]: () => config
+  })
+}
+
+type TargetNode = Pick<LGraphNode, 'id' | 'type' | 'inputs'>
+
+function createTargetNode(
+  inputWidgetName: string | undefined = TARGET_INPUT_NAME
+): TargetNode {
+  return fromPartial<TargetNode>({
+    id: 7,
+    type: TARGET_NODE_TYPE,
+    inputs: [
+      fromPartial<INodeInputSlot>({
+        widget: inputWidgetName ? { name: inputWidgetName } : undefined
+      })
+    ]
+  })
+}
+
+interface PrimitiveStub {
+  graph?: LGraph
+  outputs: INodeOutputSlot[]
+  widgets: IBaseWidget[]
+}
+
+function createPrimitiveStub(options: {
+  widget: IBaseWidget
+  slotWidget?: SlotWidget
+  targetNode?: TargetNode | null
+  hasLink?: boolean
+}): PrimitiveStub {
+  const { widget, slotWidget, hasLink = true } = options
+  const targetNode =
+    options.targetNode === undefined ? createTargetNode() : options.targetNode
+  const link = hasLink
+    ? fromPartial<LLink>({
+        target_id: targetNode?.id ?? 7,
+        target_slot: 0
+      })
+    : undefined
+
+  const stub = Object.create(PrimitiveNode.prototype) as PrimitiveStub
+  stub.graph = fromPartial<LGraph>({
+    links: link ? { 1: link } : {},
+    getNodeById: vi.fn(() => targetNode ?? null)
+  })
+  stub.outputs = [
+    fromPartial<INodeOutputSlot>({
+      links: hasLink ? [1] : [],
+      widget: slotWidget
+    })
+  ]
+  stub.widgets = [widget]
+  return stub
+}
+
+function refreshOnStub(
+  stub: PrimitiveStub,
+  defs?: Record<string, ComfyNodeDef>
+) {
+  ;(stub as unknown as PrimitiveNode).refreshComboInNode(defs)
+}
+
+function defsWithCombo(
+  values: (string | number)[]
+): Record<string, ComfyNodeDef> {
+  return {
+    [TARGET_NODE_TYPE]: fromPartial<ComfyNodeDef>({
+      input: { required: { [TARGET_INPUT_NAME]: [values, {}] } }
+    })
+  }
+}
+
+function defsWithV2Combo(
+  values: (string | number)[]
+): Record<string, ComfyNodeDef> {
+  return {
+    [TARGET_NODE_TYPE]: fromPartial<ComfyNodeDef>({
+      input: {
+        required: { [TARGET_INPUT_NAME]: ['COMBO', { options: values }] }
+      }
+    })
+  }
+}
+
+describe('PrimitiveNode.refreshComboInNode', () => {
+  it('updates combo options from the freshly passed defs', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithCombo(FRESH_OPTIONS))
+
+    expect(widget.options.values).toEqual(FRESH_OPTIONS)
+  })
+
+  it('preserves existing options when defs lookup yields nothing and slot config is missing', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({ widget, slotWidget: undefined })
+
+    refreshOnStub(stub, {})
+
+    expect(widget.options.values).toEqual(ORIGINAL_OPTIONS)
+    expect(widget.value).toBe('euler')
+  })
+
+  it('preserves existing options when slot widget GET_CONFIG returns undefined', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const slotWidget = fromPartial<SlotWidget>({
+      name: TARGET_INPUT_NAME,
+      [GET_CONFIG]: () => undefined
+    })
+    const stub = createPrimitiveStub({ widget, slotWidget })
+
+    refreshOnStub(stub)
+
+    expect(widget.options.values).toEqual(ORIGINAL_OPTIONS)
+  })
+
+  it('falls back to slot widget config when defs do not include the target node', () => {
+    const widget = createComboWidget('euler', [])
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, {})
+
+    expect(widget.options.values).toEqual(ORIGINAL_OPTIONS)
+  })
+
+  it('resets value and fires callback when current value is no longer in the new options', () => {
+    const widget = createComboWidget('removed_value', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithCombo(FRESH_OPTIONS))
+
+    expect(widget.value).toBe('euler')
+    expect(widget.callback).toHaveBeenCalledWith('euler')
+  })
+
+  it('does not change value or fire callback when current value is still valid', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithCombo(FRESH_OPTIONS))
+
+    expect(widget.value).toBe('euler')
+    expect(widget.callback).not.toHaveBeenCalled()
+  })
+
+  it('updates combo options from V2 combo specs in defs', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithV2Combo(FRESH_OPTIONS))
+
+    expect(widget.options.values).toEqual(FRESH_OPTIONS)
+  })
+
+  it('propagates an empty option list when defs return one', () => {
+    const widget = createComboWidget('euler', ORIGINAL_OPTIONS)
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([ORIGINAL_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithCombo([]))
+
+    expect(widget.options.values).toEqual([])
+    expect(widget.value).toBeUndefined()
+  })
+
+  it('falls back to input.name when the target input has no widget locator', () => {
+    const widget = createComboWidget('euler', [])
+    const targetNode = fromPartial<TargetNode>({
+      id: 9,
+      type: TARGET_NODE_TYPE,
+      inputs: [
+        fromPartial<INodeInputSlot>({
+          name: TARGET_INPUT_NAME
+        })
+      ]
+    })
+    const stub = createPrimitiveStub({ widget, targetNode })
+
+    refreshOnStub(stub, defsWithCombo(FRESH_OPTIONS))
+
+    expect(widget.options.values).toEqual(FRESH_OPTIONS)
+  })
+
+  it('skips non-combo widgets', () => {
+    const widget = fromPartial<IBaseWidget>({
+      type: 'string',
+      name: 'value',
+      value: 'hello',
+      options: { values: ORIGINAL_OPTIONS },
+      callback: vi.fn()
+    })
+    const stub = createPrimitiveStub({
+      widget,
+      slotWidget: createSlotWidget([FRESH_OPTIONS, {}])
+    })
+
+    refreshOnStub(stub, defsWithCombo(FRESH_OPTIONS))
+
+    expect(widget.options.values).toEqual(ORIGINAL_OPTIONS)
+  })
+})

src/extensions/core/widgetInputs.ts

@@ -14,6 +14,10 @@ import type {
 } from '@/lib/litegraph/src/types/widgets'
 import { assetService } from '@/platform/assets/services/assetService'
 import { createAssetWidget } from '@/platform/assets/utils/createAssetWidget'
+import {
+  getComboSpecComboOptions,
+  isComboInputSpec
+} from '@/schemas/nodeDefSchema'
 import type { ComfyNodeDef, InputSpec } from '@/schemas/nodeDefSchema'
 import { app } from '@/scripts/app'
 import {
@@ -56,21 +60,64 @@ export class PrimitiveNode extends LGraphNode {
     applyFirstWidgetValueToGraph(this, extraLinks, () => v)
   }
 
-  override refreshComboInNode() {
+  override refreshComboInNode(defs?: Record<string, ComfyNodeDef>) {
     const widget = this.widgets?.[0]
-    if (widget?.type === 'combo') {
-      // @ts-expect-error fixme ts strict error
-      widget.options.values = this.outputs[0].widget[GET_CONFIG]()[0]
+    if (widget?.type !== 'combo') return
 
-      // @ts-expect-error fixme ts strict error
-      if (!widget.options.values.includes(widget.value as string)) {
-        // @ts-expect-error fixme ts strict error
-        widget.value = widget.options.values[0]
-        widget.callback?.(widget.value)
-      }
+    const newValues = this._resolveComboValues(defs)
+    if (newValues === undefined) return
+
+    widget.options.values = newValues
+
+    if (!newValues.includes(widget.value as string | number)) {
+      widget.value = newValues[0]
+      widget.callback?.(widget.value)
     }
   }
 
+  private _resolveComboValues(
+    defs?: Record<string, ComfyNodeDef>
+  ): (string | number)[] | undefined {
+    const fromDefs = defs ? this._comboValuesFromDefs(defs) : undefined
+    if (fromDefs !== undefined) return fromDefs
+
+    const slotWidget = this.outputs?.[0]?.widget
+    const config = (
+      slotWidget?.[GET_CONFIG] as (() => InputSpec) | undefined
+    )?.()
+    if (!config || !isComboInputSpec(config)) return undefined
+
+    return getComboSpecComboOptions(config)
+  }
+
+  private _comboValuesFromDefs(
+    defs: Record<string, ComfyNodeDef>
+  ): (string | number)[] | undefined {
+    const link = this._getFirstOutputLink()
+    if (!link) return undefined
+
+    const targetNode = this.graph?.getNodeById(link.target_id)
+    const targetType = targetNode?.type
+    if (!targetType) return undefined
+
+    const targetInput = targetNode?.inputs?.[link.target_slot]
+    const inputName = targetInput?.widget?.name ?? targetInput?.name
+    if (!inputName) return undefined
+
+    const def = defs[targetType]
+    const inputSpec =
+      def?.input?.required?.[inputName] ?? def?.input?.optional?.[inputName]
+    if (!inputSpec || !isComboInputSpec(inputSpec)) return undefined
+
+    return getComboSpecComboOptions(inputSpec)
+  }
+
+  private _getFirstOutputLink(): LLink | undefined {
+    const linkId = this.outputs?.[0]?.links?.[0]
+    if (linkId == null || !this.graph) return undefined
+    return this.graph.links[linkId] ?? undefined
+  }
+
   override onAfterGraphConfigured() {
     if (this.outputs[0].links?.length && !this.widgets?.length) {
       this._onFirstConnection()