diff --git a/.cursor/rules/model-pipeline.mdc b/.cursor/rules/model-pipeline.mdc index 5139a4b..280f85e 100644 --- a/.cursor/rules/model-pipeline.mdc +++ b/.cursor/rules/model-pipeline.mdc @@ -17,3 +17,12 @@ or sampled. The symptom in front of you probably has a recorded cause; check its - Every answer goes through `reviewAnswer` before it is returned, so one turn can generate twice. The checks stay deterministic and the correction states the fix; asking the model to grade itself is the thing that does not work. +- `MODEL_ID` and `MODEL_DTYPE` are the smallest way to run this model at all, not a default nobody + measured: `node tools/verify-model.mjs` surveys every published variant and fails if that stops + being true. Do not swap in the multimodal export to gain anything but a vision encoder. +- The OPFS cache performs the download so it can resume one. Three parts of that are load bearing: + writes go through a sync access handle because a writable stream discards a partial, the `ETag` is + compared here because the Hub's CDN ignores `If-Range`, and `match` returns only a finished file — + Transformers.js also calls it as an existence check and drops the body, which strands the lock. +- Install progress comes from `setDownloadProgress`, keyed to the same file names Transformers.js + reports. Two keys for one file and the bar never reaches the end. diff --git a/.cursor/skills/verify-in-browser/SKILL.md b/.cursor/skills/verify-in-browser/SKILL.md index f83c709..dfa4d72 100644 --- a/.cursor/skills/verify-in-browser/SKILL.md +++ b/.cursor/skills/verify-in-browser/SKILL.md @@ -60,12 +60,27 @@ which `onnx/model_q4f16.onnx_data` is 448 MB. Expect roughly four minutes on a f a second on a second visit. The gate screen is the fastest read on state: it shows whether the model is installed, how much -space it occupies, whether storage is persistent, and offers **Remove model**. +space it occupies, whether storage is persistent, and offers **Remove model**. It reads `.onnx_data` +as the test of "installed", so a run that fetched only the small files still shows as not installed. + +To exercise the resume path, install with DevTools → Network → Offline switched on part way through, +or kill the tab mid-download. The cache retries three times on its own first, so leave it offline +long enough to see it give up. The gate then reads **partly downloaded** with a **Resume install** +button, `model_q4f16.onnx_data.part` and a `.part-meta` sidecar are in OPFS, and the next attempt +sends `Range: bytes=-` — visible in the Network panel as a `206`. Restarting +from zero instead means the sidecar's `ETag` no longer matched, which is the correct response to the +weights having changed upstream and worth confirming before calling it a bug. + +The progress bar during an install is reported by the cache, not by Transformers.js, because the +library is handed a file that is already on disk. A bar that stalls at half while bytes are clearly +arriving means the URL-to-filename mapping in `worker.ts` has drifted and one file is being counted +twice. To inspect the cache directly, open DevTools → Application → Storage. Weights live in the Origin Private File System under `model-cache/`, with filenames that are the download URL flattened by `cacheKeyFor` (`huggingface.co_onnx-community_…_model_q4f16.onnx`). A name ending in `.part` is an -in-flight or abandoned download and is deliberately invisible to `listCachedFiles`. +in-flight or abandoned download and is deliberately invisible to `listCachedFiles`; the matching +`.part-meta` holds the `ETag` and total size a resume is checked against. To retest an install, use **Remove model**. Clearing site data also drops the persistence grant, which is sometimes what you want to test and usually not. diff --git a/README.md b/README.md index 481f7a8..e9f9f1c 100644 --- a/README.md +++ b/README.md @@ -28,14 +28,28 @@ The model emits reasoning inside `` blocks and tool requests as JSON insi ## Installing the model -The model is **448 MB** and downloads once. Two things make it stick: +The model is **448 MB** and downloads once. Three things make it stick: 1. **Persistent storage.** Before downloading, the app calls `navigator.storage.persist()`. Without that grant the browser treats the weights as best-effort data and may evict them under storage pressure — turning a one-time download into a recurring one. Chrome grants persistence silently for installed PWAs and sufficiently engaged sites. 2. **OPFS instead of the Cache API.** Transformers.js caches downloads in the Cache API by default, but Chrome rejects the 448 MB weights file there with `Failed to execute 'put' on 'Cache': Unexpected internal error` — the download completed and then quietly vanished, so every visit re-fetched it. `src/llm/opfs-cache.ts` replaces that backend with the Origin Private File System, which is built for large binaries and streams them to disk. Downloads land under a `.part` name and are renamed only once complete, so an interrupted install can never be mistaken for a finished one. +3. **A dropped connection costs only what was left.** The `.part` file is kept, and the next attempt continues from it with `Range: bytes=-` instead of starting the 448 MB again — see [resuming an interrupted download](#resuming-an-interrupted-download). A second visit then reaches the chat in about a second instead of four minutes. -The gate screen shows whether the model is installed, how much space it occupies, whether storage is persistent, and offers a **Remove model** button to reclaim the space. +The gate screen shows whether the model is installed, how much space it occupies, whether storage is persistent, and offers a **Remove model** button to reclaim the space. A half-finished download is reported as such — `312 MB of 467 MB saved` — with a **Resume install** button, rather than counted as installed because a few of the seven files arrived. + +### Resuming an interrupted download + +Transformers.js reads a whole response into memory before handing it to a cache, so a `put`-side cache never sees a failure: at 400 MB of 448 MB there is nothing to hand over and nothing on disk. Resuming therefore has to own the fetch, and `opfsCache.match` does — it downloads the file into `.part`, retries the transfer up to three times from wherever it stopped, publishes it under the real name, and only then answers with the file. Upstream shipped the same idea for Node's filesystem cache in [transformers.js#1715](https://github.com/huggingface/transformers.js/pull/1715); the browser half is [still open](https://github.com/huggingface/transformers.js/issues/1220). + +Four details are what make it work rather than merely sound good: + +- **A partial is written through a sync access handle.** A `FileSystemWritableFileStream` buffers into a swap file that is discarded unless it is closed cleanly, so the old code's `.part` file was always empty after a failure. `createSyncAccessHandle()` writes straight to the file, and is available because the download runs in a Web Worker. +- **The entity tag is checked here, not by the server.** The natural mechanism is `If-Range`, and the Hub's CDN ignores it: a stale validator still comes back `206` with the old byte range, which would splice two different files together. So the `ETag` and total size are recorded next to the partial and compared on the next attempt; anything that does not match starts the file again. +- **A body that stops short is an unfinished download, not a shorter file.** Transformers.js sizes its buffer from `Content-Length` and zero-pads whatever never arrived, which would publish silently corrupt weights. A transfer that ends before the declared total is retried instead. +- **The download finishes before the response is returned.** Handing back a streaming body looked neater and was wrong: Transformers.js also calls `match` to ask whether a file exists and how big it is, and drops the body when it does — which left the OPFS write lock held by a reader that was never going to read. Progress therefore comes from the cache itself, reported per megabyte and translated into file names by the worker. + +`node tools/verify-model.mjs` checks the three things this needs from the host: byte ranges, a `206` that states the file's total size, and an `ETag` that CORS actually lets a script read. Installing Jarvis as a PWA (the install icon in Chrome's address bar) is what makes offline use reliable, because installed apps get persistent storage automatically. @@ -47,7 +61,7 @@ By default the weights are fetched from the Hugging Face Hub. It works well as a - **No account, no token.** The repository is public and ungated. - **Any origin may fetch it.** The Hub reflects the requesting `Origin` back in `Access-Control-Allow-Origin`, so a browser on any domain can download directly. -- **Byte ranges are supported** (`Accept-Ranges: bytes`), so downloads can resume. +- **Byte ranges are supported** (`Accept-Ranges: bytes`) and the `ETag` is exposed to scripts, which is what an interrupted download needs to continue. - **Rate limits are not a concern here.** Anonymous clients get 3,000 file requests per five minutes per IP address; one installation needs seven. - **Licensing permits redistribution.** The base model `Qwen/Qwen3.5-0.8B` is Apache-2.0, so you may mirror the weights as long as you keep the licence and attribution. @@ -75,7 +89,7 @@ Copy these seven files, keeping the `onnx/` subdirectory: | `chat_template.jinja` | < 1 MB | | **Total** | **467 MB** | -The host must send `Access-Control-Allow-Origin` for your domain and should support range requests. Those are the two things `node tools/verify-model.mjs` checks, and it reads the same two variables, so point them at your mirror and run it before deploying. **Cloudflare R2 fits well**: 467 MB sits inside the 10 GB free tier, egress is free at any volume, and a public bucket on a custom domain gives you a CDN with configurable CORS. Uploading to your own Hugging Face repository works too and takes minutes. +The host must send `Access-Control-Allow-Origin` for your domain, and should serve byte ranges with a script-readable `ETag` so an interrupted install can resume. Those are the things `node tools/verify-model.mjs` checks, and it reads the same two variables, so point them at your mirror and run it before deploying. **Cloudflare R2 fits well**: 467 MB sits inside the 10 GB free tier, egress is free at any volume, and a public bucket on a custom domain gives you a CDN with configurable CORS. Uploading to your own Hugging Face repository works too and takes minutes. GitHub Releases will not work. Release assets are served with `Access-Control-Allow-Origin: https://render.githubusercontent.com`, so a browser cannot read them. @@ -394,7 +408,23 @@ knowing even if the skill is never opened. Skills hold the detail, rules decide ## Notes on the model -`onnx-community/Qwen3.5-0.8B-Text-ONNX` is the text-only export, loaded through the standard `text-generation` pipeline with `dtype: 'q4f16'`. The multimodal build of the same model also exists, but it ships a vision encoder this app never feeds, requires the dedicated `Qwen3_5ForConditionalGeneration` class, and downloads roughly 150 MB more. +`onnx-community/Qwen3.5-0.8B-Text-ONNX` is the text-only export, loaded through the standard `text-generation` pipeline with `dtype: 'q4f16'` — which is what Transformers.js added text-only Qwen3.5 support for ([transformers.js#1602](https://github.com/huggingface/transformers.js/pull/1602)). + +### Why 448 MB is the floor + +The download is the largest thing this app asks of anyone, so it is worth saying plainly what the alternatives cost. Within Transformers.js and Qwen3.5-0.8B there are none. + +| Export | q4f16 total | Loads through | +| --------------------------------------- | ----------: | -------------------------------------------------- | +| `onnx-community/Qwen3.5-0.8B-Text-ONNX` | **448 MiB** | `pipeline('text-generation')` | +| `onnx-community/Qwen3.5-0.8B-ONNX-OPT` | 616 MiB | `Qwen3_5ForConditionalGeneration` + vision encoder | +| `onnx-community/Qwen3.5-0.8B-ONNX` | 617 MiB | `Qwen3_5ForConditionalGeneration` + vision encoder | + +Those are the only first-party ONNX conversions; every other Qwen3.5-0.8B ONNX repository on the Hub is a copy of one of them or larger. Within the export in use, `q4f16` is the smallest of the five variants published — `q4` is 526 MiB, int8 896 MiB, fp16 1.4 GiB, fp32 2.9 GiB — and INT4 is as far as ONNX Runtime Web's WebGPU backend goes. + +Two facts about the model account for the rest. Qwen3.5 has no size below 0.8B: the small series is 0.8B, 2B, 4B and 9B, so there is no smaller sibling to fall back to. And 0.8B parameters at four bits would be nearer 400 MB were it not for a 248,320-token vocabulary tied to the output layer, which is a third of the weights on its own. + +`node tools/verify-model.mjs` re-measures all of this against the Hub and fails if a smaller variant appears, so the claim above is checked rather than remembered. Tool definitions are passed straight to the pipeline via its `tools` option, added in Transformers.js v4.2, so the chat template renders them itself rather than us hand-assembling a prompt. diff --git a/src/components/ModelGate.test.tsx b/src/components/ModelGate.test.tsx index 28dccab..af9fe4e 100644 --- a/src/components/ModelGate.test.tsx +++ b/src/components/ModelGate.test.tsx @@ -65,6 +65,16 @@ describe('ModelGate', () => { expect(await screen.findByText('There may not be room for the download')).toBeInTheDocument() }) + it('offers to continue a download that stopped part way through', async () => { + stubStorage(storage({ modelCached: false, partialBytes: 300 * 1024 ** 2 })) + render({null}) + + expect(await screen.findByText('partly downloaded')).toBeInTheDocument() + expect(screen.getByText(/300 MB of 467 MB saved/)).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Resume install (167 MB left)' })).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Discard download' })).toBeInTheDocument() + }) + it('says nothing about room when the model is already installed', async () => { stubStorage(storage({ modelCached: true, quotaBytes: 1024 ** 3, usageBytes: 900 * 1024 ** 2 })) render({null}) diff --git a/src/components/ModelGate.tsx b/src/components/ModelGate.tsx index 8174091..599edc7 100644 --- a/src/components/ModelGate.tsx +++ b/src/components/ModelGate.tsx @@ -36,6 +36,10 @@ export function ModelGate({ children }: { children: ReactNode }) { // Better to say the download will not fit than to spend ten minutes finding out. const freeBytes = storage.quotaBytes - storage.usageBytes const tooLittleRoom = !installed && !hasRoomFor(storage, MODEL_DOWNLOAD_BYTES) + // An earlier attempt that died part way through is not lost work: the next one + // continues from it, so the gate offers to resume rather than to start again. + const resumeBytes = installed ? 0 : storage.partialBytes + const remainingBytes = Math.max(MODEL_DOWNLOAD_BYTES - resumeBytes, 0) return (
@@ -89,6 +93,16 @@ export function ModelGate({ children }: { children: ReactNode }) { {formatBytes(storage.modelBytes)} on disk )} + ) : resumeBytes > 0 ? ( + <> + + partly downloaded + + + {formatBytes(resumeBytes)} of {formatBytes(MODEL_DOWNLOAD_BYTES)} saved — the rest + picks up where it stopped + + ) : ( <> not installed @@ -151,15 +165,17 @@ export function ModelGate({ children }: { children: ReactNode }) {
- {installed && ( + {(installed || resumeBytes > 0) && ( )}
@@ -184,7 +200,8 @@ export function ModelGate({ children }: { children: ReactNode }) { )}

- Downloading only happens once. Afterwards the model is served from this browser. + Downloading only happens once. Afterwards the model is served from this browser, and a + transfer that is interrupted continues from where it stopped rather than starting again.

)} diff --git a/src/lib/storage.test.ts b/src/lib/storage.test.ts index 3406aa0..f3ec52d 100644 --- a/src/lib/storage.test.ts +++ b/src/lib/storage.test.ts @@ -2,6 +2,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { deleteModel, getStorageStatus, hasRoomFor, requestPersistence, type StorageStatus } from './storage' const MODEL = 'onnx-community/Qwen3.5-0.8B-Text-ONNX' +const WEIGHTS = 'onnx/model_q4f16.onnx_data' +const PREFIX = `huggingface.co_${MODEL.replace(/[^a-zA-Z0-9._-]/g, '_')}_resolve_main_` /** Stand-in for the OPFS directory handle; jsdom implements neither OPFS nor StorageManager. */ function fakeOpfs(files: Record) { @@ -43,8 +45,8 @@ let opfs: ReturnType beforeEach(() => { // Filenames mirror how the OPFS cache flattens a download URL. opfs = fakeOpfs({ - [`huggingface.co_${MODEL.replace(/[^a-zA-Z0-9._-]/g, '_')}_resolve_main_onnx_model_q4f16.onnx`]: 400, - [`huggingface.co_${MODEL.replace(/[^a-zA-Z0-9._-]/g, '_')}_resolve_main_tokenizer.json`]: 48, + [`${PREFIX}onnx_model_q4f16.onnx_data`]: 400, + [`${PREFIX}tokenizer.json`]: 48, 'huggingface.co_some-other-model_resolve_main_model.onnx': 999, }) stubNavigatorStorage(opfs) @@ -77,13 +79,13 @@ describe('requestPersistence', () => { describe('getStorageStatus', () => { it('counts only the files belonging to the requested model', async () => { - const status = await getStorageStatus(MODEL) + const status = await getStorageStatus(MODEL, WEIGHTS) expect(status.modelCached).toBe(true) expect(status.modelBytes).toBe(448) }) it('reports quota usage from the storage estimate', async () => { - const status = await getStorageStatus(MODEL) + const status = await getStorageStatus(MODEL, WEIGHTS) expect(status.usageBytes).toBe(500_000_000) expect(status.quotaBytes).toBe(2_000_000_000) }) @@ -91,9 +93,31 @@ describe('getStorageStatus', () => { it('reports an uninstalled model when nothing is cached', async () => { const empty = fakeOpfs({}) stubNavigatorStorage(empty) - const status = await getStorageStatus(MODEL) + const status = await getStorageStatus(MODEL, WEIGHTS) expect(status.modelCached).toBe(false) expect(status.modelBytes).toBe(0) + expect(status.partialBytes).toBe(0) + }) + + it('is not installed on the strength of the small files alone', async () => { + stubNavigatorStorage(fakeOpfs({ [`${PREFIX}tokenizer.json`]: 48 })) + const status = await getStorageStatus(MODEL, WEIGHTS) + expect(status.modelCached).toBe(false) + }) + + it('reports an unfinished download as bytes the next attempt starts from', async () => { + stubNavigatorStorage( + fakeOpfs({ + [`${PREFIX}tokenizer.json`]: 48, + [`${PREFIX}onnx_model_q4f16.onnx_data.part`]: 300, + [`${PREFIX}onnx_model_q4f16.onnx_data.part-meta`]: 60, + }), + ) + const status = await getStorageStatus(MODEL, WEIGHTS) + expect(status.modelCached).toBe(false) + expect(status.partialBytes).toBe(300) + // Nothing half-written is counted as installed weight. + expect(status.modelBytes).toBe(48) }) }) @@ -102,6 +126,7 @@ describe('hasRoomFor', () => { persisted: false, modelCached: false, modelBytes: 0, + partialBytes: 0, usageBytes, quotaBytes, }) @@ -125,4 +150,14 @@ describe('deleteModel', () => { await deleteModel(MODEL) expect([...opfs.store.keys()]).toEqual(['huggingface.co_some-other-model_resolve_main_model.onnx']) }) + + it('reclaims an unfinished download as well as the installed files', async () => { + const withPartial = fakeOpfs({ + [`${PREFIX}onnx_model_q4f16.onnx_data.part`]: 300, + [`${PREFIX}onnx_model_q4f16.onnx_data.part-meta`]: 60, + }) + stubNavigatorStorage(withPartial) + await deleteModel(MODEL) + expect([...withPartial.store.keys()]).toEqual([]) + }) }) diff --git a/src/lib/storage.ts b/src/lib/storage.ts index ff4719b..3c78323 100644 --- a/src/lib/storage.ts +++ b/src/lib/storage.ts @@ -1,4 +1,4 @@ -import { cacheKeyFor, clearCachedFiles, listCachedFiles } from '@/llm/opfs-cache' +import { cacheKeyFor, clearCachedFiles, listCachedFiles, listPartialFiles } from '@/llm/opfs-cache' /** * Model weights live in the Origin Private File System (see `llm/opfs-cache.ts`). @@ -12,9 +12,11 @@ import { cacheKeyFor, clearCachedFiles, listCachedFiles } from '@/llm/opfs-cache export interface StorageStatus { /** Whether the browser promised not to evict this origin's data. */ persisted: boolean - /** True once the model's files are on disk. */ + /** True once the weights are on disk, not merely some of the model's files. */ modelCached: boolean modelBytes: number + /** Bytes of an unfinished download the next attempt will continue from. */ + partialBytes: number usageBytes: number quotaBytes: number } @@ -23,6 +25,7 @@ export const EMPTY_STORAGE_STATUS: StorageStatus = { persisted: false, modelCached: false, modelBytes: 0, + partialBytes: 0, usageBytes: 0, quotaBytes: 0, } @@ -61,23 +64,35 @@ function modelFilePrefix(modelId: string): string { return cacheKeyFor(modelId) } -export async function getStorageStatus(modelId: string): Promise { +/** + * What is on disk for `modelId`. + * + * `weightsFile` decides what counts as installed. Any one of the model's seven + * files being present is not enough: a run that fetched the tokenizer and then + * lost the connection would report itself installed, and pressing Start would + * quietly begin the 448 MB download again. + */ +export async function getStorageStatus(modelId: string, weightsFile: string): Promise { if (!storageApiAvailable()) return EMPTY_STORAGE_STATUS try { - const [persisted, estimate, files] = await Promise.all([ + const [persisted, estimate, files, partials] = await Promise.all([ navigator.storage.persisted?.() ?? Promise.resolve(false), navigator.storage.estimate?.() ?? Promise.resolve({}), listCachedFiles(), + listPartialFiles(), ]) const needle = modelFilePrefix(modelId) - const modelFiles = files.filter((file) => file.name.includes(needle)) + const belongs = (name: string): boolean => name.includes(needle) + const modelFiles = files.filter((file) => belongs(file.name)) + const weightsKey = cacheKeyFor(weightsFile) return { persisted, - modelCached: modelFiles.length > 0, + modelCached: modelFiles.some((file) => file.name.endsWith(weightsKey)), modelBytes: modelFiles.reduce((sum, file) => sum + file.size, 0), + partialBytes: partials.filter((file) => belongs(file.name)).reduce((sum, file) => sum + file.size, 0), usageBytes: estimate.usage ?? 0, quotaBytes: estimate.quota ?? 0, } @@ -86,7 +101,7 @@ export async function getStorageStatus(modelId: string): Promise } } -/** Frees the weights again. The next load re-downloads them. */ +/** Frees the weights again, unfinished downloads included. The next load re-downloads them. */ export async function deleteModel(modelId: string): Promise { const needle = modelFilePrefix(modelId) await clearCachedFiles((name) => name.includes(needle)) diff --git a/src/llm/config.ts b/src/llm/config.ts index 0ff9d45..53f888e 100644 --- a/src/llm/config.ts +++ b/src/llm/config.ts @@ -1,14 +1,32 @@ /** - * Text-only export of Qwen3.5-0.8B. The multimodal build ships a vision encoder - * we never feed, needs the dedicated Qwen3_5 classes, and downloads ~150 MB more. + * Text-only export of Qwen3.5-0.8B, which is the smallest way to run this model + * through Transformers.js and the one its own text-generation support was added + * for (huggingface/transformers.js#1602). + * + * The two multimodal exports of the same weights, `-ONNX` and `-ONNX-OPT`, need + * the dedicated `Qwen3_5ForConditionalGeneration` class and a vision encoder + * this app never feeds; at q4f16 they come to 616 MiB against 448 MiB here. + * Every other Qwen3.5-0.8B ONNX repository on the Hub is a copy of one of the + * three. `node tools/verify-model.mjs` re-checks that against the Hub. */ export const MODEL_ID = 'onnx-community/Qwen3.5-0.8B-Text-ONNX' -/** INT4 weights on an fp16 graph: the smallest variant ONNX Runtime Web runs well. */ +/** + * INT4 weights on an fp16 graph: the smallest of the five variants published, + * ahead of q4 at 526 MiB, int8 at 896 MiB, fp16 at 1.4 GiB and fp32 at 2.9 GiB. + * + * 448 MiB is the floor for this model rather than a choice worth revisiting. + * Qwen3.5 has no size below 0.8B, and a 248,320-token vocabulary tied to the + * output layer is why 0.8B parameters at four bits land here instead of nearer + * 400 MB. + */ export const MODEL_DTYPE = 'q4f16' +/** The weights themselves, and so what an install is mostly waiting for. */ +export const MODEL_WEIGHTS_FILE = `onnx/model_${MODEL_DTYPE}.onnx_data` + /** Measured total of the seven files the q4f16 variant needs. */ -export const MODEL_DOWNLOAD_BYTES = 489_167_000 +export const MODEL_DOWNLOAD_BYTES = 489_174_504 /** * Where the weights come from. diff --git a/src/llm/opfs-cache.test.ts b/src/llm/opfs-cache.test.ts index 326ffbc..9781dba 100644 --- a/src/llm/opfs-cache.test.ts +++ b/src/llm/opfs-cache.test.ts @@ -1,5 +1,16 @@ -import { describe, expect, it } from 'vitest' -import { cacheKeyFor } from './opfs-cache' +import { Blob as NodeBlob } from 'node:buffer' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { + cacheKeyFor, + clearCachedFiles, + listCachedFiles, + listPartialFiles, + opfsCache, + planWrite, + setDownloadProgress, + type DownloadProgress, + type ResumeMeta, +} from './opfs-cache' describe('cacheKeyFor', () => { it('flattens a download URL into a safe filename', () => { @@ -23,3 +34,417 @@ describe('cacheKeyFor', () => { expect(cacheKeyFor('https://host/a?b=c#d')).toBe('host_a_b_c_d') }) }) + +describe('planWrite', () => { + const meta: ResumeMeta = { etag: '"abc"', total: 1000 } + + it('starts from zero when the whole file arrives', () => { + expect( + planWrite({ status: 200, etag: '"abc"', contentRange: null, contentLength: '1000' }, 600, meta), + ).toEqual({ start: 0, total: 1000 }) + }) + + it('continues the partial when the range matches it exactly', () => { + expect( + planWrite( + { status: 206, etag: '"abc"', contentRange: 'bytes 600-999/1000', contentLength: '400' }, + 600, + meta, + ), + ).toEqual({ start: 600, total: 1000 }) + }) + + it('refuses a range whose entity tag no longer matches the saved bytes', () => { + expect( + planWrite( + { status: 206, etag: '"changed"', contentRange: 'bytes 600-999/1000', contentLength: '400' }, + 600, + meta, + ), + ).toBeNull() + }) + + it('refuses a range that belongs to a differently sized file', () => { + expect( + planWrite( + { status: 206, etag: '"abc"', contentRange: 'bytes 600-1199/1200', contentLength: '600' }, + 600, + meta, + ), + ).toBeNull() + }) + + it('refuses a range that does not start where the file ends', () => { + expect( + planWrite( + { status: 206, etag: '"abc"', contentRange: 'bytes 500-999/1000', contentLength: '500' }, + 600, + meta, + ), + ).toBeNull() + }) + + it('refuses a range nobody asked for, and any other status', () => { + expect( + planWrite( + { status: 206, etag: '"abc"', contentRange: 'bytes 0-999/1000', contentLength: '1000' }, + 0, + null, + ), + ).toBeNull() + expect( + planWrite({ status: 404, etag: null, contentRange: null, contentLength: null }, 0, null), + ).toBeNull() + }) +}) + +/** + * In-memory stand-in for OPFS: jsdom implements none of it, and the resume path + * is mostly about what survives on disk between two attempts. + */ +function fakeOpfs() { + const files = new Map() + const locked = new Set() + + const bytesOf = (value: unknown): Uint8Array => + typeof value === 'string' ? new TextEncoder().encode(value) : new Uint8Array(value as ArrayBufferLike) + + function handleFor(name: string) { + const read = (): Uint8Array => files.get(name) ?? new Uint8Array() + return { + kind: 'file' as const, + name, + // Node's Blob rather than jsdom's: only that one has the `stream()` a + // Response body is read through. + getFile: async () => new NodeBlob([read()]), + move: async (next: string) => { + files.set(next, read()) + files.delete(name) + }, + createWritable: async () => { + let buffer = new Uint8Array() + return { + write: async (chunk: unknown) => { + const incoming = bytesOf(chunk instanceof Uint8Array ? chunk : chunk) + const next = new Uint8Array(buffer.length + incoming.length) + next.set(buffer) + next.set(incoming, buffer.length) + buffer = next + }, + close: async () => void files.set(name, buffer), + abort: async () => undefined, + } + }, + createSyncAccessHandle: async () => { + if (locked.has(name)) throw new Error('NoModificationAllowedError') + locked.add(name) + return { + getSize: () => read().length, + read: (buffer: Uint8Array, options?: { at?: number }) => { + const source = read() + const at = options?.at ?? 0 + const length = Math.max(0, Math.min(buffer.length, source.length - at)) + buffer.set(source.subarray(at, at + length)) + return length + }, + write: (chunk: Uint8Array, options?: { at?: number }) => { + const at = options?.at ?? 0 + const current = read() + const next = new Uint8Array(Math.max(current.length, at + chunk.length)) + next.set(current) + next.set(chunk, at) + files.set(name, next) + return chunk.length + }, + truncate: (size: number) => { + const current = read() + const next = new Uint8Array(size) + next.set(current.subarray(0, Math.min(size, current.length))) + files.set(name, next) + }, + flush: () => undefined, + close: () => void locked.delete(name), + } + }, + } + } + + const directory = { + kind: 'directory' as const, + getFileHandle: async (name: string, options?: { create?: boolean }) => { + if (!files.has(name)) { + if (!options?.create) throw new Error('NotFoundError') + files.set(name, new Uint8Array()) + } + return handleFor(name) + }, + removeEntry: async (name: string) => { + if (!files.delete(name)) throw new Error('NotFoundError') + }, + [Symbol.asyncIterator]: async function* () { + for (const name of [...files.keys()]) yield [name, handleFor(name)] as const + }, + } + + Object.defineProperty(globalThis.navigator, 'storage', { + configurable: true, + value: { + getDirectory: async () => ({ getDirectoryHandle: async () => directory }), + }, + }) + + return files +} + +const ETAG = '"sha-of-the-weights"' +const WEIGHTS: Uint8Array = new Uint8Array( + Array.from({ length: 4096 }, (_, index) => index % 251), +) + +function complete(bytes: Uint8Array, etag = ETAG): Response { + return new Response(bytes, { + status: 200, + headers: { 'content-length': String(bytes.length), etag }, + }) +} + +function tail(bytes: Uint8Array, from: number, etag = ETAG): Response { + const slice = bytes.subarray(from) + return new Response(slice, { + status: 206, + headers: { + 'content-length': String(slice.length), + 'content-range': `bytes ${from}-${bytes.length - 1}/${bytes.length}`, + etag, + }, + }) +} + +/** + * A transfer that dies part way through, the way a dropped connection does. The + * failure has to come from a later `pull`: erroring a stream discards whatever + * is still queued, so the first chunk has to be read before the break. + */ +function severed(bytes: Uint8Array, cut: number, etag = ETAG): Response { + let sent = false + const body = new ReadableStream({ + pull(controller) { + if (!sent) { + sent = true + controller.enqueue(bytes.subarray(0, cut)) + return + } + controller.error(new Error('network went away')) + }, + }) + return new Response(body, { + status: 200, + headers: { 'content-length': String(bytes.length), etag }, + }) +} + +/** A transfer that ends cleanly but short, which content-length would hide. */ +function truncated(bytes: Uint8Array, cut: number, etag = ETAG): Response { + return new Response(bytes.subarray(0, cut), { + status: 200, + headers: { 'content-length': String(bytes.length), etag }, + }) +} + +async function drain(response: Response | undefined): Promise { + if (!response) throw new Error('nothing to read') + return new Uint8Array(await response.arrayBuffer()) +} + +let files: Map +let fetchMock: ReturnType + +beforeEach(() => { + files = fakeOpfs() + fetchMock = vi.fn() + vi.stubGlobal('fetch', fetchMock) +}) + +afterEach(() => { + setDownloadProgress(null) + vi.unstubAllGlobals() + vi.restoreAllMocks() +}) + +/** + * Each test uses a URL of its own. The cache remembers which downloads it had to + * hand back to Transformers.js, and that memory is per module rather than per test. + */ +function urlFor(name: string): string { + return `https://huggingface.co/onnx-community/Model/resolve/main/${name}` +} + +describe('opfsCache downloads', () => { + it('downloads to disk and publishes under the final name', async () => { + const url = urlFor('fresh.onnx_data') + fetchMock.mockResolvedValueOnce(complete(WEIGHTS)) + + expect(await drain(await opfsCache.match(url))).toEqual(WEIGHTS) + + expect(files.get(cacheKeyFor(url))).toEqual(WEIGHTS) + expect([...files.keys()].filter((name) => name.includes('.part'))).toEqual([]) + // A first attempt asks for the file exactly as Transformers.js would. + expect(fetchMock.mock.calls[0]?.[1]).toBeUndefined() + }) + + it('reports progress, since Transformers.js never sees this download', async () => { + const url = urlFor('progress.onnx_data') + const seen: DownloadProgress[] = [] + setDownloadProgress((progress) => void seen.push(progress)) + fetchMock.mockResolvedValueOnce(complete(WEIGHTS)) + + await drain(await opfsCache.match(url)) + + expect(seen.at(-1)).toEqual({ url, loaded: WEIGHTS.length, total: WEIGHTS.length }) + }) + + it('serves an installed file without going near the network', async () => { + const url = urlFor('installed.onnx_data') + files.set(cacheKeyFor(url), WEIGHTS) + + expect(await drain(await opfsCache.match(url))).toEqual(WEIGHTS) + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('never fetches for the local paths Transformers.js also probes', async () => { + expect(await opfsCache.match('/models/onnx-community/Model/onnx/model_q4f16.onnx_data')).toBeUndefined() + expect(fetchMock).not.toHaveBeenCalled() + }) + + it('resumes from what arrived when the connection drops mid-transfer', async () => { + const url = urlFor('dropped.onnx_data') + fetchMock.mockResolvedValueOnce(severed(WEIGHTS, 1500)).mockResolvedValueOnce(tail(WEIGHTS, 1500)) + + expect(await drain(await opfsCache.match(url))).toEqual(WEIGHTS) + + // The second attempt asked only for the bytes that were still missing. + expect(fetchMock.mock.calls[1]?.[1]).toEqual({ headers: { Range: 'bytes=1500-' } }) + expect(files.get(cacheKeyFor(url))).toEqual(WEIGHTS) + expect([...files.keys()].filter((name) => name.includes('.part'))).toEqual([]) + }) + + it('keeps the partial for the next attempt when every attempt fails', async () => { + const url = urlFor('offline.onnx_data') + const key = cacheKeyFor(url) + fetchMock + .mockResolvedValueOnce(severed(WEIGHTS, 1500)) + .mockResolvedValueOnce(severed(WEIGHTS, 1500)) + .mockResolvedValueOnce(severed(WEIGHTS, 1500)) + + // Nothing loadable is on offer, so the app must not call itself installed. + expect(await opfsCache.match(url)).toBeUndefined() + expect(files.has(key)).toBe(false) + expect(files.get(`${key}.part`)).toEqual(WEIGHTS.subarray(0, 1500)) + expect(await listCachedFiles()).toEqual([]) + expect(await listPartialFiles()).toEqual([{ name: `${key}.part`, size: 1500 }]) + + fetchMock.mockResolvedValueOnce(tail(WEIGHTS, 1500)) + expect(await drain(await opfsCache.match(url))).toEqual(WEIGHTS) + expect(fetchMock.mock.calls[3]?.[1]).toEqual({ headers: { Range: 'bytes=1500-' } }) + }) + + it('starts again when the file changed upstream while a partial was on disk', async () => { + const url = urlFor('changed.onnx_data') + const key = cacheKeyFor(url) + const replacement: Uint8Array = new Uint8Array(2048).fill(7) + + // The Hub ignores If-Range, so a stale partial is answered with a 206 that + // belongs to different bytes. It has to be recognised here. + fetchMock + .mockResolvedValueOnce(severed(WEIGHTS, 900)) + .mockResolvedValueOnce(tail(replacement, 900, '"a-new-export"')) + .mockResolvedValueOnce(complete(replacement, '"a-new-export"')) + + expect(await drain(await opfsCache.match(url))).toEqual(replacement) + expect(files.get(key)).toEqual(replacement) + }) + + it('starts again when the host ignores the range and sends the whole file', async () => { + const url = urlFor('no-ranges.onnx_data') + fetchMock.mockResolvedValueOnce(severed(WEIGHTS, 700)).mockResolvedValueOnce(complete(WEIGHTS)) + + expect(await drain(await opfsCache.match(url))).toEqual(WEIGHTS) + expect(files.get(cacheKeyFor(url))).toEqual(WEIGHTS) + }) + + it('treats a body that stops short as unfinished rather than as a shorter file', async () => { + const url = urlFor('short.onnx_data') + const key = cacheKeyFor(url) + fetchMock + .mockResolvedValueOnce(truncated(WEIGHTS, 2000)) + .mockResolvedValueOnce(truncated(WEIGHTS, 2000)) + .mockResolvedValueOnce(truncated(WEIGHTS, 2000)) + + expect(await opfsCache.match(url)).toBeUndefined() + expect(files.has(key)).toBe(false) + expect(files.get(`${key}.part`)).toEqual(WEIGHTS.subarray(0, 2000)) + }) + + it('stands aside when the file cannot be reached at all', async () => { + const url = urlFor('missing.json') + fetchMock.mockResolvedValueOnce(new Response('nope', { status: 404 })) + + expect(await opfsCache.match(url)).toBeUndefined() + // One attempt only: a 404 reads the same way however often it is asked for. + expect(fetchMock).toHaveBeenCalledTimes(1) + // No empty partial left behind for the next visit to trip over. + expect([...files.keys()]).toEqual([]) + }) + + it('downloads once however many callers ask at the same time', async () => { + const url = urlFor('shared.onnx_data') + fetchMock.mockResolvedValueOnce(complete(WEIGHTS)) + + const [first, second] = await Promise.all([opfsCache.match(url), opfsCache.match(url)]) + + expect(await drain(first)).toEqual(WEIGHTS) + expect(await drain(second)).toEqual(WEIGHTS) + expect(fetchMock).toHaveBeenCalledTimes(1) + }) + + it('leaves the download to Transformers.js when the write lock is held', async () => { + const url = urlFor('locked.onnx_data') + const directory = await (await navigator.storage.getDirectory()).getDirectoryHandle('model-cache') + const partial = (await directory.getFileHandle(`${cacheKeyFor(url)}.part`, { + create: true, + })) as FileSystemFileHandle & { createSyncAccessHandle: () => Promise } + await partial.createSyncAccessHandle() + + expect(await opfsCache.match(url)).toBeUndefined() + expect(fetchMock).not.toHaveBeenCalled() + }) +}) + +describe('opfsCache.put', () => { + it('stores a file Transformers.js downloaded without touching a resumable partial', async () => { + const url = urlFor('stored.json') + const key = cacheKeyFor(url) + files.set(`${key}.part`, WEIGHTS.subarray(0, 100)) + + await opfsCache.put(url, complete(WEIGHTS)) + + expect(files.get(key)).toEqual(WEIGHTS) + expect(files.get(`${key}.part`)).toEqual(WEIGHTS.subarray(0, 100)) + }) +}) + +describe('clearCachedFiles', () => { + it('removes unfinished downloads too, since they hold the space back', async () => { + const url = urlFor('discarded.onnx_data') + const key = cacheKeyFor(url) + + fetchMock.mockResolvedValue(severed(WEIGHTS, 1200)) + expect(await opfsCache.match(url)).toBeUndefined() + files.set('huggingface.co_some-other-model_resolve_main_model.onnx', new Uint8Array(4)) + + await clearCachedFiles((name) => name.includes(cacheKeyFor('onnx-community/Model'))) + + expect([...files.keys()]).toEqual(['huggingface.co_some-other-model_resolve_main_model.onnx']) + expect(key).toContain('onnx-community_Model') + }) +}) diff --git a/src/llm/opfs-cache.ts b/src/llm/opfs-cache.ts index b0d253d..00ac3ec 100644 --- a/src/llm/opfs-cache.ts +++ b/src/llm/opfs-cache.ts @@ -1,11 +1,19 @@ /** * Model cache backed by the Origin Private File System. * - * Transformers.js defaults to the Cache API, but Chrome rejects the ~440 MB + * Transformers.js defaults to the Cache API, but Chrome rejects the ~448 MB * weights file there with "Failed to execute 'put' on 'Cache': Unexpected * internal error", so the download silently never persisted. OPFS is designed * for large binary files and streams them to disk without buffering the whole * body in memory. + * + * This backend also performs the download, which is the only place a resume can + * live. Transformers.js reads a whole response into a buffer *before* handing it + * to a cache, so by the time `put` is called every byte has already arrived and + * a connection that dropped at 400 MB left nothing behind. Owning the fetch + * means the bytes that did arrive stay on disk, and the next attempt continues + * from them with a Range request. Upstream has this for Node's filesystem cache + * (huggingface/transformers.js#1715); the browser side is still open (#1220). */ export const MODEL_CACHE_DIR = 'model-cache' @@ -13,6 +21,21 @@ export const MODEL_CACHE_DIR = 'model-cache' /** Marks an in-flight download so an interrupted write is never mistaken for a complete file. */ const PARTIAL_SUFFIX = '.part' +/** Records what a partial is a prefix of, so a later attempt can prove it still matches. */ +const META_SUFFIX = '.part-meta' + +/** Written bytes are forced to disk this often, capping what a crash costs. */ +const FLUSH_EVERY_BYTES = 16 * 1024 * 1024 + +/** Progress is reported this often rather than per chunk, which is every 64 KB. */ +const REPORT_EVERY_BYTES = 1024 * 1024 + +/** Attempts per download, each one resuming where the last stopped. */ +const ATTEMPTS = 3 + +/** Pause between attempts, long enough for a brief drop to pass. */ +const RETRY_DELAY_MS = 500 + export function cacheKeyFor(request: string): string { return request.replace(/^https?:\/\//, '').replace(/[^a-zA-Z0-9._-]/g, '_') } @@ -59,17 +82,326 @@ async function publish( await directory.removeEntry(partialName) } +/** What an unfinished download is known to be a prefix of. */ +export interface ResumeMeta { + /** Entity tag the bytes on disk came from. */ + etag: string + /** Size of the complete file. */ + total: number +} + +/** Where a response's bytes belong in the file being assembled. */ +export interface WritePlan { + /** Offset the body starts at. Zero for a full response. */ + start: number + /** Size of the complete file, or 0 when the server did not say. */ + total: number +} + +/** `bytes 1024-4095/4096` → where the body starts and how large the file is. */ +function parseContentRange(value: string | null): { start: number; total: number } | null { + const match = /^bytes (\d+)-(\d+)\/(\d+)$/.exec(value?.trim() ?? '') + if (!match) return null + return { start: Number(match[1]), total: Number(match[3]) } +} + +/** + * Whether a response can be appended to what is already on disk. + * + * A 206 is trusted only when it continues the partial exactly: same entity tag, + * same total, starting where the file ends. The check cannot be delegated to the + * server, because the Hub's CDN ignores `If-Range` — a stale validator still + * comes back as 206 with the old range, which would splice bytes from two + * different files together. A 200 is always a whole file, so it restarts. + * Anything else means this backend should stand aside. + */ +export function planWrite( + response: { + status: number + etag: string | null + contentRange: string | null + contentLength: string | null + }, + requested: number, + meta: ResumeMeta | null, +): WritePlan | null { + if (response.status === 200) { + return { start: 0, total: Number(response.contentLength ?? 0) || 0 } + } + if (response.status !== 206 || requested <= 0 || !meta) return null + + const range = parseContentRange(response.contentRange) + if (!range || range.start !== requested || range.total !== meta.total) return null + if (!response.etag || response.etag !== meta.etag) return null + return range +} + +async function readMeta(directory: FileSystemDirectoryHandle, name: string): Promise { + try { + const handle = await directory.getFileHandle(`${name}${META_SUFFIX}`) + const parsed = JSON.parse(await (await handle.getFile()).text()) as Partial + if (typeof parsed.etag !== 'string' || typeof parsed.total !== 'number' || parsed.total <= 0) return null + return { etag: parsed.etag, total: parsed.total } + } catch { + return null + } +} + +async function writeMeta( + directory: FileSystemDirectoryHandle, + name: string, + meta: ResumeMeta | null, +): Promise { + if (!meta) { + await discard(directory, `${name}${META_SUFFIX}`) + return + } + const handle = await directory.getFileHandle(`${name}${META_SUFFIX}`, { create: true }) + const writable = await handle.createWritable() + await writable.write(JSON.stringify(meta)) + await writable.close() +} + +async function discard(directory: FileSystemDirectoryHandle, ...names: string[]): Promise { + for (const name of names) await directory.removeEntry(name).catch(() => undefined) +} + +/** Bytes on disk for one file, as the app shows them while a download runs. */ +export interface DownloadProgress { + url: string + loaded: number + total: number +} + +let report: ((progress: DownloadProgress) => void) | null = null + +/** + * Where download progress goes. + * + * Transformers.js reports progress for the bodies it reads itself, and it never + * reads this one — the file is on disk by the time it is handed over. The worker + * translates these URLs into the file names the rest of the app already uses. + */ +export function setDownloadProgress(listener: ((progress: DownloadProgress) => void) | null): void { + report = listener +} + +/** + * URLs this backend cannot download, as opposed to ones it failed to. + * + * A missing sync access handle or a file another writer holds will not fix + * itself, so those are handed back to Transformers.js for good. A network + * failure is the opposite: the next attempt is the one that resumes. + */ +const unavailable = new Set() + +/** One download per URL, however many callers ask for it. */ +const inFlight = new Map>() + +function headersOf(response: Response) { + return { + status: response.status, + etag: response.headers.get('etag'), + contentRange: response.headers.get('content-range'), + contentLength: response.headers.get('content-length'), + } +} + +/** + * Fetches whatever is still missing and appends it to the partial. + * + * Returns true once every byte is on disk. A false return means the attempt is + * worth repeating: what arrived has been kept, so the next one asks for less. + * Throws only when the response was not a download at all. + */ +async function attempt( + url: string, + directory: FileSystemDirectoryHandle, + access: FileSystemSyncAccessHandle, + name: string, +): Promise { + const meta = await readMeta(directory, name) + const saved = access.getSize() + const resumeFrom = meta && saved > 0 && saved < meta.total ? saved : 0 + + // A first attempt asks for the file exactly as Transformers.js would. Only a + // continuation carries a Range, whose simple byte form needs no preflight. + let response = await fetch(url, resumeFrom > 0 ? { headers: { Range: `bytes=${resumeFrom}-` } } : undefined) + let plan = planWrite(headersOf(response), resumeFrom, meta) + + // The partial cannot be continued — the file changed upstream, or the host + // ignored the range. Ask for the whole thing and overwrite. + if (!plan && resumeFrom > 0) { + response = await fetch(url) + plan = planWrite(headersOf(response), 0, null) + } + if (!response.ok || !plan || !response.body) throw new Error(`HTTP ${response.status} for ${url}`) + + access.truncate(plan.start) + const etag = response.headers.get('etag') + await writeMeta(directory, name, plan.total > 0 && etag ? { etag, total: plan.total } : null) + + let position = plan.start + let sinceFlush = 0 + let sinceReport = 0 + const reader = response.body.getReader() + + try { + for (;;) { + const { done, value } = await reader.read() + if (done) break + access.write(value, { at: position }) + position += value.byteLength + sinceFlush += value.byteLength + sinceReport += value.byteLength + if (sinceFlush >= FLUSH_EVERY_BYTES) { + access.flush() + sinceFlush = 0 + } + if (sinceReport >= REPORT_EVERY_BYTES) { + report?.({ url, loaded: position, total: plan.total }) + sinceReport = 0 + } + } + } finally { + access.flush() + } + + // A body that stops early is an unfinished download, not a shorter file. The + // difference matters: Transformers.js sizes its buffer from the content length + // and zero-pads the rest, so publishing this would mean corrupt weights. + if (plan.total > 0 && position !== plan.total) { + report?.({ url, loaded: position, total: plan.total }) + return false + } + + report?.({ url, loaded: position, total: position }) + return true +} + +/** + * Downloads `url` into OPFS, continuing an earlier attempt where one is on disk + * and retrying the transfer a few times before giving up. + * + * Leaves the file absent rather than throwing. `match` looks for the result, so + * a failure here simply means Transformers.js downloads the file itself, as it + * did before this backend existed. + */ +async function download(url: string): Promise { + const name = cacheKeyFor(url) + const directory = await cacheDir() + const partialName = `${name}${PARTIAL_SUFFIX}` + const handle = await directory.getFileHandle(partialName, { create: true }) + + // Only a dedicated worker gets a sync access handle, and only a sync handle + // writes straight through to the file: a `FileSystemWritableFileStream` + // buffers into a swap file that is discarded unless it is closed cleanly, + // which would leave nothing to resume from. + /** Only the empty file this function just created, never a real resume point. */ + const standAside = async (): Promise => { + unavailable.add(url) + if ((await handle.getFile().catch(() => null))?.size === 0) await discard(directory, partialName) + } + + if (typeof handle.createSyncAccessHandle !== 'function') { + await standAside() + return + } + + let access: FileSystemSyncAccessHandle + try { + access = await handle.createSyncAccessHandle() + } catch { + // Another writer holds the file. Downloading it twice would be worse. + await standAside() + return + } + + let closed = false + try { + for (let attemptsLeft = ATTEMPTS; attemptsLeft > 0; attemptsLeft -= 1) { + let complete = false + try { + complete = await attempt(url, directory, access, name) + } catch (error) { + // A response that was not a download at all — a 404, a redirect to an + // error page — will read the same way next time. + if (error instanceof Error && error.message.startsWith('HTTP')) break + } + + if (complete) { + access.flush() + access.close() + closed = true + await discard(directory, name, `${name}${META_SUFFIX}`) + await publish(directory, partialName, name) + return + } + if (attemptsLeft > 1) await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY_MS)) + } + } catch { + // A failure is reported by the file's absence, not by throwing: see above. + } finally { + if (!closed) { + const saved = access.getSize() + access.flush() + access.close() + // An empty partial is not a resume point, only clutter. + if (saved === 0) await discard(directory, partialName, `${name}${META_SUFFIX}`) + } + } +} + +function once(url: string): Promise { + let pending = inFlight.get(url) + if (!pending) { + pending = download(url).finally(() => inFlight.delete(url)) + inFlight.set(url, pending) + } + return pending +} + +async function cachedResponse( + directory: FileSystemDirectoryHandle, + name: string, +): Promise { + try { + const file = await (await directory.getFileHandle(name)).getFile() + if (file.size === 0) return undefined + return new Response(file, { headers: { 'content-length': String(file.size) } }) + } catch { + return undefined + } +} + export const opfsCache = { + /** + * The cached file, downloading it first if it is not there yet. + * + * Transformers.js also calls this to ask whether a file exists and how large + * it is, and drops the body when it does — which is why the download finishes + * before anything is returned rather than streaming through the response. + */ async match(request: string): Promise { if (!opfsAvailable()) return undefined + + let directory: FileSystemDirectoryHandle try { - const handle = await (await cacheDir()).getFileHandle(cacheKeyFor(request)) - const file = await handle.getFile() - if (file.size === 0) return undefined - return new Response(file, { headers: { 'content-length': String(file.size) } }) + directory = await cacheDir() } catch { return undefined } + + const name = cacheKeyFor(request) + const cached = await cachedResponse(directory, name) + if (cached) return cached + + // Only a real URL can be fetched. Transformers.js also probes this cache + // with local paths, which are not ours to go and download. + if (unavailable.has(request) || !/^https?:\/\//.test(request)) return undefined + + await once(request) + return cachedResponse(directory, name) }, async put( @@ -81,7 +413,9 @@ export const opfsCache = { const name = cacheKeyFor(request) const directory = await cacheDir() - const partialName = `${name}${PARTIAL_SUFFIX}` + // A name of its own, so storing a file this backend did not download cannot + // collide with a partial that is being resumed. + const partialName = `${name}${PARTIAL_SUFFIX}-${crypto.randomUUID().slice(0, 8)}` const partial = await directory.getFileHandle(partialName, { create: true }) const writable = await partial.createWritable() const total = Number(response.headers.get('content-length') ?? 0) @@ -98,13 +432,15 @@ export const opfsCache = { } await writable.close() } catch (error) { + // Nothing written through a writable survives an abort, so there is no + // partial worth keeping on this path — see `download` for the one there is. await writable.abort().catch(() => undefined) - await directory.removeEntry(partialName).catch(() => undefined) + await discard(directory, partialName) throw error } // Publish under the real name only once the bytes are all on disk. - await directory.removeEntry(name).catch(() => undefined) + await discard(directory, name) await publish(directory, partialName, name) }, @@ -124,13 +460,13 @@ export interface CachedFile { size: number } -export async function listCachedFiles(): Promise { +async function listEntries(): Promise { if (!opfsAvailable()) return [] try { const directory = await cacheDir() const files: CachedFile[] = [] for await (const [name, handle] of directory as unknown as AsyncIterable<[string, FileSystemHandle]>) { - if (handle.kind !== 'file' || name.endsWith(PARTIAL_SUFFIX)) continue + if (handle.kind !== 'file') continue files.push({ name, size: (await (handle as FileSystemFileHandle).getFile()).size }) } return files @@ -139,10 +475,25 @@ export async function listCachedFiles(): Promise { } } +/** Complete files only: a partial is not something the app can load. */ +export async function listCachedFiles(): Promise { + return (await listEntries()).filter((file) => !file.name.includes(PARTIAL_SUFFIX)) +} + +/** Unfinished downloads, which a later attempt will continue rather than repeat. */ +export async function listPartialFiles(): Promise { + return (await listEntries()).filter((file) => file.name.endsWith(PARTIAL_SUFFIX)) +} + +/** + * Removes cached files whose name matches, partials included: leaving a + * half-downloaded 448 MB file behind after "Remove model" would occupy the + * space the user asked to get back. + */ export async function clearCachedFiles(predicate: (name: string) => boolean): Promise { if (!opfsAvailable()) return const directory = await cacheDir() - for (const file of await listCachedFiles()) { - if (predicate(file.name)) await directory.removeEntry(file.name).catch(() => undefined) + for (const file of await listEntries()) { + if (predicate(file.name)) await discard(directory, file.name) } } diff --git a/src/llm/worker.ts b/src/llm/worker.ts index 1ba1e1a..5ed0277 100644 --- a/src/llm/worker.ts +++ b/src/llm/worker.ts @@ -7,7 +7,7 @@ import { type TextGenerationPipeline, } from '@huggingface/transformers' import { DEFAULT_GENERATION, MODEL_DTYPE, MODEL_HOST, MODEL_ID, MODEL_PATH_TEMPLATE } from './config' -import { opfsAvailable, opfsCache } from './opfs-cache' +import { opfsAvailable, opfsCache, setDownloadProgress } from './opfs-cache' import { CLOSE_THINK, closeReasoning, splitReasoning } from './phases' import type { ChatTurn, LoadProgress, MainToWorker, WorkerToMain } from './protocol' @@ -60,6 +60,25 @@ function onProgress(event: ProgressEvent): void { } } +/** + * Where the model files sit under the host, so the cache's URLs can be reduced + * to the same names Transformers.js reports. Two keys for one file would be + * counted twice by the progress bar and it would never reach the end. + */ +const REMOTE_PREFIX = + MODEL_HOST + MODEL_PATH_TEMPLATE.replaceAll('{model}', MODEL_ID).replaceAll('{revision}', 'main') + +/** + * The cache downloads the weights itself so it can resume an interrupted + * transfer, and Transformers.js only sees the finished file — so the progress + * the user watches during an install comes from here. + */ +setDownloadProgress(({ url, loaded, total }) => { + const file = url.startsWith(REMOTE_PREFIX) ? url.slice(REMOTE_PREFIX.length) : url + progressByFile.set(file, { file, loaded, total }) + post({ type: 'progress', files: [...progressByFile.values()] }) +}) + async function load(): Promise { if (loadPromise) return loadPromise diff --git a/src/store/chat.ts b/src/store/chat.ts index b2e9d65..694a5dd 100644 --- a/src/store/chat.ts +++ b/src/store/chat.ts @@ -2,7 +2,7 @@ import { create } from 'zustand' import { runAgent } from '@/agent/loop' import { LlmClient } from '@/llm/client' import type { ChatTurn, LoadProgress } from '@/llm/protocol' -import { MODEL_ID } from '@/llm/config' +import { MODEL_ID, MODEL_WEIGHTS_FILE } from '@/llm/config' import { deleteModel, getStorageStatus, @@ -259,11 +259,14 @@ export const useChatStore = create((set, get) => { await get().setMcpServers(get().mcpServers) } catch (error) { set({ status: 'error', error: error instanceof Error ? error.message : String(error) }) + // Whatever did arrive is a resume point, and the gate offers to continue + // from it — so the figure it shows has to be the one after the failure. + void get().refreshStorage() } }, async refreshStorage() { - set({ storage: await getStorageStatus(MODEL_ID) }) + set({ storage: await getStorageStatus(MODEL_ID, MODEL_WEIGHTS_FILE) }) }, async removeModel() { diff --git a/tools/verify-model.mjs b/tools/verify-model.mjs index 875aed2..32d631c 100644 --- a/tools/verify-model.mjs +++ b/tools/verify-model.mjs @@ -29,11 +29,16 @@ const MODEL_ID = 'onnx-community/Qwen3.5-0.8B-Text-ONNX' const MODEL_HOST = process.env.VITE_MODEL_HOST || 'https://huggingface.co/' const MODEL_PATH_TEMPLATE = process.env.VITE_MODEL_PATH_TEMPLATE || '{model}/resolve/{revision}/' +const MODEL_DTYPE = 'q4f16' + +/** The weights themselves, and so the file every check below cares about most. */ +const WEIGHTS_FILE = `onnx/model_${MODEL_DTYPE}.onnx_data` + /** The seven files an install fetches, largest first. */ const MODEL_FILES = [ - 'onnx/model_q4f16.onnx_data', + WEIGHTS_FILE, 'tokenizer.json', - 'onnx/model_q4f16.onnx', + `onnx/model_${MODEL_DTYPE}.onnx`, 'tokenizer_config.json', 'chat_template.jinja', 'config.json', @@ -41,7 +46,14 @@ const MODEL_FILES = [ ] /** MODEL_DOWNLOAD_BYTES in src/llm/config.ts. Compared loosely; see below. */ -const EXPECTED_TOTAL_BYTES = 489_167_000 +const EXPECTED_TOTAL_BYTES = 489_174_504 + +/** + * The other ONNX exports of the same weights, for the size comparison below. + * Both are multimodal: they need `Qwen3_5ForConditionalGeneration` and a vision + * encoder this app never feeds, so they are listed to be measured, not used. + */ +const ALTERNATIVE_EXPORTS = ['onnx-community/Qwen3.5-0.8B-ONNX', 'onnx-community/Qwen3.5-0.8B-ONNX-OPT'] /** * The app fetches these from a page, so every request is cross-origin. Using the @@ -100,6 +112,7 @@ heading('Weight availability') console.log(` host: ${MODEL_HOST}`) let totalBytes = 0 +const fileBytes = new Map() for (const path of MODEL_FILES) { let result try { @@ -109,6 +122,7 @@ for (const path of MODEL_FILES) { continue } totalBytes += result.bytes + fileBytes.set(path, result.bytes) check( result.ok, `${path} — ${result.ok ? `${result.status}, ${result.bytes.toLocaleString()} bytes` : `HTTP ${result.status}`}`, @@ -137,6 +151,93 @@ if (failures.length > 0) { process.exit(1) } +/** + * A half-gigabyte download that restarts from zero on a dropped connection is + * one most people never finish, so `src/llm/opfs-cache.ts` continues it with a + * Range request. That needs three things from the host, none of them ours: byte + * ranges (checked above), a 206 that says which bytes it is sending, and an + * entity tag a script can read — the partial is only safe to append to if the + * remote file can be proven not to have changed. `If-Range` is not enough: the + * Hub's CDN ignores it and answers a stale validator with 206 anyway. + */ +heading('Resumable download') + +const weightsUrl = `${MODEL_HOST}${MODEL_PATH_TEMPLATE.replace('{model}', MODEL_ID).replace('{revision}', 'main')}${WEIGHTS_FILE}` +const ranged = await fetch(weightsUrl, { + headers: { origin: ORIGIN, range: 'bytes=1024-2047' }, +}) +await ranged.arrayBuffer() + +const exposed = (ranged.headers.get('access-control-expose-headers') ?? '').toLowerCase() +const contentRange = ranged.headers.get('content-range') +check(ranged.status === 206, `partial content honoured (HTTP ${ranged.status})`) +// The resume reads the whole file's size out of this header and refuses a +// partial whose total has moved, so the number in it has to be the real one. +check( + contentRange === `bytes 1024-2047/${fileBytes.get(WEIGHTS_FILE)}`, + `range states the file's total size (${contentRange})`, +) +check(Boolean(ranged.headers.get('etag')), `entity tag present (${ranged.headers.get('etag')})`) +check(exposed.includes('*') || exposed.includes('etag'), 'entity tag readable from a page, not just by curl') + +if (failures.length > 0) { + console.error('\nA broken-off install cannot be resumed against this host; it would start over.') + process.exit(1) +} + +/** + * Whether the variant the app loads is still the smallest one published. + * + * The download is the single biggest thing this app asks of anyone, so "is this + * the right export?" deserves an answer from the Hub rather than from a comment. + */ +heading('Smallest available variant') + +async function variantSizes(repo) { + const response = await fetch(`https://huggingface.co/api/models/${repo}/tree/main?recursive=true`) + if (!response.ok) throw new Error(`HTTP ${response.status}`) + const sizes = new Map() + for (const entry of await response.json()) { + if (entry.type !== 'file') continue + const match = + /^onnx\/(.+?)(?:_(fp16|fp32|quantized|q4|q4f16|int8|bnb4|uint8))?\.onnx(?:_data(?:_\d+)?)?$/.exec( + entry.path, + ) + if (!match) continue + const dtype = match[2] ?? 'fp32' + sizes.set(dtype, (sizes.get(dtype) ?? 0) + (entry.size ?? 0)) + } + return sizes +} + +try { + const sizes = await variantSizes(MODEL_ID) + const mine = sizes.get(MODEL_DTYPE) ?? 0 + for (const [dtype, bytes] of [...sizes].sort((a, b) => a[1] - b[1])) { + const mib = (bytes / 1024 / 1024).toFixed(1) + console.log(` ${dtype === MODEL_DTYPE ? '->' : ' '} ${dtype.padEnd(10)} ${mib.padStart(8)} MiB`) + } + check(mine > 0, `${MODEL_DTYPE} is published in ${MODEL_ID}`) + check( + [...sizes.values()].every((bytes) => bytes >= mine), + `${MODEL_DTYPE} is the smallest variant in this repository`, + ) + + for (const repo of ALTERNATIVE_EXPORTS) { + const alternative = await variantSizes(repo) + const total = alternative.get(MODEL_DTYPE) ?? 0 + console.log(` ${repo} at ${MODEL_DTYPE}: ${(total / 1024 / 1024).toFixed(1)} MiB`) + check(total >= mine, `${repo} is no smaller than the export in use`) + } +} catch (error) { + check(false, `variant survey failed: ${error.message}`) +} + +if (failures.length > 0) { + console.error('\nA smaller export of this model may now exist. Check before shipping the larger one.') + process.exit(1) +} + heading('Tokenizer and chat template') const tokenizer = await AutoTokenizer.from_pretrained(MODEL_ID) console.log('Tokenizer loaded.')