From 827e9a8daec4fd6ad4fb192bfe8e913bb4bc44b5 Mon Sep 17 00:00:00 2001 From: Lucas Winther Date: Thu, 29 Oct 2026 21:50:40 +0100 Subject: [PATCH] Add a Cache-API model file downloader that reports load progress and caches only complete files Co-Authored-By: Claude Opus 5.5 --- src/worker/modelFiles.ts | 121 +++++++++++++++++++++++++++++++++++++++ tests/modelFiles.test.ts | 121 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 242 insertions(+) create mode 100644 src/worker/modelFiles.ts create mode 100644 tests/modelFiles.test.ts diff --git a/src/worker/modelFiles.ts b/src/worker/modelFiles.ts new file mode 100644 index 0000000..72a1109 --- /dev/null +++ b/src/worker/modelFiles.ts @@ -0,0 +1,121 @@ +/** + * Weight-file download + cache for engine paths that don't go through + * Transformers.js's own loader (Parakeet, D-WS7a). + * + * Files land in the same Cache Storage cache Transformers.js uses + * (`transformers-cache`, D-005), keyed by their Hub URL, so the downloaded + * badge, the storage panel, and delete (src/lib/modelCache.ts) work for them + * unchanged. Progress is reported as the same `load-progress` messages the + * Whisper path forwards, so ModelLoadProgress needs no second code path. + * + * A file is `put` only once its whole body has arrived: a cancelled or + * failed download (worker terminated, network drop) leaves nothing behind + * that could later show as "available offline". + * + * Dependencies are injected so `bun test` runs this without Cache Storage. + */ + +import { TRANSFORMERS_CACHE_NAME } from "../lib/modelCache"; +import type { LoadProgressMessage } from "./protocol"; + +/** The slice of `Cache` used here. */ +export interface ModelFileCache { + match(url: string): Promise; + put(url: string, response: Response): Promise; +} + +/** The slice of `CacheStorage` used here. */ +export interface ModelFileCacheStorage { + open(name: string): Promise; +} + +export interface FetchModelFileDeps { + /** Cache Storage, or undefined where it doesn't exist (download uncached). */ + caches: ModelFileCacheStorage | undefined; + fetch: (url: string) => Promise; + onProgress: (message: LoadProgressMessage) => void; +} + +/** Fetch one model file, from cache when present, else download and cache it. */ +export async function fetchModelFile( + url: string, + file: string, + deps: FetchModelFileDeps, +): Promise { + deps.onProgress({ type: "load-progress", status: "initiate", file }); + const cache = deps.caches ? await deps.caches.open(TRANSFORMERS_CACHE_NAME) : undefined; + + const cached = await cache?.match(url); + if (cached) { + const bytes = new Uint8Array(await cached.arrayBuffer()); + deps.onProgress({ type: "load-progress", status: "done", file }); + return bytes; + } + + deps.onProgress({ type: "load-progress", status: "download", file }); + const response = await deps.fetch(url); + if (!response.ok || !response.body) { + throw new Error(`Could not download ${file} (HTTP ${response.status})`); + } + const total = Number(response.headers.get("Content-Length")) || 0; + const bytes = await readWithProgress(response.body, total, (loaded) => { + const size = total || loaded; + deps.onProgress({ + type: "load-progress", + status: "progress", + file, + loaded, + total: size, + progress: (loaded / size) * 100, + }); + }); + + if (cache) { + try { + await cache.put( + url, + new Response(bytes as Uint8Array, { + headers: { "Content-Length": String(bytes.length) }, + }), + ); + } catch { + // Quota exceeded / private mode: the model still runs this session, it + // just won't be available offline. + } + } + deps.onProgress({ type: "load-progress", status: "done", file }); + return bytes; +} + +/** + * Read a body into one buffer — preallocated when the length is known, so a + * 650 MB encoder isn't also held as a list of chunks. Reports progress at + * most once per whole percent (a large file arrives in thousands of reads). + */ +async function readWithProgress( + body: ReadableStream, + total: number, + onLoaded: (loaded: number) => void, +): Promise { + let buffer = new Uint8Array(total); + let loaded = 0; + let lastPercent = -1; + const reader = body.getReader(); + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + if (loaded + value.length > buffer.length) { + const grown = new Uint8Array(Math.max(buffer.length * 2, loaded + value.length)); + grown.set(buffer.subarray(0, loaded)); + buffer = grown; + } + buffer.set(value, loaded); + loaded += value.length; + const percent = total ? Math.floor((loaded / total) * 100) : loaded; + if (percent !== lastPercent) { + lastPercent = percent; + onLoaded(loaded); + } + } + return loaded === buffer.length ? buffer : buffer.slice(0, loaded); +} diff --git a/tests/modelFiles.test.ts b/tests/modelFiles.test.ts new file mode 100644 index 0000000..1b3bfcb --- /dev/null +++ b/tests/modelFiles.test.ts @@ -0,0 +1,121 @@ +import { describe, expect, test } from "bun:test"; +import { fetchModelFile, type ModelFileCacheStorage } from "../src/worker/modelFiles"; +import type { LoadProgressMessage } from "../src/worker/protocol"; + +// Cache-API-backed weight download for the Parakeet path (D-WS7a, D-005). + +const URL_A = "https://huggingface.co/org/model/resolve/rev/encoder.onnx"; + +function fakeCaches(initial: Record = {}) { + const stored = new Map(Object.entries(initial)); + const opened: string[] = []; + const storage: ModelFileCacheStorage = { + async open(name) { + opened.push(name); + return { + async match(url) { + const bytes = stored.get(url); + return bytes ? new Response(bytes as Uint8Array) : undefined; + }, + async put(url, response) { + stored.set(url, new Uint8Array(await response.arrayBuffer())); + }, + }; + }, + }; + return { storage, stored, opened }; +} + +function streamingFetch(bytes: Uint8Array, pieces: number, status = 200) { + const calls: string[] = []; + const fetchFn = async (url: string) => { + calls.push(url); + if (status !== 200) return new Response("nope", { status }); + const size = Math.ceil(bytes.length / pieces); + const body = new ReadableStream({ + start(controller) { + for (let i = 0; i < bytes.length; i += size) controller.enqueue(bytes.slice(i, i + size)); + controller.close(); + }, + }); + return new Response(body, { headers: { "Content-Length": String(bytes.length) } }); + }; + return { fetchFn, calls }; +} + +describe("fetchModelFile", () => { + test("a cached file is returned without touching the network", async () => { + const cached = new Uint8Array([1, 2, 3]); + const { storage, opened } = fakeCaches({ [URL_A]: cached }); + const { fetchFn, calls } = streamingFetch(new Uint8Array(), 1); + const events: LoadProgressMessage[] = []; + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: storage, + fetch: fetchFn, + onProgress: (event) => events.push(event), + }); + expect([...bytes]).toEqual([1, 2, 3]); + expect(calls).toEqual([]); + expect(opened).toEqual(["transformers-cache"]); + expect(events.map((event) => event.status)).toEqual(["initiate", "done"]); + }); + + test("a miss downloads with byte progress and stores the complete file", async () => { + const payload = new Uint8Array(1000).map((_, i) => i % 251); + const { storage, stored } = fakeCaches(); + const { fetchFn, calls } = streamingFetch(payload, 4); + const events: LoadProgressMessage[] = []; + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: storage, + fetch: fetchFn, + onProgress: (event) => events.push(event), + }); + expect(calls).toEqual([URL_A]); + expect([...bytes]).toEqual([...payload]); + expect([...stored.get(URL_A)!]).toEqual([...payload]); + const progress = events.filter((event) => event.status === "progress"); + expect(progress.length).toBe(4); + const last = progress[progress.length - 1]!; + expect(last).toMatchObject({ file: "encoder.onnx", loaded: 1000, total: 1000, progress: 100 }); + expect(events[0]!.status).toBe("initiate"); + expect(events[events.length - 1]!.status).toBe("done"); + }); + + test("an HTTP error throws a plain message and stores nothing", async () => { + const { storage, stored } = fakeCaches(); + const { fetchFn } = streamingFetch(new Uint8Array(), 1, 404); + await expect( + fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }), + ).rejects.toThrow("Could not download encoder.onnx (HTTP 404)"); + expect(stored.size).toBe(0); + }); + + test("an interrupted download stores nothing (no false 'available offline')", async () => { + const { storage, stored } = fakeCaches(); + const fetchFn = async () => + new Response( + new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(10)); + controller.error(new Error("network down")); + }, + }), + { headers: { "Content-Length": "100" } }, + ); + await expect( + fetchModelFile(URL_A, "encoder.onnx", { caches: storage, fetch: fetchFn, onProgress: () => {} }), + ).rejects.toThrow(); + expect(stored.size).toBe(0); + }); + + test("without Cache Storage it still downloads", async () => { + const { fetchFn, calls } = streamingFetch(new Uint8Array([9]), 1); + const bytes = await fetchModelFile(URL_A, "encoder.onnx", { + caches: undefined, + fetch: fetchFn, + onProgress: () => {}, + }); + expect([...bytes]).toEqual([9]); + expect(calls).toEqual([URL_A]); + }); +});