au_whisper_web/src/lib/modelCache.ts

/**
 * "Is this model already downloaded?" helper (Decisions.md D-005).
 *
 * Transformers.js v3 persists Hub downloads in the browser Cache API (Cache
 * Storage, cache name "transformers-cache"). We inspect the cached request
 * URLs and match them against catalog model ids. Used by `ModelPicker` for the
 * downloaded badge (M2) and reused by the offline indicator (M3).
 *
 * `matchCachedModels` is pure so it can be unit-tested; only
 * `getDownloadedModels` touches the `caches` global.
 */

/** Cache name used by Transformers.js v3 (`env.cacheName` default). */
export const TRANSFORMERS_CACHE_NAME = "transformers-cache";

/**
 * Pure matcher: which of `modelIds` have at least one cached ONNX weight file?
 *
 * Cached URLs look like
 * `https://huggingface.co/<modelId>/resolve/main/onnx/<file>.onnx`. We require
 * an `.onnx` hit (not just tokenizer/config JSON) so a model whose weights
 * never finished downloading does not show as available.
 */
export function matchCachedModels(
  urls: Iterable<string>,
  modelIds: readonly string[],
): Set<string> {
  const downloaded = new Set<string>();
  for (const url of urls) {
    if (!url.split("?")[0]?.endsWith(".onnx")) continue;
    for (const modelId of modelIds) {
      if (!downloaded.has(modelId) && url.includes(`/${modelId}/`)) {
        downloaded.add(modelId);
      }
    }
  }
  return downloaded;
}

/**
 * Query the browser Cache API for models with cached weights. Returns an
 * empty set where Cache Storage is unavailable (non-secure contexts, tests).
 */
export async function getDownloadedModels(
  modelIds: readonly string[],
): Promise<Set<string>> {
  if (typeof caches === "undefined") return new Set();
  try {
    const urls: string[] = [];
    // Match any transformers-owned cache, in case the name is ever versioned.
    for (const name of await caches.keys()) {
      if (!name.includes("transformers")) continue;
      const cache = await caches.open(name);
      for (const request of await cache.keys()) urls.push(request.url);
    }
    return matchCachedModels(urls, modelIds);
  } catch {
    return new Set();
  }
}

// ---------------------------------------------------------------------------
// Storage management (M6): per-model size + delete.
//
// The functions below take an injectable Cache-Storage-shaped parameter so the
// logic is unit-testable under `bun test` (no `caches` global there); the
// structural types cover exactly the subset of the Cache API we touch, and the
// real `caches` satisfies them. Scope is the transformers-cache only — the
// service worker's `whisper-runtime-*` cache (onnxruntime .wasm, shared by all
// models and versioned/cleaned by the SW itself) is deliberately not surfaced.
// ---------------------------------------------------------------------------

/** The slice of `Cache` the storage helpers use. */
export interface StorageCache {
  keys(): Promise<ReadonlyArray<{ url: string }>>;
  match(
    request: { url: string },
  ): Promise<{ blob(): Promise<{ size: number }> } | undefined>;
  delete(request: { url: string }): Promise<boolean>;
}

/** The slice of `CacheStorage` the storage helpers use. */
export interface StorageCacheStorage {
  keys(): Promise<string[]>;
  open(name: string): Promise<StorageCache>;
}

/** Per-model cache usage, for the storage management panel. */
export interface ModelStorageInfo {
  modelId: string;
  /** Number of cached files (weights + tokenizer/config aux files). */
  fileCount: number;
  /** Exact stored bytes, summed from the cached response bodies. */
  bytes: number;
}

/**
 * Pure matcher: which catalog model owns a cached URL? Same segment-boundary
 * rule as `matchCachedModels` (`/<modelId>/`), so `.en` variants never match
 * their multilingual sibling. Any file counts here — a partially downloaded
 * model still occupies storage and must be listable/deletable.
 */
export function modelIdForUrl(
  url: string,
  modelIds: readonly string[],
): string | undefined {
  return modelIds.find((modelId) => url.includes(`/${modelId}/`));
}

function defaultCacheStorage(): StorageCacheStorage | undefined {
  // The cast is sound: we only ever hand `match`/`delete` the `Request`
  // objects that `keys()` itself returned (lib.dom types the params wider).
  return typeof caches === "undefined"
    ? undefined
    : (caches as unknown as StorageCacheStorage);
}

/**
 * Measure cached storage per model: file count and exact byte size (summed
 * `blob().size` of every cached response). Models with nothing cached are
 * omitted. Deliberately lazy/heavier than `getDownloadedModels` — it reads
 * every cached body, so call it when the user opens the storage panel, not on
 * every render. Returns [] where Cache Storage is unavailable.
 */
export async function getModelStorageInfo(
  modelIds: readonly string[],
  cacheStorage: StorageCacheStorage | undefined = defaultCacheStorage(),
): Promise<ModelStorageInfo[]> {
  if (!cacheStorage) return [];
  try {
    const byModel = new Map<string, ModelStorageInfo>();
    for (const name of await cacheStorage.keys()) {
      if (!name.includes("transformers")) continue;
      const cache = await cacheStorage.open(name);
      for (const request of await cache.keys()) {
        const modelId = modelIdForUrl(request.url, modelIds);
        if (modelId === undefined) continue;
        const response = await cache.match(request);
        if (!response) continue;
        const size = (await response.blob()).size;
        const info = byModel.get(modelId) ?? { modelId, fileCount: 0, bytes: 0 };
        info.fileCount += 1;
        info.bytes += size;
        byModel.set(modelId, info);
      }
    }
    // Catalog order, so the panel lists models the same way the picker does.
    return modelIds
      .map((modelId) => byModel.get(modelId))
      .filter((info): info is ModelStorageInfo => info !== undefined);
  } catch {
    return [];
  }
}

/**
 * Delete every cached file belonging to `modelId` (weights and aux files)
 * from the transformers cache(s). Returns true if anything was deleted. The
 * weights re-download on the model's next use — nothing but cache lives here.
 */
export async function deleteModelFromCache(
  modelId: string,
  cacheStorage: StorageCacheStorage | undefined = defaultCacheStorage(),
): Promise<boolean> {
  if (!cacheStorage) return false;
  try {
    let deleted = false;
    for (const name of await cacheStorage.keys()) {
      if (!name.includes("transformers")) continue;
      const cache = await cacheStorage.open(name);
      for (const request of await cache.keys()) {
        if (modelIdForUrl(request.url, [modelId]) !== modelId) continue;
        if (await cache.delete(request)) deleted = true;
      }
    }
    return deleted;
  } catch {
    return false;
  }
}