diff --git a/src/extensions/core/groupNode.ts b/src/extensions/core/groupNode.ts index 4946a6fc8..5905bf00e 100644 --- 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) ) } } diff --git a/src/scripts/api.ts b/src/scripts/api.ts index a70b25356..70998a812 100644 --- 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 { + type: T + data: ApiCalls[T] +} + +/** Ensures workers get a fair shake. */ +type Unionize = 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 +}> + +/** Wraps all properties in {@link CustomEvent}. */ +type AsCustomEvents = { + readonly [K in keyof T]: CustomEvent +} + +/** Handles differing event and API signatures. */ +type ApiToEventType = { + [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 + +/** Dictionary of API events: `[name]: CustomEvent` */ +type ApiEvents = AsCustomEvents + +/** {@link Omit} all properties that evaluate to `never`. */ +type NeverNever = { + [K in keyof T as T[K] extends never ? never : K]: T[K] +} + +/** {@link Pick} only properties that evaluate to `never`. */ +type PickNevers = { + [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 +/** Keys (names) of API events that pass a {@link CustomEvent} `detail` object. */ +type ComplexApiEvents = keyof NeverNever + +/** EventTarget typing has no generic capability. This interface enables tsc strict. */ +export interface ComfyApi extends EventTarget { + addEventListener( + type: TEvent, + callback: ((event: ApiEvents[TEvent]) => void) | null, + options?: AddEventListenerOptions | boolean + ): void + + removeEventListener( + 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( + 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( + 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(type: T): boolean + dispatchCustomEvent( + type: T, + detail: ApiEventTypes[T] | null + ): boolean + dispatchCustomEvent( + 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)) { diff --git a/src/scripts/app.ts b/src/scripts/app.ts index eab2f434e..a400055cd 100644 --- 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) => { - 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 } diff --git a/src/scripts/changeTracker.ts b/src/scripts/changeTracker.ts index 2373b3863..1dfecc1a9 100644 --- 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. diff --git a/src/scripts/ui.ts b/src/scripts/ui.ts index e9f828bae..b5f686b6e 100644 --- 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) { diff --git a/src/stores/commandStore.ts b/src/stores/commandStore.ts index 29f8c3a17..7050707c1 100644 --- 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') } } }, diff --git a/src/types/apiTypes.ts b/src/types/apiTypes.ts index 36f4072ce..39f412a20 100644 --- 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({ diff --git a/tests-ui/tests/slow/groupNode.test.ts b/tests-ui/tests/slow/groupNode.test.ts index b056406a4..5a184c5cd 100644 --- 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')