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, GpuStatus, JobPublisher, MediaCapabilities, MediaJobContext, MediaJobInputFile, MediaJobOutputFile, MediaJobSnapshot, MediaOperationHandler, MediaRuntime, StartMediaJobInput, } from "@rakazo/adapter-kit"; import { ArtifactTooLargeError, MediaJobDeferredError, mediaRunJob, mediaRunJobKey, runContinueJob, } from "@rakazo/adapter-kit"; import type { MediaJobBlock, MediaJobStatus, MediaOperation, MediaOutput, MessageBlock, } from "@rakazo/contracts"; import { attachmentKindOf, MEDIA_JOB_ACTIVE_STATUSES, MediaOperationSchema, } from "@rakazo/contracts"; import { attachmentExtensionForMimeType, inferAttachmentMimeType } from "@rakazo/core"; import type { MediaJob, PrismaClient, ThreadEvents } from "@rakazo/db"; import { appendEventInTransaction, createThreadMessageInTransaction, Prisma, withTransactionRetry, } 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"; /** * One media operation (`generate_image`, `edit_video`, …): how to validate and label its input, * when it is available, and the handler that does the work inside a media job. */ export interface MediaOperationRegistration { operation: MediaOperation; /** Validates tool and card input. The job stores the parsed value, so it must parse again. */ inputSchema: z.ZodType; /** Card label for a job, e.g. the prompt or what the edit does. */ title(input: Input): string; /** Whether the space can run the operation now (backend configured, binaries present). */ available(context: { spaceId: string }): Promise; handler: MediaOperationHandler; /** GPU guard: the runner waits while the backend reports less free VRAM than this. */ gpu?: { minVramMb: number; status(signal?: AbortSignal): Promise }; /** Artifact ids the input reads, so starting a job refuses ones the user cannot read. */ artifactIds?(input: Input): string[]; } /** A media job that cannot start. The message is safe to show to the user and the model. */ export class MediaJobRejectedError extends Error { constructor(message: string) { super(message); this.name = "MediaJobRejectedError"; } } export interface MediaRuntimeDeps { prisma: PrismaClient; events: Pick; jobs: JobPublisher; artifacts: ArtifactStore; registrations: readonly MediaOperationRegistration[]; } export interface MediaRuntimeOptions { /** Runner lease; a crashed runner's job is recovered after it expires. */ leaseMs?: number; /** Lease renewal and cancel polling interval while a job runs. */ heartbeatMs?: number; /** Minimum time between two progress writes to the card. */ progressIntervalMs?: number; /** Deferrals (GPU busy, backend not ready) before a job fails. */ maxDeferrals?: number; /** Longest wait between two attempts of a deferred job. */ maxDeferMs?: number; /** Restarts after a runner died mid-job before the job fails. */ maxRecoveries?: number; /** Hard cap on one attempt. */ maxRunMs?: number; /** How long a handler may ignore an abort before the runner settles without it. */ abortGraceMs?: number; /** Parent of the per-job temp directories. */ tempRoot?: string; } /** The runtime plus what the background job handler and the reconciler call. */ export interface MediaRuntimeHost extends MediaRuntime { /** Runs one job (`media.run`); safe to call for any id, any number of times. */ runJob(jobId: string): Promise; /** Re-enqueues jobs whose enqueue was lost, whose runner died, or whose wait is over. */ reconcile(): Promise; } const ACTIVE_STATUSES: MediaJobStatus[] = [...MEDIA_JOB_ACTIVE_STATUSES]; const TITLE_MAX_LENGTH = 200; const ERROR_MAX_LENGTH = 500; const SUMMARY_MAX_LENGTH = 4_000; const GPU_STATUS_TIMEOUT_MS = 10_000; const GPU_RETRY_MS = 15_000; /** Retry delay while another media job holds the GPU. */ const BUSY_RETRY_MS = 5_000; /** A due job not picked up by its queue in this time is re-enqueued outside the queue. */ const STUCK_QUEUE_MS = 120_000; /** Due jobs younger than this are left to the enqueue that created them. */ const RECONCILE_GRACE_MS = 30_000; const RECONCILE_BATCH_SIZE = 100; // Shares the reconciliation lock namespace (job-reconciler.ts) with its own id. const MEDIA_CLAIM_LOCK_NAMESPACE = 1_380_019_075; const MEDIA_CLAIM_LOCK_ID = 2; /** Card errors. They are shown to the user, so they stay short and actionable. */ const UNAVAILABLE_ERROR = "This media operation is not available right now."; const INTERRUPTED_ERROR = "The job was interrupted. Try again."; const BUSY_ERROR = "The media server stayed busy. Try again later."; const TIMEOUT_ERROR = "The job took too long."; const FAILED_ERROR = "The job failed."; const MISSING_SOURCE_ERROR = "A source file is no longer available."; const TOO_LARGE_ERROR = "The result is too large to keep."; const GPU_TOO_SMALL_ERROR = "The GPU does not have enough memory for this job."; type AbortReason = "cancel" | "lease" | "abandoned" | "timeout"; function clamp(text: string, max: number): string { return text.length > max ? `${text.slice(0, max - 1)}…` : text; } function issuesText(error: z.ZodError): string { return error.issues .slice(0, 5) .map((issue) => issue.path.length ? `${issue.path.join(".")}: ${issue.message}` : issue.message, ) .join("; "); } function cardBlock(job: MediaJob, outputs: MediaOutput[]): MediaJobBlock { return { kind: "media", jobId: job.id, operation: job.operation as MediaOperation, status: job.status as MediaJobStatus, title: job.title, ...(job.progress === null ? {} : { progress: job.progress }), ...(job.error === null ? {} : { error: job.error }), outputs, }; } function mediaBlockOf(blocks: unknown, jobId: string): MediaJobBlock | undefined { if (!Array.isArray(blocks)) return undefined; return (blocks as MessageBlock[]).find( (block): block is MediaJobBlock => block.kind === "media" && block.jobId === jobId, ); } function positiveInt(value: number | undefined): number | undefined { if (value === undefined || !Number.isFinite(value)) return undefined; const rounded = Math.round(value); return rounded > 0 ? rounded : undefined; } function nonNegativeInt(value: number | undefined): number | undefined { if (value === undefined || !Number.isFinite(value)) return undefined; const rounded = Math.round(value); return rounded >= 0 ? rounded : undefined; } /** * Free VRAM a job waits for: the operation's floor, capped below the card's total so a smaller * card that can still run it (with offloading) does not wait forever. Null when the card is far * too small to try. */ export function requiredVramMb(minVramMb: number, gpu: GpuStatus): number | null { if (gpu.vramTotalMb > 0 && gpu.vramTotalMb * 2 < minVramMb) return null; if (gpu.vramTotalMb <= 0) return minVramMb; return Math.min(minVramMb, Math.floor(gpu.vramTotalMb * 0.9)); } function insideDir(dir: string, file: string): string { const resolved = path.resolve(dir, file); if (!resolved.startsWith(`${path.resolve(dir)}${path.sep}`)) { throw new Error("Media output must be a file in the job's temp directory"); } return resolved; } /** The sentence a bot is woken with when its job ends; job text is data, not instructions. */ export function mediaWakePrompt(job: { id: string; operation: string; title: string; status: string; error: string | null; summary: string; outputs: MediaOutput[]; }): string { const label = `Media job ${job.id} (${job.operation}: ${JSON.stringify(job.title)})`; if (job.status === "cancelled") { return `${label} was cancelled. Do not start it again unless the user asks.`; } if (job.status === "failed") { return `${label} failed: ${job.error ?? FAILED_ERROR} Continue your task if you can, otherwise tell the user.`; } const outputs = job.outputs.map((output) => `${output.name} (${output.artifactId})`).join(", "); return [ `${label} is done.`, job.summary ? `Result: ${clamp(job.summary, SUMMARY_MAX_LENGTH)}` : "", outputs ? `Outputs: ${outputs}.` : "", "The media card already shows the result to the user; continue your task.", ] .filter(Boolean) .join(" "); } export function createMediaRuntime( deps: MediaRuntimeDeps, options: MediaRuntimeOptions = {}, ): MediaRuntimeHost { const leaseMs = options.leaseMs ?? 90_000; const heartbeatMs = options.heartbeatMs ?? 5_000; const progressIntervalMs = options.progressIntervalMs ?? 2_000; const maxDeferrals = options.maxDeferrals ?? 30; const maxDeferMs = options.maxDeferMs ?? 60_000; const maxRecoveries = options.maxRecoveries ?? 2; const maxRunMs = options.maxRunMs ?? 2 * 60 * 60_000; const abortGraceMs = options.abortGraceMs ?? 15_000; const tempRoot = options.tempRoot ?? path.join(tmpdir(), "evoo-media"); const registry = new Map(); for (const registration of deps.registrations) { if (registry.has(registration.operation)) { throw new Error(`Media operation ${registration.operation} is registered twice`); } registry.set(registration.operation, registration); } /** Jobs this process runs, so a cancel here stops them without waiting for a heartbeat. */ const running = new Map(); const isAvailable = async ( registration: MediaOperationRegistration | undefined, spaceId: string, ): Promise => { if (!registration) return false; try { return await registration.available({ spaceId }); } catch (error) { getLogger().warn("media operation availability check failed", { operation: registration.operation, error: error instanceof Error ? error.message : String(error), }); return false; } }; const notify = async (threadId: string, seq: number | undefined) => { if (seq === undefined) return; await deps.events.notify(threadId, seq).catch((error) => { getLogger().error("media card realtime notification", error); }); }; const enqueue = async (jobId: string, availableAt?: Date) => { await deps.jobs.enqueue(mediaRunJob(jobId, availableAt)).catch((error) => { // The job row is durable; the reconciler enqueues it again. getLogger().error("media job enqueue", error); }); }; /** Rewrites the job's media block; returns the event seq when the card changed. */ const writeCard = async ( tx: Prisma.TransactionClient, job: MediaJob, appended: MediaOutput[] = [], ): Promise<{ exists: boolean; seq?: number; outputs: MediaOutput[] }> => { if (!job.messageId) return { exists: false, outputs: appended }; const message = await tx.message.findFirst({ where: { id: job.messageId, threadId: job.threadId }, select: { id: true, role: true, botId: true, blocks: true }, }); if (!message) return { exists: false, outputs: appended }; const previous = mediaBlockOf(message.blocks, job.id); const outputs = [...(previous?.outputs ?? []), ...appended]; const next = cardBlock(job, outputs); const current = Array.isArray(message.blocks) ? (message.blocks as MessageBlock[]) : []; const blocks = previous ? current.map((block) => (block === previous ? next : block)) : [...current, next]; if (JSON.stringify(blocks) === JSON.stringify(message.blocks)) { return { exists: true, outputs }; } await tx.message.update({ where: { id: message.id }, data: { blocks: blocks as Prisma.InputJsonValue }, }); const event = await appendEventInTransaction(tx, { spaceId: job.spaceId, threadId: job.threadId, botId: message.botId ?? job.botId ?? "", type: "thread.message.updated", payload: { messageId: message.id, role: message.role, blocks }, }); return { exists: true, seq: event.seq, outputs }; }; /** * Applies `data` to the job when the fence still matches and mirrors the job on its card in * the same transaction. A null token fences on "not held by a live runner". */ const updateJob = async ( jobId: string, token: string | null, data: Prisma.MediaJobUpdateManyMutationInput, appended: MediaOutput[] = [], ) => { const committed = await withTransactionRetry(() => deps.prisma.$transaction(async (tx) => { const now = new Date(); const fence: Prisma.MediaJobWhereInput = token ? { id: jobId, leaseToken: token } : { id: jobId, status: { in: ACTIVE_STATUSES }, OR: [{ leaseToken: null }, { leaseExpiresAt: { lte: now } }], }; const before = await tx.mediaJob.findFirst({ where: fence, select: { threadId: true } }); if (!before) return null; // Thread row first, like clearing a chat, so a clear and a card write cannot deadlock. await tx.thread.updateMany({ where: { id: before.threadId }, data: { nextEventSeq: { increment: 0 } }, }); const updated = await tx.mediaJob.updateMany({ where: fence, data }); if (!updated.count) return null; const job = await tx.mediaJob.findUniqueOrThrow({ where: { id: jobId } }); const card = await writeCard(tx, job, appended); return { job, card }; }), ); if (committed) await notify(committed.job.threadId, committed.card.seq); return committed; }; /** Ends a job: terminal status on the row and the card, and the bot's wake when requested. */ const settle = async ( jobId: string, token: string | null, outcome: { status: "done" | "failed" | "cancelled"; error?: string; summary?: string }, ): Promise => { const committed = await withTransactionRetry(() => deps.prisma.$transaction(async (tx) => { const now = new Date(); const fence: Prisma.MediaJobWhereInput = token ? { id: jobId, leaseToken: token } : { id: jobId, status: { in: ACTIVE_STATUSES }, OR: [{ leaseToken: null }, { leaseExpiresAt: { lte: now } }], }; const before = await tx.mediaJob.findFirst({ where: fence, select: { threadId: true } }); if (!before) return null; await tx.thread.updateMany({ where: { id: before.threadId }, data: { nextEventSeq: { increment: 0 } }, }); const updated = await tx.mediaJob.updateMany({ where: fence, data: { status: outcome.status, error: outcome.status === "failed" ? clamp(outcome.error ?? FAILED_ERROR, ERROR_MAX_LENGTH) : null, ...(outcome.status === "done" ? { progress: 1 } : {}), finishedAt: now, leaseToken: null, leaseExpiresAt: null, nextAttemptAt: null, }, }); if (!updated.count) return null; const job = await tx.mediaJob.findUniqueOrThrow({ where: { id: jobId } }); const card = await writeCard(tx, job); let wakeRunId: string | undefined; if (job.continueAfter && job.botId && card.exists) { const [bot, member] = await Promise.all([ tx.bot.findFirst({ where: { id: job.botId, spaceId: job.spaceId, userId: job.userId, archivedAt: null }, select: { id: true }, }), tx.spaceMember.findUnique({ where: { spaceId_userId: { spaceId: job.spaceId, userId: job.userId } }, select: { id: true }, }), ]); if (bot && member) { const prompt = mediaWakePrompt({ ...job, summary: outcome.summary ?? "", outputs: card.outputs, }); const task = await tx.task.create({ data: { spaceId: job.spaceId, botId: job.botId, threadId: job.threadId, userId: job.userId, prompt, status: "queued", }, }); const run = await tx.run.create({ data: { spaceId: job.spaceId, botId: job.botId, threadId: job.threadId, userId: job.userId, taskId: task.id, status: "queued", trigger: MEDIA_WAKE_TRIGGER, clientNonce: `media-wake:${job.id}`, }, select: { id: true }, }); wakeRunId = run.id; } } return { job, seq: card.seq, wakeRunId }; }), ); if (!committed) return false; await notify(committed.job.threadId, committed.seq); if (committed.wakeRunId) { // A queued run is durable; the run reconciler repairs a missed wake. await deps.jobs.enqueue(runContinueJob(committed.wakeRunId)).catch((error) => { getLogger().error("media job wake enqueue", error); }); } return true; }; /** * Takes the job for this process. At most one media job holds a live lease at a time, which * keeps the GPU to one job even when the Graphile queue is bypassed or several in-memory * drivers run. */ const claim = async (jobId: string, token: string): Promise<"claimed" | "busy" | "gone"> => deps.prisma.$transaction(async (tx) => { await tx.$executeRaw`SELECT pg_advisory_xact_lock(${MEDIA_CLAIM_LOCK_NAMESPACE}::integer, ${MEDIA_CLAIM_LOCK_ID}::integer)`; const now = new Date(); const claimable: Prisma.MediaJobWhereInput = { id: jobId, status: { in: ACTIVE_STATUSES }, AND: [ { OR: [{ leaseToken: null }, { leaseExpiresAt: { lte: now } }] }, { OR: [{ nextAttemptAt: null }, { nextAttemptAt: { lte: now } }] }, ], }; const candidate = await tx.mediaJob.findFirst({ where: claimable, select: { id: true } }); if (!candidate) return "gone"; const holder = await tx.mediaJob.findFirst({ where: { id: { not: jobId }, status: { in: ACTIVE_STATUSES }, leaseToken: { not: null }, leaseExpiresAt: { gt: now }, }, select: { id: true }, }); if (holder) return "busy"; const updated = await tx.mediaJob.updateMany({ where: claimable, data: { status: "running", leaseToken: token, leaseExpiresAt: new Date(now.getTime() + leaseMs), attempts: { increment: 1 }, nextAttemptAt: null, }, }); return updated.count ? "claimed" : "gone"; }); const waitBehindHolder = async (job: MediaJob) => { const nextAttemptAt = new Date(Date.now() + BUSY_RETRY_MS); // A stale `running` job keeps its status so the next claim still counts the crash. const committed = await updateJob(job.id, null, { ...(job.status === "running" ? {} : { status: "waiting" }), nextAttemptAt, }); if (committed) await enqueue(job.id, nextAttemptAt); }; const storeArtifactFile = async ( job: MediaJob, signal: AbortSignal, file: { path: string; name: string; mimeType: string; width?: number; height?: number; durationMs?: number; parentArtifactId?: string; }, ) => { try { return await storeArtifactFromFile(deps, { ...file, spaceId: job.spaceId, userId: job.userId, botId: job.botId, groupId: job.groupId, runId: job.runId, signal, }); } catch (error) { 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 }, select: { id: true, storageKey: true }, }); await deps.prisma.artifact.deleteMany({ where: { id: { in: rows.map((row) => row.id) } } }); const context = artifactContext(job, `media-output-cleanup:${job.id}`); for (const row of rows) { await deps.artifacts.remove(row.storageKey, context).catch(() => undefined); } }; const runClaimed = async (jobId: string, token: string, previous: MediaJob) => { const job = await deps.prisma.mediaJob.findUniqueOrThrow({ where: { id: jobId } }); const recoveries = previous.attempts - previous.deferrals; if (previous.status === "running" && recoveries > maxRecoveries) { await settle(jobId, token, { status: "failed", error: INTERRUPTED_ERROR }); return; } if (job.cancelRequested) { await settle(jobId, token, { status: "cancelled" }); return; } const registration = registry.get(job.operation as MediaOperation); if (!registration || !(await isAvailable(registration, job.spaceId))) { await settle(jobId, token, { status: "failed", error: UNAVAILABLE_ERROR }); return; } const member = await deps.prisma.spaceMember.findUnique({ where: { spaceId_userId: { spaceId: job.spaceId, userId: job.userId } }, select: { id: true }, }); const started = member ? await updateJob(jobId, token, { startedAt: job.startedAt ?? new Date() }) : null; if (!started?.card.exists) { // The conversation was cleared or the user left the space: nobody can see the result. await settle(jobId, token, { status: "cancelled" }); return; } const controller = new AbortController(); const abort = (reason: AbortReason) => { if (!controller.signal.aborted) controller.abort(reason); }; running.set(jobId, controller); const deadline = setTimeout(() => abort("timeout"), maxRunMs); deadline.unref?.(); let beating = false; const heartbeat = setInterval(() => { if (beating) return; beating = true; void deps.prisma.mediaJob .updateMany({ where: { id: jobId, leaseToken: token }, data: { leaseExpiresAt: new Date(Date.now() + leaseMs) }, }) .then(async (renewed) => { if (!renewed.count) return abort("lease"); const current = await deps.prisma.mediaJob.findUnique({ where: { id: jobId }, select: { cancelRequested: true }, }); if (current?.cancelRequested) abort("cancel"); }) .catch(() => { // A database blip keeps the job running until the lease is really gone. }) .finally(() => { beating = false; }); }, heartbeatMs); heartbeat.unref?.(); const tempDir = path.join(tempRoot, jobId); const inputs = new Map(); let lastProgress = { at: 0, value: job.progress ?? 0 }; const signal = controller.signal; const context: MediaJobContext = { jobId, spaceId: job.spaceId, userId: job.userId, botId: job.botId, input: job.input, tempDir, signal, progress: async (fraction) => { signal.throwIfAborted(); const value = Math.min(1, Math.max(0, Number.isFinite(fraction) ? fraction : 0)); const now = Date.now(); if (Math.abs(value - lastProgress.value) < 0.01 && value < 1) return; if (now - lastProgress.at < progressIntervalMs && value < 1) return; lastProgress = { at: now, value }; const written = await updateJob(jobId, token, { progress: value }); if (!written) abort("lease"); else if (!written.card.exists) abort("abandoned"); }, fetchArtifact: async (artifactId) => { signal.throwIfAborted(); const known = inputs.get(artifactId); if (known) return known; const row = await deps.prisma.artifact.findFirst({ where: { id: artifactId, spaceId: job.spaceId, userId: job.userId }, }); if (!row) throw new MediaJobRejectedError(MISSING_SOURCE_ERROR); const extension = 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()}`); await copyArtifact(job, row.storageKey, file); const input: MediaJobInputFile = { artifactId: row.id, path: file, name: row.name, mimeType: row.mimeType, kind: attachmentKindOf(row.mimeType), size: row.size, ...(row.width ? { width: row.width } : {}), ...(row.height ? { height: row.height } : {}), ...(row.durationMs !== null ? { durationMs: row.durationMs } : {}), }; inputs.set(artifactId, input); return input; }, addOutput: async (file: MediaJobOutputFile) => { signal.throwIfAborted(); const kind = attachmentKindOf(file.mimeType); const filePath = insideDir(tempDir, file.path); const parent = file.parentArtifactId ? await deps.prisma.artifact.findFirst({ where: { id: file.parentArtifactId, spaceId: job.spaceId, userId: job.userId }, select: { id: true }, }) : null; const name = path.basename(file.name).slice(0, 200) || `output${attachmentExtensionForMimeType(file.mimeType)}`; const stored: string[] = []; try { let posterArtifactId: string | undefined; if (file.posterPath) { const posterPath = insideDir(tempDir, file.posterPath); 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, signal, { path: posterPath, name: posterName, mimeType: posterMime, width: positiveInt(file.width), height: positiveInt(file.height), }); stored.push(poster.id); posterArtifactId = poster.id; } } const width = positiveInt(file.width); const height = positiveInt(file.height); const durationMs = nonNegativeInt(file.durationMs); const artifact = await storeArtifactFile(job, signal, { path: filePath, name, mimeType: file.mimeType, width, height, durationMs, parentArtifactId: parent?.id, }); stored.push(artifact.id); const output: MediaOutput = { artifactId: artifact.id, kind, mimeType: file.mimeType, name, size: artifact.size, ...(width ? { width } : {}), ...(height ? { height } : {}), ...(durationMs !== undefined ? { durationMs } : {}), ...(posterArtifactId ? { posterArtifactId } : {}), ...(parent ? { parentArtifactId: parent.id } : {}), ...(file.preset ? { preset: file.preset } : {}), }; const written = await updateJob(jobId, token, { outputArtifactIds: { push: stored } }, [ output, ]); if (!written) { abort("lease"); throw new Error("Media job lost its lease"); } if (!written.card.exists) abort("abandoned"); return output; } catch (error) { await removeArtifacts(job, stored).catch(() => undefined); throw error; } }, }; try { await rm(tempDir, { recursive: true, force: true }); await mkdir(tempDir, { recursive: true, mode: 0o700 }); if (registration.gpu) { const gpu = await registration.gpu .status(AbortSignal.any([signal, AbortSignal.timeout(GPU_STATUS_TIMEOUT_MS)])) .catch(() => undefined); const needed = gpu ? requiredVramMb(registration.gpu.minVramMb, gpu) : undefined; if (needed === null) throw new MediaJobRejectedError(GPU_TOO_SMALL_ERROR); if (gpu && needed !== undefined && gpu.vramFreeMb < needed) { throw new MediaJobDeferredError("Waiting for free GPU memory", GPU_RETRY_MS); } } const result = await Promise.race([ registration.handler.run(context), new Promise((_, reject) => { const onAbort = () => { const timer = setTimeout(() => reject(signal.reason), abortGraceMs); timer.unref?.(); }; if (signal.aborted) onAbort(); else signal.addEventListener("abort", onAbort, { once: true }); }), ]); if (signal.aborted) throw signal.reason; await settle(jobId, token, { status: "done", summary: result.summary }); } catch (error) { const reason = signal.aborted ? (signal.reason as AbortReason) : undefined; if (reason === "lease") return; if (reason === "cancel" || reason === "abandoned") { await settle(jobId, token, { status: "cancelled" }); } else if (reason === "timeout") { await settle(jobId, token, { status: "failed", error: TIMEOUT_ERROR }); } else if (error instanceof MediaJobDeferredError) { await defer(job, token, error); } else { if (!(error instanceof MediaJobRejectedError)) { getLogger().warn("media job failed", { jobId, operation: job.operation, error: error instanceof Error ? error.message : String(error), }); } const message = error instanceof Error ? error.message.split(tempDir).join("") : ""; await settle(jobId, token, { status: "failed", error: message.trim() || FAILED_ERROR }); } } finally { clearInterval(heartbeat); clearTimeout(deadline); running.delete(jobId); await rm(tempDir, { recursive: true, force: true }).catch((error) => { getLogger().error("media job temp cleanup", error); }); } }; const defer = async (job: MediaJob, token: string, error: MediaJobDeferredError) => { if (job.deferrals >= maxDeferrals) { await settle(job.id, token, { status: "failed", error: BUSY_ERROR }); return; } const delay = Math.min( maxDeferMs, Math.max(1_000, error.retryAfterMs) * 2 ** Math.min(job.deferrals, 5), ); const nextAttemptAt = new Date(Date.now() + delay); const committed = await updateJob(job.id, token, { status: "waiting", deferrals: { increment: 1 }, leaseToken: null, leaseExpiresAt: null, nextAttemptAt, }); if (committed) await enqueue(job.id, nextAttemptAt); }; return { async capabilities(spaceId) { const entries = await Promise.all( MediaOperationSchema.options.map( async (operation) => [operation, await isAvailable(registry.get(operation), spaceId)] as const, ), ); return Object.fromEntries(entries) as MediaCapabilities; }, async start(input: StartMediaJobInput) { const registration = registry.get(input.operation); if (!registration || !(await isAvailable(registration, input.spaceId))) { throw new MediaJobRejectedError(`${input.operation} is not available.`); } if (!input.botId) throw new MediaJobRejectedError("A media job needs a bot to report to."); const botId = input.botId; const parsed = registration.inputSchema.safeParse(input.input); if (!parsed.success) { throw new MediaJobRejectedError( `Invalid ${input.operation} input: ${issuesText(parsed.error)}`, ); } const artifactIds = [...new Set(registration.artifactIds?.(parsed.data) ?? [])]; if (artifactIds.length) { const readable = await deps.prisma.artifact.count({ where: { id: { in: artifactIds }, spaceId: input.spaceId, userId: input.userId }, }); if (readable !== artifactIds.length) throw new MediaJobRejectedError("Unknown artifact."); } const title = clamp( (input.title?.trim() || registration.title(parsed.data)).replace(/\s+/g, " ").trim(), TITLE_MAX_LENGTH, ); const replay = async () => { if (!input.operationKey) return null; const existing = await deps.prisma.mediaJob.findUnique({ where: { operationKey: input.operationKey }, }); if (!existing) return null; if ( existing.spaceId !== input.spaceId || existing.userId !== input.userId || existing.threadId !== input.threadId || !existing.messageId ) { throw new Error("Media job operation scope changed"); } return { jobId: existing.id, messageId: existing.messageId }; }; const replayed = await replay(); if (replayed) return replayed; let committed: { jobId: string; messageId: string; seq: number }; try { committed = await withTransactionRetry(() => deps.prisma.$transaction(async (tx) => { // Also locks the thread row, serializing with clearing the chat. const thread = await tx.thread.updateMany({ where: { id: input.threadId, spaceId: input.spaceId, userId: input.userId }, data: { nextEventSeq: { increment: 0 } }, }); if (!thread.count) { throw new MediaJobRejectedError("This conversation is no longer available."); } const job = await tx.mediaJob.create({ data: { spaceId: input.spaceId, userId: input.userId, botId, groupId: input.groupId, threadId: input.threadId, runId: input.runId, operation: input.operation, title, input: parsed.data as Prisma.InputJsonValue, continueAfter: input.continueAfter === true, operationKey: input.operationKey, }, }); const blocks: MessageBlock[] = [cardBlock(job, [])]; const message = await createThreadMessageInTransaction(tx, { threadId: input.threadId, role: "bot", blocks, botId, runId: input.runId ?? undefined, }); await tx.mediaJob.update({ where: { id: job.id }, data: { messageId: message.id } }); const event = await appendEventInTransaction(tx, { spaceId: input.spaceId, threadId: input.threadId, botId, runId: input.runId ?? undefined, type: "thread.message.created", payload: { messageId: message.id, role: "bot", blocks }, }); return { jobId: job.id, messageId: message.id, seq: event.seq }; }), ); } catch (error) { // A concurrent replay of the same tool call won the unique operation key. if (error instanceof Prisma.PrismaClientKnownRequestError && error.code === "P2002") { const existing = await replay(); if (existing) return existing; } throw error; } await notify(input.threadId, committed.seq); await enqueue(committed.jobId); return { jobId: committed.jobId, messageId: committed.messageId }; }, async get(jobId): Promise { const job = await deps.prisma.mediaJob.findUnique({ where: { id: jobId } }); if (!job) return null; const message = job.messageId ? await deps.prisma.message.findFirst({ where: { id: job.messageId, threadId: job.threadId }, select: { blocks: true }, }) : null; return { id: job.id, operation: job.operation as MediaOperation, status: job.status as MediaJobStatus, title: job.title, progress: job.progress, error: job.error, outputs: mediaBlockOf(message?.blocks, job.id)?.outputs ?? [], }; }, async cancel(jobId) { const requested = await deps.prisma.mediaJob.updateMany({ where: { id: jobId, status: { in: ACTIVE_STATUSES } }, data: { cancelRequested: true }, }); if (!requested.count) return; running.get(jobId)?.abort("cancel" satisfies AbortReason); // Not held by a live runner: end it now. A running job stops at its next heartbeat. if (await settle(jobId, null, { status: "cancelled" })) { await deps.jobs.cancel(mediaRunJobKey(jobId)).catch(() => undefined); } }, async runJob(jobId) { const previous = await deps.prisma.mediaJob.findUnique({ where: { id: jobId } }); if (!previous || !ACTIVE_STATUSES.includes(previous.status as MediaJobStatus)) return; const token = randomUUID(); const claimed = await claim(jobId, token); if (claimed === "busy") { if (previous.cancelRequested) await settle(jobId, null, { status: "cancelled" }); else await waitBehindHolder(previous); return; } if (claimed === "gone") return; await runClaimed(jobId, token, previous); }, async reconcile() { const now = Date.now(); const graceAgo = new Date(now - RECONCILE_GRACE_MS); const due = await deps.prisma.mediaJob.findMany({ where: { OR: [ { status: "queued", leaseToken: null, updatedAt: { lte: graceAgo } }, { status: "waiting", leaseToken: null, nextAttemptAt: { lte: graceAgo } }, { status: "waiting", leaseToken: null, nextAttemptAt: null, updatedAt: { lte: graceAgo }, }, { status: { in: ACTIVE_STATUSES }, leaseExpiresAt: { lte: new Date(now) } }, { status: { in: ACTIVE_STATUSES }, cancelRequested: true, leaseToken: null }, ], }, orderBy: [{ createdAt: "asc" }, { id: "asc" }], take: RECONCILE_BATCH_SIZE, }); for (const job of due) { const dueAt = job.status === "waiting" ? (job.nextAttemptAt ?? job.updatedAt) : job.leaseExpiresAt && job.leaseExpiresAt.getTime() <= now ? job.leaseExpiresAt : job.createdAt; // The run_at keeps queue order; a job stuck behind a dead worker's queue lock (Graphile // holds it for hours) runs outside the queue and relies on the claim's GPU lock. const scheduled = mediaRunJob(job.id, dueAt); const stuck = now - dueAt.getTime() > STUCK_QUEUE_MS; await deps.jobs .enqueue(stuck ? { ...scheduled, queue: undefined } : scheduled) .catch((error) => getLogger().error("media job reconciliation", error)); } }, }; } function artifactContext(job: MediaJob, operationId: string): AdapterContext { return { operationId, traceId: operationId, spaceId: job.spaceId, userId: job.userId, botId: job.botId ?? undefined, signal: new AbortController().signal, }; }