WebLLM: TOCTOU between loadEngine() resolve and acquireInference() lets a concurrent picker switch unload an engine mid-inference

[s-annam]

Summary

loadEngine / acquireInference have a time-of-check-to-time-of-use gap: an engine can be .unload()'d by a concurrent picker switch in the window between loadEngine() resolving and the caller calling acquireInference(). The in-flight-inference tracker that defers eviction (inflightInferenceCount / pendingUnload in src/lib/webllm/web-llm.ts) is exactly the mechanism meant to prevent tear-down-mid-use, but it has a hole at the boundary it guards.

Introduced in #143 (PR B of the M4 in-browser AI rewrite epic, #64).

The race

Consumer flow (src/components/features/RewriteButton.tsx:74-78, same shape in SectionRewrite.tsx):

const engine = await loadEngine(selectedModelId, onProgress); // fast path: resolved engine A
const rewritten = await rewriteBulletWithLlm(bullet, engine, modelId);
//   ^ acquireInference(modelId) is called INSIDE rewriteBulletWithLlm, AFTER loadEngine resolved

acquireInference runs at src/lib/webllm/rewrite-bullet.ts:60 (and rewrite-section.ts), inside the rewrite fn — not before the load resolves.

Interleaving:

1. Engine A already loaded. Rewrite caller hits fast-path A: loadEngine(A) returns Promise.resolve(A). The await yields to the microtask queue.
2. The picker's queued serialChain entry (also a microtask) runs evictAllExcept(B). It removes A from loadedEngines, sees inflightInferenceCount[A] === 0 (acquire not called yet), and calls A.unload() immediately.
3. Rewrite caller's continuation resumes: acquireInference(A) (too late) → A.chat.completions.create() on a torn-down engine.

The inflightInferenceCount/pendingUnload deferral only protects calls that have already acquired. The fast-path resolve → acquire boundary is unguarded.

Blast radius

Narrow. Requires a picker model-switch concurrent with an in-progress rewrite click on a different model. Recoverable: the failed inference surfaces a per-model error and the user retries (UI already supports this). serialChain keeps it from corrupting the cache. Not a data-loss or crash-the-app bug — a rare, recoverable inference failure. Deferred from #143 as a known follow-up, not a merge blocker.

Fix options

1. Acquire before the load resolves. Have consumers acquireInference(modelId) before awaiting loadEngine, paired with releaseInference in finally. Then evictAllExcept sees a positive count and parks the unload. Cost: acquire references a model id whose engine may not exist yet — evictAllExcept would need to tolerate parking an id that isn't in loadedEngines (no-op park).
2. Bracket inside loadEngine's fast path. Return a { engine, release } handle from loadEngine that has already incremented the count, so there's no caller-side gap. Changes the loadEngine signature / all call sites.
3. Re-check after acquire. After acquireInference, verify the engine is still in loadedEngines; if evicted, re-load. Adds a retry hop but no signature change.

Option 1 or 2 closes the window cleanly; 3 is a smaller patch that tolerates the race instead of eliminating it.

Acceptance

• A test reproduces the interleave: load A, start a rewrite on A, fire a picker switch to B in the same tick, assert A is not unloaded until the rewrite's releaseInference.
• Fix chosen and applied; existing web-llm cache/serialization tests stay green.

Found during review of #143.