- A handler that ignores its abort is left behind after the grace period; a late rejection from it (or from the grace timer after the job settled) was unobserved, and the worker exits on unhandled rejections. Both sides of the race are now observed and the grace timer is cleared with the job. - Ending a job removes its temp directory when no runner in this process holds it (leftovers of a runner that died here), and reconciliation sweeps directories of jobs that no longer run. - An executor test runs a media tool end to end through applyTool: shown only for operations the space can run, started with the run's scope and effect key, continue_after split off, refused when not exposed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
660 lines
24 KiB
TypeScript
660 lines
24 KiB
TypeScript
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<MediaOperationResult>;
|
|
|
|
function deferred<T = void>() {
|
|
let resolve!: (value: T) => void;
|
|
const promise = new Promise<T>((done) => {
|
|
resolve = done;
|
|
});
|
|
return { promise, resolve };
|
|
}
|
|
|
|
describePostgres("media job runtime (PostgreSQL, fake operations)", () => {
|
|
let prisma: PrismaClient;
|
|
let db: ReturnType<typeof createDb>;
|
|
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<boolean>;
|
|
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<FakeInput> = {
|
|
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<string, unknown>, extra: Record<string, unknown> = {}) =>
|
|
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<MediaJobBlock> => {
|
|
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<void>((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);
|
|
});
|
|
});
|