@qvac/sdk 0.17.1 → 0.18.0

package/dist/server/bare/plugins/llamacpp-completion/ops/kv-cache-session.js

@@ -11,36 +11,22 @@ import { getServerLogger } from '../../../../../logging/index.js';
 // the caller — typically `withRequestContext(...)`.
 const moduleLogger = getServerLogger();
 /**
- * Coordinates five KV-cache state layers:
- *
- * 1. `cachedMessageCounts` — saved message boundaries.
- * 2. `initializedCaches` — caches primed in this worker.
- * 3. On-disk `.bin` files written by the addon.
- * 4. `activeCachePaths` — per-path refs that block in-flight eviction.
- * 5. `.auto-cache-<key>` markers — SDK-generated cache ownership.
- *
- * Every turn must finish through `commitTurn` or `rollback` so all
- * inference state stays aligned, the active-path ref is released, and
- * marker metadata follows the cache directory lifecycle.
+ * What the kv-cache file on disk is known to cover, keyed by cache path.
+ * Written by `commitTurn`, read by `getSavedCount`, deleted by `rollback` /
+ * `delete` / `dropStaleSavedCount`. The same INVARIANT that existed in
+ * `kv-cache-state.ts` still holds: an entry is present only when the
+ * corresponding `.bin` file is considered trustworthy. Cancelled or
+ * zero-token turns must remove the entry so the next-turn slice doesn't read
+ * a stale boundary.
  */
-// ----- module-scoped state. The session is the single mutation point
-// for the in-memory KV-cache bookkeeping. -----
-/**
- * Number of chat messages the kv-cache file on disk is known to cover,
- * keyed by cache path. Written by `commitTurn`, read by `getSavedCount`,
- * deleted by `rollback` / `delete` / `dropStaleSavedCount`. The same
- * INVARIANT that existed in `kv-cache-state.ts` still holds: an entry is
- * present only when the corresponding `.bin` file is considered
- * trustworthy. Cancelled or zero-token turns must remove the entry so
- * the next-turn slice doesn't read a stale boundary.
- */
-const cachedMessageCounts = new Map();
+const cachedPrefixes = new Map();
 /**
  * In-memory registry of caches initialized this session. The addon
  * defers disk writes, so the absence of a `.bin` file on disk isn't
  * proof that the cache hasn't been primed in this worker process. Keyed
- * by `${modelId}:${configHash}:${cacheKey}`, so on-disk caches from
- * older worker runs still hit the lazy-load path in `beginTurn`.
+ * by the resolved cache path, so aliased keys that name one file share an
+ * entry and on-disk caches from older worker runs still hit the lazy-load
+ * path in `beginTurn`.
  */
 const initializedCaches = new Set();
 const activeCachePaths = new Map();
@@ -51,8 +37,72 @@ const AUTO_CACHE_SWEEP_INTERVAL_MS = 5 * 60 * 1000;
 let lastAutoCacheSweepMs = 0;
 let autoCacheSweepInFlight = null;
 let cacheStateLockTail = Promise.resolve();
-function initRegistryKey(modelId, configHash, cacheKey) {
-    return `${modelId}:${configHash}:${cacheKey}`;
+const cachePathLocks = new Map();
+// Internal sentinel so acquisition paths prune directories/markers only when a
+// queued lock wait was aborted, not when lock acquisition failed for another
+// reason.
+class CacheLockAbortError extends Error {
+    constructor(cause) {
+        super('cache lock wait aborted', cause !== undefined ? { cause } : undefined);
+        this.name = 'CacheLockAbortError';
+    }
+}
+// Case-fold the lock-map key so case-only path variants (e.g. "Session" vs
+// "session"), which name the SAME file on case-insensitive filesystems (default
+// macOS/Windows), serialise on one lock and can't interleave writes. Over-locks
+// case-only variants on case-sensitive filesystems, which is safe. NOTE: the
+// other per-path bookkeeping (activeCachePaths, cachedPrefixes, retention)
+// stays case-sensitive, so case-only key variants are otherwise unsupported —
+// see the KV-cache system doc.
+function lockKeyFor(cachePath) {
+    return cachePath.toLowerCase();
+}
+function releaseCachePathWriteLock(lockKey) {
+    const lock = cachePathLocks.get(lockKey);
+    if (!lock)
+        return;
+    const next = lock.waiters.shift();
+    if (next) {
+        next.grant();
+        return;
+    }
+    lock.held = false;
+    cachePathLocks.delete(lockKey);
+}
+// Returns the release fn. Aborting while queued rejects and removes the waiter.
+async function acquireCachePathWriteLock(cachePath, signal) {
+    const lockKey = lockKeyFor(cachePath);
+    let lock = cachePathLocks.get(lockKey);
+    if (!lock) {
+        lock = { held: false, waiters: [] };
+        cachePathLocks.set(lockKey, lock);
+    }
+    if (!lock.held) {
+        lock.held = true;
+        return () => releaseCachePathWriteLock(lockKey);
+    }
+    await new Promise((resolve, reject) => {
+        const waiter = {
+            grant: () => {
+                signal?.removeEventListener('abort', onAbort);
+                resolve();
+            },
+            drop: (reason) => reject(reason)
+        };
+        const onAbort = () => {
+            const index = lock.waiters.indexOf(waiter);
+            if (index !== -1)
+                lock.waiters.splice(index, 1);
+            waiter.drop(new CacheLockAbortError(signal?.reason));
+        };
+        if (signal?.aborted) {
+            onAbort();
+            return;
+        }
+        lock.waiters.push(waiter);
+        signal?.addEventListener('abort', onAbort, { once: true });
+    });
+    return () => releaseCachePathWriteLock(lockKey);
 }
 function markCachePathActive(cachePath) {
     activeCachePaths.set(cachePath, (activeCachePaths.get(cachePath) ?? 0) + 1);
@@ -67,6 +117,10 @@ function releaseCachePath(cachePath) {
     }
     activeCachePaths.set(cachePath, count - 1);
 }
+// Paths in-flight turns hold, so pruning skips a directory still in use.
+function snapshotActivePaths() {
+    return Array.from(activeCachePaths.keys());
+}
 function isCacheKeyActive(cacheKey) {
     const cacheDirectory = path.join(getKVCacheDir(), cacheKey);
     return Array.from(activeCachePaths.keys()).some((cachePath) => isCachePathWithinDirectory(cacheDirectory, cachePath));
@@ -151,15 +205,18 @@ export function createKvCacheSession(modelId, options) {
     // WeakMap so handles drop their state once the handler scope releases
     // the reference; the module-scoped maps above survive.
     const turnState = new WeakMap();
-    function makeHandle(cachePath, registryKey, autoCacheKey) {
+    function makeHandle(cachePath, autoCacheKey, releaseWriteLock = () => { }, signal) {
+        const cached = cachedPrefixes.get(cachePath);
         const handle = {
             cachePath,
-            savedCount: cachedMessageCounts.get(cachePath) ?? 0
+            savedCount: cached?.messages ?? 0,
+            toolBlockCached: cached?.toolBlock ?? false
         };
         turnState.set(handle, {
             cachePath,
-            registryKey,
             ...(autoCacheKey !== undefined && { autoCacheKey }),
+            ...(signal !== undefined && { signal }),
+            releaseWriteLock,
             committed: false,
             rolledBack: false
         });
@@ -168,19 +225,41 @@ export function createKvCacheSession(modelId, options) {
     }
     async function beginCustom(input) {
         const cachePath = await getCacheFilePath(modelId, input.configHash, input.customKey);
-        const registryKey = initRegistryKey(modelId, input.configHash, input.customKey);
-        const handle = makeHandle(cachePath, registryKey);
+        // Held for the whole turn so a same-file peer can't interleave writes.
+        // Released on commit/rollback/error; throws if aborted while queued.
+        let releaseWriteLock;
+        try {
+            releaseWriteLock = await acquireCachePathWriteLock(cachePath, input.signal);
+        }
+        catch (error) {
+            // Aborted while queued: no lock is held to release, and getCacheFilePath
+            // already mkdir'd this key's parent (a distinct dir under a case variant),
+            // so prune it before the cancellation propagates.
+            if (error instanceof CacheLockAbortError) {
+                await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
+            }
+            throw error;
+        }
+        // A turn cancelled by the time it holds the lock must not prime (native work).
+        // getCacheFilePath already mkdir'd the parent, so prune it before surfacing
+        // the cancellation the plugin rides.
+        if (input.signal?.aborted) {
+            await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
+            releaseWriteLock();
+            throw new CacheLockAbortError(input.signal.reason);
+        }
+        const handle = makeHandle(cachePath, undefined, releaseWriteLock, input.signal);
         try {
             // In-memory registry check first — the addon defers disk writes, so
             // a freshly-primed cache may not yet exist on disk. If the
             // in-memory flag isn't set, fall back to a filesystem probe so
             // caches surviving across worker restarts still hit the reuse path.
-            let exists = initializedCaches.has(registryKey);
+            let exists = initializedCaches.has(cachePath);
             if (!exists) {
                 try {
                     await fsPromises.access(cachePath);
                     exists = true;
-                    initializedCaches.add(registryKey);
+                    initializedCaches.add(cachePath);
                 }
                 catch {
                     exists = false;
@@ -188,49 +267,109 @@ export function createKvCacheSession(modelId, options) {
             }
             logCacheStatus(input.customKey, exists);
             if (!exists) {
+                // Recreate the parent dir if a same-key peer's rollback pruned it after our lock wait.
+                await fsPromises.mkdir(path.dirname(cachePath), { recursive: true });
+                // The access probe / mkdir above yielded, so a cancel may have landed
+                // since the acquire-time check — re-check right before native priming.
+                // The catch below prunes the directory just created.
+                if (input.signal?.aborted)
+                    throw new CacheLockAbortError(input.signal.reason);
                 await input.primeIfMissing(cachePath);
                 await verifyPrimedFile(cachePath, logger);
-                initializedCaches.add(registryKey);
+                initializedCaches.add(cachePath);
             }
             return handle;
         }
         catch (error) {
             releaseCachePath(cachePath);
+            await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
+            releaseWriteLock();
             throw error;
         }
     }
     async function beginAuto(input) {
-        // The pre-response cache key is derived from
-        // `history.slice(0, -1)` — `findMatchingCache` does that
-        // internally. The post-response key (used after a successful turn)
-        // is computed by the caller and passed to `commitTurn` as
-        // `targetCachePath`.
-        const setup = await withCacheStateLock(async () => {
+        // The path comes from a history lookup, so: discover it, lock it (outside
+        // the state lock, or commitTurn's rename could deadlock a waiter), then
+        // re-check existence under the lock in case a same-file peer just renamed it.
+        const discovered = await withCacheStateLock(async () => {
             const existingCache = await findMatchingCache(modelId, input.configHash, input.history);
             const cacheInfo = existingCache ?? (await getCurrentCacheInfo(modelId, input.configHash, input.history));
-            const registryKey = initRegistryKey(modelId, input.configHash, cacheInfo.cacheKey);
             return {
-                cacheExists: existingCache !== null,
                 cachePath: cacheInfo.cachePath,
-                cacheKey: cacheInfo.cacheKey,
-                registryKey,
-                handle: makeHandle(cacheInfo.cachePath, registryKey, cacheInfo.cacheKey)
+                cacheKey: cacheInfo.cacheKey
             };
         });
-        const { cacheExists, cachePath, cacheKey, registryKey, handle } = setup;
+        const { cachePath, cacheKey } = discovered;
+        let releaseWriteLock;
+        try {
+            releaseWriteLock = await acquireCachePathWriteLock(cachePath, input.signal);
+        }
+        catch (error) {
+            // Aborted while queued: no lock is held to release. Discovery already
+            // mkdir'd the parent and wrote the auto marker; clean both before the
+            // cancellation propagates.
+            if (error instanceof CacheLockAbortError) {
+                await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
+                await removeAutoCacheMarkerIfMissing(cacheKey);
+            }
+            throw error;
+        }
+        // A turn cancelled by the time it holds the lock must not prime (native work).
+        // Discovery already mkdir'd the parent and wrote the auto marker; clean both
+        // before surfacing the cancellation.
+        if (input.signal?.aborted) {
+            await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
+            await removeAutoCacheMarkerIfMissing(cacheKey);
+            releaseWriteLock();
+            throw new CacheLockAbortError(input.signal.reason);
+        }
+        let handle;
+        let cacheExists;
+        try {
+            const resolved = await withCacheStateLock(async () => {
+                // Mark the path active under the state lock so retention can't evict it.
+                const h = makeHandle(cachePath, cacheKey, releaseWriteLock, input.signal);
+                let exists = initializedCaches.has(cachePath);
+                if (!exists) {
+                    try {
+                        await fsPromises.access(cachePath);
+                        exists = true;
+                        initializedCaches.add(cachePath);
+                    }
+                    catch {
+                        exists = false;
+                    }
+                }
+                return { h, exists };
+            });
+            handle = resolved.h;
+            cacheExists = resolved.exists;
+        }
+        catch (error) {
+            releaseWriteLock();
+            throw error;
+        }
         logCacheStatus('auto', cacheExists);
         try {
             if (!cacheExists) {
+                // Recreate the parent dir if a same-file peer's rename pruned it.
+                await fsPromises.mkdir(path.dirname(cachePath), { recursive: true });
+                // The discovery / mkdir above yielded, so a cancel may have landed since
+                // the acquire-time check — re-check right before native priming. The
+                // catch below prunes the directory and auto marker.
+                if (input.signal?.aborted)
+                    throw new CacheLockAbortError(input.signal.reason);
                 await input.primeIfMissing(cachePath);
                 await verifyPrimedFile(cachePath, logger);
-                initializedCaches.add(registryKey);
+                initializedCaches.add(cachePath);
             }
             return handle;
         }
         catch (error) {
             releaseCachePath(cachePath);
-            await pruneEmptyCacheDirectories(cachePath);
+            await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
             await removeAutoCacheMarkerIfMissing(cacheKey);
+            releaseWriteLock();
             throw error;
         }
     }
@@ -255,7 +394,7 @@ export function createKvCacheSession(modelId, options) {
             // at the same path. Verify the file persisted (the addon
             // currently swallows save errors — see TODO in
             // `verifySaveAndRecord`) and record the new boundary.
-            const ok = await verifySaveAndRecord(state.cachePath, result.messageCount);
+            const ok = await verifySaveAndRecord(state.cachePath, result.messageCount, result.toolBlockCached);
             if (!ok) {
                 // The expected save didn't land — treat the turn as a rollback
                 // so the next turn re-primes cleanly.
@@ -264,51 +403,79 @@ export function createKvCacheSession(modelId, options) {
             }
             state.committed = true;
             releaseCachePath(state.cachePath);
+            state.releaseWriteLock();
             return;
         }
         const sourceCachePath = state.cachePath;
         const sourceCacheKey = state.autoCacheKey;
         const targetCacheKey = path.basename(path.dirname(path.dirname(result.targetCachePath)));
-        const renamed = await withCacheStateLock(async () => {
-            markCachePathActive(result.targetCachePath);
-            await fsPromises.mkdir(path.dirname(result.targetCachePath), { recursive: true });
-            await markAutoCacheKey(targetCacheKey);
-            if (!(await renameCacheFile(sourceCachePath, result.targetCachePath))) {
-                releaseCachePath(result.targetCachePath);
-                await pruneEmptyCacheDirectories(result.targetCachePath);
-                await removeAutoCacheMarkerIfMissing(targetCacheKey);
-                return false;
+        // Hold the target's lock across the rename AND commit verification, so a peer
+        // resolving to the same file can't observe it before its saved count is
+        // published. Abortable: a cancel while waiting for the lock rejects here.
+        const releaseTargetLock = await acquireCachePathWriteLock(result.targetCachePath, state.signal);
+        try {
+            // A cancel that landed while we waited for the target lock must not commit.
+            if (state.signal?.aborted) {
+                await runRollback(state);
+                return;
             }
-            releaseCachePath(sourceCachePath);
-            state.cachePath = result.targetCachePath;
-            state.autoCacheKey = targetCacheKey;
-            await pruneEmptyCacheDirectories(sourceCachePath);
-            if (sourceCacheKey !== undefined) {
-                await removeAutoCacheMarkerIfMissing(sourceCacheKey);
+            const renamed = await withCacheStateLock(async () => {
+                markCachePathActive(result.targetCachePath);
+                try {
+                    await fsPromises.mkdir(path.dirname(result.targetCachePath), { recursive: true });
+                    await markAutoCacheKey(targetCacheKey);
+                    if (!(await renameCacheFile(sourceCachePath, result.targetCachePath))) {
+                        releaseCachePath(result.targetCachePath);
+                        await pruneEmptyCacheDirectories(result.targetCachePath, snapshotActivePaths());
+                        await removeAutoCacheMarkerIfMissing(targetCacheKey);
+                        return false;
+                    }
+                }
+                catch (setupError) {
+                    // mkdir/markAutoCacheKey/rename threw after the target active-ref was
+                    // taken: release it (first, so it can't leak) and prune any directory
+                    // or marker created above, then propagate.
+                    releaseCachePath(result.targetCachePath);
+                    await pruneEmptyCacheDirectories(result.targetCachePath, snapshotActivePaths());
+                    await removeAutoCacheMarkerIfMissing(targetCacheKey);
+                    throw setupError;
+                }
+                releaseCachePath(sourceCachePath);
+                state.cachePath = result.targetCachePath;
+                state.autoCacheKey = targetCacheKey;
+                await pruneEmptyCacheDirectories(sourceCachePath, snapshotActivePaths());
+                if (sourceCacheKey !== undefined) {
+                    await removeAutoCacheMarkerIfMissing(sourceCacheKey);
+                }
+                cachedPrefixes.delete(sourceCachePath);
+                // state.cachePath was just reassigned to the target; clear the SOURCE
+                // entry, not the freshly-valid target.
+                initializedCaches.delete(sourceCachePath);
+                return true;
+            });
+            if (!renamed) {
+                logger.warn(`[kv-cache] Auto cache rename failed; rolling back. from=${sourceCachePath} to=${result.targetCachePath}`);
+                await runRollback(state);
+                return;
             }
-            cachedMessageCounts.delete(sourceCachePath);
-            initializedCaches.delete(state.registryKey);
-            return true;
-        });
-        if (!renamed) {
-            logger.warn(`[kv-cache] Auto cache rename failed; rolling back. from=${sourceCachePath} to=${result.targetCachePath}`);
-            await runRollback(state);
-            return;
+            const ok = await verifySaveAndRecord(result.targetCachePath, result.messageCount, result.toolBlockCached);
+            if (!ok) {
+                // Rename succeeded but the file isn't where we expected. Roll back via
+                // the target path instead of the (now-empty) source.
+                await runRollback(state);
+                return;
+            }
+            // Successful auto-rename. The handle's `cachePath` field still points at
+            // the (now-gone) source path — fine, the handle is committed and won't
+            // roll back. Future turns compute fresh paths.
+            state.committed = true;
+            releaseCachePath(state.cachePath);
+            state.releaseWriteLock();
+            scheduleAutoCacheSweep(logger);
         }
-        const ok = await verifySaveAndRecord(result.targetCachePath, result.messageCount);
-        if (!ok) {
-            // Rename succeeded but the file isn't where we expected. Roll
-            // back via the target path instead of the (now-empty) source.
-            await runRollback(state);
-            return;
+        finally {
+            releaseTargetLock();
         }
-        // Successful auto-rename. The handle's `cachePath` field still
-        // points at the (now-gone) source path — that's fine, the handle
-        // is committed and won't roll back. Future turns compute fresh
-        // paths.
-        state.committed = true;
-        releaseCachePath(state.cachePath);
-        scheduleAutoCacheSweep(logger);
     }
     async function rollback(turn) {
         const state = turnState.get(turn);
@@ -329,14 +496,17 @@ export function createKvCacheSession(modelId, options) {
         catch (unlinkError) {
             logger.warn(`[kv-cache] Failed to remove cache file during rollback; next turn may load stale KV state. path=${state.cachePath} error=${unlinkError instanceof Error ? unlinkError.message : String(unlinkError)}`);
         }
-        await pruneEmptyCacheDirectories(state.cachePath);
+        // Release before pruning so an empty parent can go; a sibling still holding
+        // the path keeps it in the active snapshot and protects the directory.
+        releaseCachePath(state.cachePath);
+        await pruneEmptyCacheDirectories(state.cachePath, snapshotActivePaths());
         if (state.autoCacheKey !== undefined) {
             await removeAutoCacheMarkerIfMissing(state.autoCacheKey);
         }
-        initializedCaches.delete(state.registryKey);
-        cachedMessageCounts.delete(state.cachePath);
-        releaseCachePath(state.cachePath);
+        initializedCaches.delete(state.cachePath);
+        cachedPrefixes.delete(state.cachePath);
         state.rolledBack = true;
+        state.releaseWriteLock();
         if (state.autoCacheKey !== undefined)
             scheduleAutoCacheSweep(logger);
     }
@@ -344,7 +514,7 @@ export function createKvCacheSession(modelId, options) {
         const state = turnState.get(turn);
         if (!state)
             return;
-        cachedMessageCounts.delete(state.cachePath);
+        cachedPrefixes.delete(state.cachePath);
     }
     return {
         beginTurn,
@@ -370,24 +540,26 @@ export function createKvCacheSession(modelId, options) {
  * Layers cleared, in order:
  *   1. On-disk: `deleteCache(...)` removes the matching directory
  *      tree (or wipes and recreates the root for `all: true`).
- *   2. `cachedMessageCounts`: prefix-cleanup by the removed directory
+ *   2. `cachedPrefixes`: prefix-cleanup by the removed directory
  *      so any per-cache count under the deleted tree is forgotten.
- *   3. `initializedCaches`: scope clear by `(kvCacheKey[, modelId])`,
- *      matching the on-disk scope.
+ *   3. `initializedCaches`: prefix-cleanup by the same removed
+ *      directory, since it is keyed by the resolved cache path.
  *
- * Concurrency with in-flight turns: this delete is wire-async with
- * respect to any turn currently holding a `TurnHandle` for the same
- * cache key. Worst case the on-disk `.bin` is removed while a turn is
- * mid-write; the turn's eventual `commitTurn(...)` then fails the
- * `verifySaveAndRecord` probe (file gone) and rolls back idempotently.
- * No coordination primitive is needed because every layer's mutation
- * is idempotent (`unlink` no-ops if missing, `Map.delete` / `Set.delete`
- * no-op on absent keys).
+ * Concurrency with in-flight turns: this delete does not take the
+ * per-cache-path write locks, so it races any turn holding a `TurnHandle`
+ * for the same key. Deleting a key that is in active use is unsupported.
+ * Each individual mutation is idempotent (`unlink` no-ops if missing,
+ * `Map.delete` / `Set.delete` no-op on absent keys), but the layers are
+ * cleared without a lock, so an interleaving that lands the delete's
+ * in-memory cleanup after a concurrent turn has already renamed/committed
+ * its file splits state: the file stays on disk while its saved-count and
+ * init flag are cleared, so the next turn sees the file, skips priming, and
+ * reports `savedCount=0`. Callers must not delete a key that is in active use.
  */
 export async function deleteKvCacheState(target) {
     if ('all' in target) {
         const removed = await deleteCacheUtil({ all: true });
-        cachedMessageCounts.clear();
+        cachedPrefixes.clear();
         initializedCaches.clear();
         // `removed` is the kv-cache root dir; surfaces it for ops
         // visibility but isn't part of the contract.
@@ -399,20 +571,20 @@ export async function deleteKvCacheState(target) {
         ...(target.modelId !== undefined && { modelId: target.modelId })
     });
     if (target.modelId === undefined) {
-        await removeAutoCacheMarker(target.kvCacheKey);
+        // Remove the marker by the root-relative key: an alias like "./<16hex>"
+        // still resolves to the canonical auto key, while a nested key such as
+        // "tenant/<16hex>" stays distinct — path.basename would wrongly collapse it
+        // onto the unrelated top-level auto marker.
+        await removeAutoCacheMarker(path.relative(getKVCacheDir(), removedPath));
     }
     // Prefix-cleanup the in-memory counts. The on-disk directory tree
     // is `{kvCacheRoot}/{kvCacheKey}[/{modelId}]/`, so every entry in
-    // `cachedMessageCounts` whose key is the removed directory itself
+    // `cachedPrefixes` whose key is the removed directory itself
     // or sits beneath it must go.
     clearCachedMessageCountsByPrefix(removedPath, path.sep);
-    // The in-memory init-set keys are
-    // `${modelId}:${configHash}:${kvCacheKey}` — clear by the user-
-    // facing kvCacheKey (and optionally narrow by modelId).
-    clearInitializedCachesByScope({
-        cacheKey: target.kvCacheKey,
-        ...(target.modelId !== undefined && { modelId: target.modelId })
-    });
+    // initializedCaches is keyed by the resolved cachePath too, so clear it by
+    // the same removed-directory prefix as cachedPrefixes above.
+    clearInitializedCachesByPrefix(removedPath, path.sep);
 }
 // ----- private helpers -----
 /**
@@ -454,7 +626,7 @@ async function verifyPrimedFile(cachePath, logger) {
     catch (statError) {
         // ENOENT is the common case here — addon prime returned without
         // calling save (most often: signal abort during prefill).
-        await pruneEmptyCacheDirectories(cachePath);
+        await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
         throw new Error(`[kv-cache] prime closure resolved but no cache file was written. path=${cachePath} cause=${statError instanceof Error ? statError.message : String(statError)}`);
     }
     if (stats.size === 0) {
@@ -463,7 +635,7 @@ async function verifyPrimedFile(cachePath, logger) {
         // primary "prime didn't persist" condition.
         try {
             await fsPromises.unlink(cachePath);
-            await pruneEmptyCacheDirectories(cachePath);
+            await pruneEmptyCacheDirectories(cachePath, snapshotActivePaths());
         }
         catch (unlinkError) {
             logger.warn(`[kv-cache] Failed to remove empty primed cache file. path=${cachePath} error=${unlinkError instanceof Error ? unlinkError.message : String(unlinkError)}`);
@@ -482,52 +654,41 @@ async function verifyPrimedFile(cachePath, logger) {
  * false), drop the `access()` probe and wrap the `model.run()` call in
  * a real try/catch that forwards the error.
  */
-async function verifySaveAndRecord(cachePath, messageCount) {
+async function verifySaveAndRecord(cachePath, messageCount, toolBlockCached) {
     try {
         await fsPromises.access(cachePath);
-        cachedMessageCounts.set(cachePath, messageCount);
+        cachedPrefixes.set(cachePath, { messages: messageCount, toolBlock: toolBlockCached });
         return true;
     }
     catch (err) {
-        cachedMessageCounts.delete(cachePath);
+        cachedPrefixes.delete(cachePath);
         logCacheSaveError(cachePath, err);
         return false;
     }
 }
 function clearCachedMessageCountsByPrefix(prefix, sep) {
     if (!prefix) {
-        cachedMessageCounts.clear();
+        cachedPrefixes.clear();
         return;
     }
-    for (const key of cachedMessageCounts.keys()) {
+    for (const key of cachedPrefixes.keys()) {
         if (key === prefix) {
-            cachedMessageCounts.delete(key);
+            cachedPrefixes.delete(key);
             continue;
         }
         if (!key.startsWith(prefix + sep))
             continue;
-        cachedMessageCounts.delete(key);
+        cachedPrefixes.delete(key);
     }
 }
-function clearInitializedCachesByScope(scope) {
-    if (scope.cacheKey === undefined && scope.modelId === undefined) {
+function clearInitializedCachesByPrefix(prefix, sep) {
+    if (!prefix) {
         initializedCaches.clear();
         return;
     }
     for (const key of initializedCaches) {
-        const firstSep = key.indexOf(':');
-        const secondSep = key.indexOf(':', firstSep + 1);
-        if (firstSep === -1 || secondSep === -1)
-            continue;
-        const entryModelId = key.slice(0, firstSep);
-        const entryCacheKey = key.slice(secondSep + 1);
-        if (scope.cacheKey !== undefined && entryCacheKey !== scope.cacheKey) {
-            continue;
-        }
-        if (scope.modelId !== undefined && entryModelId !== scope.modelId) {
-            continue;
-        }
-        initializedCaches.delete(key);
+        if (key === prefix || key.startsWith(prefix + sep))
+            initializedCaches.delete(key);
     }
 }
 // ----- test-only escape hatches -----
@@ -542,16 +703,19 @@ function clearInitializedCachesByScope(scope) {
  */
 export const __kvCacheSessionTestHooks = {
     getSavedCount(cachePath) {
-        return cachedMessageCounts.get(cachePath);
+        return cachedPrefixes.get(cachePath)?.messages;
     },
     setSavedCountForTest(cachePath, count) {
-        cachedMessageCounts.set(cachePath, count);
+        cachedPrefixes.set(cachePath, { messages: count, toolBlock: false });
+    },
+    getToolBlockCachedForTest(cachePath) {
+        return cachedPrefixes.get(cachePath)?.toolBlock ?? false;
     },
-    hasInitializedKey(modelId, configHash, cacheKey) {
-        return initializedCaches.has(initRegistryKey(modelId, configHash, cacheKey));
+    hasInitializedPath(cachePath) {
+        return initializedCaches.has(cachePath);
     },
-    markInitializedForTest(modelId, configHash, cacheKey) {
-        initializedCaches.add(initRegistryKey(modelId, configHash, cacheKey));
+    getActivePathCountForTest(cachePath) {
+        return activeCachePaths.get(cachePath) ?? 0;
     },
     getLastAutoCacheSweepMsForTest() {
         return lastAutoCacheSweepMs;
@@ -560,12 +724,13 @@ export const __kvCacheSessionTestHooks = {
         lastAutoCacheSweepMs = value;
     },
     resetForTest() {
-        cachedMessageCounts.clear();
+        cachedPrefixes.clear();
         initializedCaches.clear();
         activeCachePaths.clear();
         lastAutoCacheSweepMs = 0;
         autoCacheSweepInFlight = null;
         cacheStateLockTail = Promise.resolve();
+        cachePathLocks.clear();
     },
     waitForAutoCacheSweepForTest() {
         return autoCacheSweepInFlight ?? Promise.resolve();