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-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/adapters/src/media-settings.ts b/packages/adapters/src/media-settings.ts index b82134d..c241c47 100644 --- a/packages/adapters/src/media-settings.ts +++ b/packages/adapters/src/media-settings.ts @@ -1,10 +1,5 @@ -import { - type MediaSettings, - MediaSettingsSchema, - type MediaSettingsUpdate, - MediaSettingsUpdateSchema, - type MediaSettingsView, -} from "@rakazo/contracts"; +import type { MediaSettings, MediaSettingsUpdate, MediaSettingsView } from "@rakazo/contracts"; +import { MediaSettingsSchema, MediaSettingsUpdateSchema } from "@rakazo/contracts"; import { parseWorkflowTemplate } from "@rakazo/core/media/workflows"; import type { PrismaClient } from "@rakazo/db"; import { @@ -83,6 +78,14 @@ 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 assertValidWorkflow(label: string, workflow: string | undefined): void { if (!workflow) return; try { @@ -92,33 +95,42 @@ function assertValidWorkflow(label: string, workflow: string | undefined): void } } +function resolveSettings(previous: MediaSettings, update: MediaSettingsUpdate): MediaSettings { + const parsed = MediaSettingsUpdateSchema.parse(update); + try { + assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow); + assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow); + 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); - assertValidWorkflow("Image workflow", parsed.comfyui?.imageWorkflow); - assertValidWorkflow("Video workflow", parsed.comfyui?.videoWorkflow); - const previous = await loadMediaSettings(deps); - let comfyui: MediaSettings["comfyui"] = null; - if (parsed.comfyui) { - 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 365b97e..e1982e1 100644 --- a/packages/contracts/src/index.ts +++ b/packages/contracts/src/index.ts @@ -11,6 +11,7 @@ export * from "./local-settings.js"; export * from "./mcp.js"; export * from "./media.js"; export * from "./media-generation.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/core/src/attachments.ts b/packages/core/src/attachments.ts index b4a6759..cace390 100644 --- a/packages/core/src/attachments.ts +++ b/packages/core/src/attachments.ts @@ -1,12 +1,10 @@ +import type { AttachmentKind, AttachmentMimeType, MessageBlock } from "@rakazo/contracts"; import { ATTACHMENT_MAX_BASE64_LENGTH, ATTACHMENT_MAX_BYTES, - type AttachmentKind, - type AttachmentMimeType, attachmentKindOf, isAllowedAttachmentMimeType, isAttachmentImageMimeType, - type MessageBlock, } from "@rakazo/contracts"; export class AttachmentValidationError extends Error { diff --git a/packages/testkit/src/cli/harness.ts b/packages/testkit/src/cli/harness.ts index 5d565e9..fc38326 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", @@ -117,6 +118,7 @@ async function main() { "packages/adapters/src/job-reconciler.postgres.test.ts", "packages/adapters/src/cloud-agent.postgres.test.ts", "packages/adapters/src/media/image-operations.postgres.test.ts", + "packages/adapters/src/media-runtime.postgres.test.ts", "packages/adapters/src/builtin-memory-provider.postgres.test.ts", "packages/adapters/src/profile-consolidation.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; +}