diff --git a/packages/adapters/src/comfyui-provider.test.ts b/packages/adapters/src/comfyui-provider.test.ts index 2773c1f..4364d53 100644 --- a/packages/adapters/src/comfyui-provider.test.ts +++ b/packages/adapters/src/comfyui-provider.test.ts @@ -15,12 +15,18 @@ function setup(config: Partial = {}, 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 }; diff --git a/packages/adapters/src/comfyui-provider.ts b/packages/adapters/src/comfyui-provider.ts index e2410a1..139ac68 100644 --- a/packages/adapters/src/comfyui-provider.ts +++ b/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 - * address-checked fetch as OpenAI-compatible model servers. + * API-format workflow, poll its history, download the outputs. While a prompt runs, step + * 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 | 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) { - const graph = fillWorkflow(template, { ...values, seed: firstSeed + index }); - const outputs = await this.runPrompt(graph, job.kind, signal); - 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) } - : {}), - }); - } - job.onProgress?.((index + 1) / job.count); + // One client id per generation, so concurrent jobs each get their own websocket messages. + const clientId = randomUUID(); + const progress = new JobProgress(job.count, job.onProgress); + const stopSocket = new AbortController(); + if (job.onProgress) { + const socketSignal = signal + ? AbortSignal.any([signal, stopSocket.signal]) + : stopSocket.signal; + this.followProgress(clientId, progress, socketSignal); + } + try { + for (let index = 0; index < job.count; index += 1) { + const graph = fillWorkflow(template, { ...values, seed: firstSeed + index }); + const outputs = await this.runPrompt(graph, job.kind, clientId, signal, (promptId) => + progress.start(promptId, index), + ); + 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> { 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 { - const payload = { prompt: graph, client_id: this.clientId }; + private async submit( + graph: ComfyWorkflow, + clientId: string, + signal: AbortSignal, + ): Promise { + 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 { const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS); const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; - const headers: Record = { ...init.headers }; - if (this.apiKey) headers.Authorization = `Bearer ${this.apiKey}`; + const headers = { ...init.headers, ...this.authHeaders() }; 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 { + 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 { diff --git a/packages/adapters/src/image-generation-conformance.test.ts b/packages/adapters/src/image-generation-conformance.test.ts index 92883a5..2117f1a 100644 --- a/packages/adapters/src/image-generation-conformance.test.ts +++ b/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 MediaGenerationProvider> = { - comfyui: () => - createComfyUiProvider({ + comfyui: () => { + 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 { @@ -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); diff --git a/packages/adapters/src/pi-openai-compatible-provider.test.ts b/packages/adapters/src/pi-openai-compatible-provider.test.ts index 3594dbd..20e825e 100644 --- a/packages/adapters/src/pi-openai-compatible-provider.test.ts +++ b/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(run: () => Promise | T): Promise { + 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((resolve) => server.listen(0, "127.0.0.1", resolve)); + return { + port: (server.address() as AddressInfo).port, + upgrades, + connections, + close: () => + new Promise((resolve) => { + for (const connection of connections) connection.destroy(); + server.close(() => resolve()); + }), + }; +} + +function nextMessage(socket: GuardedWebSocket): Promise { + return new Promise((resolve) => + socket.addEventListener("message", (event) => resolve(event.data)), + ); +} + +function closeEvent(socket: GuardedWebSocket): Promise { + 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(); + } + }); +}); diff --git a/packages/adapters/src/pi-openai-compatible-provider.ts b/packages/adapters/src/pi-openai-compatible-provider.ts index 81e3b8f..614a47b 100644 --- a/packages/adapters/src/pi-openai-compatible-provider.ts +++ b/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 }, +) => 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; 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(); + 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 diff --git a/packages/adapters/src/testing/comfyui-emulator.ts b/packages/adapters/src/testing/comfyui-emulator.ts index 4da22a1..f682ab7 100644 --- a/packages/adapters/src/testing/comfyui-emulator.ts +++ b/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(); + + 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 - * tests: it validates submitted workflows like ComfyUI does, "runs" them after a few history - * polls, and returns small generated PNGs or an MP4 fixture. + * In-process substitute for a ComfyUI server's HTTP API and websocket, so the real provider runs + * in offline tests: it validates submitted workflows like ComfyUI does, "runs" them after a few + * 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(); 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( - "//", - "/", - ); - if (!url.pathname.startsWith(prefix)) return json({ error: "Not found" }, 404); - const path = url.pathname.slice(prefix.length); + const path = this.apiPath(url); + if (path === undefined) return json({ error: "Not found" }, 404); 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}`) { - return json({ error: "Unauthorized" }, 401); - } + if (!this.authorized(authorization)) 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 }) => { + 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): 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({