From e0277985f987170e0401e8324db0708a811077f6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:37:19 +0000 Subject: [PATCH 1/2] 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 27969652569bed2cdb24efa74bf71bcd6ee284e6 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 3 Oct 2026 13:47:44 +0000 Subject: [PATCH 2/2] 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; +}