Merge f/comfy-progress: live ComfyUI progress over its websocket
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
This commit is contained in:
Claude
commit
81991251a9
@ -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 };
|
||||
|
||||
@ -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<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) {
|
||||
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<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> {
|
||||
const payload = { prompt: graph, client_id: this.clientId };
|
||||
private async submit(
|
||||
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 };
|
||||
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<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 {
|
||||
|
||||
@ -12,12 +12,15 @@ import { ComfyUiEmulator } from "./testing/comfyui-emulator.js";
|
||||
type MediaGenerationProvider = ImageGenerationProvider & VideoGenerationProvider;
|
||||
|
||||
const factories: Record<string, () => 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<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);
|
||||
|
||||
@ -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();
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
* 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<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(
|
||||
"//",
|
||||
"/",
|
||||
);
|
||||
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<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({
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user