ai-chatbot/lib/video-canvas/task-worker.ts

import { resolveGroupInputs } from "./group-execution";
import { groupTaskInputs } from "./task-store";
import { executeVideoUpscale } from "./video-upscale";
import { getUpscaleCheckpoint } from "./task-recovery";
import { executeSmartEdit } from "./smart-edit/execute";
import { executeCanvasAudio } from "./audio-generation";
import { ProviderTaskUncertainError } from "./video-task-lifecycle";
import {
  getEnhanceCheckpoint,
  getVideoCheckpoint,
  getAudioCheckpoint,
  getSmartEditCheckpoint,
} from "./task-recovery";
import { requireStoryPromptResult } from "./story-task-attachment";
import {
  executeImageOperation,
  ImageOperationUncertainError,
  CanvasTaskLeaseLostError,
} from "./image-operation";
import {
  executeEnhancePOST,
  executeEnhanceGET,
} from "@/lib/canvas/operations/enhance";
import { executeExtendPOST } from "@/lib/canvas/operations/extend";
import { executeRemoveBackgroundPOST } from "@/lib/canvas/operations/rmbg";
import { executeAngleEditPOST } from "@/lib/canvas/operations/angle-edit";
import {
  checkpointCanvasGenerationTask,
  interruptCanvasGenerationTask,
  ownsCanvasTaskLease,
} from "./task-store";
import "server-only";
import { executeDirectVideoGeneration } from "@/lib/ai/tools/generate-video";
import { executeCanvasImageResponse } from "@/lib/canvas/generate-image-response";
import { executeWorkflowAnalysisResponse } from "./workflow-analysis-response";
import {
  requireAnalysisResult,
  type CanvasGenerationTask,
} from "./task-contract";
import { validateVideoReferences } from "./reference-input";
import { probeRemoteUrl } from "@/lib/net/url-access";
import {
  claimCanvasGenerationTask,
  finishCanvasGenerationTask,
  heartbeatCanvasGenerationTask,
} from "./task-store";

async function execute(task: CanvasGenerationTask, userId: string) {
  const request = task.request;
  const assertLease = async () => {
    if (!(await ownsCanvasTaskLease(task.id, task.attempts!)))
      throw new CanvasTaskLeaseLostError("任务已由新的执行进程接管");
  };
  await assertLease();
  const resume = getEnhanceCheckpoint(task);
  const videoResume = getVideoCheckpoint(task);
  const audioResume = getAudioCheckpoint(task);
  const editResume = getSmartEditCheckpoint(task);
  const upscaleResume = getUpscaleCheckpoint(task);
  if (
    (task.attempts || 1) > 1 &&
    !resume &&
    !videoResume &&
    !audioResume &&
    !editResume &&
    !upscaleResume
  )
    throw Error("缺少上游任务编号，禁止重新提交");
  const params = request.workflow
    ? resolveGroupInputs(request, await groupTaskInputs(userId, request))
    : request.params;
  if (request.kind === "video-upscale")
    return executeVideoUpscale(params, userId, task.id, {
      resume: upscaleResume || undefined,
      active: assertLease,
      save: (value) =>
        checkpointCanvasGenerationTask(task.id, value, task.attempts),
    });
  if (request.kind === "smart-edit")
    return executeSmartEdit(params, userId, {
      taskId: task.id,
      resume: editResume,
      assertActive: assertLease,
      checkpoint: (value) =>
        checkpointCanvasGenerationTask(task.id, value, task.attempts),
    });
  if (request.kind === "audio")
    return executeCanvasAudio(params, {
      resume: audioResume,
      assertActive: assertLease,
      checkpoint: (value) =>
        checkpointCanvasGenerationTask(task.id, value, task.attempts),
    });
  if (request.kind === "image-operation")
    return executeImageOperation(
      params,
      userId,
      {
        enhance: executeEnhancePOST,
        extend: executeExtendPOST,
        removeBackground: executeRemoveBackgroundPOST,
        angleEdit: executeAngleEditPOST,
        pollEnhance: executeEnhanceGET,
      },
      {
        resume,
        assertLease,
        checkpoint: (result) =>
          checkpointCanvasGenerationTask(task.id, result, task.attempts),
      }
    );
  if (request.kind === "video") {
    const checked = videoResume
      ? { images: params.images as string[] | undefined, error: undefined }
      : await validateVideoReferences(params.images, (url) =>
          probeRemoteUrl(url, { timeoutMs: 8000 })
        );
    if (checked.error) throw Error(checked.error);
    const result = await executeDirectVideoGeneration({
      ...params,
      prompt: String(params.prompt || ""),
      mode: params.mode as any,
      images: checked.images,
      taskLifecycle: {
        resume: videoResume,
        assertActive: assertLease,
        save: (checkpoint) =>
          checkpointCanvasGenerationTask(
            task.id,
            { upstreamTask: checkpoint },
            task.attempts
          ),
      },
    });
    if (!result.success || !result.videoUrl) throw Error("视频生成未返回结果");
    return {
      url: result.videoUrl,
      metadata: {
        ...result.metadata,
        temporaryVideoUrl: result.videoUrl,
        storageStatus:
          result.metadata?.permanentVideoUrl === result.videoUrl
            ? "archived"
            : "temporary",
      },
      modelId: result.modelId,
      modelRef: result.modelRef,
    };
  }
  const response =
    request.kind === "image"
      ? await executeCanvasImageResponse(
          { ...params, prompt: String(params.prompt || ""), async: false },
          userId
        )
      : await executeWorkflowAnalysisResponse(params);
  const body = await response.json();
  if (!response.ok || !body.success) throw Error(body.error || "任务执行失败");
  if (
    request.kind === "image" &&
    !body.data?.images?.some(
      (image: any) => image.originalUrl || image.thumbnailUrl
    )
  )
    throw Error("图片生成没有返回结果");
  if (request.kind === "analysis")
    requireAnalysisResult(body.data, String(params.mode));
  requireStoryPromptResult(body.data, request.attachment);
  return body.data;
}
export async function runCanvasGenerationQueue(
  userId: string,
  projectId: string
) {
  // Work is independent of the browser request. Atomic DB claiming also covers multiple processes.
  for (let i = 0; i < 40; i++) {
    const task = await claimCanvasGenerationTask(userId, projectId);
    if (!task) return;
    const timer = setInterval(
      () =>
        void heartbeatCanvasGenerationTask(task.id, task.attempts).catch(
          () => {}
        ),
      20000
    );
    try {
      await finishCanvasGenerationTask(
        task.id,
        await execute(task, userId),
        undefined,
        task.attempts
      );
    } catch (error) {
      if (error instanceof CanvasTaskLeaseLostError) return;
      if (
        error instanceof ImageOperationUncertainError ||
        error instanceof ProviderTaskUncertainError
      ) {
        await interruptCanvasGenerationTask(
          task.id,
          error.message,
          task.attempts
        );
        continue;
      }
      await finishCanvasGenerationTask(
        task.id,
        undefined,
        error instanceof Error ? error.message : "任务执行失败",
        task.attempts
      );
    } finally {
      clearInterval(timer);
    }
  }
}