From fa0986bd2d276fb6d532a62d205378c431d0c19c Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:26:59 +0000 Subject: [PATCH 1/6] feat(media): add the ComfyUI provider, default workflows and emulator Image and video generation run on the user's own GPU server. ComfyUI is the common self-hosted backend, so this adds a provider for its HTTP API that the generation operations can call. - Core: API-format workflow templates with a small placeholder format ({{prompt}}, {{width}}, {{seed}}, ...) that is validated and filled without evaluating anything. Defaults: text-to-image and image-to-image on core nodes for any SD-family checkpoint, and text/image-to-video on ComfyUI's core LTX-Video nodes; a README lists the placeholders and model files. - Adapters: createComfyUiProvider submits a workflow, polls its history, uploads input images and downloads outputs through the address-checked fetch used for other self-hosted model servers, with timeouts, capped reads and readable errors. A custom workflow from settings replaces the default it matches. Cancelling dequeues or interrupts only its own prompt. GPU status counts memory ComfyUI can free for itself as available. - Contracts: input schemas for generate_image, generate_video and edit_image. - Tests: an in-process ComfyUI emulator, a provider-neutral generation conformance suite and provider tests, all offline. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- .../adapters/src/comfyui-provider.test.ts | 227 +++++++ packages/adapters/src/comfyui-provider.ts | 596 ++++++++++++++++++ .../src/image-generation-conformance.test.ts | 111 ++++ packages/adapters/src/index.ts | 1 + .../adapters/src/testing/comfyui-emulator.ts | 393 ++++++++++++ packages/contracts/src/index.ts | 1 + .../contracts/src/media-generation.test.ts | 43 ++ packages/contracts/src/media-generation.ts | 108 ++++ packages/core/package.json | 5 + .../core/src/media/workflow-template.test.ts | 113 ++++ packages/core/src/media/workflow-template.ts | 159 +++++ packages/core/src/media/workflows/README.md | 35 + .../src/media/workflows/image-to-image.json | 45 ++ .../src/media/workflows/image-to-video.json | 86 +++ .../src/media/workflows/text-to-image.json | 41 ++ .../src/media/workflows/text-to-video.json | 77 +++ 16 files changed, 2041 insertions(+) create mode 100644 packages/adapters/src/comfyui-provider.test.ts create mode 100644 packages/adapters/src/comfyui-provider.ts create mode 100644 packages/adapters/src/image-generation-conformance.test.ts create mode 100644 packages/adapters/src/testing/comfyui-emulator.ts create mode 100644 packages/contracts/src/media-generation.test.ts create mode 100644 packages/contracts/src/media-generation.ts create mode 100644 packages/core/src/media/workflow-template.test.ts create mode 100644 packages/core/src/media/workflow-template.ts create mode 100644 packages/core/src/media/workflows/README.md create mode 100644 packages/core/src/media/workflows/image-to-image.json create mode 100644 packages/core/src/media/workflows/image-to-video.json create mode 100644 packages/core/src/media/workflows/text-to-image.json create mode 100644 packages/core/src/media/workflows/text-to-video.json diff --git a/packages/adapters/src/comfyui-provider.test.ts b/packages/adapters/src/comfyui-provider.test.ts new file mode 100644 index 0000000..2773c1f --- /dev/null +++ b/packages/adapters/src/comfyui-provider.test.ts @@ -0,0 +1,227 @@ +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import type { ComfyUiProviderConfig } from "./comfyui-provider.js"; +import { + createComfyUiProvider, + describeExecutionError, + parseGpuStatus, +} from "./comfyui-provider.js"; +import { ComfyUiEmulator } from "./testing/comfyui-emulator.js"; + +const BASE_URL = "http://127.0.0.1:8188"; +const IMAGE = { prompt: "a lighthouse", width: 512, height: 512 }; + +function setup(config: Partial = {}, emulator = new ComfyUiEmulator()) { + const provider = createComfyUiProvider({ + baseUrl: BASE_URL, + fetch: emulator.fetch, + pollIntervalMs: 1, + ...config, + }); + return { provider, emulator }; +} + +function submitted(emulator: ComfyUiEmulator, index = 0) { + const graph = emulator.prompts[index]?.graph ?? {}; + const byClass = (classType: string) => + Object.values(graph).find((node) => node.class_type === classType)?.inputs ?? {}; + return { graph, byClass }; +} + +describe("createComfyUiProvider", () => { + it("refuses metadata hosts, public hosts and keys over public http", () => { + expect(() => createComfyUiProvider({ baseUrl: "http://169.254.169.254" })).toThrow( + "blocked metadata or link-local host", + ); + expect(() => createComfyUiProvider({ baseUrl: "http://comfy.example.com" })).toThrow( + "Public model endpoints are blocked", + ); + expect(() => createComfyUiProvider({ baseUrl: "ftp://127.0.0.1" })).toThrow("http or https"); + }); + + it("validates custom workflows up front", () => { + expect(() => createComfyUiProvider({ baseUrl: BASE_URL, imageWorkflow: "{" })).toThrow( + "Workflow is not valid JSON", + ); + expect(() => + createComfyUiProvider({ baseUrl: BASE_URL, videoWorkflow: '{"nodes": [], "links": []}' }), + ).toThrow("API format"); + }); + + it("keeps a reverse-proxy path prefix and sends the API key", async () => { + const emulator = new ComfyUiEmulator({ apiKey: "comfy-key", basePath: "/comfy" }); + const { provider } = setup({ baseUrl: `${BASE_URL}/comfy/`, apiKey: "comfy-key" }, emulator); + await provider.generateImage(IMAGE); + expect(emulator.requests.map((request) => request.path)).toContain("prompt"); + expect(emulator.requests.every((request) => request.authorization === "Bearer comfy-key")).toBe( + true, + ); + const { provider: keyless } = setup({ baseUrl: `${BASE_URL}/comfy` }, emulator); + await expect(keyless.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 401"); + }); + + it("fills the default text-to-image workflow", async () => { + const { provider, emulator } = setup(); + await provider.generateImage({ ...IMAGE, width: 1001, negativePrompt: "fog", seed: 42 }); + const { byClass } = submitted(emulator); + expect(byClass("CheckpointLoaderSimple").ckpt_name).toBe("sd_xl_base_1.0.safetensors"); + expect(byClass("EmptyLatentImage")).toMatchObject({ width: 1000, height: 512 }); + expect(byClass("KSampler")).toMatchObject({ seed: 42, denoise: 1 }); + const texts = emulator.prompts[0]!.graph; + expect(Object.values(texts).map((node) => node.inputs.text)).toEqual( + expect.arrayContaining(["a lighthouse", "fog"]), + ); + }); + + it("uses the configured model, else the first checkpoint the video workflows do not use", async () => { + const emulator = new ComfyUiEmulator(); + emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors", "v1-5-pruned.safetensors"]; + emulator.checkpointListFormat = "legacy"; + await setup({}, emulator).provider.generateImage(IMAGE); + expect(submitted(emulator, 0).byClass("CheckpointLoaderSimple").ckpt_name).toBe( + "v1-5-pruned.safetensors", + ); + emulator.checkpoints.push("dreamshaper.safetensors"); + await setup({ imageModel: "dreamshaper.safetensors" }, emulator).provider.generateImage(IMAGE); + expect(submitted(emulator, 1).byClass("CheckpointLoaderSimple").ckpt_name).toBe( + "dreamshaper.safetensors", + ); + emulator.checkpoints = ["ltx-video-2b-v0.9.5.safetensors"]; + await expect(setup({}, emulator).provider.generateImage(IMAGE)).rejects.toThrow( + "ComfyUI has no image checkpoint installed", + ); + }); + + it("uploads the source image for image-to-image", async () => { + const { provider, emulator } = setup(); + const bytes = new Uint8Array(Buffer.from((await sourcePng()).bytes)); + await provider.generateImage({ + ...IMAGE, + source: { bytes, mimeType: "image/png" }, + strength: 0.3, + }); + const { byClass } = submitted(emulator); + const uploaded = String(byClass("LoadImage").image); + expect(uploaded).toMatch(/^evoo-.+\.png$/); + expect(Buffer.from(emulator.uploads.get(uploaded)!).equals(Buffer.from(bytes))).toBe(true); + expect(byClass("KSampler").denoise).toBe(0.3); + }); + + it("lets a custom workflow replace only the default it matches", async () => { + const custom = { + "1": { + class_type: "CustomTextToImage", + inputs: { text: "style: {{prompt}}", seed: "{{seed}}" }, + }, + "2": { class_type: "SaveImage", inputs: { images: ["1", 0], filename_prefix: "x" } }, + }; + const { provider, emulator } = setup({ imageWorkflow: JSON.stringify(custom) }); + await provider.generateImage(IMAGE); + expect(submitted(emulator, 0).byClass("CustomTextToImage").text).toBe("style: a lighthouse"); + await provider.generateImage({ ...IMAGE, source: await sourcePng() }); + expect(submitted(emulator, 1).byClass("LoadImage").image).toBeDefined(); + }); + + it("fills the LTX-Video workflow with an 8n+1 frame count and reports the clip", async () => { + const { provider, emulator } = setup(); + const [clip] = await provider.generateVideo({ + prompt: "a paper boat", + width: 470, + height: 836, + durationSeconds: 5, + }); + const { byClass } = submitted(emulator); + expect(byClass("EmptyLTXVLatentVideo")).toMatchObject({ width: 480, height: 832, length: 121 }); + expect(byClass("CreateVideo").fps).toBe(24); + expect(clip).toMatchObject({ + mimeType: "video/mp4", + width: 480, + height: 832, + durationMs: 5042, + }); + }); + + it("reads video from a custom workflow's gifs output and skips frame previews", async () => { + const custom = { + "1": { class_type: "TextToFrames", inputs: { text: "{{prompt}}" } }, + "2": { class_type: "PreviewImage", inputs: { images: ["1", 0] } }, + "3": { class_type: "VHS_VideoCombine", inputs: { images: ["1", 0], frame_rate: 16 } }, + }; + const { provider } = setup({ videoWorkflow: JSON.stringify(custom) }); + const clips = await provider.generateVideo({ ...IMAGE, durationSeconds: 2 }); + expect(clips.map((clip) => [clip.mimeType, clip.durationMs])).toEqual([ + ["video/mp4", undefined], + ]); + }); + + it("explains rejected workflows and failed runs", async () => { + const emulator = new ComfyUiEmulator(); + const { provider } = setup({ imageModel: "missing.safetensors" }, emulator); + await expect(provider.generateImage(IMAGE)).rejects.toThrow( + "ComfyUI rejected the workflow: CheckpointLoaderSimple: Value not in list: ckpt_name: 'missing.safetensors' not in list", + ); + const working = setup({}, emulator).provider; + emulator.failNextExecution = "CUDA out of memory. Tried to allocate 2.00 GiB\nTraceback..."; + await expect(working.generateImage(IMAGE)).rejects.toThrow( + /^ComfyUI failed in KSampler: CUDA out of memory\. Tried to allocate 2\.00 GiB$/, + ); + emulator.failNextStatus = 500; + await expect(working.generateImage(IMAGE)).rejects.toThrow("ComfyUI returned 500"); + }); + + it("reports an unreachable server", async () => { + const provider = createComfyUiProvider({ + baseUrl: BASE_URL, + fetch: async () => { + throw new TypeError("fetch failed"); + }, + }); + await expect(provider.generateImage(IMAGE)).rejects.toThrow("ComfyUI is unreachable"); + await expect(provider.gpuStatus()).rejects.toThrow("ComfyUI is unreachable"); + }); + + it("interrupts its own running prompt when the caller cancels", async () => { + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 1_000_000; + const { provider } = setup({ pollIntervalMs: 5 }, emulator); + const controller = new AbortController(); + const run = provider.generateImage({ ...IMAGE, signal: controller.signal }); + while (!emulator.requests.some((request) => request.path.startsWith("history/"))) { + await new Promise((resolve) => setTimeout(resolve, 1)); + } + controller.abort(new Error("cancelled by user")); + await expect(run).rejects.toThrow("cancelled by user"); + expect(emulator.deleted).toEqual(["prompt-1"]); + expect(emulator.interrupted).toEqual(["prompt-1"]); + }); + + 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 }; + await expect(setup({}, emulator).provider.gpuStatus()).resolves.toEqual({ + vramFreeMb: 9_000, + vramTotalMb: 12_288, + }); + emulator.gpu = null; + await expect(setup({}, emulator).provider.gpuStatus()).resolves.toBeUndefined(); + expect(parseGpuStatus({ devices: [{ type: "cuda", vram_total: 0, vram_free: 0 }] })).toBe( + undefined, + ); + }); + + it("names interrupted runs", () => { + expect(describeExecutionError([["execution_interrupted", {}]])).toBe( + "ComfyUI interrupted the run", + ); + expect(describeExecutionError(undefined)).toBe("ComfyUI failed"); + }); +}); + +async function sourcePng() { + const png = await sharp({ + create: { width: 32, height: 32, channels: 3, background: { r: 1, g: 2, b: 3 } }, + }) + .png() + .toBuffer(); + return { bytes: new Uint8Array(png), mimeType: "image/png" }; +} diff --git a/packages/adapters/src/comfyui-provider.ts b/packages/adapters/src/comfyui-provider.ts new file mode 100644 index 0000000..ec35de7 --- /dev/null +++ b/packages/adapters/src/comfyui-provider.ts @@ -0,0 +1,596 @@ +import { randomInt, randomUUID } from "node:crypto"; +import { setTimeout as sleep } from "node:timers/promises"; +import type { + AdapterDescriptor, + GeneratedMedia, + GpuStatus, + ImageGenerationCapabilities, + ImageGenerationProvider, + ImageGenerationRequest, + MediaFile, + MediaProgressListener, + VideoGenerationCapabilities, + VideoGenerationProvider, + VideoGenerationRequest, +} from "@rakazo/adapter-kit"; +import type { AttachmentKind } from "@rakazo/contracts"; +import { + ARTIFACT_MAX_BYTES_BY_KIND, + attachmentKindOf, + MEDIA_VIDEO_MAX_SECONDS, +} from "@rakazo/contracts"; +import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; +import type { ComfyWorkflow, WorkflowTemplate, WorkflowValues } from "@rakazo/core/media/workflows"; +import { + DEFAULT_WORKFLOWS, + fillWorkflow, + parseWorkflowTemplate, +} from "@rakazo/core/media/workflows"; +import { + assertAllowedOpenAiCompatibleRequestUrl, + assertHttpsForKeyedOpenAiCompatibleUrl, +} from "./openai-compatible-url.js"; +import { createOpenAiCompatibleFetch } from "./pi-openai-compatible-provider.js"; +import { readBodyCapped } from "./web-ssrf.js"; + +export const COMFYUI_PROVIDER_ID = "comfyui"; +/** Frame rate of generated video; frame counts follow LTX-Video's 8n+1 rule. */ +export const COMFYUI_VIDEO_FPS = 24; +const MAX_IMAGE_SIDE = 2048; +const REQUEST_TIMEOUT_MS = 30_000; +const TRANSFER_TIMEOUT_MS = 120_000; +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; +const MAX_JSON_BYTES = 8 * 1024 * 1024; +const MAX_ERROR_CHARS = 400; +const MIB = 1024 * 1024; +const IMAGE_STEPS = 25; +const VIDEO_STEPS = 30; +const DEFAULT_STRENGTH = 0.6; +const DEFAULT_NEGATIVE_PROMPTS = { + image: "blurry, low quality, distorted, watermark, text", + video: "worst quality, inconsistent motion, blurry, jittery, distorted", +}; +const JSON_HEADERS = { Accept: "application/json", "Content-Type": "application/json" }; + +export interface ComfyUiProviderConfig { + baseUrl: string; + apiKey?: string; + /** Checkpoint for `{{model}}`; the first installed image checkpoint when omitted. */ + imageModel?: string; + /** API-format workflow JSON that replaces the matching default image workflow. */ + imageWorkflow?: string; + /** API-format workflow JSON that replaces the matching default video workflow. */ + videoWorkflow?: string; + fetch?: typeof globalThis.fetch; + /** How long to wait between checks for a finished prompt; tests shorten it. */ + pollIntervalMs?: number; +} + +export interface ComfyUiProvider extends ImageGenerationProvider, VideoGenerationProvider { + describe(): AdapterDescriptor; + gpuStatus(signal?: AbortSignal): Promise; +} + +type MediaKind = "image" | "video"; + +interface OutputFile { + filename: string; + subfolder: string; + type: string; + mimeType: string; + kind: AttachmentKind; +} + +interface GenerationJob { + kind: MediaKind; + template: WorkflowTemplate; + prompt: string; + negativePrompt?: string; + width: number; + height: number; + seed?: number; + count: number; + source?: MediaFile; + strength?: number; + frames?: number; + signal?: AbortSignal; + onProgress?: MediaProgressListener; +} + +/** + * 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. + */ +export function createComfyUiProvider(config: ComfyUiProviderConfig): ComfyUiProvider { + return new ComfyUiClient(config); +} + +class ComfyUiClient implements ComfyUiProvider { + private readonly baseUrl: URL; + private readonly apiKey: string | undefined; + private readonly imageModel: string | undefined; + private readonly imageWorkflow: WorkflowTemplate | undefined; + private readonly videoWorkflow: WorkflowTemplate | undefined; + private readonly fetch: typeof globalThis.fetch; + private readonly pollIntervalMs: number; + private readonly clientId = randomUUID(); + + constructor(config: ComfyUiProviderConfig) { + this.baseUrl = comfyUiBaseUrl(config.baseUrl); + this.apiKey = config.apiKey?.trim() || undefined; + assertHttpsForKeyedOpenAiCompatibleUrl(this.baseUrl, this.apiKey); + this.imageModel = config.imageModel?.trim() || undefined; + this.imageWorkflow = config.imageWorkflow?.trim() + ? parseWorkflowTemplate(config.imageWorkflow) + : undefined; + this.videoWorkflow = config.videoWorkflow?.trim() + ? parseWorkflowTemplate(config.videoWorkflow) + : undefined; + this.fetch = createOpenAiCompatibleFetch(config.fetch); + this.pollIntervalMs = config.pollIntervalMs ?? POLL_INTERVAL_MS; + } + + describe() { + return { + id: COMFYUI_PROVIDER_ID, + contractVersion: "1", + adapterVersion: "0.1.0", + capabilities: { + textToImage: true, + imageToImage: true, + inpaint: false, + maxWidth: MAX_IMAGE_SIDE, + maxHeight: MAX_IMAGE_SIDE, + textToVideo: true, + imageToVideo: true, + maxDurationSeconds: MEDIA_VIDEO_MAX_SECONDS, + }, + }; + } + + async gpuStatus(signal?: AbortSignal): Promise { + return parseGpuStatus(await this.getJson("system_stats", signal)); + } + + async generateImage(request: ImageGenerationRequest): Promise { + if (request.mask) throw new Error("ComfyUI inpainting is not supported"); + return this.generate({ + ...request, + kind: "image", + template: pickTemplate( + this.imageWorkflow, + Boolean(request.source), + DEFAULT_WORKFLOWS.textToImage, + DEFAULT_WORKFLOWS.imageToImage, + ), + width: roundSide(request.width, 8), + height: roundSide(request.height, 8), + count: request.count ?? 1, + }); + } + + async generateVideo(request: VideoGenerationRequest): Promise { + const seconds = Math.min(Math.max(request.durationSeconds, 1), MEDIA_VIDEO_MAX_SECONDS); + return this.generate({ + ...request, + kind: "video", + template: pickTemplate( + this.videoWorkflow, + Boolean(request.source), + DEFAULT_WORKFLOWS.textToVideo, + DEFAULT_WORKFLOWS.imageToVideo, + ), + width: roundSide(request.width, 32), + height: roundSide(request.height, 32), + frames: 8 * Math.round((seconds * COMFYUI_VIDEO_FPS) / 8) + 1, + count: 1, + }); + } + + private async generate(job: GenerationJob): Promise { + const { template, signal } = job; + const values: WorkflowValues = { + prompt: job.prompt, + negative_prompt: job.negativePrompt ?? DEFAULT_NEGATIVE_PROMPTS[job.kind], + width: job.width, + height: job.height, + steps: job.kind === "image" ? IMAGE_STEPS : VIDEO_STEPS, + strength: job.strength ?? DEFAULT_STRENGTH, + frames: job.frames ?? 1, + fps: COMFYUI_VIDEO_FPS, + }; + if (template.placeholders.has("model")) values.model = await this.imageCheckpoint(signal); + if (template.placeholders.has("image")) { + if (!job.source) throw new Error("ComfyUI workflow needs an input image"); + values.image = await this.upload(job.source, signal); + } + const sized = template.placeholders.has("width") && template.placeholders.has("height"); + 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); + } + 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, + signal: AbortSignal | undefined, + ): 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); + try { + for (;;) { + const outputs = await this.finishedOutputs(promptId, runSignal); + if (outputs) return outputs; + await sleep(this.pollIntervalMs, undefined, { signal: runSignal }); + } + } catch (error) { + if (!runSignal.aborted) throw error; + await this.cancelPrompt(promptId); + if (signal?.aborted) throw signal.reason ?? error; + throw new Error("ComfyUI timed out"); + } + } + + private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise { + const body = JSON.stringify({ prompt: graph, client_id: this.clientId }); + return this.send("prompt", { method: "POST", headers: JSON_HEADERS, body }, signal, { + read: async (response, readSignal) => { + if (response.status === 400) { + const rejection = await readJson(response, readSignal).catch(() => undefined); + throw new Error(describeWorkflowRejection(rejection)); + } + const result = await readOkJson(response, readSignal); + const promptId = isRecord(result) ? result.prompt_id : undefined; + if (typeof promptId !== "string" || !promptId) { + throw new Error("ComfyUI did not accept the workflow"); + } + return promptId; + }, + }); + } + + private async finishedOutputs( + promptId: string, + signal: AbortSignal, + ): Promise | undefined> { + const history = await this.getJson(`history/${encodeURIComponent(promptId)}`, signal); + const entry = isRecord(history) ? history[promptId] : undefined; + if (!isRecord(entry)) return undefined; + const status = isRecord(entry.status) ? entry.status : {}; + if (status.status_str === "error") throw new Error(describeExecutionError(status.messages)); + if (status.completed === false) return undefined; + return isRecord(entry.outputs) ? entry.outputs : {}; + } + + /** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */ + private async cancelPrompt(promptId: string): Promise { + const options = { timeoutMs: CANCEL_TIMEOUT_MS, read: discardBody }; + try { + const body = JSON.stringify({ delete: [promptId] }); + await this.send("queue", { method: "POST", headers: JSON_HEADERS, body }, undefined, options); + const queue = await this.send( + "queue", + { headers: { Accept: "application/json" } }, + undefined, + { + timeoutMs: CANCEL_TIMEOUT_MS, + read: readOkJson, + }, + ); + const running = + isRecord(queue) && Array.isArray(queue.queue_running) ? queue.queue_running : []; + if (running.some((entry) => Array.isArray(entry) && entry[1] === promptId)) { + const interrupt = JSON.stringify({ prompt_id: promptId }); + await this.send( + "interrupt", + { method: "POST", headers: JSON_HEADERS, body: interrupt }, + undefined, + options, + ); + } + } catch { + // The prompt then finishes on its own; nothing waits for it. + } + } + + private async imageCheckpoint(signal: AbortSignal | undefined): Promise { + if (this.imageModel) return this.imageModel; + const info = await this.getJson("object_info/CheckpointLoaderSimple", signal); + const videoCheckpoints = new Set( + [DEFAULT_WORKFLOWS.textToVideo, DEFAULT_WORKFLOWS.imageToVideo, this.videoWorkflow].flatMap( + (template) => (template ? checkpointNames(template.graph) : []), + ), + ); + const checkpoint = parseCheckpointList(info).find((name) => !videoCheckpoints.has(name)); + if (!checkpoint) throw new Error("ComfyUI has no image checkpoint installed"); + return checkpoint; + } + + private async upload(file: MediaFile, signal: AbortSignal | undefined): Promise { + const form = new FormData(); + const extension = attachmentExtensionForMimeType(file.mimeType) || ".png"; + form.append( + "image", + new Blob([Buffer.from(file.bytes)], { type: file.mimeType }), + `evoo-${randomUUID()}${extension}`, + ); + form.append("type", "input"); + form.append("overwrite", "true"); + // Serialize here: the address-checked fetch runs on the undici package, whose FormData + // class differs from the global one. + const encoded = new Response(form); + const headers = { + Accept: "application/json", + "Content-Type": encoded.headers.get("content-type") ?? "", + }; + const body = new Uint8Array(await encoded.arrayBuffer()); + return this.send("upload/image", { method: "POST", headers, body }, signal, { + timeoutMs: TRANSFER_TIMEOUT_MS, + read: async (response, readSignal) => { + const result = await readOkJson(response, readSignal); + const name = isRecord(result) && typeof result.name === "string" ? result.name : ""; + if (!name) throw new Error("ComfyUI did not accept the image"); + const subfolder = + isRecord(result) && typeof result.subfolder === "string" ? result.subfolder : ""; + return subfolder ? `${subfolder}/${name}` : name; + }, + }); + } + + private download(file: OutputFile, signal: AbortSignal | undefined): Promise { + const query = new URLSearchParams({ + filename: file.filename, + subfolder: file.subfolder, + type: file.type, + }); + return this.send(`view?${query}`, {}, signal, { + timeoutMs: TRANSFER_TIMEOUT_MS, + read: async (response, readSignal) => { + await assertOk(response); + return readBodyCapped(response, ARTIFACT_MAX_BYTES_BY_KIND[file.kind], readSignal); + }, + }); + } + + private getJson(path: string, signal: AbortSignal | undefined): Promise { + return this.send(path, { headers: { Accept: "application/json" } }, signal, { + read: readOkJson, + }); + } + + /** One request, with the timeout covering the response body too. */ + private async send( + path: string, + init: { + method?: string; + headers?: Record; + body?: string | Uint8Array; + }, + signal: AbortSignal | undefined, + options: { + timeoutMs?: number; + read: (response: Response, signal: AbortSignal) => Promise; + }, + ): 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}`; + try { + const response = await this.fetch(new URL(path, this.baseUrl).href, { + method: init.method ?? "GET", + headers, + body: init.body, + redirect: "error", + signal: requestSignal, + }); + return await options.read(response, requestSignal); + } catch (error) { + if (signal?.aborted) throw error; + if (timeout.aborted) throw new Error("ComfyUI timed out"); + // undici reports refused or reset connections as a bare "fetch failed". + if (error instanceof TypeError && error.message === "fetch failed") { + throw new Error("ComfyUI is unreachable"); + } + throw error; + } + } +} + +function comfyUiBaseUrl(raw: string): URL { + const url = assertAllowedOpenAiCompatibleRequestUrl(raw.trim()); + url.search = ""; + url.hash = ""; + url.pathname = `${url.pathname.replace(/\/+$/, "")}/`; + return url; +} + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function truncate(text: string): string { + return text.length > MAX_ERROR_CHARS ? `${text.slice(0, MAX_ERROR_CHARS - 1)}…` : text; +} + +function roundSide(side: number, multiple: number): number { + const rounded = Math.round(side / multiple) * multiple; + return Math.min(Math.max(rounded, multiple * Math.ceil(64 / multiple)), MAX_IMAGE_SIDE); +} + +/** A custom workflow replaces the default that matches whether it takes an input image. */ +function pickTemplate( + custom: WorkflowTemplate | undefined, + withImage: boolean, + textDefault: WorkflowTemplate, + imageDefault: WorkflowTemplate, +): WorkflowTemplate { + if (custom && custom.placeholders.has("image") === withImage) return custom; + return withImage ? imageDefault : textDefault; +} + +function checkpointNames(graph: ComfyWorkflow): string[] { + return Object.values(graph).flatMap((node) => + typeof node.inputs.ckpt_name === "string" ? [node.inputs.ckpt_name] : [], + ); +} + +async function discardBody(response: Response): Promise { + await response.body?.cancel().catch(() => undefined); +} + +async function assertOk(response: Response): Promise { + if (response.ok) return; + await discardBody(response); + throw new Error(`ComfyUI returned ${response.status}`); +} + +async function readJson(response: Response, signal: AbortSignal): Promise { + const bytes = await readBodyCapped(response, MAX_JSON_BYTES, signal); + try { + return JSON.parse(new TextDecoder().decode(bytes)); + } catch { + throw new Error("ComfyUI returned invalid JSON"); + } +} + +async function readOkJson(response: Response, signal: AbortSignal): Promise { + await assertOk(response); + return readJson(response, signal); +} + +/** + * ComfyUI can evict its own models, so memory its PyTorch process holds counts as available; + * only memory other processes use (another model server on the same GPU) does not. + */ +export function parseGpuStatus(body: unknown): GpuStatus | undefined { + const devices = isRecord(body) && Array.isArray(body.devices) ? body.devices : []; + for (const device of devices) { + if (!isRecord(device) || device.type === "cpu") continue; + const total = finiteNumber(device.vram_total); + const free = finiteNumber(device.vram_free); + if (!total || free === undefined) continue; + const reserved = finiteNumber(device.torch_vram_total) ?? 0; + const reservedFree = finiteNumber(device.torch_vram_free) ?? 0; + const available = Math.min(total, Math.max(0, free - reservedFree + reserved)); + return { vramFreeMb: Math.floor(available / MIB), vramTotalMb: Math.floor(total / MIB) }; + } + return undefined; +} + +function finiteNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) ? value : undefined; +} + +/** Older ComfyUI lists combo values as `[[...names], {...}]`, newer as `["COMBO", {options}]`. */ +export function parseCheckpointList(body: unknown): string[] { + const loader = isRecord(body) ? body.CheckpointLoaderSimple : undefined; + const input = isRecord(loader) ? loader.input : undefined; + const required = isRecord(input) ? input.required : undefined; + const spec = isRecord(required) ? required.ckpt_name : undefined; + if (!Array.isArray(spec)) return []; + const names = Array.isArray(spec[0]) + ? spec[0] + : isRecord(spec[1]) && Array.isArray(spec[1].options) + ? spec[1].options + : []; + return names.filter((name): name is string => typeof name === "string" && name.length > 0); +} + +/** Saved files win over previews; video runs return clips (or animations), not frames. */ +function outputFiles(outputs: Record, kind: MediaKind): OutputFile[] { + const files: OutputFile[] = []; + for (const node of Object.values(outputs)) { + if (!isRecord(node)) continue; + for (const key of ["images", "gifs", "videos"]) { + const entries = node[key]; + if (!Array.isArray(entries)) continue; + for (const entry of entries) { + if (!isRecord(entry) || typeof entry.filename !== "string" || !entry.filename) continue; + const mimeType = inferAttachmentMimeType(entry.filename); + if (!mimeType) continue; + files.push({ + filename: entry.filename, + subfolder: typeof entry.subfolder === "string" ? entry.subfolder : "", + type: typeof entry.type === "string" ? entry.type : "output", + mimeType, + kind: attachmentKindOf(mimeType), + }); + } + } + } + const videos = files.filter((file) => file.kind === "video"); + const wanted = + kind === "image" + ? files.filter((file) => file.kind === "image") + : videos.length > 0 + ? videos + : files.filter((file) => file.mimeType === "image/gif" || file.mimeType === "image/webp"); + const saved = wanted.filter((file) => file.type === "output"); + return saved.length > 0 ? saved : wanted; +} + +export function describeWorkflowRejection(body: unknown): string { + const problems: string[] = []; + const nodeErrors = isRecord(body) && isRecord(body.node_errors) ? body.node_errors : {}; + for (const node of Object.values(nodeErrors)) { + if (!isRecord(node) || !Array.isArray(node.errors)) continue; + for (const error of node.errors) { + if (!isRecord(error)) continue; + const parts = [node.class_type, error.message, error.details]; + problems.push(parts.filter((part) => typeof part === "string" && part).join(": ")); + } + } + if (problems.length === 0 && isRecord(body) && isRecord(body.error)) { + if (typeof body.error.message === "string") problems.push(body.error.message); + } + return truncate( + problems.length > 0 + ? `ComfyUI rejected the workflow: ${problems.join("; ")}` + : "ComfyUI rejected the workflow", + ); +} + +export function describeExecutionError(messages: unknown): string { + const events = Array.isArray(messages) ? messages : []; + for (const event of events) { + if (!Array.isArray(event)) continue; + const [name, data] = event; + if (name === "execution_interrupted") return "ComfyUI interrupted the run"; + if (name !== "execution_error" || !isRecord(data)) continue; + const node = typeof data.node_type === "string" ? data.node_type : "a node"; + const message = + typeof data.exception_message === "string" + ? data.exception_message.trim().split("\n")[0] + : ""; + return truncate(`ComfyUI failed in ${node}${message ? `: ${message}` : ""}`); + } + return "ComfyUI failed"; +} diff --git a/packages/adapters/src/image-generation-conformance.test.ts b/packages/adapters/src/image-generation-conformance.test.ts new file mode 100644 index 0000000..92883a5 --- /dev/null +++ b/packages/adapters/src/image-generation-conformance.test.ts @@ -0,0 +1,111 @@ +import type { + ImageGenerationProvider, + MediaFile, + VideoGenerationProvider, +} from "@rakazo/adapter-kit"; +import { attachmentKindOf } from "@rakazo/contracts"; +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import { createComfyUiProvider } from "./comfyui-provider.js"; +import { ComfyUiEmulator } from "./testing/comfyui-emulator.js"; + +type MediaGenerationProvider = ImageGenerationProvider & VideoGenerationProvider; + +const factories: Record MediaGenerationProvider> = { + comfyui: () => + createComfyUiProvider({ + baseUrl: "http://127.0.0.1:8188", + fetch: new ComfyUiEmulator().fetch, + pollIntervalMs: 1, + }), +}; + +async function sourceImage(): Promise { + const png = await sharp({ + create: { width: 96, height: 64, channels: 3, background: { r: 10, g: 20, b: 30 } }, + }) + .png() + .toBuffer(); + return { bytes: new Uint8Array(png), mimeType: "image/png", name: "source.png" }; +} + +/** + * Offline conformance: every image/video generation provider describes its limits, returns + * the requested number of decodable images, is reproducible by seed, takes a source image, + * reports progress, returns video, and honors cancellation. + */ +for (const [name, create] of Object.entries(factories)) { + describe(`${name} media generation conformance (offline)`, () => { + it("describes the contract and its limits", () => { + const image: ImageGenerationProvider = create(); + const video: VideoGenerationProvider = create(); + expect(image.describe().contractVersion).toBe("1"); + expect(image.describe().capabilities.maxWidth).toBeGreaterThanOrEqual(1024); + expect(image.describe().capabilities.maxHeight).toBeGreaterThanOrEqual(1024); + expect(video.describe().contractVersion).toBe("1"); + expect(video.describe().capabilities.maxDurationSeconds).toBeGreaterThan(0); + }); + + it("returns the requested number of images, reproducible by seed", async () => { + const provider = create(); + const progress: number[] = []; + const request = { prompt: "a red kite", width: 256, height: 192, count: 2, seed: 7 }; + const images = await provider.generateImage({ + ...request, + onProgress: (fraction) => progress.push(fraction), + }); + expect(images).toHaveLength(2); + for (const image of images) { + expect(attachmentKindOf(image.mimeType)).toBe("image"); + const metadata = await sharp(image.bytes).metadata(); + expect([metadata.width, metadata.height]).toEqual([256, 192]); + } + expect(Buffer.from(images[0]!.bytes).equals(Buffer.from(images[1]!.bytes))).toBe(false); + 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); + }); + + it("generates from a source image", async () => { + const images = await create().generateImage({ + prompt: "the same scene at night", + width: 96, + height: 64, + source: await sourceImage(), + strength: 0.5, + }); + expect(images).toHaveLength(1); + expect(attachmentKindOf(images[0]!.mimeType)).toBe("image"); + }); + + it("returns video from text and from a start frame", async () => { + const provider = create(); + for (const source of [undefined, await sourceImage()]) { + const [clip, ...rest] = await provider.generateVideo({ + prompt: "waves rolling onto a beach", + width: 768, + height: 512, + durationSeconds: 2, + source, + }); + expect(rest).toHaveLength(0); + expect(attachmentKindOf(clip!.mimeType)).toBe("video"); + expect(clip!.bytes.byteLength).toBeGreaterThan(0); + } + }); + + it("reports GPU memory with free not above total", async () => { + const status = await create().gpuStatus?.(); + if (!status) return; + expect(status.vramTotalMb).toBeGreaterThan(0); + expect(status.vramFreeMb).toBeGreaterThanOrEqual(0); + expect(status.vramFreeMb).toBeLessThanOrEqual(status.vramTotalMb); + }); + + it("rejects when the caller cancels", async () => { + const request = { prompt: "x", width: 64, height: 64, signal: AbortSignal.abort() }; + await expect(create().generateImage(request)).rejects.toThrow(); + await expect(create().generateVideo({ ...request, durationSeconds: 1 })).rejects.toThrow(); + }); + }); +} diff --git a/packages/adapters/src/index.ts b/packages/adapters/src/index.ts index 7b4ecd0..18253dc 100644 --- a/packages/adapters/src/index.ts +++ b/packages/adapters/src/index.ts @@ -23,6 +23,7 @@ export * from "./cloud-agent-provider-env.js"; export * from "./cloud-agent-service.js"; export * from "./cloud-agent-tools.js"; export * from "./cloud-agent-tools-select.js"; +export * from "./comfyui-provider.js"; export * from "./composio-catalog-cache.js"; export * from "./composio-connector.js"; export * from "./composio-emulator.js"; diff --git a/packages/adapters/src/testing/comfyui-emulator.ts b/packages/adapters/src/testing/comfyui-emulator.ts new file mode 100644 index 0000000..4da22a1 --- /dev/null +++ b/packages/adapters/src/testing/comfyui-emulator.ts @@ -0,0 +1,393 @@ +import { createHash } from "node:crypto"; + +/** A 1 s, 32×32 H.264 MP4 (ffmpeg testsrc), small enough to inline. */ +const TINY_MP4_BASE64 = [ + "AAAAIGZ0eXBpc29tAAACAGlzb21pc28yYXZjMW1wNDEAAAM+bW9vdgAAAGxtdmhkAAAAAAAAAAAAAAAAAAAD6AAAA+gAAQAAAQAA", + "AAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAgAA", + "Ao10cmFrAAAAXHRraGQAAAADAAAAAAAAAAAAAAABAAAAAAAAA+gAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAABAAAA", + "AAAAAAAAAAAAAABAAAAAACAAAAAgAAAAAAAkZWR0cwAAABxlbHN0AAAAAAAAAAEAAAPoAAAQAAABAAAAAAIFbWRpYQAAACBtZGhk", + "AAAAAAAAAAAAAAAAAABAAAAAQABVxAAAAAAALWhkbHIAAAAAAAAAAHZpZGUAAAAAAAAAAAAAAABWaWRlb0hhbmRsZXIAAAABsG1p", + "bmYAAAAUdm1oZAAAAAEAAAAAAAAAAAAAACRkaW5mAAAAHGRyZWYAAAAAAAAAAQAAAAx1cmwgAAAAAQAAAXBzdGJsAAAAwHN0c2QA", + "AAAAAAAAAQAAALBhdmMxAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAACAAIABIAAAASAAAAAAAAAABDExhdmMgbGlieDI2NAAAAAAA", + "AAAAAAAAAAAAAAAAAAAAGP//AAAANmF2Y0MBZAAK/+EAGGdkAAqscgRJbARAAAADAEAAAAQDxIlhGAEAB2joQ4OSyLD9+PgAAAAA", + "EHBhc3AAAAABAAAAAQAAABRidHJ0AAAAAAAADAgAAAwIAAAAGHN0dHMAAAAAAAAAAQAAAAgAAAgAAAAAFHN0c3MAAAAAAAAAAQAA", + "AAEAAAAYY3R0cwAAAAAAAAABAAAACAAAEAAAAAAcc3RzYwAAAAAAAAABAAAAAQAAAAgAAAABAAAANHN0c3oAAAAAAAAAAAAAAAgA", + "AADTAAAADAAAABgAAAAdAAAAGgAAAB0AAAAfAAAAFwAAABRzdGNvAAAAAAAAAAEAAANuAAAAPXVkdGEAAAA1bWV0YQAAAAAAAAAh", + "aGRscgAAAAAAAAAAbWRpcmFwcGwAAAAAAAAAAAAAAAAIaWxzdAAAAAhmcmVlAAABiW1kYXQAAADPZYiBABP/115QL0x5O3/uan23", + "okcn80QwiyZelIq0iY0T1naTwGHVgXgsN7YBwlEikXvo+AyEjKk/afikvNFJbO7NjVpKr2f7e+xmO5rxkaelhrWKVT83gKOM5wKs", + "QoFpCKDo56fFf6IYs3U1LAHAxZY3IIAUCFca5RyJ/ibKbmkxBCXXiBuUCCuL9FhMjcqhW2misGdpwuLgknHPr/37MowWA2zcOn9B", + "OHR5mfgyt4/0x3o95hpnI+2axZFkXHhYwK55h134UUvr+r3B9b/xAAAACEGaCC2JP6KAAAAAFEGaEEvwIGTKYTe7qicJ6jU1YvCB", + "AAAAGUGaGGk8CB5MpgTfu1UOB2sZPRDVFTCp+IAAAAAWQZogiXwIH5MpgER/wf8hVV5nVt8U2QAAABlBmiipvAgf5MpgBFfFjke1", + "iXbU9gFuIojBAAAAG0GaMMn8CB/5MpgARn/KNfso3F7XaeLq3d3VjQAAABNBmjjojwIH/5MpgAEIf9P6j9Ph", +].join(""); + +export const COMFYUI_EMULATOR_MP4 = new Uint8Array(Buffer.from(TINY_MP4_BASE64, "base64")); + +const MIB = 1024 * 1024; + +type EmulatedNode = { class_type: string; inputs: Record }; +type EmulatedGraph = Record; +type FileRef = { filename: string; subfolder: string; type: string }; + +interface EmulatedPrompt { + id: string; + graph: EmulatedGraph; + polls: number; + failure?: string; + interrupted?: boolean; + /** Dropped from the queue before it ran. */ + removed?: boolean; + /** Reported as finished by the history. */ + done?: boolean; + outputs?: Record>; +} + +export interface EmulatedComfyUiRequest { + method: string; + path: string; + authorization: string | null; + body?: unknown; +} + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function json(body: unknown, status = 200): Response { + return Response.json(body, { status }); +} + +/** + * 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. + */ +export class ComfyUiEmulator { + readonly requests: EmulatedComfyUiRequest[] = []; + readonly prompts: EmulatedPrompt[] = []; + readonly uploads = new Map(); + readonly deleted: string[] = []; + readonly interrupted: string[] = []; + 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"; + /** null reports a CPU-only server. */ + gpu: { totalMb: number; freeMb: number; torchReservedMb: number; torchFreeMb: number } | null = { + totalMb: 24_576, + freeMb: 20_000, + torchReservedMb: 0, + torchFreeMb: 0, + }; + /** History polls that still report a submitted prompt as running. */ + pendingPolls = 1; + /** Answer the next request with this HTTP status. */ + failNextStatus: number | undefined; + /** Fail the next submitted prompt while it runs, with this exception message. */ + failNextExecution: string | undefined; + private readonly files = new Map(); + private nextPrompt = 1; + + /** `basePath` serves the API under a prefix, like ComfyUI behind a reverse proxy. */ + constructor(private readonly options: { apiKey?: string; basePath?: string } = {}) {} + + readonly fetch: typeof fetch = async (input, init) => { + init?.signal?.throwIfAborted(); + const url = new URL(input instanceof Request ? input.url : String(input)); + 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 body = + typeof init?.body === "string" && headers.get("content-type")?.includes("json") + ? JSON.parse(init.body) + : undefined; + this.requests.push({ method, path, authorization, body }); + if (this.failNextStatus) { + const status = this.failNextStatus; + this.failNextStatus = undefined; + return new Response("Upstream failure", { status }); + } + if (this.options.apiKey && authorization !== `Bearer ${this.options.apiKey}`) { + 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") { + return this.checkpointInfo(); + } + if (route === "POST prompt") return this.submit(body); + if (route === "GET history") return this.history(decodeURIComponent(path.slice(8))); + if (route === "GET view") return this.view(url.searchParams); + if (route === "POST upload") return this.upload(input, init); + if (route === "GET queue") return this.queue(); + if (route === "POST queue") { + for (const id of Array.isArray(body?.delete) ? body.delete : []) { + this.deleted.push(id); + // Like ComfyUI, deleting only drops pending prompts; a running one needs an interrupt. + const prompt = this.prompts.find((entry) => entry.id === id); + if (prompt && prompt !== this.unfinished()[0]) prompt.removed = true; + } + return json({}); + } + if (route === "POST interrupt") { + const id = String(body?.prompt_id ?? ""); + this.interrupted.push(id); + const running = this.unfinished()[0]; + if (running && (!id || running.id === id)) running.interrupted = true; + return json({}); + } + return json({ error: "Not found" }, 404); + }; + + private systemStats(): Response { + const gpu = this.gpu; + return json({ + system: { os: "posix", comfyui_version: "0.3.60", embedded_python: false }, + devices: [ + gpu + ? { + name: "cuda:0 Emulated GPU", + type: "cuda", + index: 0, + vram_total: gpu.totalMb * MIB, + vram_free: gpu.freeMb * MIB, + torch_vram_total: gpu.torchReservedMb * MIB, + torch_vram_free: gpu.torchFreeMb * MIB, + } + : { name: "cpu", type: "cpu", index: null, vram_total: 64 * 1024 * MIB, vram_free: 0 }, + ], + }); + } + + private checkpointInfo(): Response { + const spec = + this.checkpointListFormat === "legacy" + ? [this.checkpoints, { tooltip: "Checkpoint" }] + : ["COMBO", { options: this.checkpoints, tooltip: "Checkpoint" }]; + return json({ + CheckpointLoaderSimple: { + input: { required: { ckpt_name: spec } }, + output: ["MODEL", "CLIP", "VAE"], + name: "CheckpointLoaderSimple", + }, + }); + } + + private submit(body: unknown): Response { + const graph = isRecord(body) ? body.prompt : undefined; + if (!isRecord(graph) || typeof (body as { client_id?: unknown }).client_id !== "string") { + return json( + { error: { type: "invalid_prompt", message: "Cannot execute because prompt is invalid" } }, + 400, + ); + } + const nodeErrors: Record = {}; + const reject = (id: string, node: unknown, message: string, details: string) => { + nodeErrors[id] = { + errors: [{ type: "invalid", message, details, extra_info: {} }], + dependent_outputs: [], + class_type: isRecord(node) ? node.class_type : undefined, + }; + }; + for (const [id, node] of Object.entries(graph)) { + if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) { + reject(id, node, "Invalid node", "class_type and inputs are required"); + continue; + } + for (const [name, value] of Object.entries(node.inputs)) { + if (typeof value === "string" && value.includes("{{")) { + reject(id, node, "Unfilled placeholder", `${name}: ${value}`); + } + if (Array.isArray(value) && !isRecord(graph[String(value[0])])) { + reject(id, node, "Missing linked node", `${name}: ${String(value[0])}`); + } + } + const checkpoint = node.inputs.ckpt_name; + if (typeof checkpoint === "string" && !this.checkpoints.includes(checkpoint)) { + reject(id, node, "Value not in list", `ckpt_name: '${checkpoint}' not in list`); + } + if (node.class_type === "LoadImage" && !this.uploads.has(String(node.inputs.image))) { + reject(id, node, "Invalid image file", `image - ${String(node.inputs.image)}`); + } + } + if (Object.keys(nodeErrors).length > 0) { + return json( + { + error: { + type: "prompt_outputs_failed_validation", + message: "Prompt outputs failed validation", + }, + node_errors: nodeErrors, + }, + 400, + ); + } + const id = `prompt-${this.nextPrompt++}`; + this.prompts.push({ + id, + graph: graph as EmulatedGraph, + polls: 0, + failure: this.failNextExecution, + }); + this.failNextExecution = undefined; + return json({ prompt_id: id, number: this.nextPrompt, node_errors: {} }); + } + + private async history(id: string): Promise { + const prompt = this.prompts.find((entry) => entry.id === id); + if (!prompt || prompt.removed) return json({}); + if (!prompt.done && !prompt.interrupted && prompt.polls < this.pendingPolls) { + prompt.polls += 1; + return json({}); + } + prompt.done = true; + if (prompt.interrupted) { + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, []], + outputs: {}, + status: { + status_str: "error", + completed: false, + messages: [["execution_interrupted", { prompt_id: id, node_id: "5" }]], + }, + }, + }); + } + if (prompt.failure) { + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, []], + outputs: {}, + status: { + status_str: "error", + completed: false, + messages: [ + ["execution_start", { prompt_id: id }], + [ + "execution_error", + { + prompt_id: id, + node_id: "5", + node_type: "KSampler", + exception_message: prompt.failure, + exception_type: "RuntimeError", + traceback: ["..."], + }, + ], + ], + }, + }, + }); + } + prompt.outputs ??= await this.execute(prompt); + return json({ + [id]: { + prompt: [0, id, prompt.graph, {}, Object.keys(prompt.outputs)], + outputs: prompt.outputs, + status: { status_str: "success", completed: true, messages: [] }, + }, + }); + } + + private async execute(prompt: EmulatedPrompt): Promise>> { + const outputs: Record> = {}; + let counter = 0; + const save = (bytes: Uint8Array, extension: string, type: string): FileRef => { + counter += 1; + const ref = { + filename: `${prompt.id}_${String(counter).padStart(5, "0")}_.${extension}`, + subfolder: type === "output" ? "evoo" : "", + type, + }; + this.files.set(fileKey(ref), bytes); + return ref; + }; + for (const [id, node] of Object.entries(prompt.graph)) { + if (node.class_type === "SaveImage" || node.class_type === "PreviewImage") { + const type = node.class_type === "SaveImage" ? "output" : "temp"; + outputs[id] = { images: [save(await this.renderImage(prompt.graph), "png", type)] }; + } else if (node.class_type === "SaveVideo") { + outputs[id] = { images: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")], animated: [true] }; + } else if (node.class_type === "VHS_VideoCombine") { + outputs[id] = { gifs: [save(COMFYUI_EMULATOR_MP4, "mp4", "output")] }; + } + } + return outputs; + } + + /** A flat color derived from the seed, at the workflow's size (or its input image's). */ + private async renderImage(graph: EmulatedGraph): Promise { + const { default: sharp } = await import("sharp"); + const nodes = Object.values(graph); + const input = (name: string) => + nodes.map((node) => node.inputs[name]).find((value) => typeof value === "number"); + let width = Number(input("width") ?? 0); + let height = Number(input("height") ?? 0); + const loaded = nodes.find((node) => node.class_type === "LoadImage"); + if (loaded && !(width && height)) { + const metadata = await sharp(this.uploads.get(String(loaded.inputs.image))).metadata(); + width = metadata.width ?? 64; + height = metadata.height ?? 64; + } + const seed = String(input("seed") ?? input("noise_seed") ?? 0); + const [r = 0, g = 0, b = 0] = createHash("sha256").update(seed).digest(); + const png = await sharp({ + create: { width: width || 64, height: height || 64, channels: 3, background: { r, g, b } }, + }) + .png() + .toBuffer(); + return new Uint8Array(png); + } + + private view(params: URLSearchParams): Response { + const bytes = this.files.get( + fileKey({ + filename: params.get("filename") ?? "", + subfolder: params.get("subfolder") ?? "", + type: params.get("type") ?? "", + }), + ); + if (!bytes) return new Response("Not found", { status: 404 }); + return new Response(Buffer.from(bytes), { + headers: { "Content-Type": "application/octet-stream" }, + }); + } + + private async upload(input: RequestInfo | URL, init?: RequestInit): Promise { + const form = await new Request(String(input), { + method: "POST", + headers: init?.headers, + body: init?.body, + }).formData(); + const image = form.get("image"); + if (!(image instanceof File) || form.get("type") !== "input") { + return json({ error: "Invalid upload" }, 400); + } + this.uploads.set(image.name, new Uint8Array(await image.arrayBuffer())); + return json({ name: image.name, subfolder: "", type: "input" }); + } + + private unfinished(): EmulatedPrompt[] { + return this.prompts.filter((prompt) => !prompt.done && !prompt.removed); + } + + private queue(): Response { + const running = this.unfinished(); + return json({ + queue_running: running.slice(0, 1).map((prompt) => [0, prompt.id, prompt.graph, {}, []]), + queue_pending: running + .slice(1) + .map((prompt, index) => [index + 1, prompt.id, prompt.graph, {}, []]), + }); + } +} + +function fileKey(ref: FileRef): string { + return `${ref.type}/${ref.subfolder}/${ref.filename}`; +} diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index ba2c70b..3c000fb 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -10,6 +10,7 @@ export * from "./integration-settings.js"; export * from "./local-settings.js"; export * from "./mcp.js"; export * from "./media.js"; +export * from "./media-generation.js"; export * from "./model-tiers.js"; export * from "./openai-compatible-ui.js"; export * from "./project-goals.js"; diff --git a/packages/contracts/src/media-generation.test.ts b/packages/contracts/src/media-generation.test.ts new file mode 100644 index 0000000..4ab1a74 --- /dev/null +++ b/packages/contracts/src/media-generation.test.ts @@ -0,0 +1,43 @@ +import { describe, expect, it } from "vitest"; +import { + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, +} from "./media-generation.js"; + +describe("media generation inputs", () => { + it("fills defaults and drops unknown fields such as continue_after", () => { + expect(GenerateImageInputSchema.parse({ prompt: " a fox ", continue_after: true })).toEqual({ + prompt: "a fox", + count: 1, + }); + expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({ + prompt: "waves", + aspect: "landscape", + durationSeconds: 5, + }); + }); + + it("bounds image count, size and video duration", () => { + expect(GenerateImageInputSchema.safeParse({ prompt: "x", count: 5 }).success).toBe(false); + expect(GenerateImageInputSchema.safeParse({ prompt: "x", width: 512 }).success).toBe(false); + expect( + GenerateImageInputSchema.safeParse({ prompt: "x", width: 512, height: 4096 }).success, + ).toBe(false); + expect(GenerateVideoInputSchema.safeParse({ prompt: "x", durationSeconds: 11 }).success).toBe( + false, + ); + }); + + it("needs at least one edit and valid edit values", () => { + expect(EditImageInputSchema.safeParse({ artifactId: "a" }).success).toBe(false); + expect(EditImageInputSchema.safeParse({ artifactId: "a", rotate: 45 }).success).toBe(false); + expect(EditImageInputSchema.safeParse({ artifactId: "a", resize: {} }).success).toBe(false); + expect( + EditImageInputSchema.safeParse({ artifactId: "a", extend: { top: 4, color: "red" } }).success, + ).toBe(false); + expect( + EditImageInputSchema.parse({ artifactId: "a", resize: { aspect: "vertical" }, rotate: 90 }), + ).toEqual({ artifactId: "a", resize: { aspect: "vertical", fit: "crop" }, rotate: 90 }); + }); +}); diff --git a/packages/contracts/src/media-generation.ts b/packages/contracts/src/media-generation.ts new file mode 100644 index 0000000..6d33216 --- /dev/null +++ b/packages/contracts/src/media-generation.ts @@ -0,0 +1,108 @@ +import * as z from "zod"; +import { Id } from "./ids.js"; +import { MediaFitSchema, MediaPresetIdSchema } from "./media.js"; + +export const MEDIA_IMAGE_MAX_COUNT = 4; +export const MEDIA_VIDEO_MAX_SECONDS = 10; + +const Prompt = z.string().trim().min(1).max(4000); +const NegativePrompt = z.string().trim().max(4000).optional().describe("What to avoid."); +const Seed = z + .number() + .int() + .min(0) + .max(Number.MAX_SAFE_INTEGER) + .optional() + .describe("Reuse a seed to reproduce a result."); +const Strength = z.number().min(0).max(1).optional(); +const Aspect = MediaPresetIdSchema.describe( + "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9.", +); +const Side = z.number().int().min(64).max(2048); +const Pixels = z.number().int().min(0).max(16_384); + +export const GenerateImageInputSchema = z + .object({ + prompt: Prompt.describe("What the image shows."), + negativePrompt: NegativePrompt, + aspect: Aspect.optional(), + width: Side.optional().describe("Exact width in pixels, with height; overrides aspect."), + height: Side.optional(), + count: z.number().int().min(1).max(MEDIA_IMAGE_MAX_COUNT).default(1), + seed: Seed, + sourceArtifactId: Id.optional().describe("Image artifact to start from (image-to-image)."), + strength: Strength.describe("With a source: how far to move away from it (default 0.6)."), + }) + .refine((input) => (input.width === undefined) === (input.height === undefined), { + message: "Set width and height together", + path: ["width"], + }); +export type GenerateImageInput = z.infer; + +export const GenerateVideoInputSchema = z.object({ + prompt: Prompt.describe("What happens in the clip: subject, motion, camera."), + negativePrompt: NegativePrompt, + aspect: Aspect.default("landscape"), + durationSeconds: z.number().min(1).max(MEDIA_VIDEO_MAX_SECONDS).default(5), + seed: Seed, + sourceArtifactId: Id.optional().describe("Image artifact to use as the first frame."), +}); +export type GenerateVideoInput = z.infer; + +export const ImageFormatSchema = z.enum(["png", "jpeg", "webp"]); +export type ImageFormat = z.infer; + +/** Edits apply in this order: crop, rotate, flip, resize, extend, repaint, format. */ +export const EditImageInputSchema = z + .object({ + artifactId: Id.describe("Image artifact to edit."), + crop: z + .object({ left: Pixels, top: Pixels, width: Pixels.min(1), height: Pixels.min(1) }) + .optional() + .describe("Area to keep, in source pixels."), + rotate: z.literal([90, 180, 270]).optional().describe("Clockwise, in degrees."), + flip: z.enum(["horizontal", "vertical"]).optional(), + resize: z + .object({ + aspect: Aspect.optional().describe("Exact size of this format."), + width: z.number().int().min(1).max(8192).optional(), + height: z.number().int().min(1).max(8192).optional(), + fit: MediaFitSchema.default("crop").describe( + "How a different aspect ratio fills the size: crop, pad, or blur (blurred background).", + ), + }) + .refine((resize) => resize.aspect || resize.width || resize.height, { + message: "Set aspect, width or height", + }) + .optional() + .describe("Width or height alone keeps the aspect ratio."), + extend: z + .object({ + top: Pixels.default(0), + right: Pixels.default(0), + bottom: Pixels.default(0), + left: Pixels.default(0), + color: z + .string() + .regex(/^#[0-9a-fA-F]{6}$/) + .optional() + .describe("#RRGGBB; transparent when omitted and the format allows it."), + }) + .optional() + .describe("Borders to add, in pixels."), + prompt: Prompt.optional().describe("Repaint the image with this prompt."), + strength: Strength.describe("How far the repaint moves away from the image (default 0.6)."), + format: ImageFormatSchema.optional().describe("Default keeps the source format."), + }) + .refine( + (input) => + input.crop || + input.rotate || + input.flip || + input.resize || + input.extend || + input.prompt || + input.format, + { message: "Nothing to edit" }, + ); +export type EditImageInput = z.infer; diff --git a/packages/core/package.json b/packages/core/package.json index bcac6e5..562dfd6 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -34,6 +34,11 @@ "development": "./src/plot/plot-spec.ts", "default": "./src/plot/plot-spec.ts" }, + "./media/workflows": { + "types": "./src/media/workflow-template.ts", + "development": "./src/media/workflow-template.ts", + "default": "./src/media/workflow-template.ts" + }, "./node/desktop-runtime": { "types": "./src/node/desktop-runtime.ts", "development": "./src/node/desktop-runtime.ts", diff --git a/packages/core/src/media/workflow-template.test.ts b/packages/core/src/media/workflow-template.test.ts new file mode 100644 index 0000000..7865463 --- /dev/null +++ b/packages/core/src/media/workflow-template.test.ts @@ -0,0 +1,113 @@ +import { describe, expect, it } from "vitest"; +import { + DEFAULT_WORKFLOWS, + fillWorkflow, + parseWorkflowTemplate, + WorkflowTemplateError, +} from "./workflow-template.js"; + +const minimal = { + "1": { class_type: "CLIPTextEncode", inputs: { text: "{{prompt}}", clip: ["2", 1] } }, + "2": { class_type: "EmptyLatentImage", inputs: { width: "{{width}}", height: 512 } }, + "3": { class_type: "SaveImage", inputs: { filename_prefix: "run/{{seed}}-{{prompt}}" } }, +}; + +describe("parseWorkflowTemplate", () => { + it("accepts JSON text or parsed graphs and lists the placeholders", () => { + for (const source of [minimal, JSON.stringify(minimal)]) { + const template = parseWorkflowTemplate(source); + expect([...template.placeholders].sort()).toEqual(["prompt", "seed", "width"]); + } + }); + + it.each([ + ["not json", "Workflow is not valid JSON"], + ["[]", "Workflow must be a JSON object"], + ["{}", "Workflow must have 1 to 1000 nodes"], + [JSON.stringify({ nodes: [], links: [] }), "Workflow must be exported in ComfyUI's API format"], + [JSON.stringify({ "1": { inputs: {} } }), "Workflow node 1 needs a class_type and inputs"], + [ + JSON.stringify({ "1": { class_type: "SaveImage", inputs: { text: "plain" } } }), + "Workflow needs a {{prompt}} placeholder", + ], + [ + JSON.stringify({ + "1": { class_type: "X", inputs: { a: "{{prompt}}", b: "{{cfg}}" } }, + }), + "Unknown workflow placeholders: {{cfg}}", + ], + ])("rejects %s", (source, message) => { + expect(() => parseWorkflowTemplate(source)).toThrow(new WorkflowTemplateError(message)); + }); + + it("rejects placeholder names that only exist on Object.prototype", () => { + const source = { + "1": { + class_type: "X", + inputs: { a: "{{prompt}}", b: "{{constructor}}", c: "{{__proto__}}" }, + }, + }; + expect(() => parseWorkflowTemplate(source)).toThrow( + "Unknown workflow placeholders: {{constructor}}, {{__proto__}}", + ); + }); + + it("rejects oversized workflow text", () => { + expect(() => parseWorkflowTemplate(" ".repeat(1_000_001))).toThrow("Workflow is too large"); + }); +}); + +describe("fillWorkflow", () => { + it("replaces whole values with typed values and embedded ones with text, without mutating", () => { + const template = parseWorkflowTemplate(minimal); + const graph = fillWorkflow(template, { prompt: "a $& cat", width: 768, seed: 7 }); + expect(graph["1"]?.inputs).toEqual({ text: "a $& cat", clip: ["2", 1] }); + expect(graph["2"]?.inputs).toEqual({ width: 768, height: 512 }); + expect(graph["3"]?.inputs.filename_prefix).toBe("run/7-a $& cat"); + expect(minimal["2"].inputs.width).toBe("{{width}}"); + }); + + it("refuses missing or mistyped values", () => { + const template = parseWorkflowTemplate(minimal); + expect(() => fillWorkflow(template, { prompt: "cat", seed: 1 })).toThrow( + "Workflow value for {{width}} is missing", + ); + expect(() => fillWorkflow(template, { prompt: "cat", width: 1.5, seed: 1 })).toThrow( + "{{width}}", + ); + }); +}); + +describe("DEFAULT_WORKFLOWS", () => { + it("only image-to-* defaults need an input image", () => { + expect(DEFAULT_WORKFLOWS.textToImage.placeholders.has("image")).toBe(false); + expect(DEFAULT_WORKFLOWS.textToVideo.placeholders.has("image")).toBe(false); + expect(DEFAULT_WORKFLOWS.imageToImage.placeholders.has("image")).toBe(true); + expect(DEFAULT_WORKFLOWS.imageToVideo.placeholders.has("image")).toBe(true); + }); + + it("fill into graphs whose links point at existing nodes", () => { + const values = { + prompt: "a lighthouse", + negative_prompt: "", + model: "model.safetensors", + image: "input.png", + width: 512, + height: 512, + seed: 1, + steps: 20, + strength: 0.6, + frames: 25, + fps: 24, + }; + for (const template of Object.values(DEFAULT_WORKFLOWS)) { + const graph = fillWorkflow(template, values); + expect(JSON.stringify(graph)).not.toContain("{{"); + for (const node of Object.values(graph)) { + for (const input of Object.values(node.inputs)) { + if (Array.isArray(input)) expect(graph[String(input[0])]).toBeDefined(); + } + } + } + }); +}); diff --git a/packages/core/src/media/workflow-template.ts b/packages/core/src/media/workflow-template.ts new file mode 100644 index 0000000..a7c74e6 --- /dev/null +++ b/packages/core/src/media/workflow-template.ts @@ -0,0 +1,159 @@ +import imageToImage from "./workflows/image-to-image.json" with { type: "json" }; +import imageToVideo from "./workflows/image-to-video.json" with { type: "json" }; +import textToImage from "./workflows/text-to-image.json" with { type: "json" }; +import textToVideo from "./workflows/text-to-video.json" with { type: "json" }; + +/** + * Values a ComfyUI workflow template can ask for. A string value that is exactly `{{name}}` + * becomes the typed value (so `"{{width}}"` turns into the number 1024); `{{name}}` inside a + * longer string is replaced by its text. Templates are data: nothing is ever evaluated. + */ +export const WORKFLOW_PLACEHOLDERS = { + prompt: "string", + negative_prompt: "string", + /** Checkpoint file name: the configured image model or the first installed checkpoint. */ + model: "string", + /** Uploaded input image (image-to-image, image-to-video). */ + image: "string", + width: "integer", + height: "integer", + seed: "integer", + steps: "integer", + /** 0–1 denoise for image-to-image. */ + strength: "number", + frames: "integer", + fps: "number", +} as const satisfies Record; + +export type WorkflowPlaceholder = keyof typeof WORKFLOW_PLACEHOLDERS; +export type WorkflowValues = { + [Name in WorkflowPlaceholder]?: (typeof WORKFLOW_PLACEHOLDERS)[Name] extends "string" + ? string + : number; +}; + +/** One node of a workflow in ComfyUI's API format ("Export (API)" in ComfyUI). */ +export interface ComfyWorkflowNode { + class_type: string; + inputs: Record; +} +export type ComfyWorkflow = Record; + +export interface WorkflowTemplate { + readonly graph: ComfyWorkflow; + /** The placeholders the graph uses; `image` marks a template that needs an input image. */ + readonly placeholders: ReadonlySet; +} + +export class WorkflowTemplateError extends Error { + constructor(message: string) { + super(message); + this.name = "WorkflowTemplateError"; + } +} + +const PLACEHOLDER_PATTERN = /\{\{([a-z_]+)\}\}/g; +const WHOLE_PLACEHOLDER = /^\{\{([a-z_]+)\}\}$/; +const MAX_WORKFLOW_CHARS = 1_000_000; +const MAX_WORKFLOW_NODES = 1_000; + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === "object" && !Array.isArray(value); +} + +function isPlaceholder(name: string): name is WorkflowPlaceholder { + return Object.hasOwn(WORKFLOW_PLACEHOLDERS, name); +} + +function visitStrings(value: unknown, visit: (text: string) => void): void { + if (typeof value === "string") visit(value); + else if (Array.isArray(value)) for (const entry of value) visitStrings(entry, visit); + else if (isRecord(value)) for (const entry of Object.values(value)) visitStrings(entry, visit); +} + +/** Validates an API-format workflow (JSON text or parsed) and lists its placeholders. */ +export function parseWorkflowTemplate(source: unknown): WorkflowTemplate { + let value = source; + if (typeof source === "string") { + if (source.length > MAX_WORKFLOW_CHARS) { + throw new WorkflowTemplateError("Workflow is too large"); + } + try { + value = JSON.parse(source); + } catch { + throw new WorkflowTemplateError("Workflow is not valid JSON"); + } + } + if (!isRecord(value)) throw new WorkflowTemplateError("Workflow must be a JSON object"); + if (Array.isArray(value.nodes)) { + throw new WorkflowTemplateError("Workflow must be exported in ComfyUI's API format"); + } + const nodes = Object.entries(value); + if (nodes.length === 0 || nodes.length > MAX_WORKFLOW_NODES) { + throw new WorkflowTemplateError(`Workflow must have 1 to ${MAX_WORKFLOW_NODES} nodes`); + } + for (const [id, node] of nodes) { + if (!isRecord(node) || typeof node.class_type !== "string" || !isRecord(node.inputs)) { + throw new WorkflowTemplateError(`Workflow node ${id} needs a class_type and inputs`); + } + } + const placeholders = new Set(); + const unknown = new Set(); + visitStrings(value, (text) => { + for (const [, name = ""] of text.matchAll(PLACEHOLDER_PATTERN)) { + if (isPlaceholder(name)) placeholders.add(name); + else unknown.add(`{{${name}}}`); + } + }); + if (unknown.size > 0) { + throw new WorkflowTemplateError(`Unknown workflow placeholders: ${[...unknown].join(", ")}`); + } + if (!placeholders.has("prompt")) { + throw new WorkflowTemplateError("Workflow needs a {{prompt}} placeholder"); + } + return { graph: value as unknown as ComfyWorkflow, placeholders }; +} + +function replacePlaceholders(value: unknown, values: WorkflowValues): unknown { + if (typeof value === "string") { + const whole = WHOLE_PLACEHOLDER.exec(value)?.[1]; + if (whole) return values[whole as WorkflowPlaceholder]; + return value.replace(PLACEHOLDER_PATTERN, (_, name: string) => + String(values[name as WorkflowPlaceholder]), + ); + } + if (Array.isArray(value)) return value.map((entry) => replacePlaceholders(entry, values)); + if (isRecord(value)) { + return Object.fromEntries( + Object.entries(value).map(([key, entry]) => [key, replacePlaceholders(entry, values)]), + ); + } + return value; +} + +/** Returns a copy of the template's graph with every placeholder replaced by its typed value. */ +export function fillWorkflow(template: WorkflowTemplate, values: WorkflowValues): ComfyWorkflow { + for (const name of template.placeholders) { + const value = values[name]; + const type = WORKFLOW_PLACEHOLDERS[name]; + const valid = + type === "string" + ? typeof value === "string" + : typeof value === "number" && + Number.isFinite(value) && + (type === "number" || Number.isInteger(value)); + if (!valid) throw new WorkflowTemplateError(`Workflow value for {{${name}}} is missing`); + } + return replacePlaceholders(template.graph, values) as ComfyWorkflow; +} + +/** + * Built-in workflows. Images use core nodes only and run on any SD-family checkpoint; video + * uses ComfyUI's core LTX-Video nodes (model files in workflows/README.md). + */ +export const DEFAULT_WORKFLOWS = { + textToImage: parseWorkflowTemplate(textToImage), + imageToImage: parseWorkflowTemplate(imageToImage), + textToVideo: parseWorkflowTemplate(textToVideo), + imageToVideo: parseWorkflowTemplate(imageToVideo), +} as const; diff --git a/packages/core/src/media/workflows/README.md b/packages/core/src/media/workflows/README.md new file mode 100644 index 0000000..af97d1b --- /dev/null +++ b/packages/core/src/media/workflows/README.md @@ -0,0 +1,35 @@ +# ComfyUI workflows + +Default workflows for `generate_image`, `generate_video` and repainting in `edit_image`, in +ComfyUI's API format. A custom workflow in the Media settings replaces the matching default: +one that uses `{{image}}` replaces the image-to-image (or image-to-video) default, one without +it replaces the text-to-image (or text-to-video) default. + +## Placeholders + +A string value that is exactly `{{name}}` becomes the typed value; `{{name}}` inside a longer +string is replaced by its text. Every workflow needs `{{prompt}}`; other names are optional. + +| Placeholder | Value | +|---|---| +| `{{prompt}}`, `{{negative_prompt}}` | text | +| `{{model}}` | the configured image model, or the first installed checkpoint | +| `{{image}}` | the uploaded input image (for `LoadImage`) | +| `{{width}}`, `{{height}}` | output size in pixels | +| `{{seed}}`, `{{steps}}` | integers | +| `{{strength}}` | 0–1, image-to-image denoise | +| `{{frames}}`, `{{fps}}` | video length in frames (8n+1) and frame rate | + +To write your own, build it in ComfyUI, replace the values above with placeholders and export it +with **Export (API)**. + +## Model files + +- Images: any SD 1.5 or SDXL checkpoint in `models/checkpoints/`. +- Video (LTX-Video, core nodes, about 12 GB of VRAM): + - `models/checkpoints/ltx-video-2b-v0.9.5.safetensors` from + [Lightricks/LTX-Video](https://huggingface.co/Lightricks/LTX-Video) + - `models/text_encoders/t5xxl_fp16.safetensors` from + [comfyanonymous/flux_text_encoders](https://huggingface.co/comfyanonymous/flux_text_encoders) + +The video workflows need a ComfyUI release with the `CreateVideo` and `SaveVideo` nodes. diff --git a/packages/core/src/media/workflows/image-to-image.json b/packages/core/src/media/workflows/image-to-image.json new file mode 100644 index 0000000..e762d24 --- /dev/null +++ b/packages/core/src/media/workflows/image-to-image.json @@ -0,0 +1,45 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "{{model}}" } + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["1", 1] } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] } + }, + "4": { + "class_type": "LoadImage", + "inputs": { "image": "{{image}}" } + }, + "5": { + "class_type": "VAEEncode", + "inputs": { "pixels": ["4", 0], "vae": ["1", 2] } + }, + "6": { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["5", 0], + "seed": "{{seed}}", + "steps": "{{steps}}", + "cfg": 7, + "sampler_name": "dpmpp_2m", + "scheduler": "karras", + "denoise": "{{strength}}" + } + }, + "7": { + "class_type": "VAEDecode", + "inputs": { "samples": ["6", 0], "vae": ["1", 2] } + }, + "8": { + "class_type": "SaveImage", + "inputs": { "images": ["7", 0], "filename_prefix": "evoo/image" } + } +} diff --git a/packages/core/src/media/workflows/image-to-video.json b/packages/core/src/media/workflows/image-to-video.json new file mode 100644 index 0000000..6bf53a7 --- /dev/null +++ b/packages/core/src/media/workflows/image-to-video.json @@ -0,0 +1,86 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" } + }, + "2": { + "class_type": "CLIPLoader", + "inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["2", 0] } + }, + "4": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] } + }, + "5": { + "class_type": "LoadImage", + "inputs": { "image": "{{image}}" } + }, + "6": { + "class_type": "LTXVImgToVideo", + "inputs": { + "positive": ["3", 0], + "negative": ["4", 0], + "vae": ["1", 2], + "image": ["5", 0], + "width": "{{width}}", + "height": "{{height}}", + "length": "{{frames}}", + "batch_size": 1, + "strength": 1 + } + }, + "7": { + "class_type": "LTXVConditioning", + "inputs": { "positive": ["6", 0], "negative": ["6", 1], "frame_rate": "{{fps}}" } + }, + "8": { + "class_type": "LTXVScheduler", + "inputs": { + "steps": "{{steps}}", + "max_shift": 2.05, + "base_shift": 0.95, + "stretch": true, + "terminal": 0.1, + "latent": ["6", 2] + } + }, + "9": { + "class_type": "KSamplerSelect", + "inputs": { "sampler_name": "euler" } + }, + "10": { + "class_type": "SamplerCustom", + "inputs": { + "model": ["1", 0], + "add_noise": true, + "noise_seed": "{{seed}}", + "cfg": 3, + "positive": ["7", 0], + "negative": ["7", 1], + "sampler": ["9", 0], + "sigmas": ["8", 0], + "latent_image": ["6", 2] + } + }, + "11": { + "class_type": "VAEDecode", + "inputs": { "samples": ["10", 0], "vae": ["1", 2] } + }, + "12": { + "class_type": "CreateVideo", + "inputs": { "images": ["11", 0], "fps": "{{fps}}" } + }, + "13": { + "class_type": "SaveVideo", + "inputs": { + "video": ["12", 0], + "filename_prefix": "evoo/video", + "format": "mp4", + "codec": "h264" + } + } +} diff --git a/packages/core/src/media/workflows/text-to-image.json b/packages/core/src/media/workflows/text-to-image.json new file mode 100644 index 0000000..633cdc5 --- /dev/null +++ b/packages/core/src/media/workflows/text-to-image.json @@ -0,0 +1,41 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "{{model}}" } + }, + "2": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["1", 1] } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["1", 1] } + }, + "4": { + "class_type": "EmptyLatentImage", + "inputs": { "width": "{{width}}", "height": "{{height}}", "batch_size": 1 } + }, + "5": { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "positive": ["2", 0], + "negative": ["3", 0], + "latent_image": ["4", 0], + "seed": "{{seed}}", + "steps": "{{steps}}", + "cfg": 7, + "sampler_name": "dpmpp_2m", + "scheduler": "karras", + "denoise": 1 + } + }, + "6": { + "class_type": "VAEDecode", + "inputs": { "samples": ["5", 0], "vae": ["1", 2] } + }, + "7": { + "class_type": "SaveImage", + "inputs": { "images": ["6", 0], "filename_prefix": "evoo/image" } + } +} diff --git a/packages/core/src/media/workflows/text-to-video.json b/packages/core/src/media/workflows/text-to-video.json new file mode 100644 index 0000000..55a3b3c --- /dev/null +++ b/packages/core/src/media/workflows/text-to-video.json @@ -0,0 +1,77 @@ +{ + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": { "ckpt_name": "ltx-video-2b-v0.9.5.safetensors" } + }, + "2": { + "class_type": "CLIPLoader", + "inputs": { "clip_name": "t5xxl_fp16.safetensors", "type": "ltxv" } + }, + "3": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{prompt}}", "clip": ["2", 0] } + }, + "4": { + "class_type": "CLIPTextEncode", + "inputs": { "text": "{{negative_prompt}}", "clip": ["2", 0] } + }, + "5": { + "class_type": "LTXVConditioning", + "inputs": { "positive": ["3", 0], "negative": ["4", 0], "frame_rate": "{{fps}}" } + }, + "6": { + "class_type": "EmptyLTXVLatentVideo", + "inputs": { + "width": "{{width}}", + "height": "{{height}}", + "length": "{{frames}}", + "batch_size": 1 + } + }, + "7": { + "class_type": "LTXVScheduler", + "inputs": { + "steps": "{{steps}}", + "max_shift": 2.05, + "base_shift": 0.95, + "stretch": true, + "terminal": 0.1, + "latent": ["6", 0] + } + }, + "8": { + "class_type": "KSamplerSelect", + "inputs": { "sampler_name": "euler" } + }, + "9": { + "class_type": "SamplerCustom", + "inputs": { + "model": ["1", 0], + "add_noise": true, + "noise_seed": "{{seed}}", + "cfg": 3, + "positive": ["5", 0], + "negative": ["5", 1], + "sampler": ["8", 0], + "sigmas": ["7", 0], + "latent_image": ["6", 0] + } + }, + "10": { + "class_type": "VAEDecode", + "inputs": { "samples": ["9", 0], "vae": ["1", 2] } + }, + "11": { + "class_type": "CreateVideo", + "inputs": { "images": ["10", 0], "fps": "{{fps}}" } + }, + "12": { + "class_type": "SaveVideo", + "inputs": { + "video": ["11", 0], + "filename_prefix": "evoo/video", + "format": "mp4", + "codec": "h264" + } + } +} From e0277985f987170e0401e8324db0708a811077f6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:37:19 +0000 Subject: [PATCH 2/6] feat(media): stream job inputs and outputs, cover the job lifecycle on Postgres Media files can be hundreds of megabytes, so the runner no longer buffers them: inputs are cloned from the local artifact file when the store has one and streamed otherwise, outputs go through storeArtifactFromFile (per-kind size cap and MIME check while streaming; an oversized result fails the job with a short card error). A Postgres suite runs the runtime with fake operations: success with outputs, posters and parent versions; input, availability, artifact and thread checks before a job exists; tool-call replay; failure with temp-dir cleanup; outputs outside the temp dir and sources of another user; deferral with backoff and its budget; the GPU memory guard; cancel of a queued job and of a job running in another process; one job per lease and one job on the GPU at a time; crash recovery and its restart limit; reconciliation inside and outside the queue; the continue_after wake; a cleared conversation; the run time limit; and capabilities. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- .../src/media-runtime.postgres.test.ts | 615 ++++++++++++++++++ packages/adapters/src/media-runtime.ts | 82 ++- packages/testkit/src/cli/harness.ts | 1 + 3 files changed, 653 insertions(+), 45 deletions(-) create mode 100644 packages/adapters/src/media-runtime.postgres.test.ts diff --git a/packages/adapters/src/media-runtime.postgres.test.ts b/packages/adapters/src/media-runtime.postgres.test.ts new file mode 100644 index 0000000..8fdf406 --- /dev/null +++ b/packages/adapters/src/media-runtime.postgres.test.ts @@ -0,0 +1,615 @@ +import { randomUUID } from "node:crypto"; +import { existsSync } from "node:fs"; +import { mkdtemp, readdir, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import type { + BackgroundJob, + JobPublisher, + MediaJobContext, + MediaOperationResult, +} from "@rakazo/adapter-kit"; +import { MediaJobDeferredError, mediaRunJobKey } from "@rakazo/adapter-kit"; +import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts"; +import { clearThread, createDb, type PrismaClient } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; +import { z } from "zod"; +import { storeArtifactFromFile } from "./artifact-files.js"; +import { LocalArtifactStore } from "./artifacts.js"; +import { + createMediaRuntime, + MEDIA_WAKE_TRIGGER, + MediaJobRejectedError, + type MediaOperationRegistration, + type MediaRuntimeOptions, +} from "./media-runtime.js"; + +const describePostgres = + process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip; + +const PNG = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3, 4]); + +type FakeInput = { prompt: string; source?: string }; +type Run = (context: MediaJobContext) => Promise; + +function deferred() { + let resolve!: (value: T) => void; + const promise = new Promise((done) => { + resolve = done; + }); + return { promise, resolve }; +} + +describePostgres("media job runtime (PostgreSQL, fake operations)", () => { + let prisma: PrismaClient; + let db: ReturnType; + let root: string; + const users: string[] = []; + + beforeAll(async () => { + db = createDb(process.env.DATABASE_URL!); + prisma = db.prisma; + root = await mkdtemp(path.join(tmpdir(), "media-runtime-test-")); + }); + afterAll(async () => { + if (!db) return; + await prisma.organization.deleteMany({ where: { id: { in: users } } }); + await prisma.user.deleteMany({ where: { id: { in: users } } }); + await prisma.$disconnect(); + await db.pool.end(); + await rm(root, { recursive: true, force: true }); + }); + + async function setup( + options: { + run?: Run; + available?: () => Promise; + gpu?: MediaOperationRegistration["gpu"]; + runtime?: MediaRuntimeOptions; + } = {}, + ) { + const id = randomUUID(); + users.push(id); + await prisma.user.create({ + data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true }, + }); + await prisma.organization.create({ + data: { id, name: "Media test", slug: id, createdAt: new Date() }, + }); + await prisma.space.create({ + data: { id, organizationId: id, name: "Media test", isDefault: true }, + }); + await prisma.member.create({ + data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() }, + }); + const bot = await prisma.bot.create({ + data: { spaceId: id, userId: id, name: "Studio", color: "test-color" }, + }); + const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } }); + const task = await prisma.task.create({ + data: { + spaceId: id, + userId: id, + botId: bot.id, + threadId: thread.id, + prompt: "Make a picture", + status: "running", + }, + }); + const run = await prisma.run.create({ + data: { + spaceId: id, + userId: id, + botId: bot.id, + threadId: thread.id, + taskId: task.id, + status: "running", + trigger: "user", + }, + }); + const artifacts = new LocalArtifactStore(path.join(root, id)); + const enqueued: BackgroundJob[] = []; + const cancelled: string[] = []; + const jobs: JobPublisher = { + enqueue: async (job) => { + enqueued.push(job); + }, + cancel: async (key) => { + cancelled.push(key); + }, + close: async () => undefined, + }; + const notify = vi.fn(async () => undefined); + const handlerCalls: string[] = []; + const defaultRun: Run = async (context) => { + const input = context.input as FakeInput; + if (input.source) await context.fetchArtifact(input.source); + await context.progress(0.5); + const file = path.join(context.tempDir, "out.png"); + const poster = path.join(context.tempDir, "poster.jpg"); + await writeFile(file, PNG); + await writeFile(poster, PNG); + await context.addOutput({ + path: file, + name: "picture.png", + mimeType: "image/png", + width: 64, + height: 32, + posterPath: poster, + parentArtifactId: input.source, + }); + return { summary: `Drew ${input.prompt}.` }; + }; + const registration: MediaOperationRegistration = { + operation: "generate_image", + inputSchema: z.object({ prompt: z.string().min(1), source: z.string().optional() }), + title: (input) => input.prompt, + available: options.available ?? (async () => true), + artifactIds: (input) => (input.source ? [input.source] : []), + handler: { + operation: "generate_image", + run: async (context) => { + handlerCalls.push(context.jobId); + return (options.run ?? defaultRun)(context); + }, + }, + ...(options.gpu ? { gpu: options.gpu } : {}), + }; + const tempRoot = path.join(root, id, "tmp"); + const runtimeDeps = { + prisma, + events: { notify }, + jobs, + artifacts, + registrations: [registration as MediaOperationRegistration], + }; + const runtimeOptions: MediaRuntimeOptions = { + tempRoot, + heartbeatMs: 50, + progressIntervalMs: 0, + ...options.runtime, + }; + const runtime = createMediaRuntime(runtimeDeps, runtimeOptions); + const start = (input: Record, extra: Record = {}) => + runtime.start({ + operation: "generate_image", + input, + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + threadId: thread.id, + runId: run.id, + ...extra, + }); + const job = (jobId: string) => prisma.mediaJob.findUniqueOrThrow({ where: { id: jobId } }); + const card = async (jobId: string): Promise => { + const row = await job(jobId); + const message = await prisma.message.findUniqueOrThrow({ where: { id: row.messageId! } }); + const block = (message.blocks as MessageBlock[]).find( + (entry): entry is MediaJobBlock => entry.kind === "media", + ); + if (!block) throw new Error("media card missing"); + return block; + }; + const sourceArtifact = async () => { + const file = path.join(root, `${randomUUID()}.png`); + await writeFile(file, PNG); + return storeArtifactFromFile( + { prisma, artifacts }, + { + path: file, + name: "source.png", + mimeType: "image/png", + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + runId: null, + }, + ); + }; + return { + id, + bot, + thread, + run, + runtime, + runtimeDeps, + runtimeOptions, + enqueued, + cancelled, + notify, + handlerCalls, + tempRoot, + start, + job, + card, + sourceArtifact, + }; + } + + it("runs a job to done, stores outputs with posters and versions, and cleans its temp dir", async () => { + const ctx = await setup(); + const source = await ctx.sourceArtifact(); + const started = await ctx.start({ prompt: "a red fox", source: source.id }); + + const queued = await ctx.card(started.jobId); + expect(queued).toMatchObject({ status: "queued", title: "a red fox", outputs: [] }); + expect(ctx.enqueued.map((job) => job.name)).toEqual(["media.run"]); + expect(ctx.enqueued[0]).toMatchObject({ payload: { jobId: started.jobId }, queue: "media" }); + const message = await prisma.message.findUniqueOrThrow({ where: { id: started.messageId } }); + expect(message).toMatchObject({ role: "bot", botId: ctx.bot.id, runId: ctx.run.id }); + + await ctx.runtime.runJob(started.jobId); + + const done = await ctx.job(started.jobId); + expect(done).toMatchObject({ status: "done", progress: 1, attempts: 1, leaseToken: null }); + expect(done.finishedAt).toBeInstanceOf(Date); + const block = await ctx.card(started.jobId); + expect(block.status).toBe("done"); + expect(block.outputs).toHaveLength(1); + const [output] = block.outputs; + expect(output).toMatchObject({ + kind: "image", + mimeType: "image/png", + name: "picture.png", + size: PNG.byteLength, + width: 64, + height: 32, + parentArtifactId: source.id, + }); + expect(done.outputArtifactIds).toEqual([output!.posterArtifactId, output!.artifactId]); + const stored = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } }); + expect(stored).toMatchObject({ + spaceId: ctx.id, + userId: ctx.id, + botId: ctx.bot.id, + runId: ctx.run.id, + parentArtifactId: source.id, + width: 64, + height: 32, + }); + const events = await prisma.event.findMany({ + where: { threadId: ctx.thread.id, type: "thread.message.updated" }, + orderBy: { seq: "asc" }, + }); + const statuses = events.map( + (event) => ((event.payload as { blocks: MessageBlock[] }).blocks[0] as MediaJobBlock).status, + ); + expect(statuses[0]).toBe("running"); + expect(statuses.at(-1)).toBe("done"); + expect(ctx.notify).toHaveBeenCalled(); + expect(existsSync(path.join(ctx.tempRoot, started.jobId))).toBe(false); + // No wake without continue_after. + expect( + await prisma.run.count({ where: { trigger: MEDIA_WAKE_TRIGGER, spaceId: ctx.id } }), + ).toBe(0); + expect(await ctx.runtime.get(started.jobId)).toMatchObject({ + id: started.jobId, + status: "done", + outputs: block.outputs, + }); + }); + + it("validates input, availability and readable artifacts before creating a job", async () => { + const ctx = await setup(); + await expect(ctx.start({ prompt: "" })).rejects.toBeInstanceOf(MediaJobRejectedError); + await expect(ctx.start({ prompt: "x", source: "not-an-artifact" })).rejects.toThrow( + "Unknown artifact.", + ); + const other = await setup(); + const foreign = await other.sourceArtifact(); + await expect(ctx.start({ prompt: "x", source: foreign.id })).rejects.toThrow( + "Unknown artifact.", + ); + await expect(ctx.start({ prompt: "x" }, { threadId: other.thread.id })).rejects.toThrow( + "This conversation is no longer available.", + ); + const unavailable = await setup({ available: async () => false }); + await expect(unavailable.start({ prompt: "x" })).rejects.toThrow("not available"); + expect( + await prisma.mediaJob.count({ where: { spaceId: { in: [ctx.id, unavailable.id] } } }), + ).toBe(0); + }); + + it("returns the first job for a replayed tool call", async () => { + const ctx = await setup(); + const first = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` }); + const replay = await ctx.start({ prompt: "fox" }, { operationKey: `effect:${ctx.id}` }); + expect(replay).toEqual(first); + expect(await prisma.mediaJob.count({ where: { spaceId: ctx.id } })).toBe(1); + }); + + it("fails a job with the operation's error and still cleans up", async () => { + let tempDir = ""; + const ctx = await setup({ + run: async (context) => { + tempDir = context.tempDir; + await writeFile(path.join(context.tempDir, "partial.bin"), PNG); + throw new Error(`ComfyUI rejected the workflow (${context.tempDir}/x)`); + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + const failed = await ctx.job(jobId); + expect(failed.status).toBe("failed"); + expect(failed.error).toBe("ComfyUI rejected the workflow (/x)"); + expect((await ctx.card(jobId)).error).toBe(failed.error); + expect(existsSync(tempDir)).toBe(false); + }); + + it("refuses outputs outside the temp dir and sources of another user", async () => { + const other = await setup(); + const foreign = await other.sourceArtifact(); + const ctx = await setup({ + run: async (context) => { + await context.fetchArtifact(foreign.id); + return { summary: "" }; + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "A source file is no longer available.", + }); + + const outside = await setup({ + run: async (context) => { + await context.addOutput({ path: "/etc/hostname", name: "x.png", mimeType: "image/png" }); + return { summary: "" }; + }, + }); + const escaped = await outside.start({ prompt: "fox" }); + await outside.runtime.runJob(escaped.jobId); + expect((await outside.job(escaped.jobId)).status).toBe("failed"); + expect(await prisma.artifact.count({ where: { spaceId: outside.id } })).toBe(0); + }); + + it("defers with backoff, then fails once the deferral budget is spent", async () => { + const ctx = await setup({ + run: async () => { + throw new MediaJobDeferredError("ComfyUI queue is full", 2_000); + }, + runtime: { maxDeferrals: 2 }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + const before = Date.now(); + await ctx.runtime.runJob(jobId); + const waiting = await ctx.job(jobId); + expect(waiting).toMatchObject({ status: "waiting", deferrals: 1, leaseToken: null }); + expect(waiting.nextAttemptAt!.getTime()).toBeGreaterThanOrEqual(before + 2_000); + expect((await ctx.card(jobId)).status).toBe("waiting"); + const retry = ctx.enqueued.at(-1)!; + expect(retry).toMatchObject({ name: "media.run", payload: { jobId }, queue: "media" }); + expect(retry.availableAt?.getTime()).toBe(waiting.nextAttemptAt!.getTime()); + + // Too early: the delayed job owns the next attempt. + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(1); + + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + const second = await ctx.job(jobId); + expect(second).toMatchObject({ status: "waiting", deferrals: 2 }); + // Backoff doubles. + expect(second.nextAttemptAt!.getTime() - Date.now()).toBeGreaterThan(3_000); + + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "The media server stayed busy. Try again later.", + }); + }); + + it("waits for free GPU memory before running the handler", async () => { + let free = 1_000; + const ctx = await setup({ + gpu: { + minVramMb: 6_000, + status: async () => ({ vramFreeMb: free, vramTotalMb: 24_000 }), + }, + }); + const { jobId } = await ctx.start({ prompt: "fox" }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(0); + expect(await ctx.job(jobId)).toMatchObject({ status: "waiting", deferrals: 1 }); + + free = 20_000; + await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(1); + expect((await ctx.job(jobId)).status).toBe("done"); + }); + + it("cancels a queued job at once and a running job through the lease heartbeat", async () => { + const ctx = await setup({ + run: async (context) => { + await new Promise((resolve) => { + context.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + }); + const queued = await ctx.start({ prompt: "queued" }); + await ctx.runtime.cancel(queued.jobId); + expect(await ctx.job(queued.jobId)).toMatchObject({ status: "cancelled" }); + expect((await ctx.card(queued.jobId)).status).toBe("cancelled"); + expect(ctx.cancelled).toEqual([mediaRunJobKey(queued.jobId)]); + await ctx.runtime.runJob(queued.jobId); + expect(ctx.handlerCalls).toHaveLength(0); + + // Another process (a second runtime) asks to cancel a job this one runs. + const running = await ctx.start({ prompt: "running" }); + const runningJob = ctx.runtime.runJob(running.jobId); + await vi.waitFor(() => expect(ctx.handlerCalls).toEqual([running.jobId])); + const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions); + await otherProcess.cancel(running.jobId); + expect((await ctx.job(running.jobId)).status).toBe("running"); + await runningJob; + expect(await ctx.job(running.jobId)).toMatchObject({ status: "cancelled", leaseToken: null }); + expect((await ctx.card(running.jobId)).status).toBe("cancelled"); + }); + + it("never runs one job twice and keeps one job on the GPU at a time", async () => { + const release = deferred(); + const entered = deferred(); + const ctx = await setup({ + run: async () => { + entered.resolve(); + await release.promise; + return { summary: "ok" }; + }, + }); + const otherProcess = createMediaRuntime(ctx.runtimeDeps, ctx.runtimeOptions); + const first = await ctx.start({ prompt: "one" }); + const second = await ctx.start({ prompt: "two" }); + + const a = ctx.runtime.runJob(first.jobId); + await entered.promise; + await otherProcess.runJob(first.jobId); + await otherProcess.runJob(second.jobId); + expect(ctx.handlerCalls).toEqual([first.jobId]); + expect(await ctx.job(second.jobId)).toMatchObject({ status: "waiting", attempts: 0 }); + + release.resolve(); + await a; + await prisma.mediaJob.update({ + where: { id: second.jobId }, + data: { nextAttemptAt: new Date() }, + }); + await otherProcess.runJob(second.jobId); + expect(ctx.handlerCalls).toEqual([first.jobId, second.jobId]); + expect((await ctx.job(second.jobId)).status).toBe("done"); + }); + + it("recovers a job whose runner died and fails it after too many restarts", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }); + const crashedAt = new Date(Date.now() - 1_000); + await prisma.mediaJob.update({ + where: { id: jobId }, + data: { + status: "running", + attempts: 1, + leaseToken: "dead-runner", + leaseExpiresAt: crashedAt, + updatedAt: crashedAt, + }, + }); + ctx.enqueued.length = 0; + await ctx.runtime.reconcile(); + expect(ctx.enqueued).toEqual([ + expect.objectContaining({ name: "media.run", payload: { jobId } }), + ]); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ status: "done", attempts: 2 }); + + const doomed = await ctx.start({ prompt: "crashes" }); + await prisma.mediaJob.update({ + where: { id: doomed.jobId }, + data: { + status: "running", + attempts: 3, + leaseToken: "dead-runner", + leaseExpiresAt: crashedAt, + }, + }); + await ctx.runtime.runJob(doomed.jobId); + expect(await ctx.job(doomed.jobId)).toMatchObject({ + status: "failed", + error: "The job was interrupted. Try again.", + }); + expect(ctx.handlerCalls).toEqual([jobId]); + }); + + it("re-enqueues lost and overdue jobs, outside the queue once they look stuck", async () => { + const ctx = await setup(); + const fresh = await ctx.start({ prompt: "fresh" }); + const lost = await ctx.start({ prompt: "lost" }); + const stuck = await ctx.start({ prompt: "stuck" }); + const minuteAgo = new Date(Date.now() - 60_000); + const hourAgo = new Date(Date.now() - 3_600_000); + await prisma.mediaJob.update({ + where: { id: lost.jobId }, + data: { updatedAt: minuteAgo, createdAt: minuteAgo }, + }); + await prisma.mediaJob.update({ + where: { id: stuck.jobId }, + data: { updatedAt: hourAgo, createdAt: hourAgo }, + }); + ctx.enqueued.length = 0; + await ctx.runtime.reconcile(); + const byJob = new Map( + ctx.enqueued.map((job) => [(job.payload as { jobId: string }).jobId, job]), + ); + expect(byJob.has(fresh.jobId)).toBe(false); + expect(byJob.get(lost.jobId)).toMatchObject({ queue: "media", availableAt: minuteAgo }); + expect(byJob.get(stuck.jobId)?.queue).toBeUndefined(); + }); + + it("wakes the starting bot with the result when continue_after is set", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true }); + ctx.enqueued.length = 0; + await ctx.runtime.runJob(jobId); + const wake = await prisma.run.findFirstOrThrow({ + where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER }, + include: { task: true }, + }); + expect(wake).toMatchObject({ botId: ctx.bot.id, threadId: ctx.thread.id, status: "queued" }); + expect(wake.task.prompt).toContain(`Media job ${jobId}`); + expect(wake.task.prompt).toContain("Drew fox."); + expect(wake.task.prompt).toContain("picture.png"); + expect(ctx.enqueued).toEqual([ + expect.objectContaining({ name: "run.continue", payload: { runId: wake.id } }), + ]); + }); + + it("cancels a job whose conversation was cleared without running it", async () => { + const ctx = await setup(); + const { jobId } = await ctx.start({ prompt: "fox" }, { continueAfter: true }); + await prisma.run.update({ where: { id: ctx.run.id }, data: { status: "completed" } }); + await clearThread(prisma, { spaceId: ctx.id, botId: ctx.bot.id, threadId: ctx.thread.id }); + await ctx.runtime.runJob(jobId); + expect(ctx.handlerCalls).toHaveLength(0); + expect(await ctx.job(jobId)).toMatchObject({ status: "cancelled" }); + expect( + await prisma.run.count({ where: { spaceId: ctx.id, trigger: MEDIA_WAKE_TRIGGER } }), + ).toBe(0); + }); + + it("fails a job that outlives its run limit", async () => { + const ctx = await setup({ + run: async (context) => { + await new Promise((resolve) => setTimeout(resolve, 5_000)); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + runtime: { maxRunMs: 50, abortGraceMs: 20 }, + }); + const { jobId } = await ctx.start({ prompt: "slow" }); + await ctx.runtime.runJob(jobId); + expect(await ctx.job(jobId)).toMatchObject({ + status: "failed", + error: "The job took too long.", + }); + expect(await readdir(ctx.tempRoot)).toEqual([]); + }); + + it("reports which operations a space can run", async () => { + const ctx = await setup(); + const capabilities = await ctx.runtime.capabilities(ctx.id); + expect(capabilities.generate_image).toBe(true); + expect(capabilities.transcribe).toBe(false); + const broken = await setup({ + available: async () => { + throw new Error("probe failed"); + }, + }); + expect((await broken.runtime.capabilities(broken.id)).generate_image).toBe(false); + }); +}); diff --git a/packages/adapters/src/media-runtime.ts b/packages/adapters/src/media-runtime.ts index 681f2e9..928f08f 100644 --- a/packages/adapters/src/media-runtime.ts +++ b/packages/adapters/src/media-runtime.ts @@ -1,7 +1,11 @@ -import { createHash, randomUUID } from "node:crypto"; -import { mkdir, readFile, rm, stat, writeFile } from "node:fs/promises"; +import { randomUUID } from "node:crypto"; +import { constants, createWriteStream } from "node:fs"; +import { copyFile, mkdir, rm } from "node:fs/promises"; import { tmpdir } from "node:os"; import path from "node:path"; +import { Readable } from "node:stream"; +import { pipeline } from "node:stream/promises"; +import type { ReadableStream as NodeReadableStream } from "node:stream/web"; import type { AdapterContext, ArtifactStore, @@ -17,13 +21,13 @@ import type { StartMediaJobInput, } from "@rakazo/adapter-kit"; import { + ArtifactTooLargeError, MediaJobDeferredError, mediaRunJob, mediaRunJobKey, runContinueJob, } from "@rakazo/adapter-kit"; import { - ARTIFACT_MAX_BYTES_BY_KIND, attachmentKindOf, MEDIA_JOB_ACTIVE_STATUSES, type MediaJobBlock, @@ -33,11 +37,7 @@ import { type MediaOutput, type MessageBlock, } from "@rakazo/contracts"; -import { - attachmentExtensionForMimeType, - inferAttachmentMimeType, - validateAttachmentMimeType, -} from "@rakazo/core"; +import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; import { appendEventInTransaction, createThreadMessageInTransaction, @@ -49,6 +49,7 @@ import { } from "@rakazo/db"; import { getLogger } from "@rakazo/logging"; import type { z } from "zod"; +import { storeArtifactFromFile } from "./artifact-files.js"; /** Run trigger of the turn that wakes a bot when its media job ends (`continue_after`). */ export const MEDIA_WAKE_TRIGGER = "media"; @@ -505,6 +506,7 @@ export function createMediaRuntime( const storeArtifactFile = async ( job: MediaJob, + signal: AbortSignal, file: { path: string; name: string; @@ -515,38 +517,37 @@ export function createMediaRuntime( parentArtifactId?: string; }, ) => { - // Swap point for the streaming artifact store: until it lands, files go through put(). - const bytes = new Uint8Array(await readFile(file.path)); - const context = artifactContext(job, `media-output:${job.id}`); - const stored = await deps.artifacts.put( - { name: file.name, mimeType: file.mimeType, bytes }, - context, - ); try { - return await deps.prisma.artifact.create({ - data: { - spaceId: job.spaceId, - userId: job.userId, - botId: job.botId, - groupId: job.groupId, - runId: job.runId, - name: file.name, - mimeType: file.mimeType, - size: bytes.byteLength, - hash: createHash("sha256").update(bytes).digest("hex"), - storageKey: stored.id, - width: file.width, - height: file.height, - durationMs: file.durationMs, - parentArtifactId: file.parentArtifactId, - }, + return await storeArtifactFromFile(deps, { + ...file, + spaceId: job.spaceId, + userId: job.userId, + botId: job.botId, + groupId: job.groupId, + runId: job.runId, + signal, }); } catch (error) { - await deps.artifacts.remove(stored.id, context).catch(() => undefined); + if (error instanceof ArtifactTooLargeError) throw new MediaJobRejectedError(TOO_LARGE_ERROR); throw error; } }; + /** Copies a stored artifact to `file`: a clone of the local file when possible, else a stream. */ + const copyArtifact = async (job: MediaJob, storageKey: string, file: string) => { + const context = artifactContext(job, `media-input:${job.id}`); + const local = await deps.artifacts.localPath?.(storageKey, context).catch(() => undefined); + if (local) { + await copyFile(local, file, constants.COPYFILE_EXCL | constants.COPYFILE_FICLONE); + return; + } + const stream = await deps.artifacts.createReadStream(storageKey, context); + await pipeline( + Readable.fromWeb(stream as NodeReadableStream), + createWriteStream(file, { flags: "wx", mode: 0o600 }), + ); + }; + const removeArtifacts = async (job: MediaJob, ids: string[]) => { const rows = await deps.prisma.artifact.findMany({ where: { id: { in: ids }, spaceId: job.spaceId }, @@ -657,12 +658,7 @@ export function createMediaRuntime( attachmentExtensionForMimeType(row.mimeType) || (/^\.[a-z0-9]{1,8}$/i.test(path.extname(row.name)) ? path.extname(row.name) : ""); const file = path.join(tempDir, `input-${row.id}${extension.toLowerCase()}`); - // Swap point for the streaming artifact store (localPath / createReadStream). - const bytes = await deps.artifacts.get( - row.storageKey, - artifactContext(job, `media-input:${job.id}`), - ); - await writeFile(file, bytes, { mode: 0o600, flag: "wx" }); + await copyArtifact(job, row.storageKey, file); const input: MediaJobInputFile = { artifactId: row.id, path: file, @@ -679,12 +675,8 @@ export function createMediaRuntime( }, addOutput: async (file: MediaJobOutputFile) => { signal.throwIfAborted(); - validateAttachmentMimeType(file.mimeType); const kind = attachmentKindOf(file.mimeType); const filePath = insideDir(tempDir, file.path); - const size = (await stat(filePath)).size; - if (size > ARTIFACT_MAX_BYTES_BY_KIND[kind]) - throw new MediaJobRejectedError(TOO_LARGE_ERROR); const parent = file.parentArtifactId ? await deps.prisma.artifact.findFirst({ where: { id: file.parentArtifactId, spaceId: job.spaceId, userId: job.userId }, @@ -702,7 +694,7 @@ export function createMediaRuntime( const posterName = `${path.parse(name).name}-poster${path.extname(posterPath).toLowerCase()}`; const posterMime = inferAttachmentMimeType(posterName); if (posterMime && attachmentKindOf(posterMime) === "image") { - const poster = await storeArtifactFile(job, { + const poster = await storeArtifactFile(job, signal, { path: posterPath, name: posterName, mimeType: posterMime, @@ -716,7 +708,7 @@ export function createMediaRuntime( const width = positiveInt(file.width); const height = positiveInt(file.height); const durationMs = nonNegativeInt(file.durationMs); - const artifact = await storeArtifactFile(job, { + const artifact = await storeArtifactFile(job, signal, { path: filePath, name, mimeType: file.mimeType, diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index 0ee852c..edc60ae 100644 --- a/packages/testkit/src/cli/harness.ts +++ b/packages/testkit/src/cli/harness.ts @@ -116,6 +116,7 @@ async function main() { "packages/adapters/src/realtime.postgres.test.ts", "packages/adapters/src/job-reconciler.postgres.test.ts", "packages/adapters/src/cloud-agent.postgres.test.ts", + "packages/adapters/src/media-runtime.postgres.test.ts", "packages/adapters/src/builtin-memory-provider.postgres.test.ts", "packages/adapters/src/profile-consolidation.postgres.test.ts", ]; From 4217e4b4af74c581eff6f04a2d7de57f13e0ee76 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:42:35 +0000 Subject: [PATCH 3/6] feat(media): generate images and video on ComfyUI and edit images Bots can now make images and short clips on the user's own GPU server and rework images locally, through the media job runtime. - generate_image: prompt, aspect preset or exact size, up to 4 images, optional source image for image-to-image. Sizes target what the default models are tuned for, and sources are fitted to that size before upload, so large photos cannot exhaust GPU memory. - generate_video: prompt, aspect, duration up to 10 s, optional start frame. - edit_image: crop, rotate, flip, resize to a size or format (crop, pad or blurred background), extend, convert, in that order, with sharp; an optional prompt repaints the result on ComfyUI. Edits keep the source as parent. - Registrations declare the input artifacts so jobs refuse ones the user cannot read, and generation declares a fixed GPU memory floor (6 GB images, 12 GB video) that a custom workflow cannot lower. - Tool definitions reuse the operation input schemas plus continue_after. - Saving media settings rejects custom workflows the provider could not fill. - Tests: sharp edits on generated fixtures, handlers against the ComfyUI emulator, and a PostgreSQL suite running the operations through the media runtime (card outputs, artifact versions, foreign artifacts, GPU wait). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- packages/adapters/src/builtin-tools.ts | 31 ++ packages/adapters/src/media-operations.ts | 5 +- packages/adapters/src/media-settings.ts | 12 + .../src/media/comfyui-settings.test.ts | 28 ++ .../adapters/src/media/image-edit.test.ts | 110 ++++++ packages/adapters/src/media/image-edit.ts | 156 +++++++++ .../media/image-operations.postgres.test.ts | 172 ++++++++++ .../src/media/image-operations.test.ts | 274 +++++++++++++++ .../adapters/src/media/image-operations.ts | 315 ++++++++++++++++++ .../contracts/src/media-generation.test.ts | 1 - packages/contracts/src/media-generation.ts | 11 +- packages/testkit/src/cli/harness.ts | 1 + 12 files changed, 1107 insertions(+), 9 deletions(-) create mode 100644 packages/adapters/src/media/comfyui-settings.test.ts create mode 100644 packages/adapters/src/media/image-edit.test.ts create mode 100644 packages/adapters/src/media/image-edit.ts create mode 100644 packages/adapters/src/media/image-operations.postgres.test.ts create mode 100644 packages/adapters/src/media/image-operations.test.ts create mode 100644 packages/adapters/src/media/image-operations.ts diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index 0ff52fa..704e9b6 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -2,12 +2,25 @@ import type { ConnectorTool } from "@rakazo/adapter-kit"; import { BotSecretDestination, BotSecretName, + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, MODEL_TIERS, SecretAskPurpose, SecretHttpRequest, USER_FACT_CATEGORIES, } from "@rakazo/contracts"; import { z } from "zod"; +import { mediaToolInputSchema } from "./media-tools.js"; + +/** A media tool's parameters from its operation's input schema, plus `continue_after`. */ +function mediaParameters(schema: z.ZodType) { + const json = z.toJSONSchema(schema, { io: "input" }); + return mediaToolInputSchema({ + properties: (json.properties ?? {}) as Record, + required: json.required, + }); +} export const SUBAGENT_MODEL_TIER_DESCRIPTION = "Optional model tier: fast for quick simple work, code for programming, standard otherwise. Uses the model this space maps to the tier, or your own model when it maps none. Ignored when model_provider and model_id are set."; @@ -326,6 +339,24 @@ export const builtinAgentTools: ConnectorTool[] = [ }, }, }, + { + name: "generate_image", + description: + "Generate up to 4 images from a prompt on the user's image server, or rework a source image (image-to-image).", + inputSchema: mediaParameters(GenerateImageInputSchema), + }, + { + name: "generate_video", + description: + "Generate a short video clip from a prompt on the user's image server, optionally starting from an image.", + inputSchema: mediaParameters(GenerateVideoInputSchema), + }, + { + name: "edit_image", + description: + "Edit an image artifact into a new one: crop, rotate, flip, resize (crop, pad, or blur to fit a format), add borders, convert the format, or repaint it from a prompt. Edits apply in that order.", + inputSchema: mediaParameters(EditImageInputSchema), + }, { name: "add_mcp_server", description: diff --git a/packages/adapters/src/media-operations.ts b/packages/adapters/src/media-operations.ts index 98dbf7e..8a6d556 100644 --- a/packages/adapters/src/media-operations.ts +++ b/packages/adapters/src/media-operations.ts @@ -1,4 +1,5 @@ import type { MediaSettings } from "@rakazo/contracts"; +import { createImageOperations } from "./media/image-operations.js"; import type { MediaOperationRegistration } from "./media-runtime.js"; export interface MediaOperationDeps { @@ -10,6 +11,6 @@ export interface MediaOperationDeps { * Every media operation this deployment can run. The API and the worker build their runtime * from this one list, so both processes run the same operations. */ -export function createMediaOperations(_deps: MediaOperationDeps): MediaOperationRegistration[] { - return []; +export function createMediaOperations(deps: MediaOperationDeps): MediaOperationRegistration[] { + return [...createImageOperations(deps)]; } diff --git a/packages/adapters/src/media-settings.ts b/packages/adapters/src/media-settings.ts index d237a13..b82134d 100644 --- a/packages/adapters/src/media-settings.ts +++ b/packages/adapters/src/media-settings.ts @@ -5,6 +5,7 @@ import { MediaSettingsUpdateSchema, type MediaSettingsView, } from "@rakazo/contracts"; +import { parseWorkflowTemplate } from "@rakazo/core/media/workflows"; import type { PrismaClient } from "@rakazo/db"; import { assertAllowedOpenAiCompatibleRequestUrl, @@ -82,6 +83,15 @@ function keptKey( return new URL(previous.baseUrl).origin === new URL(baseUrl).origin ? previous.apiKey : undefined; } +function assertValidWorkflow(label: string, workflow: string | undefined): void { + if (!workflow) return; + try { + parseWorkflowTemplate(workflow); + } catch (error) { + throw new Error(`${label}: ${error instanceof Error ? error.message : String(error)}`); + } +} + /** * Validates and stores the media backends (deployment owner only; the caller authorizes). * Throws an Error with a user-facing message when a URL is not allowed. @@ -91,6 +101,8 @@ export async function saveMediaSettings( update: MediaSettingsUpdate, ): Promise { const parsed = MediaSettingsUpdateSchema.parse(update); + assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow); + assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow); const previous = await loadMediaSettings(deps); let comfyui: MediaSettings["comfyui"] = null; if (parsed.comfyui) { diff --git a/packages/adapters/src/media/comfyui-settings.test.ts b/packages/adapters/src/media/comfyui-settings.test.ts new file mode 100644 index 0000000..5026881 --- /dev/null +++ b/packages/adapters/src/media/comfyui-settings.test.ts @@ -0,0 +1,28 @@ +import { describe, expect, it } from "vitest"; +import { saveMediaSettings } from "../media-settings.js"; + +const comfyui = (workflows: { imageWorkflow?: string; videoWorkflow?: string }) => ({ + comfyui: { baseUrl: "http://127.0.0.1:8188", ...workflows }, + transcription: null, +}); + +describe("ComfyUI settings", () => { + it("refuses custom workflows the provider could not fill", async () => { + const deps = {} as Parameters[0]; + await expect( + saveMediaSettings( + deps, + comfyui({ imageWorkflow: '{"1": {"class_type": "X", "inputs": {}}}' }), + ), + ).rejects.toThrow("Image workflow: Workflow needs a {{prompt}} placeholder"); + await expect( + saveMediaSettings( + deps, + comfyui({ + videoWorkflow: + '{"1": {"class_type": "X", "inputs": {"t": "{{prompt}}", "c": "{{cfg}}"}}}', + }), + ), + ).rejects.toThrow("Video workflow: Unknown workflow placeholders: {{cfg}}"); + }); +}); diff --git a/packages/adapters/src/media/image-edit.test.ts b/packages/adapters/src/media/image-edit.test.ts new file mode 100644 index 0000000..638c7b1 --- /dev/null +++ b/packages/adapters/src/media/image-edit.test.ts @@ -0,0 +1,110 @@ +import sharp from "sharp"; +import { describe, expect, it } from "vitest"; +import { editImageGeometry, encodeImage, imageSize } from "./image-edit.js"; + +/** Left half red, right half blue. */ +async function halves(width = 40, height = 20): Promise { + const half = Math.floor(width / 2); + const blue = await sharp({ + create: { width: width - half, height, channels: 3, background: { r: 0, g: 0, b: 255 } }, + }) + .png() + .toBuffer(); + return sharp({ create: { width, height, channels: 3, background: { r: 255, g: 0, b: 0 } } }) + .composite([{ input: blue, left: half, top: 0 }]) + .png() + .toBuffer(); +} + +async function pixel(image: Uint8Array, x: number, y: number) { + const { data, info } = await sharp(image).raw().toBuffer({ resolveWithObject: true }); + const offset = (y * info.width + x) * info.channels; + return [...data.subarray(offset, offset + info.channels)]; +} + +async function size(image: Uint8Array) { + const { width, height } = await sharp(image).metadata(); + return [width, height]; +} + +describe("editImageGeometry", () => { + it("applies crop before rotate, in source coordinates", async () => { + const result = await editImageGeometry(await halves(), { + crop: { left: 20, top: 0, width: 20, height: 10 }, + rotate: 90, + }); + expect(await size(result)).toEqual([10, 20]); + expect((await pixel(result, 5, 5)).slice(0, 3)).toEqual([0, 0, 255]); + }); + + it("flips horizontally and vertically", async () => { + const flopped = await editImageGeometry(await halves(), { flip: "horizontal" }); + expect((await pixel(flopped, 0, 0)).slice(0, 3)).toEqual([0, 0, 255]); + const flipped = await editImageGeometry(await halves(), { flip: "vertical" }); + expect((await pixel(flipped, 0, 0)).slice(0, 3)).toEqual([255, 0, 0]); + }); + + it("resizes to a preset with crop, pad or blur, and keeps the aspect for one side", async () => { + const source = await halves(160, 90); + for (const fit of ["crop", "pad", "blur"] as const) { + const result = await editImageGeometry(source, { resize: { aspect: "square", fit } }); + expect(await size(result)).toEqual([1080, 1080]); + } + const padded = await editImageGeometry(source, { resize: { aspect: "vertical", fit: "pad" } }); + expect((await pixel(padded, 0, 0)).at(-1)).toBe(0); + const scaled = await editImageGeometry(source, { resize: { width: 80, fit: "crop" } }); + expect(await size(scaled)).toEqual([80, 45]); + }); + + it("extends with a color or transparency", async () => { + const colored = await editImageGeometry(await halves(), { + extend: { top: 5, right: 0, bottom: 0, left: 3, color: "#00ff00" }, + }); + expect(await size(colored)).toEqual([43, 25]); + expect((await pixel(colored, 0, 0)).slice(0, 3)).toEqual([0, 255, 0]); + const clear = await editImageGeometry(await halves(), { + extend: { top: 2, right: 2, bottom: 2, left: 2 }, + }); + expect((await pixel(clear, 0, 0)).at(-1)).toBe(0); + }); + + it("applies EXIF orientation first", async () => { + const rotatedJpeg = await sharp(await halves()) + .jpeg() + .withMetadata({ orientation: 6 }) + .toBuffer(); + expect(await imageSize(rotatedJpeg)).toEqual({ width: 20, height: 40 }); + expect(await size(await editImageGeometry(rotatedJpeg, { flip: "vertical" }))).toEqual([ + 20, 40, + ]); + }); + + it("refuses crops outside the image and oversized results", async () => { + await expect( + editImageGeometry(await halves(), { crop: { left: 30, top: 0, width: 20, height: 10 } }), + ).rejects.toThrow("The crop area is outside the 40×20 image"); + await expect( + editImageGeometry(await halves(), { + extend: { top: 9000, right: 0, bottom: 0, left: 9000 }, + }), + ).rejects.toThrow("The edited image would be too large"); + await expect( + editImageGeometry(await halves(), { resize: { width: 8192, height: 8193, fit: "crop" } }), + ).rejects.toThrow("The edited image would be too large"); + }); +}); + +describe("encodeImage", () => { + it("encodes each format and flattens transparency for JPEG", async () => { + const clear = await editImageGeometry(await halves(), { + extend: { top: 2, right: 0, bottom: 0, left: 0 }, + }); + for (const format of ["png", "jpeg", "webp"] as const) { + const encoded = await encodeImage(clear, format); + expect((await sharp(encoded.bytes).metadata()).format).toBe(format); + expect([encoded.width, encoded.height]).toEqual([40, 22]); + } + const jpeg = await encodeImage(clear, "jpeg"); + expect((await pixel(jpeg.bytes, 0, 0)).every((channel) => channel < 16)).toBe(true); + }); +}); diff --git a/packages/adapters/src/media/image-edit.ts b/packages/adapters/src/media/image-edit.ts new file mode 100644 index 0000000..e250821 --- /dev/null +++ b/packages/adapters/src/media/image-edit.ts @@ -0,0 +1,156 @@ +import type { EditImageInput, ImageFormat } from "@rakazo/contracts"; +import { MEDIA_PRESETS } from "@rakazo/contracts"; +import type sharpFactory from "sharp"; +import type { Sharp } from "sharp"; + +type SharpFactory = typeof sharpFactory; +export type ImageGeometryEdits = Pick< + EditImageInput, + "crop" | "rotate" | "flip" | "resize" | "extend" +>; + +/** Edited images stay below this many pixels (8192²). */ +const MAX_EDIT_PIXELS = 8192 * 8192; +const TRANSPARENT = { r: 0, g: 0, b: 0, alpha: 0 }; +const BLACK = { r: 0, g: 0, b: 0 }; + +let sharpModule: Promise | undefined; + +/** sharp is a native module; load it on first use so a platform without it only loses edits. */ +export function loadSharp(): Promise { + sharpModule ??= import("sharp").then((module) => module.default); + return sharpModule; +} + +export async function sharpAvailable(): Promise { + return loadSharp().then( + () => true, + () => false, + ); +} + +/** Displayed size, with EXIF orientation applied. */ +export async function imageSize( + input: string | Uint8Array, +): Promise<{ width: number; height: number }> { + const sharp = await loadSharp(); + const { autoOrient } = await sharp(input).metadata(); + return { width: autoOrient.width, height: autoOrient.height }; +} + +/** Covers exactly width×height (center crop) as PNG: generation inputs at the size asked for. */ +export async function coverImage( + input: string | Uint8Array, + width: number, + height: number, +): Promise { + const sharp = await loadSharp(); + const png = await sharp(input, { autoOrient: true }) + .resize(width, height, { fit: "cover" }) + .png() + .toBuffer(); + return new Uint8Array(png); +} + +function parseColor(hex: string | undefined) { + if (!hex) return TRANSPARENT; + const value = Number.parseInt(hex.slice(1), 16); + return { r: (value >> 16) & 255, g: (value >> 8) & 255, b: value & 255, alpha: 1 }; +} + +function assertSize(width: number, height: number): void { + if (width * height > MAX_EDIT_PIXELS) throw new Error("The edited image would be too large"); +} + +/** Runs each step on decoded pixels, so steps apply in the order given, not sharp's own order. */ +async function materialize(sharp: SharpFactory, image: Sharp) { + const { data, info } = await image.raw().toBuffer({ resolveWithObject: true }); + const { width, height, channels } = info; + return { image: sharp(data, { raw: { width, height, channels } }), width, height }; +} + +async function resizeTo( + sharp: SharpFactory, + image: Sharp, + width: number, + height: number, + fit: "crop" | "pad" | "blur", +): Promise { + if (fit === "crop") return image.resize(width, height, { fit: "cover" }); + if (fit === "pad") + return image.resize(width, height, { fit: "contain", background: TRANSPARENT }); + const source = await image.png().toBuffer(); + const background = await sharp(source) + .resize(width, height, { fit: "cover" }) + .blur(Math.max(width, height) / 40) + .toBuffer(); + const foreground = await sharp(source).resize(width, height, { fit: "inside" }).toBuffer(); + return sharp(background).composite([{ input: foreground, gravity: "centre" }]); +} + +/** + * Applies crop, rotate, flip, resize and extend in that order and returns a lossless PNG. + * EXIF orientation is applied first, so coordinates match what viewers show. + */ +export async function editImageGeometry( + input: string | Uint8Array, + edits: ImageGeometryEdits, +): Promise { + const sharp = await loadSharp(); + let current = await materialize(sharp, sharp(input, { autoOrient: true })); + const step = async (apply: (image: Sharp) => Sharp | Promise) => { + current = await materialize(sharp, await apply(current.image)); + }; + const { crop, rotate, flip, resize, extend } = edits; + if (crop) { + if (crop.left + crop.width > current.width || crop.top + crop.height > current.height) { + throw new Error(`The crop area is outside the ${current.width}×${current.height} image`); + } + await step((image) => image.extract(crop)); + } + if (rotate) await step((image) => image.rotate(rotate)); + if (flip) await step((image) => (flip === "vertical" ? image.flip() : image.flop())); + if (resize) { + const preset = resize.aspect ? MEDIA_PRESETS[resize.aspect] : undefined; + const width = preset?.width ?? resize.width; + const height = preset?.height ?? resize.height; + const scale = width ? width / current.width : height! / current.height; + assertSize(width ?? current.width * scale, height ?? current.height * scale); + await step((image) => + width && height + ? resizeTo(sharp, image, width, height, resize.fit) + : image.resize(width ?? null, height ?? null), + ); + } + if (extend) { + const { top, right, bottom, left, color } = extend; + assertSize(current.width + left + right, current.height + top + bottom); + await step((image) => + image.extend({ top, right, bottom, left, background: parseColor(color) }), + ); + } + return current.image.png().toBuffer(); +} + +export const IMAGE_FORMAT_MIME_TYPES: Record = { + png: "image/png", + jpeg: "image/jpeg", + webp: "image/webp", +}; + +/** Encodes the final image; JPEG has no transparency, so transparent areas turn black. */ +export async function encodeImage( + input: Uint8Array, + format: ImageFormat, +): Promise<{ bytes: Buffer; width: number; height: number }> { + const sharp = await loadSharp(); + const image = sharp(input); + const encoded = + format === "jpeg" + ? image.flatten({ background: BLACK }).jpeg({ quality: 90, mozjpeg: true }) + : format === "webp" + ? image.webp({ quality: 90 }) + : image.png(); + const { data, info } = await encoded.toBuffer({ resolveWithObject: true }); + return { bytes: data, width: info.width, height: info.height }; +} diff --git a/packages/adapters/src/media/image-operations.postgres.test.ts b/packages/adapters/src/media/image-operations.postgres.test.ts new file mode 100644 index 0000000..2809d4d --- /dev/null +++ b/packages/adapters/src/media/image-operations.postgres.test.ts @@ -0,0 +1,172 @@ +import { randomUUID } from "node:crypto"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { JobPublisher } from "@rakazo/adapter-kit"; +import type { MediaJobBlock, MessageBlock } from "@rakazo/contracts"; +import { ComfyUiSettingsSchema } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { createDb } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import { LocalArtifactStore } from "../artifacts.js"; +import { createMediaRuntime, MediaJobRejectedError } from "../media-runtime.js"; +import { ComfyUiEmulator } from "../testing/comfyui-emulator.js"; +import { createImageOperations } from "./image-operations.js"; + +const describePostgres = + process.env.VERIFY_DATABASE && process.env.DATABASE_URL ? describe.sequential : describe.skip; + +describePostgres("ComfyUI operations in the media runtime (PostgreSQL + ComfyUI emulator)", () => { + let db: ReturnType; + let prisma: PrismaClient; + let root: string; + const users: string[] = []; + + beforeAll(async () => { + db = createDb(process.env.DATABASE_URL!); + prisma = db.prisma; + root = await mkdtemp(join(tmpdir(), "image-operations-pg-")); + }); + + afterAll(async () => { + if (!db) return; + await prisma.mediaJob.deleteMany({ where: { userId: { in: users } } }); + await prisma.artifact.deleteMany({ where: { userId: { in: users } } }); + await prisma.organization.deleteMany({ where: { id: { in: users } } }); + await prisma.user.deleteMany({ where: { id: { in: users } } }); + await prisma.$disconnect(); + await db.pool.end(); + await rm(root, { recursive: true, force: true }); + }); + + async function setup() { + const id = randomUUID(); + users.push(id); + await prisma.user.create({ + data: { id, name: "Media test", email: `${id}@example.test`, emailVerified: true }, + }); + await prisma.organization.create({ + data: { id, name: "Media test", slug: id, createdAt: new Date() }, + }); + await prisma.space.create({ data: { id, organizationId: id, name: "Media", isDefault: true } }); + await prisma.member.create({ + data: { id, organizationId: id, userId: id, role: "owner", createdAt: new Date() }, + }); + const bot = await prisma.bot.create({ + data: { spaceId: id, userId: id, name: "Studio", color: "test-color" }, + }); + const thread = await prisma.thread.create({ data: { spaceId: id, userId: id, botId: bot.id } }); + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 0; + const jobs: JobPublisher = { + enqueue: async () => undefined, + cancel: async () => undefined, + close: async () => undefined, + }; + const runtime = createMediaRuntime( + { + prisma, + events: { notify: async () => undefined }, + jobs, + artifacts: new LocalArtifactStore(join(root, id, "artifacts")), + registrations: createImageOperations({ + loadSettings: async () => ({ + comfyui: ComfyUiSettingsSchema.parse({ baseUrl: "http://127.0.0.1:8188" }), + transcription: null, + }), + fetch: emulator.fetch, + }), + }, + { tempRoot: join(root, id, "jobs") }, + ); + const scope = { + spaceId: id, + userId: id, + botId: bot.id, + groupId: null, + threadId: thread.id, + runId: null, + }; + return { emulator, runtime, scope }; + } + + async function card(messageId: string, jobId: string): Promise { + const message = await prisma.message.findUniqueOrThrow({ where: { id: messageId } }); + return (message.blocks as MessageBlock[]).find( + (block): block is MediaJobBlock => block.kind === "media" && block.jobId === jobId, + ); + } + + it("generates images onto the card and edits one into a linked version", async () => { + const { runtime, scope } = await setup(); + const generated = await runtime.start({ + ...scope, + operation: "generate_image", + input: { prompt: "a lighthouse at dusk", aspect: "landscape", count: 2, seed: 3 }, + }); + await runtime.runJob(generated.jobId); + const job = await runtime.get(generated.jobId); + expect(job).toMatchObject({ status: "done", title: "a lighthouse at dusk" }); + expect(job?.outputs.map((output) => [output.mimeType, output.width, output.height])).toEqual([ + ["image/png", 1344, 768], + ["image/png", 1344, 768], + ]); + expect(await card(generated.messageId, generated.jobId)).toMatchObject({ + status: "done", + outputs: job?.outputs, + }); + + const source = job!.outputs[0]!.artifactId; + const edited = await runtime.start({ + ...scope, + operation: "edit_image", + input: { artifactId: source, resize: { aspect: "vertical", fit: "blur" }, format: "jpeg" }, + }); + await runtime.runJob(edited.jobId); + const edit = await runtime.get(edited.jobId); + expect(edit).toMatchObject({ status: "done", title: "9:16 · JPEG" }); + const [output] = edit!.outputs; + expect(output).toMatchObject({ mimeType: "image/jpeg", width: 1080, height: 1920 }); + const artifact = await prisma.artifact.findUniqueOrThrow({ where: { id: output!.artifactId } }); + expect(artifact).toMatchObject({ parentArtifactId: source, width: 1080, height: 1920 }); + }); + + it("refuses source artifacts the user cannot read", async () => { + const { runtime, scope } = await setup(); + const other = await setup(); + const generated = await other.runtime.start({ + ...other.scope, + operation: "generate_image", + input: { prompt: "private" }, + }); + await other.runtime.runJob(generated.jobId); + const foreign = (await other.runtime.get(generated.jobId))!.outputs[0]!.artifactId; + await expect( + runtime.start({ + ...scope, + operation: "generate_video", + input: { prompt: "animate it", sourceArtifactId: foreign }, + }), + ).rejects.toThrow(MediaJobRejectedError); + await expect( + runtime.start({ + ...scope, + operation: "edit_image", + input: { artifactId: foreign, rotate: 90 }, + }), + ).rejects.toThrow(MediaJobRejectedError); + }); + + it("waits while the GPU lacks free memory", async () => { + const { emulator, runtime, scope } = await setup(); + emulator.gpu = { totalMb: 24_576, freeMb: 4_000, torchReservedMb: 0, torchFreeMb: 0 }; + const video = await runtime.start({ + ...scope, + operation: "generate_video", + input: { prompt: "rain on a window" }, + }); + await runtime.runJob(video.jobId); + expect(await runtime.get(video.jobId)).toMatchObject({ status: "waiting", outputs: [] }); + expect(emulator.prompts).toHaveLength(0); + }); +}); diff --git a/packages/adapters/src/media/image-operations.test.ts b/packages/adapters/src/media/image-operations.test.ts new file mode 100644 index 0000000..e028552 --- /dev/null +++ b/packages/adapters/src/media/image-operations.test.ts @@ -0,0 +1,274 @@ +import { copyFile, mkdtemp, rm, stat, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { MediaJobContext, MediaJobOutputFile } from "@rakazo/adapter-kit"; +import { MediaJobDeferredError } from "@rakazo/adapter-kit"; +import type { ComfyUiSettings, MediaOutput } from "@rakazo/contracts"; +import { attachmentKindOf, ComfyUiSettingsSchema } from "@rakazo/contracts"; +import sharp from "sharp"; +import { afterEach, describe, expect, it } from "vitest"; +import { ComfyUiEmulator } from "../testing/comfyui-emulator.js"; +import { + createImageOperations, + editImageTitle, + fileStem, + generationSize, + IMAGE_GENERATION_MIN_VRAM_MB, + VIDEO_GENERATION_MIN_VRAM_MB, +} from "./image-operations.js"; + +const dirs: string[] = []; +afterEach(async () => { + await Promise.all(dirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +async function tempDir(): Promise { + const dir = await mkdtemp(join(tmpdir(), "image-operations-")); + dirs.push(dir); + return dir; +} + +async function fixture(name: string, width: number, height: number): Promise { + const path = join(await tempDir(), name); + const image = sharp({ + create: { width, height, channels: 3, background: { r: 200, g: 120, b: 40 } }, + }); + await writeFile(path, await (name.endsWith(".jpg") ? image.jpeg() : image.png()).toBuffer()); + return path; +} + +function setup( + settings: ComfyUiSettings | null = ComfyUiSettingsSchema.parse({ + baseUrl: "http://127.0.0.1:8188", + }), +) { + const emulator = new ComfyUiEmulator(); + emulator.pendingPolls = 0; + const operations = createImageOperations({ + loadSettings: async () => ({ comfyui: settings, transcription: null }), + fetch: emulator.fetch, + }); + const operation = (name: string) => operations.find((entry) => entry.operation === name)!; + return { emulator, operations, operation }; +} + +async function run( + operation: ReturnType["operation"], + name: string, + input: unknown, + artifacts: Record = {}, +) { + const dir = await tempDir(); + const outputs: MediaJobOutputFile[] = []; + const progress: number[] = []; + const context: MediaJobContext = { + jobId: "job-1", + spaceId: "space-1", + userId: "user-1", + botId: null, + input, + tempDir: dir, + signal: new AbortController().signal, + progress: async (fraction) => { + progress.push(fraction); + }, + fetchArtifact: async (artifactId) => { + const source = artifacts[artifactId]; + if (!source) throw new Error("Artifact not found"); + const name = source.split("/").at(-1)!; + const path = join(dir, `input-${name}`); + await copyFile(source, path); + const mimeType = name.endsWith(".jpg") + ? "image/jpeg" + : name.endsWith(".mp4") + ? "video/mp4" + : "image/png"; + return { + artifactId, + path, + name, + mimeType, + kind: attachmentKindOf(mimeType), + size: (await stat(path)).size, + }; + }, + addOutput: async (file) => { + outputs.push(file); + const output: MediaOutput = { + artifactId: `output-${outputs.length}`, + kind: attachmentKindOf(file.mimeType), + mimeType: file.mimeType, + name: file.name, + size: (await stat(file.path)).size, + }; + return output; + }, + }; + const result = await operation(name).handler.run(context); + return { result, outputs, progress }; +} + +describe("image operations", () => { + it("are available by configuration and guard the GPU", async () => { + const configured = setup(); + expect(configured.operations.map((entry) => entry.operation)).toEqual([ + "generate_image", + "generate_video", + "edit_image", + ]); + for (const entry of configured.operations) { + await expect(entry.available({ spaceId: "space-1" })).resolves.toBe(true); + } + const generateImage = configured.operation("generate_image"); + expect(generateImage.gpu?.minVramMb).toBe(IMAGE_GENERATION_MIN_VRAM_MB); + expect(configured.operation("generate_video").gpu?.minVramMb).toBe( + VIDEO_GENERATION_MIN_VRAM_MB, + ); + await expect(generateImage.gpu?.status()).resolves.toEqual({ + vramFreeMb: 20_000, + vramTotalMb: 24_576, + }); + const unconfigured = setup(null); + await expect( + unconfigured.operation("generate_image").available({ spaceId: "s" }), + ).resolves.toBe(false); + await expect(unconfigured.operation("edit_image").available({ spaceId: "s" })).resolves.toBe( + true, + ); + await expect(run(unconfigured.operation, "generate_image", { prompt: "x" })).rejects.toThrow( + "ComfyUI is not configured", + ); + }); + + it("generates images at the aspect asked for", async () => { + const { operation } = setup(); + const { result, outputs, progress } = await run(operation, "generate_image", { + prompt: "A red kite, über the sea!", + aspect: "vertical", + count: 2, + seed: 11, + continue_after: true, + }); + expect( + outputs.map(({ name, mimeType, width, height }) => [name, mimeType, width, height]), + ).toEqual([ + ["a-red-kite-uber-the-sea-1.png", "image/png", 768, 1344], + ["a-red-kite-uber-the-sea-2.png", "image/png", 768, 1344], + ]); + expect(progress.at(-1)).toBe(1); + expect(result.summary).toBe("Generated 2 images with seed 11."); + }); + + it("starts from a source image at the source's aspect and links the result to it", async () => { + const { operation, emulator } = setup(); + const photo = await fixture("photo.jpg", 300, 200); + const { outputs } = await run( + operation, + "generate_image", + { prompt: "night", sourceArtifactId: "photo-1" }, + { "photo-1": photo }, + ); + const [upload] = emulator.uploads.values(); + expect((await sharp(upload!).metadata()).width).toBe(1280); + expect(outputs[0]).toMatchObject({ width: 1280, height: 832, parentArtifactId: "photo-1" }); + const clip = await fixture("clip.mp4", 8, 8); + await expect( + run(operation, "generate_image", { prompt: "x", sourceArtifactId: "c" }, { c: clip }), + ).rejects.toThrow("The source must be an image"); + }); + + it("generates video in the preset's aspect", async () => { + const { operation, emulator } = setup(); + const { outputs, result } = await run(operation, "generate_video", { + prompt: "paper boat on a river", + aspect: "vertical", + }); + expect(outputs).toEqual([ + expect.objectContaining({ + name: "paper-boat-on-a-river.mp4", + mimeType: "video/mp4", + width: 480, + height: 832, + durationMs: 5042, + }), + ]); + expect(emulator.prompts[0]?.graph["6"]?.inputs).toMatchObject({ width: 480, height: 832 }); + expect(result.summary).toMatch(/^Generated a 5 s video with seed \d+\.$/); + }); + + it("edits locally in order and keeps the source as parent", async () => { + const { operation, emulator } = setup(null); + const photo = await fixture("photo.jpg", 400, 300); + const { outputs, result } = await run( + operation, + "edit_image", + { + artifactId: "photo-1", + crop: { left: 0, top: 0, width: 300, height: 300 }, + rotate: 90, + resize: { aspect: "vertical", fit: "blur" }, + format: "webp", + }, + { "photo-1": photo }, + ); + expect(outputs).toEqual([ + expect.objectContaining({ + name: "photo-edit.webp", + mimeType: "image/webp", + width: 1080, + height: 1920, + parentArtifactId: "photo-1", + }), + ]); + expect((await sharp(outputs[0]!.path).metadata()).format).toBe("webp"); + expect(result.summary).toBe("Edited photo.jpg into photo-edit.webp (1080×1920)."); + expect(emulator.requests).toHaveLength(0); + await expect( + run(operation, "edit_image", { artifactId: "photo-1", prompt: "snow" }, { "photo-1": photo }), + ).rejects.toThrow("ComfyUI is not configured"); + }); + + it("repaints an edit on ComfyUI at the edited size, deferring while the GPU is busy", async () => { + const { operation, emulator } = setup(); + const photo = await fixture("photo.png", 300, 200); + const input = { + artifactId: "photo-1", + crop: { left: 0, top: 0, width: 250, height: 200 }, + prompt: "snow", + }; + const { outputs } = await run(operation, "edit_image", input, { "photo-1": photo }); + expect(outputs[0]).toMatchObject({ name: "photo-edit.png", width: 250, height: 200 }); + const [upload] = emulator.uploads.values(); + expect(await sharp(upload!).metadata()).toMatchObject({ width: 1152, height: 896 }); + emulator.gpu = { totalMb: 24_576, freeMb: 2_000, torchReservedMb: 0, torchFreeMb: 0 }; + await expect(run(operation, "edit_image", input, { "photo-1": photo })).rejects.toBeInstanceOf( + MediaJobDeferredError, + ); + }); +}); + +describe("helpers", () => { + it("sizes generations near the pixel budget", () => { + expect(generationSize(1, 1024 * 1024, 64)).toEqual({ width: 1024, height: 1024 }); + expect(generationSize(16 / 9, 768 * 512, 32)).toEqual({ width: 832, height: 480 }); + expect(generationSize(21 / 9, 1024 * 1024, 64)).toEqual({ width: 1536, height: 640 }); + }); + + it("names files and titles edits briefly", () => { + expect(fileStem(" Ünïcode — Straße 42 at night, very long prompt here ", "image")).toBe( + "unicode-strasse-42-at-night-very", + ); + expect(fileStem("夜の海", "video")).toBe("video"); + expect( + editImageTitle({ + artifactId: "a", + crop: { left: 0, top: 0, width: 1, height: 1 }, + rotate: 90, + resize: { aspect: "vertical", fit: "crop" }, + format: "png", + }), + ).toBe("crop · 90° · 9:16 · PNG"); + expect(editImageTitle({ artifactId: "a", resize: { width: 800, fit: "crop" } })).toBe("800×…"); + expect(editImageTitle({ artifactId: "a", prompt: "make it snow" })).toBe("make it snow"); + }); +}); diff --git a/packages/adapters/src/media/image-operations.ts b/packages/adapters/src/media/image-operations.ts new file mode 100644 index 0000000..68c08aa --- /dev/null +++ b/packages/adapters/src/media/image-operations.ts @@ -0,0 +1,315 @@ +import { randomInt } from "node:crypto"; +import { writeFile } from "node:fs/promises"; +import { join } from "node:path"; +import type { + GeneratedMedia, + GpuStatus, + MediaJobContext, + MediaJobInputFile, + MediaOperationResult, +} from "@rakazo/adapter-kit"; +import { MediaJobDeferredError } from "@rakazo/adapter-kit"; +import type { + EditImageInput, + GenerateImageInput, + GenerateVideoInput, + ImageFormat, + MediaPresetId, +} from "@rakazo/contracts"; +import { + attachmentKindOf, + EditImageInputSchema, + GenerateImageInputSchema, + GenerateVideoInputSchema, + MEDIA_PRESETS, +} from "@rakazo/contracts"; +import { attachmentExtensionForMimeType } from "@rakazo/core"; +import type { ComfyUiProvider } from "../comfyui-provider.js"; +import { createComfyUiProvider } from "../comfyui-provider.js"; +import type { MediaOperationDeps } from "../media-operations.js"; +import type { MediaOperationRegistration } from "../media-runtime.js"; +import { + coverImage, + editImageGeometry, + encodeImage, + IMAGE_FORMAT_MIME_TYPES, + imageSize, + sharpAvailable, +} from "./image-edit.js"; + +/** Free GPU memory a generation needs before it starts; fixed, so a custom workflow cannot lower it. */ +export const IMAGE_GENERATION_MIN_VRAM_MB = 6_000; +export const VIDEO_GENERATION_MIN_VRAM_MB = 12_000; +/** Pixel budgets the default workflows are tuned for (SDXL-class images, LTX-Video clips). */ +const IMAGE_PIXELS = 1024 * 1024; +const VIDEO_PIXELS = 768 * 512; +const GPU_RETRY_MS = 60_000; + +export interface ImageOperationsDeps extends MediaOperationDeps { + /** Tests point the provider at an emulator. */ + fetch?: typeof globalThis.fetch; +} + +type ProviderSource = () => Promise; + +/** A size close to `pixels` with this aspect ratio, in multiples of `multiple`. */ +export function generationSize(aspect: number, pixels: number, multiple: number) { + return { + width: Math.max(multiple, Math.round(Math.sqrt(pixels * aspect) / multiple) * multiple), + height: Math.max(multiple, Math.round(Math.sqrt(pixels / aspect) / multiple) * multiple), + }; +} + +function presetAspect(preset: MediaPresetId): number { + return MEDIA_PRESETS[preset].width / MEDIA_PRESETS[preset].height; +} + +/** File name stem from a prompt: its first words in lowercase ASCII. */ +export function fileStem(text: string, fallback: string): string { + const words = text + .normalize("NFKD") + .replace(/\p{M}/gu, "") + .toLowerCase() + .replace(/ß/g, "ss") + .match(/[a-z0-9]+/g); + return words?.slice(0, 6).join("-").slice(0, 48).replace(/-+$/, "") || fallback; +} + +/** Short card title for an edit without a prompt, such as "crop · 9:16 · PNG". */ +export function editImageTitle(input: EditImageInput): string { + if (input.prompt) return input.prompt; + const { resize } = input; + const parts = [ + input.crop && "crop", + input.rotate && `${input.rotate}°`, + input.flip && "flip", + resize && + (resize.aspect + ? MEDIA_PRESETS[resize.aspect].aspect + : `${resize.width ?? "…"}×${resize.height ?? "…"}`), + input.extend && "extend", + input.format?.toUpperCase(), + ]; + return parts.filter(Boolean).join(" · "); +} + +function reportProgress(context: MediaJobContext) { + return (fraction: number) => { + context.progress(fraction).catch(() => undefined); + }; +} + +async function requireProvider(source: ProviderSource): Promise { + const provider = await source(); + if (!provider) throw new Error("ComfyUI is not configured"); + return provider; +} + +/** Defers instead of starting a generation the GPU has no room for right now. */ +async function assertGpuReady( + provider: ComfyUiProvider, + minVramMb: number, + signal: AbortSignal, +): Promise { + const status: GpuStatus | undefined = await provider.gpuStatus(signal); + if (status && status.vramFreeMb < Math.min(minVramMb, status.vramTotalMb)) { + throw new MediaJobDeferredError("Waiting for free GPU memory", GPU_RETRY_MS); + } +} + +async function fetchImage( + context: MediaJobContext, + artifactId: string, +): Promise { + const file = await context.fetchArtifact(artifactId); + if (file.kind !== "image") throw new Error("The source must be an image"); + return { ...file, ...(await imageSize(file.path)) }; +} + +async function addGeneratedOutputs( + context: MediaJobContext, + results: GeneratedMedia[], + stem: string, + parentArtifactId: string | undefined, +): Promise { + for (const [index, result] of results.entries()) { + const extension = attachmentExtensionForMimeType(result.mimeType); + const name = `${stem}${results.length > 1 ? `-${index + 1}` : ""}${extension}`; + const path = join(context.tempDir, name); + await writeFile(path, result.bytes); + const size = + attachmentKindOf(result.mimeType) === "image" + ? await imageSize(path) + : { width: result.width, height: result.height }; + await context.addOutput({ + path, + name, + mimeType: result.mimeType, + ...size, + ...(result.durationMs !== undefined ? { durationMs: result.durationMs } : {}), + ...(parentArtifactId ? { parentArtifactId } : {}), + }); + } +} + +async function generateImage( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = GenerateImageInputSchema.parse(context.input); + const provider = await requireProvider(providers); + const source = input.sourceArtifactId + ? await fetchImage(context, input.sourceArtifactId) + : undefined; + const aspect = input.aspect + ? presetAspect(input.aspect) + : source + ? source.width / source.height + : 1; + const size = + input.width && input.height + ? { width: input.width, height: input.height } + : generationSize(aspect, IMAGE_PIXELS, 64); + const seed = input.seed ?? randomInt(2 ** 48 - 1); + const images = await provider.generateImage({ + prompt: input.prompt, + negativePrompt: input.negativePrompt, + ...size, + count: input.count, + seed, + source: source && { + bytes: await coverImage(source.path, size.width, size.height), + mimeType: "image/png", + }, + strength: input.strength, + signal: context.signal, + onProgress: reportProgress(context), + }); + await addGeneratedOutputs(context, images, fileStem(input.prompt, "image"), source?.artifactId); + const count = images.length === 1 ? "1 image" : `${images.length} images`; + return { summary: `Generated ${count} with seed ${seed}.` }; +} + +async function generateVideo( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = GenerateVideoInputSchema.parse(context.input); + const provider = await requireProvider(providers); + const source = input.sourceArtifactId + ? await fetchImage(context, input.sourceArtifactId) + : undefined; + const aspect = input.aspect + ? presetAspect(input.aspect) + : source + ? source.width / source.height + : presetAspect("landscape"); + const size = generationSize(aspect, VIDEO_PIXELS, 32); + const seed = input.seed ?? randomInt(2 ** 48 - 1); + const clips = await provider.generateVideo({ + prompt: input.prompt, + negativePrompt: input.negativePrompt, + ...size, + durationSeconds: input.durationSeconds, + seed, + source: source && { + bytes: await coverImage(source.path, size.width, size.height), + mimeType: "image/png", + }, + signal: context.signal, + onProgress: reportProgress(context), + }); + await addGeneratedOutputs(context, clips, fileStem(input.prompt, "video"), source?.artifactId); + return { summary: `Generated a ${input.durationSeconds} s video with seed ${seed}.` }; +} + +function defaultFormat(mimeType: string): ImageFormat { + if (mimeType === "image/jpeg") return "jpeg"; + if (mimeType === "image/webp") return "webp"; + return "png"; +} + +async function editImage( + context: MediaJobContext, + providers: ProviderSource, +): Promise { + const input = EditImageInputSchema.parse(context.input); + const source = await fetchImage(context, input.artifactId); + let png: Uint8Array = await editImageGeometry(source.path, input); + if (input.prompt) { + const provider = await requireProvider(providers); + await assertGpuReady(provider, IMAGE_GENERATION_MIN_VRAM_MB, context.signal); + const { width, height } = await imageSize(png); + // Repaint near the model's native size, then return to the edited size. + const size = generationSize(width / height, IMAGE_PIXELS, 64); + const [repainted] = await provider.generateImage({ + prompt: input.prompt, + ...size, + source: { bytes: await coverImage(png, size.width, size.height), mimeType: "image/png" }, + strength: input.strength, + signal: context.signal, + onProgress: reportProgress(context), + }); + if (!repainted) throw new Error("ComfyUI produced no image"); + png = await coverImage(repainted.bytes, width, height); + } + const format = input.format ?? defaultFormat(source.mimeType); + const encoded = await encodeImage(png, format); + const mimeType = IMAGE_FORMAT_MIME_TYPES[format]; + const stem = source.name.replace(/\.[^.]+$/, "") || "image"; + const name = `${stem}-edit${attachmentExtensionForMimeType(mimeType)}`; + const path = join(context.tempDir, name); + await writeFile(path, encoded.bytes); + await context.addOutput({ + path, + name, + mimeType, + width: encoded.width, + height: encoded.height, + parentArtifactId: source.artifactId, + }); + return { summary: `Edited ${source.name} into ${name} (${encoded.width}×${encoded.height}).` }; +} + +/** + * Media operations for image and video generation on ComfyUI and for image edits with sharp. + * Generation is available when ComfyUI is configured; edits whenever sharp loads (repainting + * an edit also needs ComfyUI). + */ +export function createImageOperations(deps: ImageOperationsDeps): MediaOperationRegistration[] { + const providers: ProviderSource = async () => { + const { comfyui } = await deps.loadSettings(); + return comfyui ? createComfyUiProvider({ ...comfyui, fetch: deps.fetch }) : null; + }; + const comfyUiConfigured = async () => Boolean((await deps.loadSettings()).comfyui); + const status = async (signal?: AbortSignal) => (await providers())?.gpuStatus(signal); + const sourceIds = (input: { sourceArtifactId?: string }) => + input.sourceArtifactId ? [input.sourceArtifactId] : []; + const image: MediaOperationRegistration = { + operation: "generate_image", + inputSchema: GenerateImageInputSchema, + title: (input) => input.prompt, + available: comfyUiConfigured, + handler: { operation: "generate_image", run: (context) => generateImage(context, providers) }, + gpu: { minVramMb: IMAGE_GENERATION_MIN_VRAM_MB, status }, + artifactIds: sourceIds, + }; + const video: MediaOperationRegistration = { + operation: "generate_video", + inputSchema: GenerateVideoInputSchema, + title: (input) => input.prompt, + available: comfyUiConfigured, + handler: { operation: "generate_video", run: (context) => generateVideo(context, providers) }, + gpu: { minVramMb: VIDEO_GENERATION_MIN_VRAM_MB, status }, + artifactIds: sourceIds, + }; + const edit: MediaOperationRegistration = { + operation: "edit_image", + inputSchema: EditImageInputSchema, + title: editImageTitle, + available: sharpAvailable, + handler: { operation: "edit_image", run: (context) => editImage(context, providers) }, + artifactIds: (input) => [input.artifactId], + }; + return [image, video, edit] as MediaOperationRegistration[]; +} diff --git a/packages/contracts/src/media-generation.test.ts b/packages/contracts/src/media-generation.test.ts index 4ab1a74..32f9b78 100644 --- a/packages/contracts/src/media-generation.test.ts +++ b/packages/contracts/src/media-generation.test.ts @@ -13,7 +13,6 @@ describe("media generation inputs", () => { }); expect(GenerateVideoInputSchema.parse({ prompt: "waves" })).toEqual({ prompt: "waves", - aspect: "landscape", durationSeconds: 5, }); }); diff --git a/packages/contracts/src/media-generation.ts b/packages/contracts/src/media-generation.ts index 6d33216..2e61c42 100644 --- a/packages/contracts/src/media-generation.ts +++ b/packages/contracts/src/media-generation.ts @@ -15,9 +15,8 @@ const Seed = z .optional() .describe("Reuse a seed to reproduce a result."); const Strength = z.number().min(0).max(1).optional(); -const Aspect = MediaPresetIdSchema.describe( - "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9.", -); +const ASPECTS = "vertical 9:16, landscape 16:9, square 1:1, portrait 4:5, classic 4:3, cinema 21:9"; +const aspect = (fallback: string) => MediaPresetIdSchema.describe(`${ASPECTS}. ${fallback}`); const Side = z.number().int().min(64).max(2048); const Pixels = z.number().int().min(0).max(16_384); @@ -25,7 +24,7 @@ export const GenerateImageInputSchema = z .object({ prompt: Prompt.describe("What the image shows."), negativePrompt: NegativePrompt, - aspect: Aspect.optional(), + aspect: aspect("Default square, or the source image's aspect ratio.").optional(), width: Side.optional().describe("Exact width in pixels, with height; overrides aspect."), height: Side.optional(), count: z.number().int().min(1).max(MEDIA_IMAGE_MAX_COUNT).default(1), @@ -42,7 +41,7 @@ export type GenerateImageInput = z.infer; export const GenerateVideoInputSchema = z.object({ prompt: Prompt.describe("What happens in the clip: subject, motion, camera."), negativePrompt: NegativePrompt, - aspect: Aspect.default("landscape"), + aspect: aspect("Default landscape, or the start frame's aspect ratio.").optional(), durationSeconds: z.number().min(1).max(MEDIA_VIDEO_MAX_SECONDS).default(5), seed: Seed, sourceArtifactId: Id.optional().describe("Image artifact to use as the first frame."), @@ -64,7 +63,7 @@ export const EditImageInputSchema = z flip: z.enum(["horizontal", "vertical"]).optional(), resize: z .object({ - aspect: Aspect.optional().describe("Exact size of this format."), + aspect: aspect("Resizes to the format's exact size.").optional(), width: z.number().int().min(1).max(8192).optional(), height: z.number().int().min(1).max(8192).optional(), fit: MediaFitSchema.default("crop").describe( diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index 0ee852c..5d565e9 100644 --- a/packages/testkit/src/cli/harness.ts +++ b/packages/testkit/src/cli/harness.ts @@ -116,6 +116,7 @@ async function main() { "packages/adapters/src/realtime.postgres.test.ts", "packages/adapters/src/job-reconciler.postgres.test.ts", "packages/adapters/src/cloud-agent.postgres.test.ts", + "packages/adapters/src/media/image-operations.postgres.test.ts", "packages/adapters/src/builtin-memory-provider.postgres.test.ts", "packages/adapters/src/profile-consolidation.postgres.test.ts", ]; From 3b97805f205753259914304f0f36fe9c495a06a4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:45:50 +0000 Subject: [PATCH 4/6] refactor(media): share the ComfyUI JSON request helpers Queue, interrupt and prompt requests now go through one JSON helper, so the cancel path reads like the rest of the client. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- packages/adapters/src/comfyui-provider.ts | 73 +++++++++++------------ 1 file changed, 35 insertions(+), 38 deletions(-) diff --git a/packages/adapters/src/comfyui-provider.ts b/packages/adapters/src/comfyui-provider.ts index ec35de7..e2410a1 100644 --- a/packages/adapters/src/comfyui-provider.ts +++ b/packages/adapters/src/comfyui-provider.ts @@ -75,6 +75,7 @@ export interface ComfyUiProvider extends ImageGenerationProvider, VideoGeneratio } type MediaKind = "image" | "video"; +type ReadBody = (response: Response, signal: AbortSignal) => Promise; interface OutputFile { filename: string; @@ -259,20 +260,18 @@ class ComfyUiClient implements ComfyUiProvider { } private async submit(graph: ComfyWorkflow, signal: AbortSignal): Promise { - const body = JSON.stringify({ prompt: graph, client_id: this.clientId }); - return this.send("prompt", { method: "POST", headers: JSON_HEADERS, body }, signal, { - read: async (response, readSignal) => { - if (response.status === 400) { - const rejection = await readJson(response, readSignal).catch(() => undefined); - throw new Error(describeWorkflowRejection(rejection)); - } - const result = await readOkJson(response, readSignal); - const promptId = isRecord(result) ? result.prompt_id : undefined; - if (typeof promptId !== "string" || !promptId) { - throw new Error("ComfyUI did not accept the workflow"); - } - return promptId; - }, + const payload = { prompt: graph, client_id: this.clientId }; + return this.postJson("prompt", payload, signal, async (response, readSignal) => { + if (response.status === 400) { + const rejection = await readJson(response, readSignal).catch(() => undefined); + throw new Error(describeWorkflowRejection(rejection)); + } + const result = await readOkJson(response, readSignal); + const promptId = isRecord(result) ? result.prompt_id : undefined; + if (typeof promptId !== "string" || !promptId) { + throw new Error("ComfyUI did not accept the workflow"); + } + return promptId; }); } @@ -291,29 +290,20 @@ class ComfyUiClient implements ComfyUiProvider { /** Best effort: drops the prompt from the queue, or interrupts it when it is ours and running. */ private async cancelPrompt(promptId: string): Promise { - const options = { timeoutMs: CANCEL_TIMEOUT_MS, read: discardBody }; try { - const body = JSON.stringify({ delete: [promptId] }); - await this.send("queue", { method: "POST", headers: JSON_HEADERS, body }, undefined, options); - const queue = await this.send( + await this.postJson( "queue", - { headers: { Accept: "application/json" } }, + { delete: [promptId] }, undefined, - { - timeoutMs: CANCEL_TIMEOUT_MS, - read: readOkJson, - }, + discardBody, + CANCEL_TIMEOUT_MS, ); + const queue = await this.getJson("queue", undefined, CANCEL_TIMEOUT_MS); const running = isRecord(queue) && Array.isArray(queue.queue_running) ? queue.queue_running : []; if (running.some((entry) => Array.isArray(entry) && entry[1] === promptId)) { - const interrupt = JSON.stringify({ prompt_id: promptId }); - await this.send( - "interrupt", - { method: "POST", headers: JSON_HEADERS, body: interrupt }, - undefined, - options, - ); + const interrupt = { prompt_id: promptId }; + await this.postJson("interrupt", interrupt, undefined, discardBody, CANCEL_TIMEOUT_MS); } } catch { // The prompt then finishes on its own; nothing waits for it. @@ -379,10 +369,20 @@ class ComfyUiClient implements ComfyUiProvider { }); } - private getJson(path: string, signal: AbortSignal | undefined): Promise { - return this.send(path, { headers: { Accept: "application/json" } }, signal, { - read: readOkJson, - }); + private getJson(path: string, signal: AbortSignal | undefined, timeoutMs?: number) { + const init = { headers: { Accept: "application/json" } }; + return this.send(path, init, signal, { timeoutMs, read: readOkJson }); + } + + private postJson( + path: string, + payload: unknown, + signal: AbortSignal | undefined, + read: ReadBody, + timeoutMs?: number, + ): Promise { + const init = { method: "POST", headers: JSON_HEADERS, body: JSON.stringify(payload) }; + return this.send(path, init, signal, { timeoutMs, read }); } /** One request, with the timeout covering the response body too. */ @@ -394,10 +394,7 @@ class ComfyUiClient implements ComfyUiProvider { body?: string | Uint8Array; }, signal: AbortSignal | undefined, - options: { - timeoutMs?: number; - read: (response: Response, signal: AbortSignal) => Promise; - }, + options: { timeoutMs?: number; read: ReadBody }, ): Promise { const timeout = AbortSignal.timeout(options.timeoutMs ?? REQUEST_TIMEOUT_MS); const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; From 66a257b08392d2a6fd5d6161ac05a3966d491e25 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:46:34 +0000 Subject: [PATCH 5/6] fix(media): describe media tools by what the model can rely on The image server is the deployment's, not necessarily the user's, and repainting an edit only works when generation is available. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- packages/adapters/src/builtin-tools.ts | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/adapters/src/builtin-tools.ts b/packages/adapters/src/builtin-tools.ts index 704e9b6..df1e679 100644 --- a/packages/adapters/src/builtin-tools.ts +++ b/packages/adapters/src/builtin-tools.ts @@ -342,19 +342,19 @@ export const builtinAgentTools: ConnectorTool[] = [ { name: "generate_image", description: - "Generate up to 4 images from a prompt on the user's image server, or rework a source image (image-to-image).", + "Generate up to 4 images from a prompt on the self-hosted image server, or rework a source image (image-to-image).", inputSchema: mediaParameters(GenerateImageInputSchema), }, { name: "generate_video", description: - "Generate a short video clip from a prompt on the user's image server, optionally starting from an image.", + "Generate a short video clip from a prompt on the self-hosted image server, optionally starting from an image.", inputSchema: mediaParameters(GenerateVideoInputSchema), }, { name: "edit_image", description: - "Edit an image artifact into a new one: crop, rotate, flip, resize (crop, pad, or blur to fit a format), add borders, convert the format, or repaint it from a prompt. Edits apply in that order.", + "Edit an image artifact into a new one: crop, rotate, flip, resize (crop, pad, or blur to fit a format), add borders, convert the format, or repaint it from a prompt (only when generate_image is available). Edits apply in that order.", inputSchema: mediaParameters(EditImageInputSchema), }, { From 27969652569bed2cdb24efa74bf71bcd6ee284e6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:47:44 +0000 Subject: [PATCH 6/6] feat(media): media tools, media RPCs and owner-only media settings API Bots and users need to start and follow media jobs, and the deployment owner needs to point the deployment at their GPU and transcription servers. - Media tools are named like their operations. A run sees one only when the space can run that operation (MediaRuntime.capabilities); the per-bot allowlist still applies. applyTool hands every media tool to one dispatcher (media-tools.ts) that splits off continue_after, starts the job under the tool call's effect key (a replay returns the same job) and returns at once; refusals (invalid input, unknown artifact, unavailable) go back to the model as tool errors. - media.jobs.start starts an export, transcription or image edit from a card: the conversation is resolved like every thread procedure, the card's bot (or the group's lead) posts the new card, and the runtime refuses source artifacts outside the actor's space. media.jobs.get and media.jobs.cancel only see the actor's own jobs. - media.settings.get/update are deployment-owner only, never return keys and answer invalid URLs with a 400 that says what is wrong. - createApp accepts media operations as a test seam. - Tests: a Postgres suite over the RPCs (owner-only settings, encryption, key retention per server, URL policy, start/get/cancel isolation between users, card jobs end to end) and unit tests for tool exposure and dispatch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u --- apps/api/src/app.ts | 11 +- apps/api/src/media-jobs.ts | 118 ++++++ apps/api/src/router.ts | 30 ++ packages/adapters/src/executor-tools.ts | 30 +- packages/adapters/src/executor.ts | 20 ++ packages/adapters/src/media-settings.ts | 53 ++- packages/adapters/src/media-tools.test.ts | 125 +++++++ packages/contracts/src/index.ts | 1 + packages/contracts/src/media-jobs.ts | 36 ++ packages/contracts/src/rpc.ts | 18 + packages/testkit/src/cli/harness.ts | 1 + packages/testkit/src/media.postgres.test.ts | 378 ++++++++++++++++++++ 12 files changed, 788 insertions(+), 33 deletions(-) create mode 100644 apps/api/src/media-jobs.ts create mode 100644 packages/adapters/src/media-tools.test.ts create mode 100644 packages/contracts/src/media-jobs.ts create mode 100644 packages/testkit/src/media.postgres.test.ts diff --git a/apps/api/src/app.ts b/apps/api/src/app.ts index 2a9e5c1..d3ff4e7 100644 --- a/apps/api/src/app.ts +++ b/apps/api/src/app.ts @@ -15,6 +15,7 @@ import type { ComposioProvider, ConnectorRegistry, DestinationEmulator, + MediaOperationRegistration, RemoteConnectorDependencies, } from "@rakazo/adapters"; import { @@ -143,6 +144,8 @@ export async function createApp( email?: TransactionalEmailProvider; remoteConnectors?: RemoteConnectorDependencies; logger?: Logger; + /** Test seam: media operations instead of the configured ones. */ + mediaOperations?: MediaOperationRegistration[]; } = {}, ): Promise { const { @@ -155,6 +158,7 @@ export async function createApp( email: emailOverride, remoteConnectors, logger: loggerOverride, + mediaOperations, ...envOverrides } = overrides; const env = { ...loadEnv(process.env), ...envOverrides }; @@ -384,9 +388,9 @@ export async function createApp( events, jobs, artifacts, - registrations: createMediaOperations({ - loadSettings: () => loadMediaSettings({ prisma, secrets }), - }), + registrations: + mediaOperations ?? + createMediaOperations({ loadSettings: () => loadMediaSettings({ prisma, secrets }) }), }); const executor = createRunExecutor({ prisma, @@ -468,6 +472,7 @@ export async function createApp( const router = createRouter({ cloudAgent, + media, prisma, events, auth, diff --git a/apps/api/src/media-jobs.ts b/apps/api/src/media-jobs.ts new file mode 100644 index 0000000..4a289e8 --- /dev/null +++ b/apps/api/src/media-jobs.ts @@ -0,0 +1,118 @@ +import { ORPCError } from "@orpc/server"; +import type { MediaRuntime } from "@rakazo/adapter-kit"; +import { + type EncryptedSecretStore, + loadMediaSettings, + MediaJobRejectedError, + MediaSettingsInvalidError, + mediaSettingsView, + saveMediaSettings, +} from "@rakazo/adapters"; +import type { Actor, MediaJobStartInput, MediaSettingsUpdate } from "@rakazo/contracts"; +import type { PrismaClient } from "@rakazo/db"; +import { resolveThreadTarget } from "./thread-target.js"; + +export interface MediaRouteDeps { + prisma: PrismaClient; + secrets: EncryptedSecretStore; + media?: MediaRuntime; +} + +/** Media backends run on the deployment owner's own server, so only they configure them. */ +function requireDeploymentOwner(actor: Actor) { + if (!actor.isDeploymentOwner) throw new ORPCError("FORBIDDEN"); +} + +function requireMedia(deps: MediaRouteDeps): MediaRuntime { + if (!deps.media) throw new ORPCError("NOT_IMPLEMENTED"); + return deps.media; +} + +export async function getMediaSettings(deps: MediaRouteDeps, actor: Actor) { + requireDeploymentOwner(actor); + return mediaSettingsView(await loadMediaSettings(deps)); +} + +export async function updateMediaSettings( + deps: MediaRouteDeps, + actor: Actor, + input: MediaSettingsUpdate, +) { + requireDeploymentOwner(actor); + try { + return await saveMediaSettings(deps, input); + } catch (error) { + if (error instanceof MediaSettingsInvalidError) { + throw new ORPCError("BAD_REQUEST", { message: error.message }); + } + throw error; + } +} + +/** + * Starts an operation a user picked on a card (export, transcribe, edit) in the same conversation. + * The thread is authorized like every thread procedure; the runtime checks that the operation's + * source artifacts belong to the actor, and the job can only ever read the actor's artifacts. + */ +export async function startMediaJob(deps: MediaRouteDeps, actor: Actor, input: MediaJobStartInput) { + const media = requireMedia(deps); + const target = await resolveThreadTarget(deps.prisma, actor, input); + const source = await deps.prisma.message.findFirst({ + where: { id: input.messageId, threadId: target.threadId }, + select: { botId: true }, + }); + if (!source) throw new ORPCError("NOT_FOUND"); + const bots = + target.kind === "bot" + ? [target.botId] + : [ + ...(target.leadBotId && target.memberBotIds.includes(target.leadBotId) + ? [target.leadBotId] + : []), + ...target.memberBotIds, + ]; + const botId = source.botId && bots.includes(source.botId) ? source.botId : bots[0]; + if (!botId) throw new ORPCError("NOT_FOUND"); + try { + return await media.start({ + operation: input.operation, + input: input.input, + spaceId: actor.spaceId, + userId: actor.userId, + botId, + groupId: target.kind === "group" ? target.groupId : null, + threadId: target.threadId, + runId: null, + }); + } catch (error) { + if (error instanceof MediaJobRejectedError) { + throw new ORPCError("BAD_REQUEST", { message: error.message }); + } + throw error; + } +} + +async function ownedJobId(deps: MediaRouteDeps, actor: Actor, jobId: string): Promise { + const job = await deps.prisma.mediaJob.findFirst({ + where: { id: jobId, spaceId: actor.spaceId, userId: actor.userId }, + select: { id: true }, + }); + if (!job) throw new ORPCError("NOT_FOUND"); + return job.id; +} + +export async function getMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) { + const media = requireMedia(deps); + const snapshot = await media.get(await ownedJobId(deps, actor, jobId)); + if (!snapshot) throw new ORPCError("NOT_FOUND"); + return snapshot; +} + +export async function cancelMediaJob(deps: MediaRouteDeps, actor: Actor, jobId: string) { + const media = requireMedia(deps); + const id = await ownedJobId(deps, actor, jobId); + await media.cancel(id); + const snapshot = await media.get(id); + if (!snapshot) throw new ORPCError("NOT_FOUND"); + return snapshot; +} diff --git a/apps/api/src/router.ts b/apps/api/src/router.ts index 2baf21e..4255c3d 100644 --- a/apps/api/src/router.ts +++ b/apps/api/src/router.ts @@ -6,6 +6,7 @@ import type { ArtifactStore, ConnectorCatalogItem, JobPublisher, + MediaRuntime, MemoryStore, SandboxProvider, } from "@rakazo/adapter-kit"; @@ -166,6 +167,13 @@ import { renameMainConversation, } from "./main-conversation.js"; import { buildMcpUpdateMaterial } from "./mcp-material.js"; +import { + cancelMediaJob, + getMediaJob, + getMediaSettings, + startMediaJob, + updateMediaSettings, +} from "./media-jobs.js"; import { disconnectMemoryProvider, getMemoryProviderConfig, @@ -446,6 +454,7 @@ function mcpAssignmentDto(row: { export interface RouterDeps { cloudAgent?: CloudAgentConnection | null; + media?: MediaRuntime; prisma: PrismaClient; events: ThreadEvents; auth: Auth; @@ -3481,6 +3490,27 @@ export function createRouter(deps: RouterDeps) { return { ok: true as const }; }), }, + media: { + settings: { + get: authed.media.settings.get.handler(({ context }) => + getMediaSettings(deps, context.actor), + ), + update: authed.media.settings.update.handler(({ context, input }) => + updateMediaSettings(deps, context.actor, input), + ), + }, + jobs: { + start: authed.media.jobs.start.handler(({ context, input }) => + startMediaJob(deps, context.actor, input), + ), + get: authed.media.jobs.get.handler(({ context, input }) => + getMediaJob(deps, context.actor, input.jobId), + ), + cancel: authed.media.jobs.cancel.handler(({ context, input }) => + cancelMediaJob(deps, context.actor, input.jobId), + ), + }, + }, connections: { catalog: authed.connections.catalog.handler(async ({ context, input }) => { const adapterContext = connectionContext( diff --git a/packages/adapters/src/executor-tools.ts b/packages/adapters/src/executor-tools.ts index 2e819b8..1045aa2 100644 --- a/packages/adapters/src/executor-tools.ts +++ b/packages/adapters/src/executor-tools.ts @@ -1,6 +1,7 @@ -import type { ConnectorTool } from "@rakazo/adapter-kit"; +import type { ConnectorTool, MediaCapabilities } from "@rakazo/adapter-kit"; import { agentConnectionTools, builtinAgentTools } from "./builtin-tools.js"; import { selectCloudAgentTools } from "./cloud-agent-tools-select.js"; +import { selectMediaTools } from "./media-tools.js"; import { selectMemoryTools } from "./memory-tools.js"; import { filterImageReturningComputerTools } from "./model-vision.js"; import { @@ -81,6 +82,8 @@ export function selectBuiltinToolsForRun(options: { trigger: string; semanticMemoryEnabled: boolean; cloudAgentEnabled?: boolean; + /** Media tools appear only for operations the space can run. */ + mediaCapabilities?: MediaCapabilities; messagingChannelRun: boolean; /** Cross-owner agent connections only exist for chat-linked bots. */ agentConnectionsEnabled?: boolean; @@ -88,21 +91,24 @@ export function selectBuiltinToolsForRun(options: { foreignAgentRun?: boolean; bot: BuiltinToolPolicy; }) { - const selected = selectCloudAgentTools( - selectMemoryTools( - filterBuiltinToolsForRun( - filterBuiltinToolsForThread( - filterPageBrowserTools( - filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed), - options.pageBrowserAllowed ?? options.graphicalToolsAllowed, + const selected = selectMediaTools( + selectCloudAgentTools( + selectMemoryTools( + filterBuiltinToolsForRun( + filterBuiltinToolsForThread( + filterPageBrowserTools( + filterImageReturningComputerTools(builtinAgentTools, options.graphicalToolsAllowed), + options.pageBrowserAllowed ?? options.graphicalToolsAllowed, + ), + options.groupId, ), - options.groupId, + options.trigger, ), - options.trigger, + options.semanticMemoryEnabled, ), - options.semanticMemoryEnabled, + Boolean(options.cloudAgentEnabled), ), - Boolean(options.cloudAgentEnabled), + options.mediaCapabilities, ).filter( (tool) => !(options.foreignAgentRun && tool.name === "update_user_profile") && diff --git a/packages/adapters/src/executor.ts b/packages/adapters/src/executor.ts index 5cb9181..31b094d 100644 --- a/packages/adapters/src/executor.ts +++ b/packages/adapters/src/executor.ts @@ -283,6 +283,7 @@ import { needsOAuthProbe, parseMcpServerToolArgs, } from "./mcp-server-tool.js"; +import { executeMediaTool, isMediaToolName } from "./media-tools.js"; import { loadAgentMemoryContext } from "./memory-context.js"; import type { MemoryProviderResolver } from "./memory-provider-factory.js"; import { selectMemoryTools } from "./memory-tools.js"; @@ -1477,6 +1478,7 @@ export function createRunExecutor(deps: ExecutorDeps) { trigger: run.trigger, semanticMemoryEnabled, cloudAgentEnabled: cloudAgentsEnabled(cloudAgent, run.spaceId), + mediaCapabilities: await deps.media?.capabilities(run.spaceId).catch(() => undefined), messagingChannelRun, agentConnectionsEnabled: hasMessagingIdentity, foreignAgentRun, @@ -2598,6 +2600,24 @@ export function createRunExecutor(deps: ExecutorDeps) { return computerScreenToolResult(() => tool(browser, computer, context, args), finish); } + if (isMediaToolName(name)) { + return finish( + await executeMediaTool( + deps.media, + { + spaceId: run.spaceId, + userId: run.userId, + botId: bot.id, + groupId: thread.groupId, + threadId: run.threadId, + runId: run.id, + operationKey: effectKey, + }, + name, + args, + ), + ); + } if (name.startsWith("cloud_agent_")) { return finish( await executeCloudAgentTool( diff --git a/packages/adapters/src/media-settings.ts b/packages/adapters/src/media-settings.ts index d237a13..325c45c 100644 --- a/packages/adapters/src/media-settings.ts +++ b/packages/adapters/src/media-settings.ts @@ -82,31 +82,48 @@ function keptKey( return new URL(previous.baseUrl).origin === new URL(baseUrl).origin ? previous.apiKey : undefined; } +/** A media setting the deployment owner has to correct; the message says what is wrong. */ +export class MediaSettingsInvalidError extends Error { + constructor(message: string) { + super(message); + this.name = "MediaSettingsInvalidError"; + } +} + +function resolveSettings(previous: MediaSettings, update: MediaSettingsUpdate): MediaSettings { + const parsed = MediaSettingsUpdateSchema.parse(update); + try { + let comfyui: MediaSettings["comfyui"] = null; + if (parsed.comfyui) { + const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl); + const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl); + assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); + comfyui = { ...parsed.comfyui, baseUrl, apiKey }; + } + let transcription: MediaSettings["transcription"] = null; + if (parsed.transcription) { + const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl); + const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl); + assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); + transcription = { ...parsed.transcription, baseUrl, apiKey }; + } + return MediaSettingsSchema.parse({ comfyui, transcription }); + } catch (error) { + throw new MediaSettingsInvalidError( + error instanceof Error ? error.message : "These media settings are not valid", + ); + } +} + /** * Validates and stores the media backends (deployment owner only; the caller authorizes). - * Throws an Error with a user-facing message when a URL is not allowed. + * Throws MediaSettingsInvalidError when a URL is not allowed. */ export async function saveMediaSettings( deps: MediaSettingsDeps, update: MediaSettingsUpdate, ): Promise { - const parsed = MediaSettingsUpdateSchema.parse(update); - const previous = await loadMediaSettings(deps); - let comfyui: MediaSettings["comfyui"] = null; - if (parsed.comfyui) { - const baseUrl = normalizeComfyUiBaseUrl(parsed.comfyui.baseUrl); - const apiKey = keptKey(parsed.comfyui.apiKey, previous.comfyui, baseUrl); - assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); - comfyui = { ...parsed.comfyui, baseUrl, apiKey }; - } - let transcription: MediaSettings["transcription"] = null; - if (parsed.transcription) { - const baseUrl = normalizeTranscriptionBaseUrl(parsed.transcription.baseUrl); - const apiKey = keptKey(parsed.transcription.apiKey, previous.transcription, baseUrl); - assertHttpsForKeyedOpenAiCompatibleUrl(new URL(baseUrl), apiKey); - transcription = { ...parsed.transcription, baseUrl, apiKey }; - } - const settings = MediaSettingsSchema.parse({ comfyui, transcription }); + const settings = resolveSettings(await loadMediaSettings(deps), update); if (!settings.comfyui && !settings.transcription) { await deps.prisma.integrationProviderConfig.deleteMany({ where: { id: MEDIA_SETTINGS_ID } }); return mediaSettingsView(settings); diff --git a/packages/adapters/src/media-tools.test.ts b/packages/adapters/src/media-tools.test.ts new file mode 100644 index 0000000..9af7370 --- /dev/null +++ b/packages/adapters/src/media-tools.test.ts @@ -0,0 +1,125 @@ +import type { MediaCapabilities, MediaRuntime, StartMediaJobInput } from "@rakazo/adapter-kit"; +import { describe, expect, it, vi } from "vitest"; +import { selectBuiltinToolsForRun } from "./executor-tools.js"; +import { MediaJobRejectedError } from "./media-runtime.js"; +import { + executeMediaTool, + isMediaToolName, + MEDIA_CONTINUE_AFTER_PROPERTY, + mediaToolInputSchema, + selectMediaTools, +} from "./media-tools.js"; + +const none: MediaCapabilities = { + generate_image: false, + generate_video: false, + edit_image: false, + edit_video: false, + export_video: false, + transcribe: false, +}; + +const scope = { + spaceId: "space-1", + userId: "user-1", + botId: "bot-1", + groupId: null, + threadId: "thread-1", + runId: "run-1", + operationKey: "effect-1", +}; + +function runtime(start: MediaRuntime["start"]): MediaRuntime { + return { + capabilities: async () => none, + start, + get: async () => null, + cancel: async () => undefined, + }; +} + +describe("media tools", () => { + it("names tools after their operations", () => { + expect(isMediaToolName("generate_image")).toBe(true); + expect(isMediaToolName("transcribe")).toBe(true); + expect(isMediaToolName("render_plot")).toBe(false); + }); + + it("shows a media tool only when the space can run its operation", () => { + const tools = [{ name: "generate_image" }, { name: "transcribe" }, { name: "web_search" }]; + expect(selectMediaTools(tools, undefined).map((tool) => tool.name)).toEqual(["web_search"]); + expect(selectMediaTools(tools, { ...none, transcribe: true }).map((tool) => tool.name)).toEqual( + ["transcribe", "web_search"], + ); + }); + + it("leaves non-media tools of a run untouched", () => { + const base = { + graphicalToolsAllowed: true, + groupId: null, + trigger: "user", + semanticMemoryEnabled: false, + messagingChannelRun: false, + bot: { kind: "specialist", restrictBuiltinTools: false, allowedBuiltinTools: [] }, + }; + const names = (mediaCapabilities?: MediaCapabilities) => + selectBuiltinToolsForRun({ ...base, mediaCapabilities }) + .map((tool) => tool.name) + .filter((name) => !isMediaToolName(name)); + expect(names({ ...none, generate_image: true, transcribe: true })).toEqual(names(undefined)); + expect(selectBuiltinToolsForRun(base).some((tool) => isMediaToolName(tool.name))).toBe(false); + }); + + it("adds continue_after to every media tool schema", () => { + expect( + mediaToolInputSchema({ + properties: { prompt: { type: "string" } }, + required: ["prompt"], + }), + ).toEqual({ + type: "object", + properties: { prompt: { type: "string" }, ...MEDIA_CONTINUE_AFTER_PROPERTY }, + required: ["prompt"], + }); + }); + + it("starts the job with the tool scope and returns at once", async () => { + const start = vi.fn(async (_input: StartMediaJobInput) => ({ + jobId: "job-1", + messageId: "message-1", + })); + const result = await executeMediaTool(runtime(start), scope, "generate_image", { + prompt: "a fox", + continue_after: true, + }); + expect(start).toHaveBeenCalledWith({ + operation: "generate_image", + input: { prompt: "a fox" }, + ...scope, + continueAfter: true, + }); + expect(result).toMatchObject({ jobId: "job-1", status: "queued" }); + expect((result as { note: string }).note).toContain("woken"); + + await executeMediaTool(runtime(start), scope, "generate_image", { prompt: "a fox" }); + expect(start.mock.calls[1]?.[0]).toMatchObject({ continueAfter: false }); + }); + + it("returns refusals to the model and rethrows other failures", async () => { + const refused = runtime(async () => { + throw new MediaJobRejectedError("Invalid generate_image input: prompt: Required"); + }); + expect(await executeMediaTool(refused, scope, "generate_image", {})).toEqual({ + error: "Invalid generate_image input: prompt: Required", + }); + expect(await executeMediaTool(undefined, scope, "transcribe", {})).toEqual({ + error: "transcribe is not available.", + }); + const broken = runtime(async () => { + throw new Error("database down"); + }); + await expect(executeMediaTool(broken, scope, "transcribe", {})).rejects.toThrow( + "database down", + ); + }); +}); diff --git a/packages/contracts/src/index.ts b/packages/contracts/src/index.ts index 48aa488..9ed1f63 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -10,6 +10,7 @@ export * from "./integration-settings.js"; export * from "./local-settings.js"; export * from "./mcp.js"; export * from "./media.js"; +export * from "./media-jobs.js"; export * from "./media-settings.js"; export * from "./model-tiers.js"; export * from "./openai-compatible-ui.js"; diff --git a/packages/contracts/src/media-jobs.ts b/packages/contracts/src/media-jobs.ts new file mode 100644 index 0000000..477a6fe --- /dev/null +++ b/packages/contracts/src/media-jobs.ts @@ -0,0 +1,36 @@ +import * as z from "zod"; +import { Id } from "./ids.js"; +import { MediaJobStatusSchema, MediaOperationSchema, MediaOutputSchema } from "./media.js"; + +/** Operations a user starts from a media card (export, transcribe, edit); generation goes through a bot. */ +export const MediaCardOperationSchema = z.enum(["edit_image", "export_video", "transcribe"]); +export type MediaCardOperation = z.infer; + +export const MediaJobStartInputSchema = z + .object({ + /** The conversation, like other thread procedures: a bot chat or a group. */ + botId: Id.optional(), + groupId: Id.optional(), + /** The message the action came from; its bot posts the new card. */ + messageId: Id, + operation: MediaCardOperationSchema, + /** The operation's input; validated by the operation. */ + input: z.record(z.string(), z.unknown()), + }) + .refine((input) => Boolean(input.botId) !== Boolean(input.groupId), { + message: "Pass botId or groupId", + }); +export type MediaJobStartInput = z.infer; + +export const MediaJobSnapshotSchema = z.object({ + id: Id, + operation: MediaOperationSchema, + status: MediaJobStatusSchema, + title: z.string(), + progress: z.number().min(0).max(1).nullable(), + error: z.string().nullable(), + outputs: z.array(MediaOutputSchema), +}); +export type MediaJobSnapshotDto = z.infer; + +export const MediaJobStartOutputSchema = z.object({ jobId: Id, messageId: Id }); diff --git a/packages/contracts/src/rpc.ts b/packages/contracts/src/rpc.ts index 6640d4c..30de55a 100644 --- a/packages/contracts/src/rpc.ts +++ b/packages/contracts/src/rpc.ts @@ -82,6 +82,12 @@ import { IntegrationProviderConfigSchema, IntegrationSetupStateSchema, } from "./integration-settings.js"; +import { + MediaJobSnapshotSchema, + MediaJobStartInputSchema, + MediaJobStartOutputSchema, +} from "./media-jobs.js"; +import { MediaSettingsUpdateSchema, MediaSettingsViewSchema } from "./media-settings.js"; import { ModelTierMappingSchema, SetModelTierInput } from "./model-tiers.js"; import { CreateProjectGoalInput, @@ -690,6 +696,18 @@ export const appContract = { get: oc.output(IntegrationSetupStateSchema), save: oc.input(IntegrationProviderConfigSchema).output(z.object({ ok: z.literal(true) })), }, + media: { + /** Deployment-wide media backends; deployment owner only. */ + settings: { + get: oc.output(MediaSettingsViewSchema), + update: oc.input(MediaSettingsUpdateSchema).output(MediaSettingsViewSchema), + }, + jobs: { + start: oc.input(MediaJobStartInputSchema).output(MediaJobStartOutputSchema), + get: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema), + cancel: oc.input(z.object({ jobId: Id })).output(MediaJobSnapshotSchema), + }, + }, connections: { catalog: oc .input(z.object({ query: z.string().optional(), connectorId: z.string().optional() })) diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index edc60ae..497a72d 100644 --- a/packages/testkit/src/cli/harness.ts +++ b/packages/testkit/src/cli/harness.ts @@ -106,6 +106,7 @@ async function main() { "packages/testkit/src/connections.test.ts", "packages/testkit/src/bot-secrets.test.ts", "packages/testkit/src/delegation.postgres.test.ts", + "packages/testkit/src/media.postgres.test.ts", "packages/testkit/src/user-profile.postgres.test.ts", "packages/db/src/space-membership.postgres.test.ts", "packages/db/src/main-agent.postgres.test.ts", diff --git a/packages/testkit/src/media.postgres.test.ts b/packages/testkit/src/media.postgres.test.ts new file mode 100644 index 0000000..f5de0c6 --- /dev/null +++ b/packages/testkit/src/media.postgres.test.ts @@ -0,0 +1,378 @@ +import { mkdtempSync, rmSync } from "node:fs"; +import { writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import { ComposioEmulator, type MediaOperationRegistration } from "@rakazo/adapters"; +import { + type MediaJobBlock, + MediaOutputSchema, + type MediaSettingsView, + type MessageBlock, +} from "@rakazo/contracts"; +import { createThreadMessage } from "@rakazo/db"; +import { afterAll, beforeAll, describe, expect, it } from "vitest"; +import type { createApp } from "../../../apps/api/src/app.ts"; +import { sessionCookieHeader } from "./index.js"; + +type App = { request: (input: string, init?: RequestInit) => Response | Promise }; +type AppHandles = Awaited>; + +process.env.WAKEUP_DRIVER = "memory"; +process.env.SANDBOX_PROVIDER = "fake"; +process.env.AGENT_RUNTIME = "scripted"; + +const hasDb = process.env.VERIFY_DATABASE === "1" && Boolean(process.env.DATABASE_URL); +const describeWithDatabase = hasDb ? describe.sequential : describe.skip; + +const PNG_BASE64 = Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 1, 2, 3]).toString( + "base64", +); + +// testkit has no zod of its own; the contracts' schemas carry it. +const SourceInputSchema = MediaOutputSchema.pick({ artifactId: true }); + +/** Fake card operations: an export that copies its source, and a transcription that waits. */ +function fakeOperations(): MediaOperationRegistration[] { + const exportVideo: MediaOperationRegistration<{ artifactId: string }> = { + operation: "export_video", + inputSchema: SourceInputSchema, + title: () => "Export", + available: async () => true, + artifactIds: (input) => [input.artifactId], + handler: { + operation: "export_video", + run: async (context) => { + const { artifactId } = context.input as { artifactId: string }; + const source = await context.fetchArtifact(artifactId); + const out = path.join(context.tempDir, "export.png"); + await writeFile(out, Buffer.from(PNG_BASE64, "base64")); + await context.addOutput({ + path: out, + name: "export.png", + mimeType: "image/png", + parentArtifactId: source.artifactId, + preset: "square", + }); + return { summary: "Exported." }; + }, + }, + }; + const transcribe: MediaOperationRegistration<{ artifactId: string }> = { + operation: "transcribe", + inputSchema: SourceInputSchema, + title: () => "Transcript", + available: async () => true, + artifactIds: (input) => [input.artifactId], + handler: { + operation: "transcribe", + run: async (context) => { + await new Promise((resolve) => { + context.signal.addEventListener("abort", () => resolve(), { once: true }); + }); + context.signal.throwIfAborted(); + return { summary: "" }; + }, + }, + }; + return [exportVideo, transcribe] as MediaOperationRegistration[]; +} + +describeWithDatabase("media RPCs (PostgreSQL)", () => { + let handles: AppHandles; + let app: App; + const stamp = `${Date.now()}-${Math.random().toString(36).slice(2)}`; + const dataDir = mkdtempSync(path.join(tmpdir(), "rakazo-media-")); + + beforeAll(async () => { + const { createApp } = await import("../../../apps/api/src/app.ts"); + handles = await createApp({ + databaseUrl: process.env.DATABASE_URL!, + dataDir, + sandboxProvider: "fake", + agentRuntime: "scripted", + wakeupDriver: "memory", + signupsEnabled: "true", + composio: new ComposioEmulator(), + mediaOperations: fakeOperations(), + }); + app = handles.app; + }); + + afterAll(async () => { + await handles?.prisma.integrationProviderConfig.deleteMany({ where: { id: "media" } }); + await handles?.stop(); + rmSync(dataDir, { recursive: true, force: true }); + }); + + async function user(label: string) { + const cookie = await signup(app, `${label}-${stamp}@rakazo.test`, label); + const me = await rpc<{ userId: string; spaceId: string }>(app, cookie, "me"); + const bot = await rpc<{ id: string }>(app, cookie, "bots/create", { + name: `${label} bot`, + title: "", + description: "", + instructions: "", + notifyOnFinish: false, + }); + const thread = await handles.prisma.thread.findUniqueOrThrow({ where: { botId: bot.id } }); + const card = await createThreadMessage(handles.prisma, { + threadId: thread.id, + role: "bot", + botId: bot.id, + blocks: [{ kind: "text", text: "Here is your clip." }], + }); + const artifact = await rpc<{ id: string }>(app, cookie, "artifacts/create", { + botId: bot.id, + name: "clip.png", + mimeType: "image/png", + contentBase64: PNG_BASE64, + }); + return { cookie, me, bot, thread, card, artifact }; + } + + async function waitForStatus(jobId: string, status: string) { + const start = Date.now(); + while (Date.now() - start < 20_000) { + const job = await handles.prisma.mediaJob.findUnique({ where: { id: jobId } }); + if (job?.status === status) return job; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + throw new Error(`media job ${jobId} never reached ${status}`); + } + + it("lets only the deployment owner read and change media settings, and never returns keys", async () => { + const owner = await user("media-owner"); + const other = await user("media-other"); + await handles.prisma.deploymentSettings.upsert({ + where: { id: "default" }, + create: { id: "default", ownerUserId: owner.me.userId }, + update: { ownerUserId: owner.me.userId }, + }); + + const update = { + comfyui: { + baseUrl: "http://127.0.0.1:8188/", + apiKey: "comfy-test-key", + imageModel: "sd_xl_base_1.0.safetensors", + imageWorkflow: '{"3":{"class_type":"KSampler","inputs":{}}}', + videoWorkflow: "", + }, + transcription: { baseUrl: "http://127.0.0.1:8080", model: "whisper-1", apiKey: "" }, + }; + expect((await raw(app, other.cookie, "media/settings/get", {})).status).toBe(403); + expect((await raw(app, other.cookie, "media/settings/update", update)).status).toBe(403); + expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe( + 0, + ); + + const saved = await rpc(app, owner.cookie, "media/settings/update", update); + expect(saved).toEqual({ + comfyui: { + baseUrl: "http://127.0.0.1:8188", + hasApiKey: true, + imageModel: "sd_xl_base_1.0.safetensors", + imageWorkflow: update.comfyui.imageWorkflow, + }, + transcription: { baseUrl: "http://127.0.0.1:8080/v1", model: "whisper-1", hasApiKey: false }, + }); + expect(JSON.stringify(saved)).not.toContain("comfy-test-key"); + expect(await rpc(app, owner.cookie, "media/settings/get")).toEqual(saved); + const row = await handles.prisma.integrationProviderConfig.findUniqueOrThrow({ + where: { id: "media" }, + }); + expect(row.ciphertext).not.toContain("comfy-test-key"); + expect(row.ciphertext).not.toContain("127.0.0.1"); + + // An omitted key stays with the same server, never follows the settings elsewhere. + const { apiKey: _key, ...withoutKey } = update.comfyui; + const kept = await rpc(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: withoutKey, + }); + expect(kept.comfyui?.hasApiKey).toBe(true); + const moved = await rpc(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...withoutKey, baseUrl: "http://127.0.0.1:8189" }, + }); + expect(moved.comfyui?.hasApiKey).toBe(false); + + for (const baseUrl of ["http://169.254.169.254", "http://user:pass@127.0.0.1:8188"]) { + const refused = await raw(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...update.comfyui, baseUrl }, + }); + expect(refused.status, baseUrl).toBe(400); + } + const badWorkflow = await raw(app, owner.cookie, "media/settings/update", { + ...update, + comfyui: { ...update.comfyui, imageWorkflow: "[1, 2" }, + }); + expect(badWorkflow.status).toBe(400); + + const previous = process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + try { + delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + expect( + ( + await raw(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "http://transcribe.example.com" }, + }) + ).status, + ).toBe(400); + process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = "1"; + // A key is never sent in cleartext to a public host. + expect( + ( + await raw(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "http://transcribe.example.com", apiKey: "secret" }, + }) + ).status, + ).toBe(400); + const https = await rpc(app, owner.cookie, "media/settings/update", { + comfyui: null, + transcription: { baseUrl: "https://transcribe.example.com", apiKey: "secret" }, + }); + expect(https.transcription).toMatchObject({ + baseUrl: "https://transcribe.example.com/v1", + hasApiKey: true, + }); + } finally { + if (previous === undefined) delete process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC; + else process.env.RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC = previous; + } + + expect( + await rpc(app, owner.cookie, "media/settings/update", { comfyui: null, transcription: null }), + ).toEqual({ comfyui: null, transcription: null }); + expect(await handles.prisma.integrationProviderConfig.count({ where: { id: "media" } })).toBe( + 0, + ); + }); + + it("starts, reads and cancels card jobs only for the actor's own threads and files", async () => { + const alice = await user("media-alice"); + const mallory = await user("media-mallory"); + + const started = await rpc<{ jobId: string; messageId: string }>( + app, + alice.cookie, + "media/jobs/start", + { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: alice.artifact.id }, + }, + ); + const job = await waitForStatus(started.jobId, "done"); + expect(job).toMatchObject({ + spaceId: alice.me.spaceId, + userId: alice.me.userId, + botId: alice.bot.id, + threadId: alice.thread.id, + runId: null, + continueAfter: false, + }); + const snapshot = await rpc<{ status: string; outputs: MediaJobBlock["outputs"] }>( + app, + alice.cookie, + "media/jobs/get", + { jobId: started.jobId }, + ); + expect(snapshot.status).toBe("done"); + expect(snapshot.outputs).toEqual([ + expect.objectContaining({ + name: "export.png", + parentArtifactId: alice.artifact.id, + preset: "square", + }), + ]); + const card = await handles.prisma.message.findUniqueOrThrow({ + where: { id: started.messageId }, + }); + expect(card).toMatchObject({ threadId: alice.thread.id, botId: alice.bot.id, role: "bot" }); + expect((card.blocks as MessageBlock[])[0]).toMatchObject({ kind: "media", status: "done" }); + + // Another user can neither read nor cancel the job, nor start one in Alice's conversation + // or on her files. + expect( + (await raw(app, mallory.cookie, "media/jobs/get", { jobId: started.jobId })).status, + ).toBe(404); + expect( + (await raw(app, mallory.cookie, "media/jobs/cancel", { jobId: started.jobId })).status, + ).toBe(404); + for (const body of [ + { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: mallory.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: alice.card.id, + operation: "export_video", + input: { artifactId: mallory.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: mallory.card.id, + operation: "export_video", + input: { artifactId: alice.artifact.id }, + }, + { + botId: mallory.bot.id, + messageId: mallory.card.id, + operation: "generate_image", + input: { prompt: "anything" }, + }, + ]) { + const response = await raw(app, mallory.cookie, "media/jobs/start", body); + expect(response.status, JSON.stringify(body)).toBeGreaterThanOrEqual(400); + } + expect(await handles.prisma.mediaJob.count({ where: { userId: mallory.me.userId } })).toBe(0); + + const slow = await rpc<{ jobId: string }>(app, alice.cookie, "media/jobs/start", { + botId: alice.bot.id, + messageId: alice.card.id, + operation: "transcribe", + input: { artifactId: alice.artifact.id }, + }); + await waitForStatus(slow.jobId, "running"); + const cancelled = await rpc<{ status: string }>(app, alice.cookie, "media/jobs/cancel", { + jobId: slow.jobId, + }); + expect(["running", "cancelled"]).toContain(cancelled.status); + await waitForStatus(slow.jobId, "cancelled"); + }); +}); + +async function signup(app: App, email: string, name: string) { + const res = await app.request("/api/auth/sign-up/email", { + method: "POST", + headers: { "content-type": "application/json", origin: "http://127.0.0.1:5173" }, + body: JSON.stringify({ email, password: "password12", name }), + }); + if (res.status >= 400) throw new Error(`signup failed ${res.status}: ${await res.text()}`); + return sessionCookieHeader(res); +} + +function raw(app: App, cookie: string, proc: string, body: unknown = {}) { + return app.request(`/rpc/${proc}`, { + method: "POST", + headers: { "content-type": "application/json", cookie, origin: "http://127.0.0.1:5173" }, + body: JSON.stringify({ json: body }), + }); +} + +async function rpc(app: App, cookie: string, proc: string, body: unknown = {}): Promise { + const res = await raw(app, cookie, proc, body); + const text = await res.text(); + const parsed = JSON.parse(text) as { json?: T; error?: { message?: string } }; + if (res.status >= 400 || parsed.error) { + throw new Error(`${proc} ${res.status}: ${parsed.error?.message ?? text}`); + } + return parsed.json as T; +}