ai-chatbot/lib/video-canvas/video-upscale-remote.ts

import { createHash } from "node:crypto";
import { createWriteStream } from "node:fs";
import { Readable, Transform } from "node:stream";
import { pipeline } from "node:stream/promises";
import path from "node:path";
import { z } from "zod";
import type { VideoUpscaleInput } from "./video-upscale-contract";
import { ProviderTaskUncertainError } from "./video-task-lifecycle";

export type UpscaleCheckpoint = { id: string; endpointHash: string };
const sha = (value: string) => createHash("sha256").update(value).digest("hex");
export function upscaleConfig() {
  const value = process.env.VIDEO_UPSCALE_URL;
  if (!value) return null;
  const url = new URL(value);
  if (
    !["http:", "https:"].includes(url.protocol) ||
    url.username ||
    url.password ||
    url.search ||
    url.hash
  )
    throw Error("视频超分服务地址配置无效");
  const base = value.replace(/\/$/, ""),
    token = process.env.VIDEO_UPSCALE_TOKEN;
  if (!token || token.length < 32) throw Error("视频超分服务认证未配置");
  return { base, token };
}
const statusSchema = z.object({
  id: z.string().regex(/^[a-f0-9]{64}$/),
  status: z.enum(["queued", "running", "succeeded", "failed"]),
  stage: z.string().max(300),
  error: z.string().max(500).optional(),
  result: z
    .object({
      width: z.number().int().positive().max(3840),
      height: z.number().int().positive().max(3840),
      durationSeconds: z.number().positive().max(120.2),
      fps: z.number().positive().max(60),
      frames: z.number().int().positive().max(7200),
      audio: z.boolean(),
    })
    .optional(),
});
export async function remoteUpscale(
  plan: VideoUpscaleInput,
  userId: string,
  taskId: string,
  dir: string,
  lifecycle: {
    resume?: UpscaleCheckpoint;
    active: () => Promise<void>;
    save: (checkpoint: UpscaleCheckpoint, stage: string) => Promise<void>;
  }
) {
  const config = upscaleConfig();
  if (!config)
    throw Error(
      lifecycle.resume
        ? "远程超分任务的服务配置已移除，不能回退本机重复超分"
        : "视频高清服务尚未接通"
    );
  const checkpoint = lifecycle.resume || {
    id: sha("video-upscale-v1:" + userId + ":" + taskId),
    endpointHash: sha(config.base),
  };
  if (
    !/^[a-f0-9]{64}$/.test(checkpoint.id) ||
    checkpoint.endpointHash !== sha(config.base)
  )
    throw Error("超分任务所属服务器已变化，请恢复原服务配置后继续");
  await lifecycle.save(checkpoint, "提交算力服务器");
  const headers = {
    authorization: "Bearer " + config.token,
    "content-type": "application/json",
  };
  async function request(route: string, body?: unknown) {
    await lifecycle.active();
    let response: Response;
    try {
      response = await fetch(config!.base + route, {
        method: body ? "POST" : "GET",
        headers,
        body: body ? JSON.stringify(body) : undefined,
        signal: AbortSignal.timeout(20000),
        redirect: "error",
        cache: "no-store",
      });
    } catch {
      throw new ProviderTaskUncertainError(
        "算力服务器连接中断，任务已保存，可继续查询"
      );
    }
    if ([429, 500, 502, 503, 504].includes(response.status))
      throw new ProviderTaskUncertainError(
        "算力服务器繁忙，任务已保存等待继续"
      );
    if (!response.ok)
      throw Error(
        `超分服务请求失败（${response.status}），请检查服务状态或任务参数`
      );
    let data: unknown;
    try {
      data = await response.json();
    } catch {
      throw new ProviderTaskUncertainError("超分服务响应中断，将查询同一任务");
    }
    const result = statusSchema.parse(data);
    if (result.id !== checkpoint.id)
      throw Error("超分服务返回了不匹配的任务编号");
    return result;
  }
  let status = await request("/v1/upscales", {
    id: checkpoint.id,
    tenant: sha(userId),
    plan,
  });
  const deadline = Date.now() + 8 * 60000;
  while (status.status === "queued" || status.status === "running") {
    await lifecycle.save(checkpoint, status.stage);
    if (Date.now() > deadline)
      throw new ProviderTaskUncertainError("远程超分仍在运行，已保存任务进度");
    await new Promise((resolve) => setTimeout(resolve, 3000));
    status = await request("/v1/upscales/" + checkpoint.id);
  }
  if (status.status === "failed")
    throw Error(status.error || "算力服务器超分失败");
  if (!status.result) throw Error("超分服务未返回成片参数");
  await lifecycle.save(checkpoint, "读取远程成片");
  async function download(kind: string, destination: string, maxBytes: number) {
    await lifecycle.active();
    try {
      const response = await fetch(
        config!.base + "/v1/upscales/" + checkpoint.id + "/" + kind,
        {
          headers,
          signal: AbortSignal.timeout(180000),
          redirect: "error",
          cache: "no-store",
        }
      );
      if (
        !response.ok ||
        !response.body ||
        Number(response.headers.get("content-length")) > maxBytes
      )
        throw Error("远程成片不可读取或超过处理上限");
      let size = 0;
      const limit = new Transform({
        transform(chunk, _encoding, callback) {
          size += chunk.length;
          callback(
            size > maxBytes ? Error("远程成片超过处理上限") : null,
            chunk
          );
        },
      });
      await pipeline(
        Readable.fromWeb(response.body as any),
        limit,
        createWriteStream(destination, { mode: 0o600 })
      );
    } catch {
      throw new ProviderTaskUncertainError(
        "成片已生成但读取中断，将继续读取同一任务"
      );
    }
  }
  const file = path.join(dir, "result.mp4");
  await download("video", file, 300 * 1024 * 1024);

  return {
    file,
    width: status.result.width,
    height: status.result.height,
    duration: status.result.durationSeconds,
    fps: status.result.fps,
    frames: status.result.frames,
    audio: status.result.audio,
  };
}