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 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
This commit is contained in:
Claude
parent
24c9245af1
commit
e0277985f9
615
packages/adapters/src/media-runtime.postgres.test.ts
Normal file
615
packages/adapters/src/media-runtime.postgres.test.ts
Normal file
@ -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<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("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,
|
||||
},
|
||||
});
|
||||
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);
|
||||
});
|
||||
});
|
||||
@ -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<Uint8Array>),
|
||||
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,
|
||||
|
||||
@ -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",
|
||||
];
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user