Add zod schema for TaskItem (#59)

* Add zod schema for TaskItem

* nit

src/scripts/api.ts

@@ -1,3 +1,6 @@
+import { HistoryTaskItem, PendingTaskItem, RunningTaskItem } from "/types/apiTypes";
+
+
 interface QueuePromptRequestBody {
 	client_id: string;
 	// Mapping from node id to node info + input values
@@ -268,7 +271,7 @@ class ComfyApi extends EventTarget {
 	 * Gets the current state of the queue
 	 * @returns The currently running and queued items
 	 */
-	async getQueue() {
+	async getQueue(): Promise<{ Running: RunningTaskItem[], Pending: PendingTaskItem[] }>{
 		try {
 			const res = await this.fetchApi("/queue");
 			const data = await res.json();
@@ -290,7 +293,7 @@ class ComfyApi extends EventTarget {
 	 * Gets the prompt execution history
 	 * @returns Prompt history including node outputs
 	 */
-	async getHistory(max_items = 200) {
+	async getHistory(max_items: number = 200): Promise<{History: HistoryTaskItem[]}> {
 		try {
 			const res = await this.fetchApi(`/history?max_items=${max_items}`);
 			return { History: Object.values(await res.json()) };

src/scripts/ui.ts

@@ -3,6 +3,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";
 
 export const ComfyDialog = _ComfyDialog;
 
@@ -221,9 +222,9 @@ class ComfyList {
 					textContent: section,
 				}),
 				$el("div.comfy-list-items", [
-					...(this.#reverse ? items[section].reverse() : items[section]).map((item) => {
+					...(this.#reverse ? items[section].reverse() : items[section]).map((item: TaskItem) => {
 						// Allow items to specify a custom remove action (e.g. for interrupt current prompt)
-						const removeAction = item.remove || {
+						const removeAction = "remove" in item ? item.remove : {
 							name: "Delete",
 							cb: () => api.deleteItem(this.#type, item.prompt[1]),
 						};
@@ -232,7 +233,7 @@ class ComfyList {
 								textContent: "Load",
 								onclick: async () => {
 									await app.loadGraphData(item.prompt[3].extra_pnginfo.workflow, true, false);
-									if (item.outputs) {
+									if ("outputs" in item) {
 										app.nodeOutputs = item.outputs;
 									}
 								},

src/types/apiTypes.ts

@@ -0,0 +1,117 @@
+import { z } from "zod";
+import { zComfyWorkflow } from "./comfyWorkflow";
+
+const zNodeId = z.number();
+const zNodeType = z.string();
+const zQueueIndex = z.number();
+const zPromptId = z.string();
+
+const zPromptItem = z.object({
+    inputs: z.record(z.string(), z.any()),
+    class_type: zNodeType,
+});
+
+const zPrompt = z.array(zPromptItem);
+
+const zExtraPngInfo = z.object({
+    workflow: zComfyWorkflow,
+}).passthrough();
+
+const zExtraData = z.object({
+    extra_pnginfo: zExtraPngInfo,
+    client_id: z.string(),
+});
+const zOutputsToExecute = z.array(zNodeId);
+
+const zExecutionStartMessage = z.tuple([
+    z.literal("execution_start"),
+    z.object({
+        prompt_id: zPromptId,
+    }),
+]);
+
+const zExecutionCachedMessage = z.tuple([
+    z.literal("execution_cached"),
+    z.object({
+        prompt_id: zPromptId,
+        nodes: z.array(zNodeId),
+    }),
+]);
+
+const zExecutionInterruptedMessage = z.tuple([
+    z.literal("execution_interrupted"),
+    z.object({
+        // InterruptProcessingException
+        prompt_id: zPromptId,
+        node_id: zNodeId,
+        node_type: zNodeType,
+        executed: z.array(zNodeId),
+    }),
+]);
+
+const zExecutionErrorMessage = z.tuple([
+    z.literal("execution_error"),
+    z.object({
+        prompt_id: zPromptId,
+        node_id: zNodeId,
+        node_type: zNodeType,
+        executed: z.array(zNodeId),
+
+        exception_message: z.string(),
+        exception_type: z.string(),
+        traceback: z.string(),
+        current_inputs: z.any(),
+        current_outputs: z.any(),
+    }),
+]);
+
+const zStatusMessage = z.union([
+    zExecutionStartMessage,
+    zExecutionCachedMessage,
+    zExecutionInterruptedMessage,
+    zExecutionErrorMessage,
+]);
+
+const zStatus = z.object({
+    status_str: z.enum(["success", "error"]),
+    completed: z.boolean(),
+    messages: z.array(zStatusMessage),
+});
+
+// TODO: this is a placeholder
+const zOutput = z.any();
+
+const zTaskPrompt = z.tuple([
+    zQueueIndex,
+    zPromptId,
+    zPrompt,
+    zExtraData,
+    zOutputsToExecute,
+]);
+
+const zRunningTaskItem = z.object({
+    prompt: zTaskPrompt,
+    remove: z.object({
+        name: z.literal("Cancel"),
+        cb: z.function(),
+    }),
+});
+
+const zPendingTaskItem = z.object({
+    prompt: zTaskPrompt,
+});
+
+const zHistoryTaskItem = z.object({
+    prompt: zTaskPrompt,
+    status: zStatus.optional(),
+    outputs: z.record(zNodeId, zOutput),
+});
+
+const zTaskItem = z.union([zRunningTaskItem, zPendingTaskItem, zHistoryTaskItem]);
+
+export type RunningTaskItem = z.infer<typeof zRunningTaskItem>;
+export type PendingTaskItem = z.infer<typeof zPendingTaskItem>;
+export type HistoryTaskItem = z.infer<typeof zHistoryTaskItem>;
+export type TaskItem = z.infer<typeof zTaskItem>;
+
+// TODO: validate `/history` `/queue` API endpoint responses.

src/types/comfyWorkflow.ts

@@ -90,7 +90,7 @@ const zExtra = z.object({
     info: zInfo.optional(),
 }).passthrough();
 
-const zComfyWorkflow = z.object({
+export const zComfyWorkflow = z.object({
     last_node_id: z.number(),
     last_link_id: z.number(),
     nodes: z.array(zComfyNode),