import { randomUUID } from "node:crypto"; import { existsSync } from "node:fs"; import { mkdir, 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 type { PrismaClient } from "@rakazo/db"; import { clearThread, createDb } 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 type { MediaOperationRegistration, MediaRuntimeOptions } from "./media-runtime.js"; import { createMediaRuntime, MEDIA_WAKE_TRIGGER, MediaJobRejectedError } 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("runs on a card smaller than the floor once it is nearly free, and fails on far too small", async () => { let free = 5_000; const ctx = await setup({ gpu: { minVramMb: 12_000, status: async () => ({ vramFreeMb: free, vramTotalMb: 8_192 }), }, }); const { jobId } = await ctx.start({ prompt: "fox" }); await ctx.runtime.runJob(jobId); expect(await ctx.job(jobId)).toMatchObject({ status: "waiting", deferrals: 1 }); free = 7_600; await prisma.mediaJob.update({ where: { id: jobId }, data: { nextAttemptAt: new Date() } }); await ctx.runtime.runJob(jobId); expect((await ctx.job(jobId)).status).toBe("done"); const tiny = await setup({ gpu: { minVramMb: 12_000, status: async () => ({ vramFreeMb: 4_000, vramTotalMb: 4_096 }), }, }); const small = await tiny.start({ prompt: "fox" }); await tiny.runtime.runJob(small.jobId); expect(tiny.handlerCalls).toHaveLength(0); expect(await tiny.job(small.jobId)).toMatchObject({ status: "failed", error: "The GPU does not have enough memory for this job.", deferrals: 0, }); }); 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, }, }); // The dead runner's scratch files are still on disk. await mkdir(path.join(ctx.tempRoot, doomed.jobId), { recursive: true }); await writeFile(path.join(ctx.tempRoot, doomed.jobId, "partial.mp4"), PNG); await ctx.runtime.runJob(doomed.jobId); expect(await ctx.job(doomed.jobId)).toMatchObject({ status: "failed", error: "The job was interrupted. Try again.", }); expect(existsSync(path.join(ctx.tempRoot, doomed.jobId))).toBe(false); expect(ctx.handlerCalls).toEqual([jobId]); }); it("sweeps temp directories of jobs that no longer run", async () => { const ctx = await setup(); const active = await ctx.start({ prompt: "active" }); const ended = await ctx.start({ prompt: "ended" }); await ctx.runtime.cancel(ended.jobId); for (const name of [active.jobId, ended.jobId, "unknown-job"]) { await mkdir(path.join(ctx.tempRoot, name), { recursive: true }); } await ctx.runtime.reconcile(); expect((await readdir(ctx.tempRoot)).sort()).toEqual([active.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); }); });