Add TS types - API (#1736)

* nit * Add TS types - API events * Replace all API event emits with type-safe variants * Add missing API type * nit * Remove test code, nit
2026-02-08 00:50:05 +00:00 · 2024-11-30 05:15:25 +11:00
parent 0bf30e7621
commit 2017b9016b
8 changed files with 225 additions and 112 deletions
--- a/src/extensions/core/groupNode.ts
+++ b/src/extensions/core/groupNode.ts
@@ -1102,10 +1102,9 @@ export class GroupNodeHandler {
        const innerNodeIndex = this.innerNodes?.findIndex((n) => n.id == id)
        if (innerNodeIndex > -1) {
          this.node.runningInternalNodeId = innerNodeIndex
-          api.dispatchEvent(
-            new CustomEvent(type, {
-              detail: getEvent(detail, this.node.id + '', this.node)
-            })
+          api.dispatchCustomEvent(
+            type,
+            getEvent(detail, `${this.node.id}`, this.node)
          )
        }
      }
--- a/src/scripts/api.ts
+++ b/src/scripts/api.ts
@@ -1,19 +1,26 @@
-import type { ComfyWorkflowJSON } from '@/types/comfyWorkflow'
-import {
-  type HistoryTaskItem,
-  type PendingTaskItem,
-  type RunningTaskItem,
-  type ComfyNodeDef,
-  type EmbeddingsResponse,
-  type ExtensionsResponse,
-  type PromptResponse,
-  type SystemStats,
-  type User,
-  type Settings,
-  type UserDataFullInfo,
-  validateComfyNodeDef,
-  LogsRawResponse
+import type { ComfyWorkflowJSON, NodeId } from '@/types/comfyWorkflow'
+import type {
+  HistoryTaskItem,
+  PendingTaskItem,
+  RunningTaskItem,
+  ComfyNodeDef,
+  EmbeddingsResponse,
+  ExtensionsResponse,
+  PromptResponse,
+  SystemStats,
+  User,
+  Settings,
+  UserDataFullInfo,
+  LogsRawResponse,
+  ExecutingWsMessage,
+  ExecutedWsMessage,
+  ProgressWsMessage,
+  ExecutionStartWsMessage,
+  ExecutionErrorWsMessage,
+  StatusWsMessage,
+  StatusWsMessageStatus
 } from '@/types/apiTypes'
+import { validateComfyNodeDef } from '@/types/apiTypes'
 import axios from 'axios'

 interface QueuePromptRequestBody {
@@ -30,6 +37,105 @@ interface QueuePromptRequestBody {
  number?: number
 }

+/** Dictionary of Frontend-generated API calls */
+interface FrontendApiCalls {
+  graphChanged: ComfyWorkflowJSON
+  promptQueued: { number: number; batchCount: number }
+  graphCleared: never
+  reconnecting: never
+  reconnected: never
+}
+
+/** Dictionary of calls originating from ComfyUI core */
+interface BackendApiCalls {
+  progress: ProgressWsMessage
+  executing: ExecutingWsMessage
+  executed: ExecutedWsMessage
+  status: StatusWsMessage
+  execution_start: ExecutionStartWsMessage
+  execution_success: never
+  execution_error: ExecutionErrorWsMessage
+  execution_cached: never
+  logs: never
+  /** Mr Blob Preview, I presume? */
+  b_preview: Blob
+}
+
+/** Dictionary of all api calls */
+interface ApiCalls extends BackendApiCalls, FrontendApiCalls {}
+
+/** Used to create a discriminating union on type value. */
+interface ApiMessage<T extends keyof ApiCalls> {
+  type: T
+  data: ApiCalls[T]
+}
+
+/** Ensures workers get a fair shake. */
+type Unionize<T> = T[keyof T]
+
+/**
+ *  Discriminated union of generic, i.e.:
+ * ```ts
+ * // Convert
+ * type ApiMessageUnion = ApiMessage<'status' | 'executing' | ...>
+ * // To
+ * type ApiMessageUnion = ApiMessage<'status'> | ApiMessage<'executing'> | ...
+ * ```
+ */
+type ApiMessageUnion = Unionize<{
+  [Key in keyof ApiCalls]: ApiMessage<Key>
+}>
+
+/** Wraps all properties in {@link CustomEvent}. */
+type AsCustomEvents<T> = {
+  readonly [K in keyof T]: CustomEvent<T[K]>
+}
+
+/** Handles differing event and API signatures. */
+type ApiToEventType<T = ApiCalls> = {
+  [K in keyof T]: K extends 'status'
+    ? StatusWsMessageStatus
+    : K extends 'executing'
+      ? NodeId
+      : T[K]
+}
+
+/** Dictionary of types used in the detail for a custom event */
+type ApiEventTypes = ApiToEventType<ApiCalls>
+
+/** Dictionary of API events: `[name]: CustomEvent<Type>` */
+type ApiEvents = AsCustomEvents<ApiEventTypes>
+
+/** {@link Omit} all properties that evaluate to `never`. */
+type NeverNever<T> = {
+  [K in keyof T as T[K] extends never ? never : K]: T[K]
+}
+
+/** {@link Pick} only properties that evaluate to `never`. */
+type PickNevers<T> = {
+  [K in keyof T as T[K] extends never ? K : never]: T[K]
+}
+
+/** Keys (names) of API events that _do not_ pass a {@link CustomEvent} `detail` object. */
+type SimpleApiEvents = keyof PickNevers<ApiEventTypes>
+/** Keys (names) of API events that pass a {@link CustomEvent} `detail` object. */
+type ComplexApiEvents = keyof NeverNever<ApiEventTypes>
+
+/** EventTarget typing has no generic capability.  This interface enables tsc strict. */
+export interface ComfyApi extends EventTarget {
+  addEventListener<TEvent extends keyof ApiEvents>(
+    type: TEvent,
+    callback: ((event: ApiEvents[TEvent]) => void) | null,
+    options?: AddEventListenerOptions | boolean
+  ): void
+
+  removeEventListener<TEvent extends keyof ApiEvents>(
+    type: TEvent,
+    callback: ((event: ApiEvents[TEvent]) => void) | null,
+    options?: EventListenerOptions | boolean
+  ): void
+}
+
 export class ComfyApi extends EventTarget {
  #registered = new Set()
  api_host: string
@@ -92,15 +198,51 @@ export class ComfyApi extends EventTarget {
    return fetch(this.apiURL(route), options)
  }

-  addEventListener(
-    type: string,
-    callback: any,
-    options?: AddEventListenerOptions
+  addEventListener<TEvent extends keyof ApiEvents>(
+    type: TEvent,
+    callback: ((event: ApiEvents[TEvent]) => void) | null,
+    options?: AddEventListenerOptions | boolean
  ) {
-    super.addEventListener(type, callback, options)
+    // Type assertion: strictFunctionTypes.  So long as we emit events in a type-safe fashion, this is safe.
+    super.addEventListener(type, callback as EventListener, options)
    this.#registered.add(type)
  }

+  removeEventListener<TEvent extends keyof ApiEvents>(
+    type: TEvent,
+    callback: ((event: ApiEvents[TEvent]) => void) | null,
+    options?: EventListenerOptions | boolean
+  ): void {
+    super.removeEventListener(type, callback as EventListener, options)
+  }
+
+  /**
+   * Dispatches a custom event.
+   * Provides type safety for the contravariance issue with EventTarget (last checked TS 5.6).
+   * @param type The type of event to emit
+   * @param detail The detail property used for a custom event ({@link CustomEventInit.detail})
+   */
+  dispatchCustomEvent<T extends SimpleApiEvents>(type: T): boolean
+  dispatchCustomEvent<T extends ComplexApiEvents>(
+    type: T,
+    detail: ApiEventTypes[T] | null
+  ): boolean
+  dispatchCustomEvent<T extends keyof ApiEventTypes>(
+    type: T,
+    detail?: ApiEventTypes[T]
+  ): boolean {
+    const event =
+      detail === undefined
+        ? new CustomEvent(type)
+        : new CustomEvent(type, { detail })
+    return super.dispatchEvent(event)
+  }
+
+  /** @deprecated Use {@link dispatchCustomEvent}. */
+  dispatchEvent(event: never): boolean {
+    return super.dispatchEvent(event)
+  }
+
  /**
   * Poll status  for colab and other things that don't support websockets.
   */
@@ -108,10 +250,10 @@ export class ComfyApi extends EventTarget {
    setInterval(async () => {
      try {
        const resp = await this.fetchApi('/prompt')
-        const status = await resp.json()
-        this.dispatchEvent(new CustomEvent('status', { detail: status }))
+        const status = (await resp.json()) as StatusWsMessageStatus
+        this.dispatchCustomEvent('status', status)
      } catch (error) {
-        this.dispatchEvent(new CustomEvent('status', { detail: null }))
+        this.dispatchCustomEvent('status', null)
      }
    }, 1000)
  }
@@ -138,7 +280,7 @@ export class ComfyApi extends EventTarget {
    this.socket.addEventListener('open', () => {
      opened = true
      if (isReconnect) {
-        this.dispatchEvent(new CustomEvent('reconnected'))
+        this.dispatchCustomEvent('reconnected')
      }
    })

@@ -155,8 +297,8 @@ export class ComfyApi extends EventTarget {
        this.#createSocket(true)
      }, 300)
      if (opened) {
-        this.dispatchEvent(new CustomEvent('status', { detail: null }))
-        this.dispatchEvent(new CustomEvent('reconnecting'))
+        this.dispatchCustomEvent('status', null)
+        this.dispatchCustomEvent('reconnecting')
      }
    })

@@ -182,9 +324,7 @@ export class ComfyApi extends EventTarget {
              const imageBlob = new Blob([buffer.slice(4)], {
                type: imageMime
              })
-              this.dispatchEvent(
-                new CustomEvent('b_preview', { detail: imageBlob })
-              )
+              this.dispatchCustomEvent('b_preview', imageBlob)
              break
            default:
              throw new Error(
@@ -192,7 +332,7 @@ export class ComfyApi extends EventTarget {
              )
          }
        } else {
-          const msg = JSON.parse(event.data)
+          const msg = JSON.parse(event.data) as ApiMessageUnion
          switch (msg.type) {
            case 'status':
              if (msg.data.sid) {
@@ -201,31 +341,32 @@ export class ComfyApi extends EventTarget {
                window.name = clientId // use window name so it isnt reused when duplicating tabs
                sessionStorage.setItem('clientId', clientId) // store in session storage so duplicate tab can load correct workflow
              }
-              this.dispatchEvent(
-                new CustomEvent('status', { detail: msg.data.status })
-              )
+              this.dispatchCustomEvent('status', msg.data.status ?? null)
              break
            case 'executing':
-              this.dispatchEvent(
-                new CustomEvent('executing', {
-                  detail: msg.data.display_node || msg.data.node
-                })
+              this.dispatchCustomEvent(
+                'executing',
+                msg.data.display_node || msg.data.node
              )
              break
+            case 'execution_start':
+            case 'execution_error':
            case 'progress':
            case 'executed':
-            case 'execution_start':
+            case 'graphChanged':
+            case 'promptQueued':
+            case 'b_preview':
+              this.dispatchCustomEvent(msg.type, msg.data)
+              break
            case 'execution_success':
-            case 'execution_error':
            case 'execution_cached':
            case 'logs':
-              this.dispatchEvent(
-                new CustomEvent(msg.type, { detail: msg.data })
-              )
+              this.dispatchCustomEvent(msg.type)
              break
            default:
              if (this.#registered.has(msg.type)) {
-                this.dispatchEvent(
+                // Fallback for custom types - calls super direct.
+                super.dispatchEvent(
                  new CustomEvent(msg.type, { detail: msg.data })
                )
              } else if (!this.reportedUnknownMessageTypes.has(msg.type)) {
--- a/src/scripts/app.ts
+++ b/src/scripts/app.ts
@@ -24,13 +24,12 @@ import {
  type NodeId,
  validateComfyWorkflow
 } from '../types/comfyWorkflow'
-import { ComfyNodeDef, StatusWsMessageStatus } from '@/types/apiTypes'
+import type { ComfyNodeDef } from '@/types/apiTypes'
 import { adjustColor, ColorAdjustOptions } from '@/utils/colorUtil'
 import { ComfyAppMenu } from './ui/menu/index'
 import { getStorageValue } from './utils'
 import { ComfyWorkflow } from '@/stores/workflowStore'
 import {
-  LGraphGroup,
  LGraphCanvas,
  LGraph,
  LGraphNode,
@@ -139,8 +138,10 @@ export class ComfyApp {
  // x, y, scale
  zoom_drag_start: [number, number, number] | null
  lastNodeErrors: any[] | null
-  lastExecutionError: { node_id: number } | null
-  progress: { value: number; max: number } | null
+  /** @type {ExecutionErrorWsMessage} */
+  lastExecutionError: { node_id?: NodeId } | null
+  /** @type {ProgressWsMessage} */
+  progress: { value?: number; max?: number } | null
  configuringGraph: boolean
  ctx: CanvasRenderingContext2D
  bodyTop: HTMLElement
@@ -1542,12 +1543,9 @@ export class ComfyApp {
   * Handles updates from the API socket
   */
  #addApiUpdateHandlers() {
-    api.addEventListener(
-      'status',
-      ({ detail }: CustomEvent<StatusWsMessageStatus>) => {
-        this.ui.setStatus(detail)
-      }
-    )
+    api.addEventListener('status', ({ detail }) => {
+      this.ui.setStatus(detail)
+    })

    api.addEventListener('progress', ({ detail }) => {
      this.progress = detail
@@ -2547,9 +2545,7 @@ export class ComfyApp {
    } finally {
      this.#processingQueue = false
    }
-    api.dispatchEvent(
-      new CustomEvent('promptQueued', { detail: { number, batchCount } })
-    )
+    api.dispatchCustomEvent('promptQueued', { number, batchCount })
    return !this.lastNodeErrors
  }

--- a/src/scripts/changeTracker.ts
+++ b/src/scripts/changeTracker.ts
@@ -83,9 +83,7 @@ export class ChangeTracker {
  }

  updateModified() {
-    api.dispatchEvent(
-      new CustomEvent('graphChanged', { detail: this.activeState })
-    )
+    api.dispatchCustomEvent('graphChanged', this.activeState)

    // Get the workflow from the store as ChangeTracker is raw object, i.e.
    // `this.workflow` is not reactive.
--- a/src/scripts/ui.ts
+++ b/src/scripts/ui.ts
@@ -4,7 +4,7 @@ import { ComfyDialog as _ComfyDialog } from './ui/dialog'
 import { toggleSwitch } from './ui/toggleSwitch'
 import { ComfySettingsDialog } from './ui/settings'
 import { ComfyApp, app } from './app'
-import { TaskItem } from '@/types/apiTypes'
+import { TaskItem, type StatusWsMessageStatus } from '@/types/apiTypes'
 import { showSettingsDialog } from '@/services/dialogService'
 import { useSettingStore } from '@/stores/settingStore'
 import { useCommandStore } from '@/stores/commandStore'
@@ -611,7 +611,7 @@ export class ComfyUI {
            app.clean()
            app.graph.clear()
            app.resetView()
-            api.dispatchEvent(new CustomEvent('graphCleared'))
+            api.dispatchCustomEvent('graphCleared')
          }
        }
      }),
@@ -641,10 +641,11 @@ export class ComfyUI {

    this.restoreMenuPosition = dragElement(this.menuContainer, this.settings)

+    // @ts-expect-error
    this.setStatus({ exec_info: { queue_remaining: 'X' } })
  }

-  setStatus(status) {
+  setStatus(status: StatusWsMessageStatus | null) {
    this.queueSize.textContent =
      'Queue size: ' + (status ? status.exec_info.queue_remaining : 'ERR')
    if (status) {
--- a/src/stores/commandStore.ts
+++ b/src/stores/commandStore.ts
@@ -204,7 +204,7 @@ export const useCommandStore = defineStore('command', () => {
        ) {
          app.clean()
          app.graph.clear()
-          api.dispatchEvent(new CustomEvent('graphCleared'))
+          api.dispatchCustomEvent('graphCleared')
        }
      }
    },
--- a/src/types/apiTypes.ts
+++ b/src/types/apiTypes.ts
@@ -30,7 +30,8 @@ const zStatusWsMessageStatus = z.object({
 })

 const zStatusWsMessage = z.object({
-  status: zStatusWsMessageStatus.nullable().optional()
+  status: zStatusWsMessageStatus.nullish(),
+  sid: z.string().nullish()
 })

 const zProgressWsMessage = z.object({
--- a/tests-ui/tests/slow/groupNode.test.ts
+++ b/tests-ui/tests/slow/groupNode.test.ts
@@ -362,7 +362,7 @@ describe('group node', () => {
    const { ez, graph, app } = await start()
    const nodes = createDefaultWorkflow(ez, graph)

-    let reroutes: EzNode[] = []
+    const reroutes: EzNode[] = []
    let prevNode = nodes.ckpt
    for (let i = 0; i < 5; i++) {
      const reroute = ez.Reroute()
@@ -568,37 +568,23 @@ describe('group node', () => {

    const { api } = await import('../../../src/scripts/api')

-    api.dispatchEvent(new CustomEvent('execution_start', {}))
-    api.dispatchEvent(
-      new CustomEvent('executing', { detail: `${nodes.save.id}` })
-    )
+    api.dispatchCustomEvent('execution_start', undefined)
+    api.dispatchCustomEvent('executing', `${nodes.save.id}`)
    // Event should be forwarded to group node id
    expect(group.node['imgs']).toBeFalsy()
-    api.dispatchEvent(
-      new CustomEvent('executed', {
-        detail: {
-          node: `${nodes.save.id}`,
-          display_node: `${nodes.save.id}`,
-          output: {
-            images: [
-              {
-                filename: 'test.png',
-                type: 'output'
-              }
-            ]
-          }
-        }
-      })
-    )
+    api.dispatchCustomEvent('executed', {
+      node: `${nodes.save.id}`,
+      display_node: `${nodes.save.id}`,
+      output: {
+        images: [{ filename: 'test.png', type: 'output' }]
+      }
+    })

    // Trigger paint
    group.node.onDrawBackground?.(app.canvas.ctx, app.canvas.canvas)

    expect(group.node['images']).toEqual([
-      {
-        filename: 'test.png',
-        type: 'output'
-      }
+      { filename: 'test.png', type: 'output' }
    ])

    // Reload
@@ -610,26 +596,17 @@ describe('group node', () => {
    group.node['getInnerNodes']()

    // Check it works for internal node ids
-    api.dispatchEvent(new CustomEvent('execution_start', {}))
-    api.dispatchEvent(new CustomEvent('executing', { detail: `${group.id}:5` }))
+    api.dispatchCustomEvent('execution_start', undefined)
+    api.dispatchCustomEvent('executing', `${group.id}:5`)
    // Event should be forwarded to group node id
    expect(group.node['imgs']).toBeFalsy()
-    api.dispatchEvent(
-      new CustomEvent('executed', {
-        detail: {
-          node: `${group.id}:5`,
-          display_node: `${group.id}:5`,
-          output: {
-            images: [
-              {
-                filename: 'test2.png',
-                type: 'output'
-              }
-            ]
-          }
-        }
-      })
-    )
+    api.dispatchCustomEvent('executed', {
+      node: `${group.id}:5`,
+      display_node: `${group.id}:5`,
+      output: {
+        images: [{ filename: 'test2.png', type: 'output' }]
+      }
+    })

    // Trigger paint
    group.node.onDrawBackground?.(app.canvas.ctx, app.canvas.canvas)
@@ -791,7 +768,7 @@ describe('group node', () => {
    })

    expect(dialogShow).toBeCalledTimes(1)
-    // @ts-expect-error
+    // @ts-expect-error Mocked
    const call = dialogShow.mock.calls[0][0].innerHTML
    expect(call).toContain('the following node types were not found')
    expect(call).toContain('NotKSampler')