From a83024c35ccc346acd722804a5f5b4e7a0f24140 Mon Sep 17 00:00:00 2001
From: syntaxbullet <syntaxbullet@protonmail.com>
Date: Sun, 5 Jul 2026 15:47:08 +0200
Subject: [PATCH] feat: add configurable history timeout and polling interval
 for ComfyUI API

---
 app/comfy.ts                | 17 ++++++++++++++---
 commands/generation.test.ts |  2 ++
 commands/generation.ts      |  9 ++++++++-
 3 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/app/comfy.ts b/app/comfy.ts
index 4abe3ea..ce21560 100644
--- a/app/comfy.ts
+++ b/app/comfy.ts
@@ -49,6 +49,8 @@ type ComfyObjectInfo = {
 };
 
 const comfyBaseUrl = process.env.COMFYUI_URL ?? "http://127.0.0.1:8188";
+const comfyHistoryTimeoutMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_TIMEOUT_MS, 20 * 60 * 1000);
+const comfyHistoryPollIntervalMs = parsePositiveInteger(process.env.COMFYUI_HISTORY_POLL_INTERVAL_MS, 1000);
 const defaultModels: Record<GenerateArchitecture, string> = {
   sdxl: "sd_xl_base_1.0.safetensors",
   "z-image": "z_image_bf16.safetensors",
@@ -170,15 +172,18 @@ async function uploadDataUrl(dataUrl: string, filename: string) {
 }
 
 async function waitForHistory(promptId: string) {
-  for (let attempt = 0; attempt < 240; attempt++) {
+  const startedAt = Date.now();
+  let attempts = 0;
+  while (Date.now() - startedAt < comfyHistoryTimeoutMs) {
     const response = await fetch(`${comfyBaseUrl}/history/${promptId}`);
+    attempts += 1;
     if (response.ok) {
       const history = await response.json() as Record<string, unknown>;
       if (history[promptId]) return history[promptId];
     }
-    await Bun.sleep(500);
+    await Bun.sleep(comfyHistoryPollIntervalMs);
   }
-  throw new Error(`Timed out waiting for ComfyUI prompt ${promptId}`);
+  throw new Error(`Timed out waiting for ComfyUI prompt ${promptId} after ${Math.round(comfyHistoryTimeoutMs / 1000)} seconds and ${attempts} checks`);
 }
 
 export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
@@ -413,6 +418,12 @@ function unique<T>(values: T[]) {
   return Array.from(new Set(values));
 }
 
+function parsePositiveInteger(value: string | undefined, fallback: number) {
+  if (!value) return fallback;
+  const parsed = Number(value);
+  return Number.isFinite(parsed) && parsed > 0 ? Math.round(parsed) : fallback;
+}
+
 function nodeErrorsMessage(nodeErrors: unknown): string | undefined {
   if (!nodeErrors) return undefined;
   if (Array.isArray(nodeErrors) && nodeErrors.length === 0) return undefined;
diff --git a/commands/generation.test.ts b/commands/generation.test.ts
index 7d80d40..a9074f2 100644
--- a/commands/generation.test.ts
+++ b/commands/generation.test.ts
@@ -33,6 +33,7 @@ describe("generation commands", () => {
     expect(next.document.assets.find((asset) => asset.id === "generated-asset")?.source).toBe("generated-source");
     expect(next.document.artboards[0]?.layers[0]?.id).toBe("generated-layer");
     expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["generated-layer"] });
+    expect(next.editor.generation).toEqual({ candidates: [], selectedCandidateId: undefined });
   });
 
   test("replaces source asset pixels for inpaint candidates", () => {
@@ -46,6 +47,7 @@ describe("generation commands", () => {
     expect(next.document.assets.find((asset) => asset.id === "source-asset")?.source).toBe("composited-source");
     expect(next.document.assets.find((asset) => asset.id === "source-asset")?.mimeType).toBe("image/png");
     expect(next.editor.selection).toEqual({ artboardId: "a1", layerIds: ["source-layer"] });
+    expect(next.editor.generation).toEqual({ candidates: [], selectedCandidateId: undefined });
   });
 });
 
diff --git a/commands/generation.ts b/commands/generation.ts
index 16c575a..b010412 100644
--- a/commands/generation.ts
+++ b/commands/generation.ts
@@ -3,7 +3,7 @@ import type { ImageDocument } from "@core/document";
 import type { ArtboardId, AssetId, LayerId } from "@core/id";
 import type { ImageLayer } from "@core/image-layer";
 import type { Layer } from "@core/layer";
-import type { GenerationCandidate } from "@editor/state";
+import type { GenerationCandidate, GenerationState } from "@editor/state";
 import type { Command } from "./command";
 import { commandIds } from "./ids";
 
@@ -146,6 +146,7 @@ export const generationApplyCandidateAsLayerCommand: Command<GenerationApplyCand
       document: insertLayerAtTop({ ...state.document, assets: [...state.document.assets, asset] }, candidate.placement.artboardId, layer),
       editor: {
         ...state.editor,
+        generation: clearCommittedGenerationPreview(state.editor.generation),
         selection: { artboardId: candidate.placement.artboardId, layerIds: [layer.id] },
       },
     };
@@ -170,6 +171,7 @@ export const generationReplaceCandidatePixelsCommand: Command<GenerationReplaceC
       },
       editor: {
         ...state.editor,
+        generation: clearCommittedGenerationPreview(state.editor.generation),
         selection: { artboardId: targetLayerLocation.artboardId, layerIds: [candidate.inpaint.targetLayerId] },
       },
     };
@@ -215,3 +217,8 @@ function findLayerInTree(layers: readonly Layer[], layerId: LayerId): Layer | un
   }
   return undefined;
 }
+
+function clearCommittedGenerationPreview(generation: GenerationState): GenerationState {
+  if (generation.candidates.length === 0 && !generation.selectedCandidateId) return generation;
+  return { candidates: [], selectedCandidateId: undefined };
+}