src/worker/modelCache.ts

// Tải trọng số model (≈1,84 GB) về OPFS của trình duyệt:
//  • tải song song 6 luồng, mỗi đoạn tối đa 32 MiB bằng HTTP Range, thử lại tối đa 3 lần
//  • kiểm tra SHA-256 từng file trước khi ghi "biên nhận" (verified.json)
//  • khoá `navigator.locks` để hai tab không tải cùng lúc
//  • trả về một đối tượng cache có `match()` để transformers.js đọc thẳng từ OPFS

import { createSHA256 } from 'hash-wasm';
import { MODEL_FILES, MODEL_REVISION, MODEL_TOTAL_BYTES, type ModelFile } from './manifest';

const CHUNK = 32 * 1024 ** 2;
const PARALLEL_DOWNLOADS = 6;
const STORAGE_DIR = 'pi-minicpm5-' + MODEL_REVISION;

export interface DownloadProgress {
  phase: 'download';
  loaded: number;
  total: number;
  cached: number;
  file?: string;
}

export interface OpenOptions {
  signal: AbortSignal;
  cachedOnly?: boolean;
  onProgress?: (p: DownloadProgress) => void;
}

export interface ModelCacheHandle {
  cachedBytes: number;
  match(request: string | Request): Promise<Response | undefined>;
  put(): Promise<void>;
}

type Receipts = Record<string, { sha256: string; modified: number }>;

const encode = (name: string) => encodeURIComponent(name);

const sleep = (ms: number, signal?: AbortSignal) =>
  new Promise<void>((resolve, reject) => {
    signal?.throwIfAborted();
    const onAbort = () => {
      clearTimeout(timer);
      reject(signal!.reason);
    };
    const timer = setTimeout(() => {
      signal?.removeEventListener('abort', onAbort);
      resolve();
    }, ms);
    signal?.addEventListener('abort', onAbort, { once: true });
  });

/** Tải một đoạn byte [start, end] (dùng Range nếu file lớn), có thử lại. */
async function fetchRange(url: string, start: number, end: number, total: number, signal: AbortSignal) {
  for (let attempt = 0; ; attempt++) {
    try {
      const ranged = total > CHUNK;
      const res = await fetch(url, {
        headers: ranged ? { Range: `bytes=${start}-${end}` } : {},
        signal: AbortSignal.any([signal, AbortSignal.timeout(90_000)]),
        cache: 'no-store',
      });
      if (!res.ok) throw new Error(`Download returned HTTP ${res.status}`);
      if (ranged && res.status !== 206) {
        await res.body?.cancel();
        throw new Error('The model host did not honor byte ranges.');
      }
      const contentRange = res.headers.get('Content-Range');
      if (contentRange && contentRange !== `bytes ${start}-${end}/${total}`) {
        await res.body?.cancel();
        throw new Error('Incorrect download range.');
      }
      const bytes = new Uint8Array(await res.arrayBuffer());
      if (bytes.length !== end - start + 1) throw new Error('Incomplete download chunk.');
      return bytes;
    } catch (err) {
      signal.throwIfAborted();
      if (attempt === 2) throw err;
      await sleep(500 * 2 ** attempt, signal);
    }
  }
}

async function scan(create: boolean) {
  if (!navigator.storage?.getDirectory) {
    throw new Error(
      'This app needs browser file storage (OPFS) to cache the 1.84 GB model. Open it in a regular browser tab.',
    );
  }
  const root = await (await navigator.storage.getDirectory()).getDirectoryHandle(STORAGE_DIR, { create });
  const receiptsFile = await root.getFileHandle('verified.json', { create });
  let receipts: Receipts = {};
  try {
    receipts = JSON.parse(await (await receiptsFile.getFile()).text());
  } catch {
    receipts = {};
  }
  if (!receipts || typeof receipts !== 'object' || Array.isArray(receipts)) receipts = {};

  let cached = 0;
  const pending: Array<[string, ModelFile]> = [];
  for (const [name, info] of Object.entries(MODEL_FILES)) {
    try {
      const file = await (await root.getFileHandle(encode(name))).getFile();
      const receipt = receipts[name];
      if (file.size === info.bytes && receipt?.sha256 === info.sha256 && receipt.modified === file.lastModified) {
        cached += info.bytes;
        continue;
      }
    } catch (err) {
      if ((err as DOMException).name !== 'NotFoundError') throw err;
    }
    pending.push([name, info]);
  }
  return { root, receiptsFile, receipts, pending, cached };
}

/** True nếu toàn bộ file model đã có trong OPFS và đã được xác minh. */
export async function isModelCached(): Promise<boolean> {
  try {
    return (await scan(false)).pending.length === 0;
  } catch {
    return false;
  }
}

function cacheMiss() {
  return Object.assign(new Error('The saved model is incomplete. Confirm the download to restore it.'), {
    code: 'MODEL_CACHE_MISS',
  });
}

async function openInner(baseUrl: string, { signal, cachedOnly = false, onProgress = () => {} }: OpenOptions) {
  let state: Awaited<ReturnType<typeof scan>>;
  try {
    state = await scan(!cachedOnly);
  } catch (err) {
    if (cachedOnly && (err as DOMException).name === 'NotFoundError') throw cacheMiss();
    throw err;
  }
  const { root, receiptsFile, receipts, pending, cached } = state;
  if (cachedOnly && pending.length) throw cacheMiss();

  let loaded = cached;
  onProgress({ phase: 'download', loaded, total: MODEL_TOTAL_BYTES, cached });

  if (pending.length) {
    const estimate = await navigator.storage.estimate();
    if (estimate.quota && estimate.quota - (estimate.usage ?? 0) < MODEL_TOTAL_BYTES - loaded + CHUNK) {
      throw new Error('Not enough browser storage for the model. Free at least 2 GB and retry.');
    }
  }

  const internal = new AbortController();
  const combined = AbortSignal.any([signal, internal.signal]);

  // Ghi biên nhận tuần tự để tránh hai lần ghi chồng nhau.
  let receiptChain: Promise<void> = Promise.resolve();
  const saveReceipts = () => {
    const json = JSON.stringify(receipts);
    const next = receiptChain.then(async () => {
      const w = await receiptsFile.createWritable();
      await w.write(json);
      await w.close();
    });
    receiptChain = next.catch(() => {});
    return next;
  };

  async function worker() {
    while (pending.length) {
      combined.throwIfAborted();
      const [name, info] = pending.shift()!;
      const handle = await root.getFileHandle(encode(name), { create: true });
      const writable = await handle.createWritable();
      try {
        const hasher = await createSHA256();
        hasher.init();
        for (let start = 0; start < info.bytes; start += CHUNK) {
          const end = Math.min(start + CHUNK, info.bytes) - 1;
          const chunk = await fetchRange(baseUrl + name, start, end, info.bytes, combined);
          combined.throwIfAborted();
          hasher.update(chunk);
          await writable.write(chunk);
          loaded += chunk.length;
          onProgress({ phase: 'download', file: name, loaded, total: MODEL_TOTAL_BYTES, cached });
        }
        if (hasher.digest('hex') !== info.sha256) throw new Error('SHA-256 verification failed: ' + name);
        await writable.close();
        receipts[name] = { sha256: info.sha256, modified: (await handle.getFile()).lastModified };
        await saveReceipts();
      } catch (err) {
        await writable.abort().catch(() => {});
        internal.abort(err);
        throw err;
      }
    }
  }

  const results = await Promise.allSettled(Array.from({ length: PARALLEL_DOWNLOADS }, worker));
  const failed = results.find((r): r is PromiseRejectedResult => r.status === 'rejected');
  if (failed) throw failed.reason;
  await receiptChain;

  const handle: ModelCacheHandle = {
    cachedBytes: cached,
    async match(request) {
      const url = typeof request === 'string' ? request : request.url;
      if (!url.startsWith(baseUrl)) return undefined;
      const name = url.slice(baseUrl.length);
      if (!Object.hasOwn(MODEL_FILES, name)) return undefined;
      const file = await (await root.getFileHandle(encode(name))).getFile();
      return new Response(file, { headers: { 'Content-Length': String(file.size) } });
    },
    async put() {
      /* chỉ đọc: file được ghi bởi bộ tải ở trên */
    },
  };
  return handle;
}

/** Mở (và nếu cần thì tải) cache model. Dùng khoá để chỉ một tab tải tại một thời điểm. */
export async function openModelCache(baseUrl: string, options: OpenOptions): Promise<ModelCacheHandle> {
  if (!navigator.locks) return openInner(baseUrl, options);
  return navigator.locks.request('pi-minicpm5-download-' + MODEL_REVISION, { signal: options.signal }, () =>
    openInner(baseUrl, options),
  ) as Promise<ModelCacheHandle>;
}