vfxreview/lib/render-pipeline/jobs.ts

import { Prisma, RenderJobType } from "@prisma/client";
import { db } from "@/lib/db";
import { PipelineError } from "./errors";
import { applyTransition } from "./transitions";
import { getConfigBoolean, getConfigNumber } from "./config";
import { machineCanClaim } from "./machines";
import { reapExpiredLeases } from "./reaper";
import { buildPreviewManifest, createPreviewJob } from "./preview";

const JOB_TYPES: RenderJobType[] = ["AE_RENDER", "PREVIEW_ONLY", "DELIVERY_BUILD"];

/**
 * E8 — atomically claim the next queued job (§6.4). `FOR UPDATE SKIP LOCKED`
 * makes double-claims impossible under concurrent workers. All scheduling
 * policy (availability windows, Render Now, urgent priority) is enforced
 * here server-side — workers poll dumbly (§7.10).
 */
export async function claimNextJob(machineId: string, types?: string[]) {
  const machine = await db.machine.findUnique({ where: { id: machineId } });
  if (!machine) throw new PipelineError(404, "Machine not found — register first (POST /api/ext/workers/register)");

  // Defensive sweep on each claim (studio decision 18.1-Q8)
  await reapExpiredLeases().catch((err) => console.error("[render-pipeline] claim-time reap failed", err));

  const claimTypes = (types?.length ? types : ["AE_RENDER", "PREVIEW_ONLY", "DELIVERY_BUILD"]).filter(
    (t): t is RenderJobType => (JOB_TYPES as string[]).includes(t)
  );
  if (claimTypes.length === 0) throw new PipelineError(422, "No valid job types requested");

  const gate = machineCanClaim(machine, new Date());
  if (gate === "none") return null;

  const [leaseSeconds, urgentThreshold] = await Promise.all([
    getConfigNumber("render.leaseSeconds"),
    getConfigNumber("render.urgentPriorityThreshold"),
  ]);

  const urgentFilter =
    gate === "urgent-only" ? Prisma.sql`AND priority <= ${urgentThreshold}` : Prisma.empty;

  const claimed = await db.$queryRaw<{ id: string }[]>(Prisma.sql`
    UPDATE "render_jobs"
    SET status = 'CLAIMED',
        "machineId" = ${machineId},
        "claimedAt" = now(),
        "leaseExpiresAt" = now() + (${leaseSeconds}::int * interval '1 second'),
        "updatedAt" = now()
    WHERE id = (
      SELECT id FROM "render_jobs"
      WHERE status = 'QUEUED'
        AND type::text = ANY(${claimTypes}::text[])
        ${urgentFilter}
      ORDER BY priority ASC, "createdAt" ASC
      FOR UPDATE SKIP LOCKED
      LIMIT 1
    )
    RETURNING id
  `);
  if (claimed.length === 0) return null;

  const job = await db.renderJob.findUniqueOrThrow({
    where: { id: claimed[0].id },
    include: { export: { select: { id: true, status: true, versionString: true } } },
  });

  // T2: QUEUED → RENDERING
  if (job.export && job.export.status === "QUEUED") {
    await db.$transaction(async (tx) => {
      await applyTransition(tx, job.export!, "RENDERING", {
        type: "WORKER",
        id: machineId,
        note: `Claimed by ${machine.name} (attempt ${job.attempt})`,
      });
    });
  }

  // Delivery packages have their own small lifecycle on the same queue.
  if (job.deliveryId) {
    await db.deliveryPackage.updateMany({
      where: { id: job.deliveryId, status: "QUEUED" },
      data: { status: "BUILDING" },
    });
  }

  return {
    job: {
      id: job.id,
      type: job.type,
      attempt: job.attempt,
      maxAttempts: job.maxAttempts,
      priority: job.priority,
      exportId: job.exportId,
      deliveryId: job.deliveryId,
      leaseExpiresAt: job.leaseExpiresAt,
      manifest: job.manifest,
    },
  };
}

async function getJobForReport(jobId: string, machineId: string) {
  const job = await db.renderJob.findUnique({
    where: { id: jobId },
    include: { export: { select: { id: true, status: true } } },
  });
  if (!job) throw new PipelineError(404, "Render job not found");
  // Zombie-worker guard (§7.4): reports from a machine that doesn't hold the lease are rejected
  if (job.machineId !== machineId) {
    throw new PipelineError(409, "Job is not held by this machine (lease reassigned?)");
  }
  return job;
}

/** E9 — progress report; renews the lease. */
export async function reportProgress(
  jobId: string,
  machineId: string,
  body: { progress?: number; currentFrame?: number; totalFrames?: number; etaSeconds?: number; logTail?: string }
) {
  const job = await getJobForReport(jobId, machineId);

  if (job.status === "CANCELLED") {
    return { ok: true, cancelRequested: true, leaseExpiresAt: job.leaseExpiresAt };
  }
  if (["COMPLETED", "FAILED", "EXPIRED"].includes(job.status)) {
    return { ok: true, alreadyApplied: true, cancelRequested: false, leaseExpiresAt: job.leaseExpiresAt };
  }

  const leaseSeconds = await getConfigNumber("render.leaseSeconds");
  const leaseExpiresAt = new Date(Date.now() + leaseSeconds * 1000);
  await db.renderJob.update({
    where: { id: job.id },
    data: {
      status: "RUNNING",
      startedAt: job.startedAt ?? new Date(),
      leaseExpiresAt,
      ...(body.progress != null ? { progress: Math.min(1, Math.max(0, body.progress)) } : {}),
      ...(body.currentFrame != null ? { currentFrame: body.currentFrame } : {}),
      ...(body.totalFrames != null ? { totalFrames: body.totalFrames } : {}),
      ...(body.etaSeconds != null ? { etaSeconds: body.etaSeconds } : {}),
      ...(body.logTail != null ? { logTail: body.logTail.slice(-20000) } : {}),
    },
  });
  return { ok: true, leaseExpiresAt, cancelRequested: false };
}

/** E10 — failure report. Auto-requeues while attempts remain and the error is retryable. */
export async function reportFail(
  jobId: string,
  machineId: string,
  body: {
    stage?: string;
    exitCode?: number;
    errorMessage?: string;
    logTail?: string;
    logFileKey?: string;
    retryable?: boolean;
  }
) {
  return db.$transaction(async (tx) => {
    const job = await tx.renderJob.findUnique({
      where: { id: jobId },
      include: { export: { select: { id: true, status: true } } },
    });
    if (!job) throw new PipelineError(404, "Render job not found");
    if (job.machineId !== machineId) throw new PipelineError(409, "Job is not held by this machine (lease reassigned?)");

    if (job.status === "FAILED") {
      return { job: { id: job.id, status: job.status }, export: job.export ? { id: job.export.id, status: job.export.status } : null, autoRequeued: false, alreadyApplied: true };
    }
    if (job.status === "CANCELLED") {
      // worker acknowledging a cancel — record the wind-down, stay CANCELLED
      await tx.renderJob.update({ where: { id: job.id }, data: { finishedAt: job.finishedAt ?? new Date(), logTail: body.logTail?.slice(-20000) ?? job.logTail } });
      return { job: { id: job.id, status: "CANCELLED" }, export: job.export ? { id: job.export.id, status: job.export.status } : null, autoRequeued: false, alreadyApplied: true };
    }
    if (job.status === "COMPLETED") {
      throw new PipelineError(409, "Job already reported complete — contradictory fail report rejected");
    }

    await tx.renderJob.update({
      where: { id: job.id },
      data: {
        status: "FAILED",
        finishedAt: new Date(),
        exitCode: body.exitCode ?? null,
        errorMessage: body.errorMessage?.slice(0, 2000) ?? null,
        logTail: body.logTail?.slice(-20000) ?? job.logTail,
        logFileKey: body.logFileKey ?? job.logFileKey,
      },
    });

    if (!job.export) {
      return { job: { id: job.id, status: "FAILED" }, export: null, autoRequeued: false };
    }

    const retryable = body.retryable !== false;
    const autoRequeue = retryable && job.attempt < job.maxAttempts;
    const isPreview = job.type === "PREVIEW_ONLY";

    if (autoRequeue) {
      const next = await tx.renderJob.create({
        data: {
          type: job.type,
          exportId: job.exportId,
          attempt: job.attempt + 1,
          maxAttempts: job.maxAttempts,
          priority: job.priority,
          manifest: job.manifest as Prisma.InputJsonValue,
        },
      });
      const note = `${body.stage ?? "RENDER"} failed (${body.errorMessage ?? "no message"}) — auto-requeued as attempt ${next.attempt}`;

      if (isPreview) {
        // The export stays in GENERATING_PREVIEW across preview attempts —
        // the EXRs are untouched, only the preview build is repeated.
        await tx.exportEvent.create({
          data: {
            exportId: job.export.id,
            fromStatus: job.export.status,
            toStatus: job.export.status,
            actorType: "SYSTEM",
            actorId: machineId,
            note,
          },
        });
        return { job: { id: job.id, status: "FAILED" }, export: { id: job.export.id, status: job.export.status }, autoRequeued: true, nextJob: { id: next.id, attempt: next.attempt } };
      }

      const exp = await applyTransition(tx, job.export, "QUEUED", { type: "SYSTEM", id: machineId, note });
      return { job: { id: job.id, status: "FAILED" }, export: { id: exp.id, status: exp.status }, autoRequeued: true, nextJob: { id: next.id, attempt: next.attempt } };
    }

    const exp = await applyTransition(tx, job.export, isPreview ? "PREVIEW_FAILED" : "RENDER_FAILED", {
      type: "WORKER",
      id: machineId,
      note: `${body.stage ?? "RENDER"} failed: ${body.errorMessage ?? "no message"}${retryable ? " (max attempts reached)" : " (not retryable)"}`,
    });
    return { job: { id: job.id, status: "FAILED" }, export: { id: exp.id, status: exp.status }, autoRequeued: false };
  });
}

/**
 * E11 — aerender finished OK.
 *
 * Hands off to validation (§8) when it is enabled: the worker then runs the
 * checks and reports them via E12, which queues the preview stage. With
 * validation disabled the old direct handoff to the preview stage is used, so
 * the kill switch restores the pre-Phase-3 flow exactly.
 */
export async function completeRender(
  jobId: string,
  machineId: string,
  body: { renderSeconds?: number; logFileKey?: string; logTail?: string; exrFileCount?: number; exrTotalBytes?: number }
) {
  return db.$transaction(async (tx) => {
    const job = await tx.renderJob.findUnique({
      where: { id: jobId },
      include: { export: { select: { id: true, status: true } } },
    });
    if (!job) throw new PipelineError(404, "Render job not found");
    if (job.machineId !== machineId) throw new PipelineError(409, "Job is not held by this machine (lease reassigned?)");

    if (job.status === "COMPLETED") {
      return { job: { id: job.id, status: job.status }, export: job.export ? { id: job.export.id, status: job.export.status } : null, alreadyApplied: true };
    }
    if (["FAILED", "CANCELLED", "EXPIRED"].includes(job.status)) {
      throw new PipelineError(409, `Job is ${job.status} — contradictory complete report rejected`);
    }

    // With validation to follow, the job deliberately stays open and keeps its
    // lease: a worker that dies mid-validation is then reaped and retried like
    // any other stalled job, instead of leaving the Export stuck in VALIDATING.
    const validationEnabled = job.export ? await getConfigBoolean("validation.enabled") : false;

    await tx.renderJob.update({
      where: { id: job.id },
      data: {
        status: validationEnabled ? "RUNNING" : "COMPLETED",
        progress: 1,
        ...(validationEnabled ? {} : { finishedAt: new Date() }),
        exitCode: 0,
        logFileKey: body.logFileKey ?? job.logFileKey,
        logTail: body.logTail?.slice(-20000) ?? job.logTail,
      },
    });

    let exportResult: { id: string; status: string } | null = null;
    let previewJob: { id: string; type: string } | null = null;

    if (job.export) {
      const exportRow = await tx.export.findUniqueOrThrow({ where: { id: job.export.id } });
      const stats = {
        ...(body.exrFileCount != null ? { exrFileCount: body.exrFileCount } : {}),
        ...(body.exrTotalBytes != null ? { exrTotalBytes: BigInt(Math.round(body.exrTotalBytes)) } : {}),
      };

      // The EXR render is only the first stage. Validation comes next when
      // enabled; the worker reports its results via E12 on the same job.
      if (validationEnabled) {
        const exp = await applyTransition(tx, job.export, "VALIDATING", {
          type: "WORKER",
          id: machineId,
          note: `Render complete in ${body.renderSeconds ?? "?"}s — validating`,
        }, stats);
        return {
          job: { id: job.id, status: "RUNNING" },
          export: { id: exp.id, status: exp.status },
          previewJob: null,
          validationRequired: true,
          alreadyApplied: false,
        };
      }

      // Validation disabled: hand off to the preview stage (headless AE rebuild
      // → delivery MOV + review MP4) exactly as before Phase 3.
      const previewEnabled = await getConfigBoolean("preview.enabled");
      if (previewEnabled) {
        const manifest = await buildPreviewManifest(exportRow);
        const created = await createPreviewJob(tx, exportRow, manifest, job.maxAttempts, job.priority);
        previewJob = { id: created.id, type: created.type };
        const exp = await applyTransition(tx, job.export, "GENERATING_PREVIEW", {
          type: "WORKER",
          id: machineId,
          note: `Render complete in ${body.renderSeconds ?? "?"}s — queued preview build`,
        }, stats);
        exportResult = { id: exp.id, status: exp.status };
      } else {
        const exp = await applyTransition(tx, job.export, "READY_FOR_QC", {
          type: "WORKER",
          id: machineId,
          note: `Render complete in ${body.renderSeconds ?? "?"}s (preview stage disabled)`,
        }, stats);
        exportResult = { id: exp.id, status: exp.status };
      }
    }
    return { job: { id: job.id, status: "COMPLETED" }, export: exportResult, previewJob, alreadyApplied: false };
  });
}

/** E15 — cancel a single job from the web UI (delegates to export-level cancel when linked). */
export async function getJobDetail(jobId: string) {
  const job = await db.renderJob.findUnique({
    where: { id: jobId },
    include: {
      machine: { select: { id: true, name: true } },
      export: { include: { shot: { select: { shotCode: true } } } },
    },
  });
  if (!job) throw new PipelineError(404, "Render job not found");
  return job;
}