Spaces:
Running
Running
| import { useCallback, useEffect, useState } from "react"; | |
| import { MLC_REPO, mlcAppConfig, mlcUrls, mlcVariant } from "../runtime/mlc"; | |
| import { MODEL_IDS, type ModelId, type ContextPreset } from "../types"; | |
| import { | |
| clearLoadedModel, | |
| markLoadedModel, | |
| readLoadedModels, | |
| type LoadedModels, | |
| } from "../lib/modelCacheRecord"; | |
| export interface StorageUsage { | |
| usage: number | null; | |
| quota: number | null; | |
| } | |
| export interface ModelCacheState { | |
| /** Per quant, from WebLLM's own cache check; null while still unknown. */ | |
| cached: Record<ModelId, boolean | null>; | |
| /** Per quant, those that have completed a load on this device. */ | |
| loadedHere: LoadedModels; | |
| storage: StorageUsage; | |
| /** The quant currently being cleared, if any. */ | |
| clearing: ModelId | null; | |
| /** True while a full (all-quants) cache purge is in flight. */ | |
| clearingAll: boolean; | |
| refresh(): Promise<void>; | |
| clear(id: ModelId): Promise<void>; | |
| /** Wipe every WebLLM weight/config entry for this app, including orphans. */ | |
| clearAll(): Promise<void>; | |
| /** Record a successful load, then re-check the cache. */ | |
| markLoaded(id: ModelId): void; | |
| } | |
| const UNKNOWN = Object.fromEntries(MODEL_IDS.map((id) => [id, null])) as Record< | |
| ModelId, | |
| boolean | null | |
| >; | |
| async function readStorage(): Promise<StorageUsage> { | |
| try { | |
| const est = await navigator.storage?.estimate?.(); | |
| return { usage: est?.usage ?? null, quota: est?.quota ?? null }; | |
| } catch { | |
| return { usage: null, quota: null }; | |
| } | |
| } | |
| /** | |
| * Delete every cache entry across all of this origin's caches whose URL | |
| * matches. WebLLM's `deleteModelAllInfoInCache` is index-driven: it only | |
| * removes the shards listed in `tensor-cache.json`. Shards left behind by an | |
| * interrupted download - or by a browser-evicted index - become orphans it can | |
| * never reach, and they are what keep `storage.estimate().usage` high long | |
| * after the UI reports "not downloaded". Matching on the URL is the only way | |
| * to reach those bytes. | |
| */ | |
| async function purgeCaches(match: (url: string) => boolean): Promise<void> { | |
| if (typeof caches === "undefined") return; | |
| const names = await caches.keys(); | |
| await Promise.all( | |
| names.map(async (name) => { | |
| const cache = await caches.open(name); | |
| const hits = (await cache.keys()).filter((req) => match(req.url)); | |
| await Promise.all(hits.map((req) => cache.delete(req))); | |
| }), | |
| ); | |
| } | |
| /** Purge everything cached for one model's subdir: weights, wasm, tokenizer, index. */ | |
| function purgeModelCache(id: ModelId): Promise<void> { | |
| const base = mlcUrls(id).model; // "…/resolve/main/<subdir>", no trailing slash | |
| return purgeCaches((u) => u === base || u.startsWith(base + "/")); | |
| } | |
| /** Purge every WebLLM artifact for this app, across all quants at once. */ | |
| function purgeAllModelCaches(): Promise<void> { | |
| const marker = `${MLC_REPO}/resolve/main/`; | |
| return purgeCaches((u) => u.includes(marker)); | |
| } | |
| /** | |
| * Reports, for every quant rather than only the selected one, whether its | |
| * weights are still in the browser cache - plus how much storage the origin is | |
| * actually using, which is the figure that explains a cache that went away. | |
| * | |
| * Two independent signals are kept, because WebLLM's alone is not trustworthy | |
| * on its own: `hasModelInCache` needs *every* param shard present, so a browser | |
| * that evicted Cache Storage under pressure reports "not downloaded" for a | |
| * quant this device has loaded many times. Our own record tells those apart. | |
| */ | |
| export function useModelCache(context: ContextPreset): ModelCacheState { | |
| const [cached, setCached] = useState<Record<ModelId, boolean | null>>(UNKNOWN); | |
| const [loadedHere, setLoadedHere] = useState<LoadedModels>(() => readLoadedModels()); | |
| const [storage, setStorage] = useState<StorageUsage>({ usage: null, quota: null }); | |
| const [clearing, setClearing] = useState<ModelId | null>(null); | |
| const [clearingAll, setClearingAll] = useState(false); | |
| const refresh = useCallback(async () => { | |
| try { | |
| const { hasModelInCache } = await import("@mlc-ai/web-llm"); | |
| const results = await Promise.all( | |
| MODEL_IDS.map(async (id) => { | |
| try { | |
| const hit = await hasModelInCache( | |
| mlcVariant(id).cacheLabel, | |
| mlcAppConfig(id, context) | |
| ); | |
| return [id, hit] as const; | |
| } catch { | |
| return [id, null] as const; | |
| } | |
| }) | |
| ); | |
| setCached(Object.fromEntries(results) as Record<ModelId, boolean | null>); | |
| } catch { | |
| setCached({ ...UNKNOWN }); | |
| } | |
| setLoadedHere(readLoadedModels()); | |
| setStorage(await readStorage()); | |
| }, [context]); | |
| const clear = useCallback( | |
| async (id: ModelId) => { | |
| setClearing(id); | |
| try { | |
| const { deleteModelAllInfoInCache } = await import("@mlc-ai/web-llm"); | |
| await deleteModelAllInfoInCache(mlcVariant(id).cacheLabel, mlcAppConfig(id, context)); | |
| } catch { | |
| /* nothing cached to remove */ | |
| } | |
| // WebLLM's delete only removes what its index lists; also purge any shards | |
| // still sitting under this model's subdir so usage actually drops. | |
| await purgeModelCache(id).catch(() => { | |
| /* best-effort */ | |
| }); | |
| // Forget our own record too: after an explicit clear the honest status is | |
| // "not downloaded", not "evicted". | |
| setLoadedHere(clearLoadedModel(id)); | |
| await refresh(); | |
| setClearing(null); | |
| }, | |
| [context, refresh] | |
| ); | |
| const clearAll = useCallback(async () => { | |
| setClearingAll(true); | |
| try { | |
| await purgeAllModelCaches(); | |
| // Forget our own load records so every quant reads "not downloaded". | |
| let rec = readLoadedModels(); | |
| for (const id of MODEL_IDS) rec = clearLoadedModel(id); | |
| setLoadedHere(rec); | |
| } catch { | |
| /* storage unavailable; the usage re-read below reflects reality */ | |
| } | |
| await refresh(); | |
| setClearingAll(false); | |
| }, [refresh]); | |
| const markLoaded = useCallback( | |
| (id: ModelId) => { | |
| setLoadedHere(markLoadedModel(id)); | |
| void refresh(); | |
| }, | |
| [refresh] | |
| ); | |
| useEffect(() => { | |
| void refresh(); | |
| }, [refresh]); | |
| return { cached, loadedHere, storage, clearing, clearingAll, refresh, clear, clearAll, markLoaded }; | |
| } | |