ai-chatbot/lib/video-canvas/image-operation.ts

import { z } from "zod";

export const imageOperationSchema = z
  .object({
    operation: z.enum(["enhance", "extend", "remove-bg", "angle-edit"]),
    imageUrl: z
      .string()
      .url()
      .refine((value) => /^https?:\/\//i.test(value), "请使用已上传的图片地址"),
    prompt: z.string().max(16000).optional(),
    scaleFactor: z.number().finite().min(1).max(8).optional(),
    xScale: z.number().finite().positive().max(4).optional(),
    yScale: z.number().finite().positive().max(4).optional(),
    horizontalAngle: z.number().finite().min(0).max(360).optional(),
    verticalAngle: z.number().finite().min(-30).max(90).optional(),
    zoom: z.number().finite().min(0).max(10).optional(),
  })
  .superRefine((input, ctx) => {
    if (
      input.operation === "angle-edit" &&
      [input.horizontalAngle, input.verticalAngle, input.zoom].some((value) => value === undefined)
    )
      ctx.addIssue({ code: z.ZodIssueCode.custom, message: "角度参数不完整" });
  });
export type ImageOperation = z.infer<typeof imageOperationSchema>;
export class ImageOperationUncertainError extends Error {}
export class CanvasTaskLeaseLostError extends Error {}
type Handler = (
  request: Request,
  userId: string,
  onTaskCreated?: (taskId: string) => Promise<void>,
) => Promise<Response>;
export async function executeImageOperation(
  params: unknown,
  userId: string,
  handlers: {
    enhance: Handler;
    extend: Handler;
    removeBackground: Handler;
    angleEdit: Handler;
    pollEnhance: Handler;
  },
  options: {
    checkpoint?: (data: Record<string, unknown>) => Promise<void>;
    wait?: (ms: number) => Promise<void>;
    maxPolls?: number;
    resume?: { operation: "enhance"; taskId: string; sourceImageUrl: string } | null;
    assertLease?: () => Promise<void>;
  } = {},
) {
  const input = imageOperationSchema.parse(params);
  const payload =
    input.operation === "extend"
      ? {
          image_url: input.imageUrl,
          prompt: input.prompt || "图像拓展处理",
          includeBase64: false,
          xScale: input.xScale,
          yScale: input.yScale,
        }
      : input.operation === "angle-edit"
        ? {
            image_url: input.imageUrl,
            horizontal_angle: input.horizontalAngle,
            vertical_angle: input.verticalAngle,
            zoom: input.zoom,
          }
        : { imageUrl: input.imageUrl, scaleFactor: input.scaleFactor };
  const handler = {
    enhance: handlers.enhance,
    extend: handlers.extend,
    "remove-bg": handlers.removeBackground,
    "angle-edit": handlers.angleEdit,
  }[input.operation];
  const checkpoint = async (taskId: string) => {
    await options.assertLease?.();
    await options.checkpoint?.({
      upstreamTask: { operation: "enhance", taskId, sourceImageUrl: input.imageUrl },
    });
  };
  await options.assertLease?.();
  if (
    options.resume &&
    (input.operation !== "enhance" ||
      options.resume.operation !== "enhance" ||
      !options.resume.taskId?.trim() ||
      options.resume.sourceImageUrl !== input.imageUrl)
  )
    throw Error("已保存的高清任务与原请求不匹配，已停止恢复");
  let response = options.resume
    ? Response.json({ processing: true, taskId: options.resume.taskId }, { status: 202 })
    : await handler(
        new Request("http://canvas.internal/operation", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(payload),
        }),
        userId,
        checkpoint,
      );
  let body = await response.json();
  if (response.status === 202 || body.processing) {
    if (input.operation !== "enhance" || !body.taskId)
      throw new ImageOperationUncertainError("图片任务已受理但缺少可查询的任务标识，请勿重复提交");
    const taskId = String(body.taskId);
    try {
      await checkpoint(taskId);
    } catch (error) {
      if (error instanceof CanvasTaskLeaseLostError) throw error;
      throw new ImageOperationUncertainError("高清任务已受理，但任务标识保存失败，请勿重复提交");
    }
    const query = new URLSearchParams({ taskId, sourceImageUrl: input.imageUrl });
    for (let i = 0; i < (options.maxPolls ?? 30); i++) {
      await (options.wait || ((ms) => new Promise((resolve) => setTimeout(resolve, ms))))(1200);
      await options.assertLease?.();
      try {
        response = await handlers.pollEnhance(
          new Request(`http://canvas.internal/operation?${query}`),
          userId,
        );
        body = await response.json();
      } catch {
        throw new ImageOperationUncertainError(
          "高清任务查询连接中断，已保留任务标识，请勿重复提交",
        );
      }
      if (!response.ok && body.terminal) throw Error(body.detail || body.error || "高清任务失败");
      if (!response.ok)
        throw new ImageOperationUncertainError(body.error || "高清任务暂时无法查询，请勿重复提交");
      if (response.status !== 202 && !body.processing) break;
    }
    if (response.status === 202 || body.processing)
      throw new ImageOperationUncertainError(
        "高清任务仍在上游处理中，已保留任务标识，请勿重复提交",
      );
  }
  const url = body.original_url || body.image_url;
  if (!response.ok || !body.success || typeof url !== "string" || !/^https?:\/\//i.test(url))
    throw Error(body.detail || body.error || "图片处理没有返回有效结果");
  return { images: [{ originalUrl: url, thumbnailUrl: body.thumbnail_url || url }] };
}