Hook up more of the API

src/lib/api.ts

@@ -1,3 +1,8 @@
+import type { Progress, SerializedPrompt, SerializedPromptOutput, SerializedPromptOutputs } from "./components/ComfyApp";
+import type TypedEmitter from "typed-emitter";
+import EventEmitter from "events";
+import type { GalleryOutput } from "./nodes/ComfyWidgetNodes";
+
 type PromptRequestBody = {
     client_id: string,
     prompt: any,
@@ -8,27 +13,68 @@ type PromptRequestBody = {
 
 export type QueueItemType = "queue" | "history";
 
-export type ComfyAPIQueueStatus = {
-    exec_info: {
-        queue_remaining: number | "X";
-    }
+export type ComfyAPIStatusExecInfo = {
+    queueRemaining: number | "X";
 }
 
-export default class ComfyAPI extends EventTarget {
-    private registered: Set<string> = new Set<string>();
+export type ComfyAPIStatusResponse = {
+    execInfo?: ComfyAPIStatusExecInfo,
+    error?: string
+}
+
+export type ComfyAPIQueueResponse = {
+    running: ComfyAPIHistoryItem[],
+    pending: ComfyAPIHistoryItem[],
+    error?: string
+}
+
+export type NodeID = string;
+export type PromptID = string; // UUID
+
+export type ComfyAPIHistoryItem = [
+    number,   // prompt number
+    PromptID,
+    SerializedPrompt,
+    any,      // extra data
+    NodeID[]  // good outputs
+]
+
+export type ComfyAPIPromptResponse = {
+    promptID?: PromptID,
+    error?: string
+}
+
+export type ComfyAPIHistoryEntry = {
+    prompt: ComfyAPIHistoryItem,
+    outputs: SerializedPromptOutputs
+}
+
+export type ComfyAPIHistoryResponse = {
+    history: Record<PromptID, ComfyAPIHistoryEntry>,
+    error?: string
+}
+
+type ComfyAPIEvents = {
+    status: (status: ComfyAPIStatusResponse | null, error?: Error | null) => void,
+    progress: (progress: Progress) => void,
+    reconnecting: () => void,
+    reconnected: () => void,
+    executing: (promptID: PromptID | null, runningNodeID: NodeID | null) => void,
+    executed: (promptID: PromptID, nodeID: NodeID, output: SerializedPromptOutput) => void,
+    execution_cached: (promptID: PromptID, nodes: NodeID[]) => void,
+    execution_error: (promptID: PromptID, message: string) => void,
+}
+
+export default class ComfyAPI {
+    private eventBus: TypedEmitter<ComfyAPIEvents> = new EventEmitter() as TypedEmitter<ComfyAPIEvents>;
 
     socket: WebSocket | null = null;
     clientId: string | null = null;
     hostname: string | null = null;
     port: number | null = 8188;
 
-    constructor() {
-        super();
-    }
-
-    override addEventListener(type: string, callback: EventListenerOrEventListenerObject | null, options?: AddEventListenerOptions | boolean) {
-        super.addEventListener(type, callback, options);
-        this.registered.add(type);
+    addEventListener<E extends keyof ComfyAPIEvents>(type: E, callback: ComfyAPIEvents[E]) {
+        this.eventBus.addListener(type, callback);
     }
 
     /**
@@ -39,9 +85,9 @@ export default class ComfyAPI extends EventTarget {
             try {
                 const resp = await fetch(this.getBackendUrl() + "/prompt");
                 const status = await resp.json();
-                this.dispatchEvent(new CustomEvent("status", { detail: status }));
+                this.eventBus.emit("status", { execInfo: { queueRemaining: status.exec_info.queue_remaining } });
             } catch (error) {
-                this.dispatchEvent(new CustomEvent("status", { detail: null }));
+                this.eventBus.emit("status", { error: error.toString() });
             }
         }, 1000);
     }
@@ -77,7 +123,7 @@ export default class ComfyAPI extends EventTarget {
         this.socket.addEventListener("open", () => {
             opened = true;
             if (isReconnect) {
-                this.dispatchEvent(new CustomEvent("reconnected"));
+                this.eventBus.emit("reconnected");
             }
         });
 
@@ -94,8 +140,8 @@ export default class ComfyAPI extends EventTarget {
                 this.createSocket(true);
             }, 300);
             if (opened) {
-                this.dispatchEvent(new CustomEvent("status", { detail: null }));
-                this.dispatchEvent(new CustomEvent("reconnecting"));
+                this.eventBus.emit("status", null);
+                this.eventBus.emit("reconnecting");
             }
         });
 
@@ -108,29 +154,25 @@ export default class ComfyAPI extends EventTarget {
                             this.clientId = msg.data.sid;
                             sessionStorage["Comfy.SessionId"] = this.clientId;
                         }
-                        this.dispatchEvent(new CustomEvent("status", { detail: msg.data.status }));
+                        this.eventBus.emit("status", msg.data.status);
                         break;
                     case "progress":
-                        this.dispatchEvent(new CustomEvent("progress", { detail: msg.data }));
+                        this.eventBus.emit("progress", msg.data as Progress);
                         break;
                     case "executing":
-                        this.dispatchEvent(new CustomEvent("executing", { detail: msg.data }));
+                        this.eventBus.emit("executing", msg.data.prompt_id, msg.data.node);
                         break;
                     case "executed":
-                        this.dispatchEvent(new CustomEvent("executed", { detail: msg.data }));
+                        this.eventBus.emit("executed", msg.data.prompt_id, msg.data.node, msg.data.output);
                         break;
                     case "execution_cached":
-                        this.dispatchEvent(new CustomEvent("execution_cached", { detail: msg.data }));
+                        this.eventBus.emit("execution_cached", msg.data.prompt_id, msg.data.nodes);
                         break;
                     case "execution_error":
-                        this.dispatchEvent(new CustomEvent("execution_error", { detail: msg.data }));
+                        this.eventBus.emit("execution_error", msg.data.prompt_id, msg.data.message);
                         break;
                     default:
-                        if (this.registered.has(msg.type)) {
-                            this.dispatchEvent(new CustomEvent(msg.type, { detail: msg.data }));
-                        } else {
-                            throw new Error("Unknown message type");
-                        }
+                        throw new Error(`Unknown message type: ${msg.type} ${msg}`);
                 }
             } catch (error) {
                 console.warn("Unhandled message:", event.data);
@@ -149,27 +191,27 @@ export default class ComfyAPI extends EventTarget {
      * Gets a list of extension urls
      * @returns An array of script urls to import
      */
-    async getExtensions() {
-        const resp = await fetch(this.getBackendUrl() + `/extensions`, { cache: "no-store" });
-        return await resp.json();
+    async getExtensions(): Promise<any> {
+        return fetch(this.getBackendUrl() + `/extensions`, { cache: "no-store" })
+            .then(resp => resp.json())
     }
 
     /**
      * Gets a list of embedding names
      * @returns An array of script urls to import
      */
-    async getEmbeddings() {
-        const resp = await fetch(this.getBackendUrl() + "/embeddings", { cache: "no-store" });
-        return await resp.json();
+    async getEmbeddings(): Promise<any> {
+        return fetch(this.getBackendUrl() + "/embeddings", { cache: "no-store" })
+            .then(resp => resp.json())
     }
 
     /**
      * Loads node object definitions for the graph
      * @returns The node definitions
      */
-    async getNodeDefs() {
-        const resp = await fetch(this.getBackendUrl() + "/object_info", { cache: "no-store" });
-        return await resp.json();
+    async getNodeDefs(): Promise<any> {
+        return fetch(this.getBackendUrl() + "/object_info", { cache: "no-store" })
+            .then(resp => resp.json())
     }
 
     /**
@@ -177,11 +219,11 @@ export default class ComfyAPI extends EventTarget {
      * @param {number} number The index at which to queue the prompt, passing -1 will insert the prompt at the front of the queue
      * @param {object} prompt The prompt data to queue
      */
-    async queuePrompt(number: number, { output, workflow }) {
+    async queuePrompt(number: number, { output, workflow }, extra_data: any): Promise<ComfyAPIPromptResponse> {
         const body: PromptRequestBody = {
             client_id: this.clientId,
             prompt: output,
-            extra_data: { extra_pnginfo: { workflow } },
+            extra_data,
             front: false,
             number: number
         };
@@ -192,67 +234,52 @@ export default class ComfyAPI extends EventTarget {
             body.number = number;
         }
 
-        const res = await fetch(this.getBackendUrl() + "/prompt", {
+        let postBody = null;
+        try {
+            postBody = JSON.stringify(body)
+        }
+        catch (error) {
+            return Promise.reject({ error })
+        }
+
+        return fetch(this.getBackendUrl() + "/prompt", {
             method: "POST",
             headers: {
                 "Content-Type": "application/json",
             },
-            body: JSON.stringify(body),
-        });
-
-        if (res.status !== 200) {
-            throw {
-                response: await res.text(),
-            };
-        }
-    }
-
-    /**
-     * Loads a list of items (queue or history)
-     * @param {string} type The type of items to load, queue or history
-     * @returns The items of the specified type grouped by their status
-     */
-    async getItems(type: QueueItemType) {
-        if (type === "queue") {
-            return this.getQueue();
-        }
-        return this.getHistory();
+            body: postBody
+        })
+            .then(res => res.json())
+            .then(raw => { return { promptID: raw.prompt_id } })
+            .catch(res => { throw res.text() })
+            .catch(error => { return { error } })
     }
 
     /**
      * Gets the current state of the queue
      * @returns The currently running and queued items
      */
-    async getQueue() {
-        try {
-            const res = await fetch(this.getBackendUrl() + "/queue");
-            const data = await res.json();
-            return {
-                // Running action uses a different endpoint for cancelling
-                Running: data.queue_running.map((prompt) => ({
-                    prompt,
-                    remove: { name: "Cancel", cb: () => this.interrupt() },
-                })),
-                Pending: data.queue_pending.map((prompt) => ({ prompt })),
-            };
-        } catch (error) {
-            console.error(error);
-            return { Running: [], Pending: [], error };
-        }
+    async getQueue(): Promise<ComfyAPIQueueResponse> {
+        return fetch(this.getBackendUrl() + "/queue")
+            .then(res => res.json())
+            .then(data => {
+                return {
+                    running: data.queue_running,
+                    pending: data.queue_pending,
+                }
+            })
+            .catch(error => { return { running: [], pending: [], error } })
     }
 
     /**
      * Gets the prompt execution history
      * @returns Prompt history including node outputs
      */
-    async getHistory() {
-        try {
-            const res = await fetch(this.getBackendUrl() + "/history");
-            return { History: Object.values(await res.json()) };
-        } catch (error) {
-            console.error(error);
-            return { History: [], error };
-        }
+    async getHistory(): Promise<ComfyAPIHistoryResponse> {
+        return fetch(this.getBackendUrl() + "/history")
+            .then(res => res.json())
+            .then(history => { return { history } })
+            .catch(error => { return { history: {}, error } })
     }
 
     /**
@@ -260,18 +287,21 @@ export default class ComfyAPI extends EventTarget {
      * @param {*} type The endpoint to post to
      * @param {*} body Optional POST data
      */
-    private async postItem(type: string, body: any) {
+    private async postItem(type: QueueItemType, body: any): Promise<Response> {
         try {
-            await fetch("/" + type, {
-                method: "POST",
-                headers: {
-                    "Content-Type": "application/json",
-                },
-                body: body ? JSON.stringify(body) : undefined,
-            });
-        } catch (error) {
-            console.error(error);
+            body = body ? JSON.stringify(body) : body
         }
+        catch (error) {
+            return Promise.reject(error)
+        }
+
+        return fetch("/" + type, {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: body
+        });
     }
 
     /**
@@ -279,22 +309,22 @@ export default class ComfyAPI extends EventTarget {
      * @param {string} type The type of item to delete, queue or history
      * @param {number} id The id of the item to delete
      */
-    async deleteItem(type: string, id: number) {
-        await this.postItem(type, { delete: [id] });
+    async deleteItem(type: QueueItemType, id: number): Promise<Response> {
+        return this.postItem(type, { delete: [id] });
     }
 
     /**
      * Clears the specified list
      * @param {string} type The type of list to clear, queue or history
      */
-    async clearItems(type: string) {
-        await this.postItem(type, { clear: true });
+    async clearItems(type: QueueItemType): Promise<Response> {
+        return this.postItem(type, { clear: true });
     }
 
     /**
      * Interrupts the execution of the running prompt
      */
-    async interrupt() {
-        await this.postItem("interrupt", null);
+    async interrupt(): Promise<Response> {
+        return fetch("/interrupt", { method: "POST" });
     }
 }