refactor: extract addAfterConfigureHandler to graphConfigureUtil

Move private method from ComfyApp to a standalone exported utility function.

Eliminates unsafe 'as any' cast in tests by making the function directly importable.

Amp-Thread-ID: https://ampcode.com/threads/T-019c9bda-0293-7528-9b0a-f72444199cbe
Co-authored-by: Amp <amp@ampcode.com>

Author: DrJKL

src/scripts/app.test.ts

@@ -6,11 +6,7 @@ import type {
   LGraphNode
 } from '@/lib/litegraph/src/litegraph'
 import { ComfyApp } from './app'
-import {
-  createNode,
-  fixLinkInputSlots,
-  hasLegacyLinkInputSlotMismatch
-} from '@/utils/litegraphUtil'
+import { createNode } from '@/utils/litegraphUtil'
 import {
   pasteAudioNode,
   pasteAudioNodes,
@@ -26,9 +22,7 @@ vi.mock('@/utils/litegraphUtil', () => ({
   isImageNode: vi.fn(),
   isVideoNode: vi.fn(),
   isAudioNode: vi.fn(),
-  executeWidgetsCallback: vi.fn(),
-  fixLinkInputSlots: vi.fn(),
-  hasLegacyLinkInputSlotMismatch: vi.fn()
+  executeWidgetsCallback: vi.fn()
 }))
 
 vi.mock('@/composables/usePaste', () => ({
@@ -146,37 +140,6 @@ describe('ComfyApp', () => {
     })
   })
 
-  describe('addAfterConfigureHandler', () => {
-    function createConfigureGraph() {
-      return {
-        nodes: [],
-        onConfigure: vi.fn()
-      } as unknown as LGraph
-    }
-
-    it('runs legacy slot repair when mismatch is detected', () => {
-      vi.mocked(hasLegacyLinkInputSlotMismatch).mockReturnValue(true)
-      const graph = createConfigureGraph()
-
-      ;(app as any).addAfterConfigureHandler(graph)
-      graph.onConfigure?.({} as never)
-
-      expect(hasLegacyLinkInputSlotMismatch).toHaveBeenCalledWith(graph)
-      expect(fixLinkInputSlots).toHaveBeenCalledWith(graph)
-    })
-
-    it('skips legacy slot repair when no mismatch is present', () => {
-      vi.mocked(hasLegacyLinkInputSlotMismatch).mockReturnValue(false)
-      const graph = createConfigureGraph()
-
-      ;(app as any).addAfterConfigureHandler(graph)
-      graph.onConfigure?.({} as never)
-
-      expect(hasLegacyLinkInputSlotMismatch).toHaveBeenCalledWith(graph)
-      expect(fixLinkInputSlots).not.toHaveBeenCalled()
-    })
-  })
-
   describe('handleAudioFileList', () => {
     it('should create audio nodes and select them', async () => {
       const mockNode1 = createMockNode({ id: 1, type: 'LoadAudio' })

src/scripts/app.ts

@@ -5,8 +5,7 @@ import { reactive, unref } from 'vue'
 import { shallowRef } from 'vue'
 
 import { useCanvasPositionConversion } from '@/composables/element/useCanvasPositionConversion'
-import { layoutStore } from '@/renderer/core/layout/store/layoutStore'
-import { flushScheduledSlotLayoutSync } from '@/renderer/extensions/vueNodes/composables/useSlotElementTracking'
+import { addAfterConfigureHandler } from '@/utils/graphConfigureUtil'
 
 import { st, t } from '@/i18n'
 import type { IContextMenuValue } from '@/lib/litegraph/src/interfaces'
@@ -87,8 +86,6 @@ import {
 import {
   executeWidgetsCallback,
   createNode,
-  fixLinkInputSlots,
-  hasLegacyLinkInputSlotMismatch,
   isImageNode,
   isVideoNode
 } from '@/utils/litegraphUtil'
@@ -792,37 +789,6 @@ export class ComfyApp {
     }
   }
 
-  private addAfterConfigureHandler(graph: LGraph) {
-    const { onConfigure } = graph
-    graph.onConfigure = function (...args) {
-      // Set pending sync flag to suppress link rendering until slots are synced
-      if (LiteGraph.vueNodesMode) {
-        layoutStore.setPendingSlotSync(true)
-      }
-
-      try {
-        if (hasLegacyLinkInputSlotMismatch(this)) fixLinkInputSlots(this)
-
-        // Fire callbacks before the onConfigure, this is used by widget inputs to setup the config
-        triggerCallbackOnAllNodes(this, 'onGraphConfigured')
-
-        const r = onConfigure?.apply(this, args)
-
-        // Fire after onConfigure, used by primitives to generate widget using input nodes config
-        triggerCallbackOnAllNodes(this, 'onAfterGraphConfigured')
-
-        return r
-      } finally {
-        // Flush pending slot layout syncs to fix link alignment after undo/redo
-        // Using finally ensures links aren't permanently suppressed if an error occurs
-        if (LiteGraph.vueNodesMode) {
-          flushScheduledSlotLayoutSync()
-          app.canvas?.setDirty(true, true)
-        }
-      }
-    }
-  }
-
   /**
    * Set up the app on the page
    */
@@ -861,7 +827,7 @@ export class ComfyApp {
       }
     })
 
-    this.addAfterConfigureHandler(graph)
+    addAfterConfigureHandler(graph, () => this.canvas)
 
     this.rootGraphInternal = graph
     this.canvas = new LGraphCanvas(canvasEl, graph)

src/utils/graphConfigureUtil.test.ts

@@ -0,0 +1,70 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+
+import type { LGraph } from '@/lib/litegraph/src/litegraph'
+import {
+  fixLinkInputSlots,
+  hasLegacyLinkInputSlotMismatch
+} from '@/utils/litegraphUtil'
+
+import { addAfterConfigureHandler } from './graphConfigureUtil'
+
+vi.mock('@/utils/litegraphUtil', () => ({
+  fixLinkInputSlots: vi.fn(),
+  hasLegacyLinkInputSlotMismatch: vi.fn()
+}))
+
+vi.mock('@/utils/graphTraversalUtil', () => ({
+  triggerCallbackOnAllNodes: vi.fn()
+}))
+
+vi.mock('@/renderer/core/layout/store/layoutStore', () => ({
+  layoutStore: { setPendingSlotSync: vi.fn() }
+}))
+
+vi.mock(
+  '@/renderer/extensions/vueNodes/composables/useSlotElementTracking',
+  () => ({
+    flushScheduledSlotLayoutSync: vi.fn()
+  })
+)
+
+function createConfigureGraph(): LGraph {
+  return {
+    nodes: [],
+    onConfigure: vi.fn()
+  } as Partial<LGraph> as LGraph
+}
+
+describe('addAfterConfigureHandler', () => {
+  beforeEach(() => {
+    vi.clearAllMocks()
+  })
+
+  it('runs legacy slot repair when mismatch is detected', () => {
+    vi.mocked(hasLegacyLinkInputSlotMismatch).mockReturnValue(true)
+    const graph = createConfigureGraph()
+
+    addAfterConfigureHandler(graph, () => undefined)
+    graph.onConfigure!.call(
+      graph,
+      {} as Parameters<NonNullable<LGraph['onConfigure']>>[0]
+    )
+
+    expect(hasLegacyLinkInputSlotMismatch).toHaveBeenCalledWith(graph)
+    expect(fixLinkInputSlots).toHaveBeenCalledWith(graph)
+  })
+
+  it('skips legacy slot repair when no mismatch is present', () => {
+    vi.mocked(hasLegacyLinkInputSlotMismatch).mockReturnValue(false)
+    const graph = createConfigureGraph()
+
+    addAfterConfigureHandler(graph, () => undefined)
+    graph.onConfigure!.call(
+      graph,
+      {} as Parameters<NonNullable<LGraph['onConfigure']>>[0]
+    )
+
+    expect(hasLegacyLinkInputSlotMismatch).toHaveBeenCalledWith(graph)
+    expect(fixLinkInputSlots).not.toHaveBeenCalled()
+  })
+})

src/utils/graphConfigureUtil.ts

@@ -0,0 +1,42 @@
+import type { LGraph, LGraphCanvas } from '@/lib/litegraph/src/litegraph'
+import { LiteGraph } from '@/lib/litegraph/src/litegraph'
+import { layoutStore } from '@/renderer/core/layout/store/layoutStore'
+import { flushScheduledSlotLayoutSync } from '@/renderer/extensions/vueNodes/composables/useSlotElementTracking'
+import { triggerCallbackOnAllNodes } from '@/utils/graphTraversalUtil'
+import {
+  fixLinkInputSlots,
+  hasLegacyLinkInputSlotMismatch
+} from '@/utils/litegraphUtil'
+
+/**
+ * Wraps graph.onConfigure to add legacy slot repair,
+ * node configure callbacks, and layout sync flushing.
+ */
+export function addAfterConfigureHandler(
+  graph: LGraph,
+  getCanvas: () => LGraphCanvas | undefined
+) {
+  const { onConfigure } = graph
+  graph.onConfigure = function (...args) {
+    if (LiteGraph.vueNodesMode) {
+      layoutStore.setPendingSlotSync(true)
+    }
+
+    try {
+      if (hasLegacyLinkInputSlotMismatch(this)) fixLinkInputSlots(this)
+
+      triggerCallbackOnAllNodes(this, 'onGraphConfigured')
+
+      const r = onConfigure?.apply(this, args)
+
+      triggerCallbackOnAllNodes(this, 'onAfterGraphConfigured')
+
+      return r
+    } finally {
+      if (LiteGraph.vueNodesMode) {
+        flushScheduledSlotLayoutSync()
+        getCanvas()?.setDirty(true, true)
+      }
+    }
+  }
+}