ai-chatbot/lib/canvas/download-preparation.ts

import "server-only";
import { createHash } from "node:crypto";
import { cacheImageBytes, readImageBytes, type ImageBytes } from "@/lib/storage/image-byte-cache";

export type DownloadPreparationInput = {
  userId: string;
  urls: string[];
  requestOrigin: string;
  cookie: string;
};
type Result =
  | { status: "preparing" }
  | { status: "ready"; asset: ImageBytes }
  | { status: "failed"; error: string };
type Job = {
  userId: string;
  input: DownloadPreparationInput;
  status: "queued" | "running" | "ready" | "failed";
  updatedAt: number;
  asset?: ImageBytes;
};
const FAILURE_MESSAGE = "原图暂时无法读取，图片仍保留在画布中，请稍后重试";
const MAX_MEMORY_BYTES = 128 * 1024 * 1024;

// A download is independent of node autosave and of the archive worker's batch.
// Polls share one preparation; only bytes are cached, never cookies or signed URLs.
export function createDownloadPreparationQueue(deps: {
  load: (input: DownloadPreparationInput, signal: AbortSignal) => Promise<ImageBytes>;
  read: (identity: string) => Promise<ImageBytes | null>;
  write: (identity: string, asset: ImageBytes) => Promise<unknown>;
  schedule: (run: () => Promise<void>) => void;
  now?: () => number;
  timeoutMs?: number;
  concurrency?: number;
}) {
  const jobs = new Map<string, Job>();
  const now = deps.now || Date.now;
  const concurrency = deps.concurrency || 4;
  let active = 0;

  const prune = () => {
    let bytes = 0;
    for (const [key, job] of jobs) {
      if ((job.status === "ready" && now() - job.updatedAt > 10 * 60_000) ||
          (job.status === "failed" && now() - job.updatedAt > 15_000)) jobs.delete(key);
      else bytes += job.asset?.buffer.length || 0;
    }
    // Completed jobs are in insertion order; keep the newest fallback bytes if
    // the optional disk cache is full. Never evict a running preparation.
    for (const [key, job] of jobs) {
      if (bytes <= MAX_MEMORY_BYTES) break;
      if (job.status === "ready") { bytes -= job.asset?.buffer.length || 0; jobs.delete(key); }
    }
  };
  const drain = async (): Promise<void> => {
    const running: Promise<void>[] = [];
    for (const [key, job] of jobs) {
      if (active >= concurrency) break;
      if (job.status !== "queued") continue;
      active++;
      job.status = "running";
      running.push((async () => {
        const controller = new AbortController();
        let timer: ReturnType<typeof setTimeout> | undefined;
        try {
          const asset = await Promise.race([
            deps.load(job.input, controller.signal),
            new Promise<never>((_, reject) => {
              timer = setTimeout(() => {
                controller.abort();
                reject(new Error("Download preparation deadline"));
              }, deps.timeoutMs || 90_000);
            }),
          ]);
          if (asset.contentType !== "image/png" || !asset.buffer.length || asset.buffer.length > MAX_MEMORY_BYTES) {
            throw new Error("Invalid prepared original");
          }
          // Keep a bounded in-memory fallback until persistence is confirmed.
          job.asset = asset;
          await deps.write(key, asset).catch(() => undefined);
          const stored = await deps.read(key).catch(() => null);
          if (stored?.contentType === "image/png" && stored.buffer.length) job.asset = undefined;
          job.status = "ready";
        } catch {
          job.status = "failed";
        } finally {
          if (timer) clearTimeout(timer);
          // Do not retain authentication data after a preparation finishes.
          job.input = { ...job.input, cookie: "", urls: [] };
          job.updatedAt = now();
          active--;
          prune();
        }
      })().then(() => drain()));
    }
    await Promise.all(running);
  };

  return {
    async request(input: DownloadPreparationInput): Promise<Result> {
      const key = "canvas-download:v2:" + createHash("sha256")
        .update(JSON.stringify([input.userId, input.requestOrigin, input.cookie, input.urls])).digest("hex");
      const cached = await deps.read(key).catch(() => null);
      if (cached?.contentType === "image/png") return { status: "ready", asset: cached };
      prune();
      const current = jobs.get(key);
      if (current?.status === "ready" && current.asset) return { status: "ready", asset: current.asset };
      if (current?.status === "failed") return { status: "failed", error: FAILURE_MESSAGE };
      if (current && current.status !== "ready") return { status: "preparing" };
      if (jobs.size >= 64) {
        for (const [oldKey, job] of jobs) {
          if (job.status === "ready" || job.status === "failed") jobs.delete(oldKey);
          if (jobs.size < 64) break;
        }
      }
      if (jobs.size >= 64 || [...jobs.values()].filter(job => job.userId === input.userId &&
          (job.status === "queued" || job.status === "running")).length >= 4) {
        return { status: "failed", error: "正在准备的原图较多，请等待当前下载完成后再试" };
      }
      jobs.set(key, { userId: input.userId, input, status: "queued", updatedAt: now() });
      deps.schedule(drain);
      return { status: "preparing" };
    },
  };
}

const state = globalThis as typeof globalThis & {
  __canvasDownloadPreparationV2?: ReturnType<typeof createDownloadPreparationQueue>;
};

export function getDownloadPreparationQueue(load: Parameters<typeof createDownloadPreparationQueue>[0]["load"],
  after: (run: () => Promise<void>) => void) {
  return (state.__canvasDownloadPreparationV2 ??= createDownloadPreparationQueue({
    load,
    read: readImageBytes,
    write: cacheImageBytes,
    schedule(run) {
      // Also run when invoked outside a live Next request. The queue itself
      // deduplicates both callbacks and limits concurrent work.
      const timer = setTimeout(() => { void run(); }, 1_000);
      timer.unref?.();
      try { after(async () => { clearTimeout(timer); await run(); }); }
      catch { clearTimeout(timer); setImmediate(() => { void run(); }); }
    },
  }));
}