4e88a9dfe6
- A handler that ignores its abort is left behind after the grace period; a late rejection from it (or from the grace timer after the job settled) was unobserved, and the worker exits on unhandled rejections. Both sides of the race are now observed and the grace timer is cleared with the job. - Ending a job removes its temp directory when no runner in this process holds it (leftovers of a runner that died here), and reconciliation sweeps directories of jobs that no longer run. - An executor test runs a media tool end to end through applyTool: shown only for operations the space can run, started with the run's scope and effect key, continue_after split off, refused when not exposed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01MRtGG9ARfqqhma3t85Ro2u
1089 lines
41 KiB
TypeScript
1089 lines
41 KiB
TypeScript
import { randomUUID } from "node:crypto";
|
|
import { constants, createWriteStream } from "node:fs";
|
|
import { copyFile, mkdir, readdir, 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<Input = unknown> {
|
|
operation: MediaOperation;
|
|
/** Validates tool and card input. The job stores the parsed value, so it must parse again. */
|
|
inputSchema: z.ZodType<Input>;
|
|
/** 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<boolean>;
|
|
handler: MediaOperationHandler;
|
|
/** GPU guard: the runner waits while the backend reports less free VRAM than this. */
|
|
gpu?: { minVramMb: number; status(signal?: AbortSignal): Promise<GpuStatus | undefined> };
|
|
/** 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<ThreadEvents, "notify">;
|
|
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<void>;
|
|
/** Re-enqueues jobs whose enqueue was lost, whose runner died, or whose wait is over. */
|
|
reconcile(): Promise<void>;
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
/**
|
|
* Runs the operation, giving up `graceMs` after an abort if the handler ignores the signal. A
|
|
* late rejection of either side is observed, so an abandoned handler cannot crash the process.
|
|
*/
|
|
async function runHandler(
|
|
handler: MediaOperationHandler,
|
|
context: MediaJobContext,
|
|
graceMs: number,
|
|
) {
|
|
const work = Promise.resolve().then(() => handler.run(context));
|
|
work.catch(() => undefined);
|
|
let timer: ReturnType<typeof setTimeout> | undefined;
|
|
const onAbort = () => {
|
|
timer = setTimeout(() => giveUp(context.signal.reason), graceMs);
|
|
timer.unref?.();
|
|
};
|
|
let giveUp: (reason: unknown) => void = () => undefined;
|
|
const abandoned = new Promise<never>((_, reject) => {
|
|
giveUp = reject;
|
|
});
|
|
abandoned.catch(() => undefined);
|
|
if (context.signal.aborted) onAbort();
|
|
else context.signal.addEventListener("abort", onAbort, { once: true });
|
|
try {
|
|
return await Promise.race([work, abandoned]);
|
|
} finally {
|
|
clearTimeout(timer);
|
|
context.signal.removeEventListener("abort", onAbort);
|
|
}
|
|
}
|
|
|
|
/** 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<MediaOperation, MediaOperationRegistration>();
|
|
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<string, AbortController>();
|
|
|
|
const isAvailable = async (
|
|
registration: MediaOperationRegistration | undefined,
|
|
spaceId: string,
|
|
): Promise<boolean> => {
|
|
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);
|
|
});
|
|
};
|
|
|
|
const removeTempDir = async (jobId: string) => {
|
|
await rm(path.join(tempRoot, jobId), { recursive: true, force: true }).catch((error) => {
|
|
getLogger().error("media job temp cleanup", error);
|
|
});
|
|
};
|
|
|
|
/** Removes temp directories of jobs that no longer run (their runner died on this host). */
|
|
const sweepTempDirs = async () => {
|
|
const names = await readdir(tempRoot).catch(() => [] as string[]);
|
|
const candidates = names.filter((name) => !running.has(name));
|
|
if (!candidates.length) return;
|
|
const active = await deps.prisma.mediaJob.findMany({
|
|
where: { id: { in: candidates }, status: { in: ACTIVE_STATUSES } },
|
|
select: { id: true },
|
|
});
|
|
const keep = new Set(active.map((job) => job.id));
|
|
for (const name of candidates) if (!keep.has(name)) await removeTempDir(name);
|
|
};
|
|
|
|
/** 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<boolean> => {
|
|
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;
|
|
// Leftovers of a runner that died on this host; a live runner's directory is its own.
|
|
if (!running.has(jobId)) await removeTempDir(jobId);
|
|
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<Uint8Array>),
|
|
createWriteStream(file, { flags: "wx", mode: 0o600 }),
|
|
);
|
|
};
|
|
|
|
const removeArtifacts = async (job: MediaJob, ids: string[]) => {
|
|
const rows = await deps.prisma.artifact.findMany({
|
|
where: { id: { in: ids }, spaceId: job.spaceId },
|
|
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<string, MediaJobInputFile>();
|
|
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 runHandler(registration.handler, context, abortGraceMs);
|
|
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 removeTempDir(jobId);
|
|
}
|
|
};
|
|
|
|
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<MediaJobSnapshot | null> {
|
|
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() {
|
|
await sweepTempDirs();
|
|
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,
|
|
};
|
|
}
|