feat(media): live ComfyUI progress over its websocket

A video generation runs a single ComfyUI prompt for minutes, and the media
card only advanced after each prompt finished, so it sat still until the
clip was done.

The ComfyUI provider now follows ComfyUI's standard /ws stream while its
prompts run and maps step progress to the job's overall fraction across
all prompts: never backwards, and short of each prompt's share until its
outputs are downloaded. The history poll stays the source of truth for
completion and errors; a socket that cannot connect, drops, or stays
silent leaves generation exactly as before, and it is closed on finish,
failure, cancel and abort. Each generation uses its own client id, so
concurrent jobs do not take over each other's messages.

The socket goes through a new createOpenAiCompatibleWebSocket next to the
guarded fetch: same URL and key checks, same resolved-address lookup,
ws/wss derived from http/https, the same Authorization header, a message
size cap, and a hard drop when a server never finishes the closing
handshake. It uses the undici package's WebSocket, so no dependency was
added. A provider given a replaced fetch without a matching socket opens
none, so tests and emulators never reach the network.

The ComfyUI emulator now serves the websocket and reports progress for
running prompts; provider, guard and conformance tests cover ordering,
multi-prompt mapping, unavailable and failing sockets, stray messages,
closing on cancel, and the address checks against a loopback server.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u

packages/adapters/src/comfyui-provider.test.ts

@@ -15,12 +15,18 @@ function setup(config: Partial<ComfyUiProviderConfig> = {}, emulator = new Comfy
   const provider = createComfyUiProvider({
     baseUrl: BASE_URL,
     fetch: emulator.fetch,
+    webSocket: emulator.openWebSocket,
     pollIntervalMs: 1,
     ...config,
   });
   return { provider, emulator };
 }
 
+/** Progress rounded so float sums compare exactly. */
+function rounded(progress: number[]): number[] {
+  return progress.map((fraction) => Number(fraction.toFixed(4)));
+}
+
 function submitted(emulator: ComfyUiEmulator, index = 0) {
   const graph = emulator.prompts[index]?.graph ?? {};
   const byClass = (classType: string) =>
@@ -195,6 +201,119 @@ describe("createComfyUiProvider", () => {
     expect(emulator.interrupted).toEqual(["prompt-1"]);
   });
 
+  it("reports live step progress in order, short of done until the outputs are in", async () => {
+    const { provider, emulator } = setup();
+    emulator.pendingPolls = 4;
+    const progress: number[] = [];
+    await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
+    expect(progress).toEqual([0.25, 0.5, 0.75, 0.99, 1]);
+    const [socket] = emulator.sockets;
+    expect(socket?.url.href).toBe(
+      `ws://127.0.0.1:8188/ws?clientId=${emulator.prompts[0]?.clientId}`,
+    );
+    expect(socket?.closedByClient).toBe(true);
+  });
+
+  it("maps progress across the prompts of a job", async () => {
+    const { provider, emulator } = setup();
+    emulator.pendingPolls = 2;
+    const progress: number[] = [];
+    await provider.generateImage({
+      ...IMAGE,
+      count: 2,
+      onProgress: (fraction) => progress.push(fraction),
+    });
+    expect(rounded(progress)).toEqual([0.25, 0.495, 0.5, 0.75, 0.995, 1]);
+    expect(emulator.sockets).toHaveLength(1);
+  });
+
+  it("sends the key and keeps the path prefix on the websocket", async () => {
+    const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" });
+    const { provider } = setup(
+      { baseUrl: "https://127.0.0.1:8188/comfy", apiKey: "comfy-key" },
+      emulator,
+    );
+    const progress: number[] = [];
+    await provider.generateImage({ ...IMAGE, onProgress: (fraction) => progress.push(fraction) });
+    const socket = emulator.sockets[0]!;
+    expect(`${socket.url.origin}${socket.url.pathname}`).toBe("wss://127.0.0.1:8188/comfy/ws");
+    expect(socket.authorization).toBe("Bearer comfy-key");
+    expect(progress).toEqual([0.99, 1]);
+  });
+
+  it("generates as before when the websocket is unavailable", async () => {
+    const emulator = new ComfyUiEmulator();
+    emulator.webSocketAvailable = false;
+    const progress: number[] = [];
+    const onProgress = (fraction: number) => progress.push(fraction);
+    await setup({}, emulator).provider.generateImage({ ...IMAGE, onProgress });
+    const throwing = setup({
+      webSocket: () => {
+        throw new Error("no websocket");
+      },
+    });
+    await throwing.provider.generateImage({ ...IMAGE, onProgress });
+    // A substituted fetch without a matching socket opens none, rather than a real connection.
+    const fetchOnly = new ComfyUiEmulator();
+    await setup({ webSocket: undefined }, fetchOnly).provider.generateImage({
+      ...IMAGE,
+      onProgress,
+    });
+    expect(progress).toEqual([1, 1, 1]);
+    expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3]);
+    expect(fetchOnly.sockets).toEqual([]);
+  });
+
+  it("ignores stray messages, a socket that drops mid-run and a failing listener", async () => {
+    const { provider, emulator } = setup();
+    emulator.pendingPolls = 3;
+    const progress: number[] = [];
+    await provider.generateImage({
+      ...IMAGE,
+      onProgress: (fraction) => {
+        progress.push(fraction);
+        const socket = emulator.sockets[0]!;
+        socket.push("not json");
+        socket.push({ type: "progress", data: { value: "2", max: 3, prompt_id: "prompt-1" } });
+        socket.push({ type: "progress", data: { value: 2, max: 0, prompt_id: "prompt-1" } });
+        socket.push({ type: "progress", data: { value: 3, max: 3, prompt_id: "prompt-9" } });
+        socket.drop();
+      },
+    });
+    expect(rounded(progress)).toEqual([0.3333, 1]);
+    const failing = setup().provider;
+    await expect(
+      failing.generateImage({
+        ...IMAGE,
+        onProgress: () => {
+          throw new Error("listener failed");
+        },
+      }),
+    ).resolves.toHaveLength(1);
+  });
+
+  it("closes the websocket when the caller cancels or the run fails", async () => {
+    const emulator = new ComfyUiEmulator();
+    emulator.pendingPolls = 1_000_000;
+    const { provider } = setup({ pollIntervalMs: 5 }, emulator);
+    const controller = new AbortController();
+    const onProgress = () => undefined;
+    const run = provider.generateImage({ ...IMAGE, signal: controller.signal, onProgress });
+    while (!emulator.requests.some((request) => request.path.startsWith("history/"))) {
+      await new Promise((resolve) => setTimeout(resolve, 1));
+    }
+    expect(emulator.sockets[0]?.readyState).toBe(1);
+    controller.abort(new Error("cancelled by user"));
+    expect(emulator.sockets[0]?.closedByClient).toBe(true);
+    await expect(run).rejects.toThrow("cancelled by user");
+    emulator.pendingPolls = 1;
+    emulator.failNextExecution = "CUDA out of memory";
+    await expect(provider.generateImage({ ...IMAGE, onProgress })).rejects.toThrow(
+      "CUDA out of memory",
+    );
+    expect(emulator.sockets.map((socket) => socket.readyState)).toEqual([3, 3]);
+  });
+
   it("reports GPU memory other processes leave to ComfyUI", async () => {
     const emulator = new ComfyUiEmulator();
     emulator.gpu = { totalMb: 12_288, freeMb: 3_000, torchReservedMb: 7_000, torchFreeMb: 1_000 };

packages/adapters/src/comfyui-provider.ts

@@ -30,7 +30,11 @@ import {
   assertAllowedOpenAiCompatibleRequestUrl,
   assertHttpsForKeyedOpenAiCompatibleUrl,
 } from "./openai-compatible-url.js";
-import { createOpenAiCompatibleFetch } from "./pi-openai-compatible-provider.js";
+import type { OpenWebSocket } from "./pi-openai-compatible-provider.js";
+import {
+  createOpenAiCompatibleFetch,
+  createOpenAiCompatibleWebSocket,
+} from "./pi-openai-compatible-provider.js";
 import { readBodyCapped } from "./web-ssrf.js";
 
 export const COMFYUI_PROVIDER_ID = "comfyui";
@@ -43,6 +47,8 @@ const CANCEL_TIMEOUT_MS = 5_000;
 const IMAGE_PROMPT_TIMEOUT_MS = 10 * 60_000;
 const VIDEO_PROMPT_TIMEOUT_MS = 30 * 60_000;
 const POLL_INTERVAL_MS = 1_000;
+/** Live progress stops short of a finished prompt, which only counts once its outputs are in. */
+const RUNNING_PROMPT_MAX_PROGRESS = 0.99;
 const MAX_JSON_BYTES = 8 * 1024 * 1024;
 const MAX_ERROR_CHARS = 400;
 const MIB = 1024 * 1024;
@@ -65,6 +71,11 @@ export interface ComfyUiProviderConfig {
   /** API-format workflow JSON that replaces the matching default video workflow. */
   videoWorkflow?: string;
   fetch?: typeof globalThis.fetch;
+  /**
+   * Opens the progress websocket; tests pass the emulator's. A replaced `fetch` without one gets
+   * no live progress, so the socket never takes a different route than the requests.
+   */
+  webSocket?: OpenWebSocket;
   /** How long to wait between checks for a finished prompt; tests shorten it. */
   pollIntervalMs?: number;
 }
@@ -103,8 +114,9 @@ interface GenerationJob {
 
 /**
  * Image and video generation on a self-hosted ComfyUI server through its HTTP API: submit an
- * API-format workflow, poll its history, download the outputs. Requests go through the same
+ * API-format workflow, poll its history, download the outputs. While a prompt runs, step
- * address-checked fetch as OpenAI-compatible model servers.
+ * progress comes from ComfyUI's websocket when it is reachable. Requests and the websocket go
+ * through the same address checks as OpenAI-compatible model servers.
  */
 export function createComfyUiProvider(config: ComfyUiProviderConfig): ComfyUiProvider {
   return new ComfyUiClient(config);
@@ -117,8 +129,8 @@ class ComfyUiClient implements ComfyUiProvider {
   private readonly imageWorkflow: WorkflowTemplate | undefined;
   private readonly videoWorkflow: WorkflowTemplate | undefined;
   private readonly fetch: typeof globalThis.fetch;
+  private readonly openSocket: ReturnType<typeof createOpenAiCompatibleWebSocket> | undefined;
   private readonly pollIntervalMs: number;
-  private readonly clientId = randomUUID();
 
   constructor(config: ComfyUiProviderConfig) {
     this.baseUrl = comfyUiBaseUrl(config.baseUrl);
@@ -132,6 +144,10 @@ class ComfyUiClient implements ComfyUiProvider {
       ? parseWorkflowTemplate(config.videoWorkflow)
       : undefined;
     this.fetch = createOpenAiCompatibleFetch(config.fetch);
+    this.openSocket =
+      config.fetch && !config.webSocket
+        ? undefined
+        : createOpenAiCompatibleWebSocket(config.webSocket);
     this.pollIntervalMs = config.pollIntervalMs ?? POLL_INTERVAL_MS;
   }
 
@@ -213,38 +229,73 @@ class ComfyUiClient implements ComfyUiProvider {
     const timed = template.placeholders.has("frames") && template.placeholders.has("fps");
     const firstSeed = job.seed ?? randomInt(2 ** 48 - 1);
     const results: GeneratedMedia[] = [];
-    for (let index = 0; index < job.count; index += 1) {
+    // One client id per generation, so concurrent jobs each get their own websocket messages.
-      const graph = fillWorkflow(template, { ...values, seed: firstSeed + index });
+    const clientId = randomUUID();
-      const outputs = await this.runPrompt(graph, job.kind, signal);
+    const progress = new JobProgress(job.count, job.onProgress);
-      const files = outputFiles(outputs, job.kind);
+    const stopSocket = new AbortController();
-      if (files.length === 0) throw new Error(`ComfyUI produced no ${job.kind}`);
+    if (job.onProgress) {
-      for (const file of files) {
+      const socketSignal = signal
-        results.push({
+        ? AbortSignal.any([signal, stopSocket.signal])
-          bytes: await this.download(file, signal),
+        : stopSocket.signal;
-          mimeType: file.mimeType,
+      this.followProgress(clientId, progress, socketSignal);
-          name: file.filename,
+    }
-          ...(sized && file.kind === job.kind ? { width: job.width, height: job.height } : {}),
+    try {
-          ...(timed && file.kind === "video"
+      for (let index = 0; index < job.count; index += 1) {
-            ? { durationMs: Math.round(((values.frames ?? 1) / COMFYUI_VIDEO_FPS) * 1000) }
+        const graph = fillWorkflow(template, { ...values, seed: firstSeed + index });
-            : {}),
+        const outputs = await this.runPrompt(graph, job.kind, clientId, signal, (promptId) =>
-        });
+          progress.start(promptId, index),
-      }
+        );
-      job.onProgress?.((index + 1) / job.count);
+        const files = outputFiles(outputs, job.kind);
+        if (files.length === 0) throw new Error(`ComfyUI produced no ${job.kind}`);
+        for (const file of files) {
+          results.push({
+            bytes: await this.download(file, signal),
+            mimeType: file.mimeType,
+            name: file.filename,
+            ...(sized && file.kind === job.kind ? { width: job.width, height: job.height } : {}),
+            ...(timed && file.kind === "video"
+              ? { durationMs: Math.round(((values.frames ?? 1) / COMFYUI_VIDEO_FPS) * 1000) }
+              : {}),
+          });
+        }
+        progress.finish(index);
+      }
+      return results;
+    } finally {
+      stopSocket.abort();
+    }
+  }
+
+  /**
+   * Best effort: listens on ComfyUI's websocket for step progress. The history poll stays the
+   * source of truth, so a socket that fails, closes or stays silent changes nothing else.
+   */
+  private followProgress(clientId: string, progress: JobProgress, signal: AbortSignal): void {
+    if (!this.openSocket) return;
+    const url = new URL("ws", this.baseUrl);
+    url.searchParams.set("clientId", clientId);
+    try {
+      const socket = this.openSocket(url.href, { headers: this.authHeaders(), signal });
+      socket.addEventListener("message", (event) => progress.receive(event.data));
+    } catch {
+      // Without the socket, progress still advances as each prompt finishes.
     }
-    return results;
   }
 
   /** Submits a prompt and waits for its outputs; a cancelled or timed-out prompt is dequeued. */
   private async runPrompt(
     graph: ComfyWorkflow,
     kind: MediaKind,
+    clientId: string,
     signal: AbortSignal | undefined,
+    onQueued: (promptId: string) => void,
   ): Promise<Record<string, unknown>> {
     const deadline = AbortSignal.timeout(
       kind === "image" ? IMAGE_PROMPT_TIMEOUT_MS : VIDEO_PROMPT_TIMEOUT_MS,
     );
     const runSignal = signal ? AbortSignal.any([signal, deadline]) : deadline;
-    const promptId = await this.submit(graph, runSignal);
+    const promptId = await this.submit(graph, clientId, runSignal);
+    onQueued(promptId);
     try {
       for (;;) {
         const outputs = await this.finishedOutputs(promptId, runSignal);
@@ -259,8 +310,12 @@ class ComfyUiClient implements ComfyUiProvider {
     }
   }
 
-  private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise<string> {
+  private async submit(
-    const payload = { prompt: graph, client_id: this.clientId };
+    graph: ComfyWorkflow,
+    clientId: string,
+    signal: AbortSignal,
+  ): Promise<string> {
+    const payload = { prompt: graph, client_id: clientId };
     return this.postJson("prompt", payload, signal, async (response, readSignal) => {
       if (response.status === 400) {
         const rejection = await readJson(response, readSignal).catch(() => undefined);
@@ -398,8 +453,7 @@ class ComfyUiClient implements ComfyUiProvider {
   ): Promise<T> {
     const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS);
     const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout;
-    const headers: Record<string, string> = { ...init.headers };
+    const headers = { ...init.headers, ...this.authHeaders() };
-    if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`;
     try {
       const response = await this.fetch(new URL(path, this.baseUrl).href, {
         method: init.method ?? "GET",
@@ -419,6 +473,65 @@ class ComfyUiClient implements ComfyUiProvider {
       throw error;
     }
   }
+
+  private authHeaders(): Record<string, string> {
+    return this.apiKey ? { Authorization: `Bearer ${this.apiKey}` } : {};
+  }
+}
+
+/**
+ * Maps prompt progress to the job's overall fraction across its prompts: never backwards, and
+ * short of each prompt's share until its outputs are downloaded. Only ComfyUI's `progress`
+ * messages for the running prompt count; the history poll decides completion and errors.
+ */
+class JobProgress {
+  private promptId: string | undefined;
+  private index = 0;
+  private reported = 0;
+
+  constructor(
+    private readonly count: number,
+    private readonly listener: MediaProgressListener | undefined,
+  ) {}
+
+  start(promptId: string, index: number): void {
+    this.promptId = promptId;
+    this.index = index;
+  }
+
+  finish(index: number): void {
+    this.promptId = undefined;
+    this.report((index + 1) / this.count);
+  }
+
+  /** A websocket message; anything malformed or about another prompt is ignored. */
+  receive(data: unknown): void {
+    if (typeof data !== "string" || !this.promptId) return;
+    let message: unknown;
+    try {
+      message = JSON.parse(data);
+    } catch {
+      return;
+    }
+    if (!isRecord(message) || message.type !== "progress" || !isRecord(message.data)) return;
+    const { value, max, prompt_id: promptId } = message.data;
+    if (promptId !== this.promptId) return;
+    const done = finiteNumber(value);
+    const total = finiteNumber(max);
+    if (done === undefined || total === undefined || total <= 0) return;
+    const fraction = Math.min(Math.max(done / total, 0), RUNNING_PROMPT_MAX_PROGRESS);
+    this.report((this.index + fraction) / this.count);
+  }
+
+  private report(fraction: number): void {
+    if (fraction <= this.reported) return;
+    this.reported = fraction;
+    try {
+      this.listener?.(fraction);
+    } catch {
+      // A failing progress listener must not fail the generation.
+    }
+  }
 }
 
 function comfyUiBaseUrl(raw: string): URL {

packages/adapters/src/image-generation-conformance.test.ts

@@ -12,12 +12,15 @@ import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";
 type MediaGenerationProvider = ImageGenerationProvider & VideoGenerationProvider;
 
 const factories: Record<string, () => MediaGenerationProvider> = {
-  comfyui: () =>
+  comfyui: () => {
-    createComfyUiProvider({
+    const emulator = new ComfyUiEmulator();
+    return createComfyUiProvider({
       baseUrl: "http://127.0.0.1:8188",
-      fetch: new ComfyUiEmulator().fetch,
+      fetch: emulator.fetch,
+      webSocket: emulator.openWebSocket,
       pollIntervalMs: 1,
-    }),
+    });
+  },
 };
 
 async function sourceImage(): Promise<MediaFile> {
@@ -61,6 +64,9 @@ for (const [name, create] of Object.entries(factories)) {
         expect([metadata.width, metadata.height]).toEqual([256, 192]);
       }
       expect(Buffer.from(images[0]!.bytes).equals(Buffer.from(images[1]!.bytes))).toBe(false);
+      // Progress never goes backwards and reaches 1 only once every output is in.
+      expect(progress).toEqual([...progress].sort((a, b) => a - b));
+      expect(progress.filter((fraction) => fraction >= 1)).toEqual([1]);
       expect(progress.at(-1)).toBe(1);
       const again = await provider.generateImage({ ...request, count: 1 });
       expect(Buffer.from(again[0]!.bytes).equals(Buffer.from(images[0]!.bytes))).toBe(true);

packages/adapters/src/pi-openai-compatible-provider.test.ts

@@ -1,12 +1,18 @@
+import { createHash } from "node:crypto";
+import { createServer } from "node:http";
+import type { AddressInfo, Socket } from "node:net";
 import { builtinModels } from "@earendil-works/pi-ai/providers/all";
 import { OPENAI_COMPATIBLE_PROVIDER_ID } from "@rakazo/contracts";
-import { describe, expect, it } from "vitest";
+import { Agent } from "undici";
+import { describe, expect, it, vi } from "vitest";
 import { buildModelConnectPlaintext } from "./model-connect.js";
 import { listPiCatalog } from "./pi-models.js";
 import { parseModelSecret, secretValuesToRedact, serializeModelSecret } from "./pi-oauth.js";
+import type { GuardedWebSocket } from "./pi-openai-compatible-provider.js";
 import {
   createOpenAiCompatibleFetch,
   createOpenAiCompatibleLookup,
+  createOpenAiCompatibleWebSocket,
   OPENAI_COMPATIBLE_CATALOG_MODEL_ID,
   openAiCompatibleCatalogProvider,
   prepareOpenAiCompatibleConnect,
@@ -369,3 +375,178 @@ describe("openai-compatible provider", () => {
     expect(entries.length).toBeGreaterThan(0);
   });
 });
+
+async function withPublicHosts<T>(run: () => Promise<T> | T): Promise<T> {
+  const previous = process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC;
+  process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = "1";
+  try {
+    return await run();
+  } finally {
+    if (previous === undefined) delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC;
+    else process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = previous;
+  }
+}
+
+/** A loopback server that accepts websocket handshakes and sends one text message. */
+async function webSocketServer(options: { answerClose: boolean }) {
+  const server = createServer();
+  const upgrades: { url?: string; authorization?: string }[] = [];
+  const connections: Socket[] = [];
+  server.on("upgrade", (request, socket: Socket) => {
+    upgrades.push({ url: request.url, authorization: request.headers.authorization });
+    connections.push(socket);
+    const accept = createHash("sha1")
+      .update(`${request.headers["sec-websocket-key"]}258EAFA5-E914-47DA-95CA-C5AB0DC85B11`)
+      .digest("base64");
+    socket.write(
+      [
+        "HTTP/1.1 101 Switching Protocols",
+        "Upgrade: websocket",
+        "Connection: Upgrade",
+        `Sec-WebSocket-Accept: ${accept}`,
+        "",
+        "",
+      ].join("\r\n"),
+    );
+    const payload = Buffer.from('{"type":"status"}');
+    socket.write(Buffer.concat([Buffer.from([0x81, payload.length]), payload]));
+    socket.on("data", (chunk: Buffer) => {
+      const isClose = ((chunk[0] ?? 0) & 0x0f) === 0x8;
+      if (isClose && options.answerClose) socket.end(Buffer.from([0x88, 0]));
+    });
+  });
+  await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
+  return {
+    port: (server.address() as AddressInfo).port,
+    upgrades,
+    connections,
+    close: () =>
+      new Promise<void>((resolve) => {
+        for (const connection of connections) connection.destroy();
+        server.close(() => resolve());
+      }),
+  };
+}
+
+function nextMessage(socket: GuardedWebSocket): Promise<unknown> {
+  return new Promise((resolve) =>
+    socket.addEventListener("message", (event) => resolve(event.data)),
+  );
+}
+
+function closeEvent(socket: GuardedWebSocket): Promise<void> {
+  return new Promise((resolve) => socket.addEventListener("close", resolve));
+}
+
+describe("openai-compatible websocket", () => {
+  it("refuses what the fetch guard refuses, before connecting", async () => {
+    const opened: string[] = [];
+    const open = createOpenAiCompatibleWebSocket((url) => {
+      opened.push(url);
+      throw new Error("should not connect");
+    });
+    const { signal } = new AbortController();
+    expect(() => open("http://169.254.169.254/ws", { signal })).toThrow(
+      "blocked metadata or link-local host",
+    );
+    expect(() => open("http://comfy.example.com/ws", { signal })).toThrow(
+      "Public model endpoints are blocked",
+    );
+    expect(() => open("ftp://127.0.0.1/ws", { signal })).toThrow("http or https");
+    await withPublicHosts(() => {
+      const headers = { Authorization: "Bearer secret" };
+      expect(() => open("http://comfy.example.com/ws", { headers, signal })).toThrow(
+        "must use HTTPS",
+      );
+    });
+    expect(() => open("http://127.0.0.1/ws", { signal: AbortSignal.abort() })).toThrow();
+    expect(opened).toEqual([]);
+  });
+
+  it("derives ws or wss from the server URL and connects through its own Agent", () => {
+    const seen: { url: string; dispatcher: unknown }[] = [];
+    const controller = new AbortController();
+    const open = createOpenAiCompatibleWebSocket((url, init) => {
+      seen.push({ url, dispatcher: init.dispatcher });
+      return { addEventListener: () => undefined, close: () => undefined };
+    });
+    open("http://127.0.0.1:8188/ws?clientId=a", { signal: controller.signal });
+    open("https://127.0.0.1:8188/comfy/ws", { signal: controller.signal });
+    controller.abort();
+    expect(seen.map((entry) => entry.url)).toEqual([
+      "ws://127.0.0.1:8188/ws?clientId=a",
+      "wss://127.0.0.1:8188/comfy/ws",
+    ]);
+    expect(seen.every((entry) => entry.dispatcher instanceof Agent)).toBe(true);
+  });
+
+  it("resolves through the address check, sends the headers and closes on abort", async () => {
+    const server = await webSocketServer({ answerClose: true });
+    try {
+      const resolved: string[] = [];
+      const open = createOpenAiCompatibleWebSocket(undefined, async (hostname) => {
+        resolved.push(hostname);
+        return [{ address: "127.0.0.1", family: 4 }];
+      });
+      const controller = new AbortController();
+      const socket = open(`http://localhost:${server.port}/ws?clientId=a`, {
+        headers: { Authorization: "Bearer local-key" },
+        signal: controller.signal,
+      });
+      await expect(nextMessage(socket)).resolves.toBe('{"type":"status"}');
+      expect(resolved).toEqual(["localhost"]);
+      expect(server.upgrades).toEqual([
+        { url: "/ws?clientId=a", authorization: "Bearer local-key" },
+      ]);
+      const closed = closeEvent(socket);
+      controller.abort();
+      await closed;
+    } finally {
+      await server.close();
+    }
+  });
+
+  it("does not connect to a public hostname that resolves to a private address", async () => {
+    const server = await webSocketServer({ answerClose: true });
+    try {
+      const resolved: string[] = [];
+      await withPublicHosts(async () => {
+        const open = createOpenAiCompatibleWebSocket(undefined, async (hostname) => {
+          resolved.push(hostname);
+          return [{ address: "127.0.0.1", family: 4 }];
+        });
+        const socket = open(`http://comfy.example.test:${server.port}/ws`, {
+          signal: new AbortController().signal,
+        });
+        await closeEvent(socket);
+      });
+      expect(resolved).toEqual(["comfy.example.test"]);
+      expect(server.upgrades).toEqual([]);
+    } finally {
+      await server.close();
+    }
+  });
+
+  it("drops the connection when the server never finishes the closing handshake", async () => {
+    const server = await webSocketServer({ answerClose: false });
+    try {
+      const controller = new AbortController();
+      const socket = createOpenAiCompatibleWebSocket()(`http://127.0.0.1:${server.port}/ws`, {
+        signal: controller.signal,
+      });
+      await nextMessage(socket);
+      const closed = closeEvent(socket);
+      const dropped = new Promise((resolve) => server.connections[0]?.on("end", resolve));
+      vi.useFakeTimers({ toFake: ["setTimeout"] });
+      try {
+        controller.abort();
+        vi.advanceTimersByTime(5_000);
+      } finally {
+        vi.useRealTimers();
+      }
+      await Promise.all([closed, dropped]);
+    } finally {
+      await server.close();
+    }
+  });
+});

packages/adapters/src/pi-openai-compatible-provider.ts

@@ -1,4 +1,5 @@
 import { lookup } from "node:dns/promises";
+import type { Socket } from "node:net";
 import { isIP, type LookupFunction } from "node:net";
 import {
   createProvider,
@@ -9,7 +10,8 @@ import {
 } from "@earendil-works/pi-ai";
 import { openAICompletionsApi } from "@earendil-works/pi-ai/api/openai-completions.lazy";
 import { DEFAULT_MODEL_CONTEXT_WINDOW, DEFAULT_MODEL_MAX_TOKENS } from "@rakazo/contracts";
-import { Agent } from "undici";
+import type { Dispatcher } from "undici";
+import { Agent, buildConnector, WebSocket as UndiciWebSocket } from "undici";
 import { declaredVisionModelIds, inputModalities } from "./model-modalities.js";
 import {
   createAddressCheckedLookup,
@@ -45,6 +47,8 @@ const MAX_MODEL_IDS = 500;
 const MAX_MODEL_ID_LENGTH = 256;
 
 const OPENAI_COMPAT_BASE = "http://127.0.0.1:1/v1";
+const MAX_WEBSOCKET_MESSAGE_BYTES = 8 * 1024 * 1024;
+const WEBSOCKET_CLOSE_GRACE_MS = 5_000;
 const resolveHostname: ResolveHostname = (hostname) =>
   lookup(hostname, { all: true, verbatim: true });
 
@@ -244,6 +248,77 @@ async function closeDispatcherWithResponse(
   });
 }
 
+/** The part of a WHATWG WebSocket that guarded callers use; tests substitute emulated sockets. */
+export interface GuardedWebSocket {
+  addEventListener(type: "message", listener: (event: { data: unknown }) => void): void;
+  addEventListener(type: "close", listener: () => void): void;
+  close(): void;
+}
+
+export type OpenWebSocket = (
+  url: string,
+  init: { dispatcher: Dispatcher; headers: Record<string, string> },
+) => GuardedWebSocket;
+
+const openUndiciWebSocket: OpenWebSocket = (url, init) => new UndiciWebSocket(url, init);
+
+/**
+ * The websocket counterpart of `createOpenAiCompatibleFetch`: takes the server's http(s) URL,
+ * applies the same URL, key and resolved-address checks, and connects over ws(s) through an
+ * Agent that caps message size. Aborting `signal` closes the socket, and drops the connection
+ * when the server does not finish the closing handshake.
+ */
+export function createOpenAiCompatibleWebSocket(
+  open: OpenWebSocket = openUndiciWebSocket,
+  resolve: ResolveHostname = resolveHostname,
+) {
+  return (
+    rawUrl: string,
+    init: { headers?: Record<string, string>; signal: AbortSignal },
+  ): GuardedWebSocket => {
+    init.signal.throwIfAborted();
+    const url = assertAllowedOpenAiCompatibleRequestUrl(rawUrl);
+    const headers = init.headers ?? {};
+    if (headersCarryAuthorization(headers)) assertHttpsForKeyedOpenAiCompatibleUrl(url, "present");
+    const connect = buildConnector({ lookup: createOpenAiCompatibleLookup(url, resolve) });
+    const connections = new Set<Socket>();
+    const dispatcher = new Agent({
+      connect: (options, callback) =>
+        connect(options, (...result) => {
+          if (result[1]) connections.add(result[1]);
+          callback(...result);
+        }),
+      webSocket: { maxPayloadSize: MAX_WEBSOCKET_MESSAGE_BYTES },
+    });
+    let released = false;
+    const release = () => {
+      if (released) return;
+      released = true;
+      for (const connection of connections) connection.destroy();
+      dispatcher.destroy().catch(() => undefined);
+    };
+    url.protocol = url.protocol === "https:" ? "wss:" : "ws:";
+    let socket: GuardedWebSocket;
+    try {
+      socket = open(url.href, { dispatcher, headers });
+    } catch (error) {
+      release();
+      throw error;
+    }
+    socket.addEventListener("close", release);
+    init.signal.addEventListener(
+      "abort",
+      () => {
+        if (released) return;
+        socket.close();
+        setTimeout(release, WEBSOCKET_CLOSE_GRACE_MS).unref();
+      },
+      { once: true },
+    );
+    return socket;
+  };
+}
+
 /** Always-visible catalog provider with a placeholder model entry. */
 export function openAiCompatibleCatalogProvider(): Provider {
   // The vision gate resolves against this catalog rather than the per-run

packages/adapters/src/testing/comfyui-emulator.ts

@@ -31,6 +31,8 @@ type FileRef = { filename: string; subfolder: string; type: string };
 
 interface EmulatedPrompt {
   id: string;
+  /** The `client_id` it was submitted with; its websocket messages go to that client. */
+  clientId: string;
   graph: EmulatedGraph;
   polls: number;
   failure?: string;
@@ -57,10 +59,58 @@ function json(body: unknown, status = 200): Response {
   return Response.json(body, { status });
 }
 
+type SocketListener = (event: { data: unknown }) => void;
+
+/** A client's websocket to the emulator; `push` and `drop` act as the server. */
+export class EmulatedComfyUiSocket {
+  /** 0 connecting, 1 open, 2 closing, 3 closed, as in `WebSocket.readyState`. */
+  readyState = 0;
+  closedByClient = false;
+  private readonly listeners = new Map<string, SocketListener[]>();
+
+  constructor(
+    readonly url: URL,
+    readonly authorization: string | null,
+  ) {}
+
+  addEventListener(type: "message", listener: (event: { data: unknown }) => void): void;
+  addEventListener(type: "close", listener: () => void): void;
+  addEventListener(type: string, listener: SocketListener): void {
+    this.listeners.set(type, [...(this.listeners.get(type) ?? []), listener]);
+  }
+
+  close(): void {
+    if (this.readyState >= 2) return;
+    this.closedByClient = true;
+    this.readyState = 2;
+    // The server answers the closing handshake.
+    queueMicrotask(() => this.drop());
+  }
+
+  open(): void {
+    if (this.readyState === 0) this.readyState = 1;
+  }
+
+  /** Sends a text message to the client while the socket is open. */
+  push(message: unknown): void {
+    if (this.readyState !== 1) return;
+    const data = typeof message === "string" ? message : JSON.stringify(message);
+    for (const listener of this.listeners.get("message") ?? []) listener({ data });
+  }
+
+  /** Ends the connection from the server side, like a refused handshake or a lost network. */
+  drop(): void {
+    if (this.readyState === 3) return;
+    this.readyState = 3;
+    for (const listener of this.listeners.get("close") ?? []) listener({ data: undefined });
+  }
+}
+
 /**
- * In-process substitute for a ComfyUI server's HTTP API, so the real provider runs in offline
+ * In-process substitute for a ComfyUI server's HTTP API and websocket, so the real provider runs
- * tests: it validates submitted workflows like ComfyUI does, "runs" them after a few history
+ * in offline tests: it validates submitted workflows like ComfyUI does, "runs" them after a few
- * polls, and returns small generated PNGs or an MP4 fixture.
+ * history polls while reporting their progress, and returns small generated PNGs or an MP4
+ * fixture.
  */
 export class ComfyUiEmulator {
   readonly requests: EmulatedComfyUiRequest[] = [];
@@ -68,6 +118,9 @@ export class ComfyUiEmulator {
   readonly uploads = new Map<string, Uint8Array>();
   readonly deleted: string[] = [];
   readonly interrupted: string[] = [];
+  readonly sockets: EmulatedComfyUiSocket[] = [];
+  /** false refuses websocket connections, like a proxy that does not forward them. */
+  webSocketAvailable = true;
   checkpoints = ["sd_xl_base_1.0.safetensors", "ltx-video-2b-v0.9.5.safetensors"];
   /** `options` answers like current ComfyUI (`["COMBO", {options}]`), `legacy` like older ones. */
   checkpointListFormat: "options" | "legacy" = "options";
@@ -96,12 +149,8 @@ export class ComfyUiEmulator {
     const method = init?.method ?? "GET";
     const headers = new Headers(init?.headers);
     const authorization = headers.get("authorization");
-    const prefix = `/${(this.options.basePath ?? "").replace(/^\/+|\/+$/g, "")}/`.replace(
+    const path = this.apiPath(url);
-      "//",
+    if (path === undefined) return json({ error: "Not found" }, 404);
-      "/",
-    );
-    if (!url.pathname.startsWith(prefix)) return json({ error: "Not found" }, 404);
-    const path = url.pathname.slice(prefix.length);
     const body =
       typeof init?.body === "string" && headers.get("content-type")?.includes("json")
         ? JSON.parse(init.body)
@@ -112,9 +161,7 @@ export class ComfyUiEmulator {
       this.failNextStatus = undefined;
       return new Response("Upstream failure", { status });
     }
-    if (this.options.apiKey && authorization !== `Bearer ${this.options.apiKey}`) {
+    if (!this.authorized(authorization)) return json({ error: "Unauthorized" }, 401);
-      return json({ error: "Unauthorized" }, 401);
-    }
     const route = `${method} ${path.split("/")[0]}`;
     if (route === "GET system_stats") return this.systemStats();
     if (route === "GET object_info" && path === "object_info/CheckpointLoaderSimple") {
@@ -144,6 +191,51 @@ export class ComfyUiEmulator {
     return json({ error: "Not found" }, 404);
   };
 
+  /**
+   * Opens a websocket to the emulator in place of the network. Like ComfyUI's `/ws`, it sends a
+   * status message on connect and the progress of running prompts submitted with its client id.
+   */
+  readonly openWebSocket = (url: string, init: { headers: Record<string, string> }) => {
+    const target = new URL(url);
+    const authorization = new Headers(init.headers).get("authorization");
+    const socket = new EmulatedComfyUiSocket(target, authorization);
+    this.sockets.push(socket);
+    const path = this.apiPath(target);
+    this.requests.push({ method: "GET", path: path ?? target.pathname, authorization });
+    const accepted =
+      this.webSocketAvailable &&
+      path === "ws" &&
+      this.authorized(authorization) &&
+      target.searchParams.has("clientId");
+    queueMicrotask(() => {
+      if (!accepted) return socket.drop();
+      socket.open();
+      const status = { exec_info: { queue_remaining: this.unfinished().length } };
+      socket.push({ type: "status", data: { status, sid: target.searchParams.get("clientId") } });
+    });
+    return socket;
+  };
+
+  private apiPath(url: URL): string | undefined {
+    const prefix = `/${(this.options.basePath ?? "").replace(/^\/+|\/+$/g, "")}/`.replace(
+      "//",
+      "/",
+    );
+    return url.pathname.startsWith(prefix) ? url.pathname.slice(prefix.length) : undefined;
+  }
+
+  private authorized(authorization: string | null): boolean {
+    return !this.options.apiKey || authorization === `Bearer ${this.options.apiKey}`;
+  }
+
+  private notify(prompt: EmulatedPrompt, type: string, data: Record<string, unknown>): void {
+    for (const socket of this.sockets) {
+      if (socket.url.searchParams.get("clientId") === prompt.clientId) {
+        socket.push({ type, data: { ...data, prompt_id: prompt.id } });
+      }
+    }
+  }
+
   private systemStats(): Response {
     const gpu = this.gpu;
     return json({
@@ -230,6 +322,7 @@ export class ComfyUiEmulator {
     const id = `prompt-${this.nextPrompt++}`;
     this.prompts.push({
       id,
+      clientId: (body as { client_id: string }).client_id,
       graph: graph as EmulatedGraph,
       polls: 0,
       failure: this.failNextExecution,
@@ -242,9 +335,24 @@ export class ComfyUiEmulator {
     const prompt = this.prompts.find((entry) => entry.id === id);
     if (!prompt || prompt.removed) return json({});
     if (!prompt.done && !prompt.interrupted && prompt.polls < this.pendingPolls) {
+      // Each poll while it runs, the sampler advances one of `pendingPolls` steps.
       prompt.polls += 1;
+      const node = Object.entries(prompt.graph).find(([, entry]) =>
+        entry.class_type.includes("Sampler"),
+      )?.[0];
+      this.notify(prompt, "executing", { node: node ?? null });
+      this.notify(prompt, "progress", { value: prompt.polls, max: this.pendingPolls, node });
       return json({});
     }
+    if (!prompt.done) {
+      this.notify(prompt, "executing", { node: null });
+      const outcome = prompt.interrupted
+        ? "execution_interrupted"
+        : prompt.failure
+          ? "execution_error"
+          : "execution_success";
+      this.notify(prompt, outcome, {});
+    }
     prompt.done = true;
     if (prompt.interrupted) {
       return json({