Fix UI crash when selecting broken node + TS fixes (#3859)

src/components/graph/DomWidgets.vue

@@ -16,12 +16,12 @@ import { computed, watch } from 'vue'
 
 import DomWidget from '@/components/graph/widgets/DomWidget.vue'
 import { useChainCallback } from '@/composables/functional/useChainCallback'
-import { type DomWidgetState, useDomWidgetStore } from '@/stores/domWidgetStore'
+import { useDomWidgetStore } from '@/stores/domWidgetStore'
 import { useCanvasStore } from '@/stores/graphStore'
 
 const domWidgetStore = useDomWidgetStore()
-const widgetStates = computed(
-  () => Array.from(domWidgetStore.widgetStates.values()) as DomWidgetState[]
+const widgetStates = computed(() =>
+  Array.from(domWidgetStore.widgetStates.values())
 )
 
 const updateWidgets = () => {

src/components/graph/GraphCanvas.vue

@@ -293,7 +293,7 @@ onMounted(async () => {
   workspaceStore.spinner = true
   // ChangeTracker needs to be initialized before setup, as it will overwrite
   // some listeners of litegraph canvas.
-  ChangeTracker.init(comfyApp)
+  ChangeTracker.init()
   await loadCustomNodesI18n()
   try {
     await settingStore.loadSettingValues()

src/components/graph/selectionToolbox/ExecuteButton.vue

@@ -36,7 +36,7 @@ const buttonHovered = ref(false)
 const selectedOutputNodes = computed(
   () =>
     canvasStore.selectedItems.filter(
-      (item) => isLGraphNode(item) && item.constructor.nodeData.output_node
+      (item) => isLGraphNode(item) && item.constructor.nodeData?.output_node
     ) as LGraphNode[]
 )
 
@@ -45,7 +45,7 @@ const isDisabled = computed(() => selectedOutputNodes.value.length === 0)
 function outputNodeStokeStyle(this: LGraphNode) {
   if (
     this.selected &&
-    this.constructor.nodeData.output_node &&
+    this.constructor.nodeData?.output_node &&
     buttonHovered.value
   ) {
     return { color: 'orange', lineWidth: 2, padding: 10 }

src/composables/useCoreCommands.ts

@@ -320,7 +320,7 @@ export function useCoreCommands(): ComfyCommand[] {
       function: async () => {
         const batchCount = useQueueSettingsStore().batchCount
         const queueNodeIds = getSelectedNodes()
-          .filter((node) => node.constructor.nodeData.output_node)
+          .filter((node) => node.constructor.nodeData?.output_node)
           .map((node) => node.id)
         if (queueNodeIds.length === 0) {
           toastStore.add({

src/extensions/core/groupNode.ts

@@ -797,7 +797,7 @@ export class GroupNodeConfig {
 
 export class GroupNodeHandler {
   node: LGraphNode
-  groupData
+  groupData: any
   innerNodes: any
 
   constructor(node: LGraphNode) {

src/extensions/core/nodeTemplates.ts

@@ -401,6 +401,7 @@ app.registerExtension({
                   // @ts-expect-error
                   data.groupNodes = {}
                 }
+                if (nodeData == null) throw new TypeError('nodeData is not set')
                 // @ts-expect-error
                 data.groupNodes[nodeData.name] = groupData
                 // @ts-expect-error

src/extensions/core/uploadAudio.ts

@@ -117,7 +117,10 @@ app.registerExtension({
           node.addDOMWidget(inputName, /* name=*/ 'audioUI', audio)
         audioUIWidget.serialize = false
 
-        const isOutputNode = node.constructor.nodeData.output_node
+        const { nodeData } = node.constructor
+        if (nodeData == null) throw new TypeError('nodeData is null')
+
+        const isOutputNode = nodeData.output_node
         if (isOutputNode) {
           // Hide the audio widget when there is no audio initially.
           audioUIWidget.element.classList.add('empty-audio-widget')

src/schemas/comfyWorkflowSchema.ts

@@ -216,6 +216,7 @@ const zComfyNode = z
 
 const zGroup = z
   .object({
+    id: z.number().optional(),
     title: z.string(),
     bounding: z.tuple([z.number(), z.number(), z.number(), z.number()]),
     color: z.string().optional(),

src/scripts/changeTracker.ts

@@ -10,6 +10,7 @@ import { ComfyWorkflow, useWorkflowStore } from '@/stores/workflowStore'
 
 import { api } from './api'
 import type { ComfyApp } from './app'
+import { app } from './app'
 
 function clone<T>(obj: T): T {
   return JSON.parse(JSON.stringify(obj))
@@ -36,11 +37,6 @@ export class ChangeTracker {
   ds?: { scale: number; offset: [number, number] }
   nodeOutputs?: Record<string, any>
 
-  static app?: ComfyApp
-  get app(): ComfyApp {
-    return ChangeTracker.app!
-  }
-
   constructor(
     /**
      * The workflow that this change tracker is tracking
@@ -68,18 +64,18 @@ export class ChangeTracker {
 
   store() {
     this.ds = {
-      scale: this.app.canvas.ds.scale,
-      offset: [this.app.canvas.ds.offset[0], this.app.canvas.ds.offset[1]]
+      scale: app.canvas.ds.scale,
+      offset: [app.canvas.ds.offset[0], app.canvas.ds.offset[1]]
     }
   }
 
   restore() {
     if (this.ds) {
-      this.app.canvas.ds.scale = this.ds.scale
-      this.app.canvas.ds.offset = this.ds.offset
+      app.canvas.ds.scale = this.ds.scale
+      app.canvas.ds.offset = this.ds.offset
     }
     if (this.nodeOutputs) {
-      this.app.nodeOutputs = this.nodeOutputs
+      app.nodeOutputs = this.nodeOutputs
     }
   }
 
@@ -105,10 +101,8 @@ export class ChangeTracker {
   }
 
   checkState() {
-    if (!this.app.graph || this.changeCount) return
-    // @ts-expect-error zod type issue on ComfyWorkflowJSON. ComfyWorkflowJSON
-    // is stricter than LiteGraph's serialisation schema.
-    const currentState = clone(this.app.graph.serialize()) as ComfyWorkflowJSON
+    if (!app.graph || this.changeCount) return
+    const currentState = clone(app.graph.serialize()) as ComfyWorkflowJSON
     if (!this.activeState) {
       this.activeState = currentState
       return
@@ -132,7 +126,7 @@ export class ChangeTracker {
       target.push(this.activeState)
       this.restoringState = true
       try {
-        await this.app.loadGraphData(prevState, false, false, this.workflow, {
+        await app.loadGraphData(prevState, false, false, this.workflow, {
           showMissingModelsDialog: false,
           showMissingNodesDialog: false,
           checkForRerouteMigration: false
@@ -189,13 +183,11 @@ export class ChangeTracker {
     }
   }
 
-  static init(app: ComfyApp) {
+  static init() {
     const getCurrentChangeTracker = () =>
       useWorkflowStore().activeWorkflow?.changeTracker
     const checkState = () => getCurrentChangeTracker()?.checkState()
 
-    ChangeTracker.app = app
-
     let keyIgnored = false
     window.addEventListener(
       'keydown',
@@ -237,7 +229,7 @@ export class ChangeTracker {
           if (await changeTracker.undoRedo(e)) return
 
           // If our active element is some type of input then handle changes after they're done
-          if (ChangeTracker.bindInput(app, bindInputEl)) return
+          if (ChangeTracker.bindInput(bindInputEl)) return
           logger.debug('checkState on keydown')
           changeTracker.checkState()
         })
@@ -339,7 +331,7 @@ export class ChangeTracker {
     })
   }
 
-  static bindInput(_app: ComfyApp, activeEl: Element | null): boolean {
+  static bindInput(activeEl: Element | null): boolean {
     if (
       !activeEl ||
       activeEl.tagName === 'CANVAS' ||

src/stores/domWidgetStore.ts

@@ -2,14 +2,14 @@
  * Stores all DOM widgets that are used in the canvas.
  */
 import { defineStore } from 'pinia'
-import { markRaw, ref } from 'vue'
+import { type Raw, markRaw, ref } from 'vue'
 
 import type { PositionConfig } from '@/composables/element/useAbsolutePosition'
 import type { BaseDOMWidget } from '@/scripts/domWidget'
 
 export interface DomWidgetState extends PositionConfig {
   // Raw widget instance
-  widget: BaseDOMWidget<object | string>
+  widget: Raw<BaseDOMWidget<object | string>>
   visible: boolean
   readonly: boolean
   zIndex: number
@@ -23,7 +23,7 @@ export const useDomWidgetStore = defineStore('domWidget', () => {
     widget: BaseDOMWidget<V>
   ) => {
     widgetStates.value.set(widget.id, {
-      widget: markRaw(widget) as unknown as BaseDOMWidget<object | string>,
+      widget: markRaw(widget) as unknown as Raw<BaseDOMWidget<object | string>>,
       visible: true,
       readonly: false,
       zIndex: 0,

src/stores/nodeDefStore.ts

@@ -292,6 +292,8 @@ export const useNodeDefStore = defineStore('nodeDef', () => {
   }
   function fromLGraphNode(node: LGraphNode): ComfyNodeDefImpl | null {
     // Frontend-only nodes don't have nodeDef
+    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+    // @ts-ignore Optional chaining used in index
     return nodeDefsByName.value[node.constructor?.nodeData?.name] ?? null
   }
 

src/types/litegraph-augmentation.d.ts

@@ -29,6 +29,12 @@ declare module '@comfyorg/litegraph/dist/types/widgets' {
      * The minimum size of the node if the widget is present.
      */
     minNodeSize?: Size
+
+    /** If the widget is advanced, this will be set to true. */
+    advanced?: boolean
+
+    /** If the widget is hidden, this will be set to true. */
+    hidden?: boolean
   }
 
   interface IBaseWidget {
@@ -60,7 +66,7 @@ declare module '@comfyorg/litegraph' {
     type?: string
     comfyClass: string
     title: string
-    nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2
+    nodeData?: ComfyNodeDefV1 & ComfyNodeDefV2 & { [key: symbol]: unknown }
     category?: string
     new (): T
   }

src/utils/executionUtil.ts

@@ -202,6 +202,5 @@ export const graphToPrompt = async (
     output = newOutput
   }
 
-  // @ts-expect-error Convert ISerializedGraph to ComfyWorkflowJSON
   return { workflow: workflow as ComfyWorkflowJSON, output }
 }

tests-ui/tests/comfyWorkflow.test.ts

@@ -8,61 +8,63 @@ const WORKFLOW_DIR = 'tests-ui/workflows'
 
 describe('parseComfyWorkflow', () => {
   it('parses valid workflow', async () => {
-    fs.readdirSync(WORKFLOW_DIR).forEach(async (file) => {
+    for await (const file of fs.readdirSync(WORKFLOW_DIR)) {
       if (file.endsWith('.json')) {
         const data = fs.readFileSync(`${WORKFLOW_DIR}/${file}`, 'utf-8')
-        expect(await validateComfyWorkflow(JSON.parse(data))).not.toBeNull()
+        await expect(
+          validateComfyWorkflow(JSON.parse(data))
+        ).resolves.not.toBeNull()
       }
-    })
+    }
   })
 
   it('workflow.nodes', async () => {
     const workflow = JSON.parse(JSON.stringify(defaultGraph))
     workflow.nodes = undefined
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     workflow.nodes = null
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     workflow.nodes = []
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
   })
 
   it('workflow.version', async () => {
     const workflow = JSON.parse(JSON.stringify(defaultGraph))
     workflow.version = undefined
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     workflow.version = '1.0.1' // Invalid format (string)
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     // 2018-2024 schema: 0.4
     workflow.version = 0.4
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
   })
 
   it('workflow.extra', async () => {
     const workflow = JSON.parse(JSON.stringify(defaultGraph))
     workflow.extra = undefined
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     workflow.extra = null
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     workflow.extra = {}
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     workflow.extra = { foo: 'bar' } // Should accept extra fields.
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
   })
 
   it('workflow.nodes.pos', async () => {
     const workflow = JSON.parse(JSON.stringify(defaultGraph))
     workflow.nodes[0].pos = [1, 2, 3]
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     workflow.nodes[0].pos = [1, 2]
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     // Should automatically transform the legacy format object to array.
     workflow.nodes[0].pos = { '0': 3, '1': 4 }
@@ -97,13 +99,13 @@ describe('parseComfyWorkflow', () => {
   it('workflow.nodes.widget_values', async () => {
     const workflow = JSON.parse(JSON.stringify(defaultGraph))
     workflow.nodes[0].widgets_values = ['foo', 'bar']
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     workflow.nodes[0].widgets_values = 'foo'
-    expect(await validateComfyWorkflow(workflow)).toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
 
     workflow.nodes[0].widgets_values = undefined
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
 
     // The object format of widgets_values is used by VHS nodes to perform
     // dynamic widgets display.
@@ -126,7 +128,7 @@ describe('parseComfyWorkflow', () => {
         'INT' // Data type
       ]
     ]
-    expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+    await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
   })
 
   describe('workflow.nodes.properties.aux_id', () => {
@@ -137,7 +139,7 @@ describe('parseComfyWorkflow', () => {
     it.each(validAuxIds)('valid aux_id: %s', async (aux_id) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.aux_id = aux_id
-      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
     })
     const invalidAuxIds = [
       'invalid spaces in username/repo',
@@ -148,7 +150,7 @@ describe('parseComfyWorkflow', () => {
     it.each(invalidAuxIds)('invalid aux_id: %s', async (aux_id) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.aux_id = aux_id
-      expect(await validateComfyWorkflow(workflow)).toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
     })
   })
 
@@ -157,14 +159,14 @@ describe('parseComfyWorkflow', () => {
     it.each(validCnrIds)('valid cnr_id: %s', async (cnr_id) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.cnr_id = cnr_id
-      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
     })
 
     const invalidCnrIds = ['invalid cnr-id', 'invalid^cnr-id', 'invalid cnr id']
     it.each(invalidCnrIds)('invalid cnr_id: %s', async (cnr_id) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.cnr_id = cnr_id
-      expect(await validateComfyWorkflow(workflow)).toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
     })
   })
 
@@ -186,7 +188,7 @@ describe('parseComfyWorkflow', () => {
     it.each(validVersionStrings)('valid version: %s', async (ver) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.ver = ver
-      expect(await validateComfyWorkflow(workflow)).not.toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.not.toBeNull()
     })
 
     const invalidVersionStrings = [
@@ -200,7 +202,7 @@ describe('parseComfyWorkflow', () => {
     it.each(invalidVersionStrings)('invalid version: %s', async (ver) => {
       const workflow = JSON.parse(JSON.stringify(defaultGraph))
       workflow.nodes[0].properties.ver = ver
-      expect(await validateComfyWorkflow(workflow)).toBeNull()
+      await expect(validateComfyWorkflow(workflow)).resolves.toBeNull()
     })
   })
 })