mlx-bun 0.5.0 → 1.0.0

package/src/prompt-cache.ts

@@ -1,682 +0,0 @@
-import { type CheckpointAttachment, materializeCheckpoint, attachmentBytes, cloneAttachments, disposeAttachments } from "./backends/mlx/checkpoint-state";
-import { leaseCacheState, minimumReusableOffset } from "./backends/mlx/state-views";
-import { withResource, cleanupFailure, disposeResources, ownResource } from "./engine/resources";
-// Byte-capped LRU prompt cache — the RAM tier of the Layer-0 KV store.
-//
-// The mlx-lm lesson (PLAN.md): a count-capped cache of multi-GB KV
-// entries is an OOM footgun. Ours accounts bytes (sum of KV array bytes
-// per entry) and evicts least-recently-used until under the cap.
-//
-// Usage pattern (single generation queue, so take/put is race-free):
-//   const hit = cache.take(promptIds);   // longest-prefix match
-//   generate(model, promptIds, { cache: hit?.caches ?? fresh })
-//   cache.put([...promptIds, ...generated], caches, ns, hit?.retain);
-// take() lends zero-copy views and leaves every RAM donor in place, including
-// recurrent and method state. SSD persistence runs independently of RAM reuse.
-// put() supersedes same-ns prefix-ancestors when the new entry is trimmable.
-//
-// TIERING (Layer 0, unified-engine plan): when a ColdTier is attached,
-// take() itself runs the two-tier dance — RAM peek vs cold find, restore
-// (streamed copy) + trim when the cold tier holds a strictly longer
-// prefix — so EVERY consumer (the serial lane, the batch scheduler, future
-// prefix sharing) gets SSD restores through the same take()/put() it
-// already calls. Eviction spills to the tier (the #spill hook); idle
-// entries demote to it when policy permits. TieredPromptCache owns queued
-// persistence and promotes restored prefixes into RAM.
-
-import type { Cache } from "./model/gemma4";
-import type { PrefixCache, PrefixCacheHit } from "./contracts/prefix-cache";
-import { LruRetention, type RetentionPolicy, type RetentionCandidate } from "./storage/retention-policy";
-import { cloneKvCaches } from "./kv-store";
-
-/** Reference-counted release: wraps an entry's `retain` (e.g. an mmap
- *  unmap) so DONOR + CLONES can each hold a share — the underlying hook
- *  runs once, after the LAST holder releases. Prefix sharing hands out
- *  zero-copy VIEWS of a donor's buffers; if the donor's unmap ran while a
- *  clone still referenced the mapped pages, the clone would read freed
- *  memory. */
-function makeSharedRetain(retain: (() => void) | undefined): { acquire(): () => void } {
-  let count = 0;
-  let released = false;
-  return {
-    acquire() {
-      count++;
-      let mine = false;
-      return () => {
-        if (mine) return; // idempotent per share
-        mine = true;
-        count--;
-        if (count === 0 && !released) {
-          released = true;
-          retain?.();
-        }
-      };
-    },
-  };
-}
-
-/** The cold (SSD) tier the RAM cache tiers over. Structural — this module
- *  never imports ssd-cache (which imports us); the server binds
- *  SsdCacheStore + the model into this shape (src/server.ts). Lookup and
- *  restore failures return null; store returns false on failure. Operations
- *  that require persistence retain their donor if storing fails. */
-export interface ColdTier {
-  /** Longest stored usable prefix for prompt/ns — index-only, no I/O.
-   *  `handle` is the tier's opaque entry token, passed back to restore. */
-  find(prompt: number[], ns: string): { prefixLen: number; handle: unknown } | null;
-  /** Restore GPU-visible state. The default SSD backend copies into owned
-   *  arrays. An optional backing release runs after consumers dispose their
-   *  arrays. Null on any failure. */
-  restore(handle: unknown): { tokens: number[]; caches: Cache[]; attachments?: CheckpointAttachment[]; retain: () => void } | null;
-  findExact?(tokens: number[], ns: string): { prefixLen: number; handle: unknown } | null;
-  restoreAsync?(handle: unknown): Promise<ReturnType<ColdTier["restore"]>>;
-  /** Tensor objects have no target layers; restore with their own schema. */
-  restoreObjectAsync?(handle: unknown): Promise<ReturnType<ColdTier["restore"]>>;
-  /** Borrow live caches and finish persistence before returning. A false
-   *  result prevents a required demotion. */
-  store(tokens: number[], caches: Cache[], ns: string, attachments?: CheckpointAttachment[]): boolean | void;
-}
-
-/** Spill sink for eviction/demotion (write-behind, 2026-07-06).
- *
- *  `spillOwned` (preferred, NON-BLOCKING): the cache hands the sink an
- *  entry it OWNS — zero-copy CLONES of the caches (made via the same
- *  injectable cloner take() uses, BEFORE the donor is disposed; entries
- *  are immutable so the clones stay consistent forever) plus a copy of
- *  the tokens. The sink must dispose entry.caches when done, on BOTH
- *  fulfill and reject paths (the server chains storeAsync -> dispose on
- *  ssdWriteChain). A sink that throws SYNCHRONOUSLY has not taken
- *  ownership — the cache disposes the clones and the spill degrades to
- *  a no-op; eviction/demotion never unwinds.
- *
- *  `spillSync` (fallback): called with the LIVE entry about to be
- *  disposed; must finish reading before returning. This is the old
- *  contract — a full synchronous serialize+write under the generation
- *  lock — kept for tests and simple embedders. */
-export interface SpillSink {
-  spillOwned?: (entry: PromptCacheEntry) => void;
-  spillSync?: (entry: PromptCacheEntry) => void;
-}
-
-export interface PromptCacheEntry extends PrefixCacheHit<Cache[], CheckpointAttachment[]> {
-  /** Namespace key — adapter spec for LoRA requests ("" = base model).
-   *  KV computed under one adapter must never seed another's prefill. */
-  ns: string;
-  /** Release hook for entries restored from the SSD cold tier: unmaps the
-   *  backing COW file mapping. MUST run only after `caches` are disposed
-   *  (they may alias file pages); every dispose path here honors that, and
-   *  callers who take() an entry carry the thunk with it. */
-  retain?: () => void;
-}
-
-export function cacheBytes(caches: Cache[]): number {
-  let total = 0;
-  for (const c of caches) {
-    total += withResource(leaseCacheState(c), (state) => state.reduce((bytes, a) => bytes + a.nbytes, 0));
-  }
-  return total;
-}
-
-/** Length of the longest common prefix of two token sequences. */
-export function commonPrefixLength(a: number[], b: number[]): number {
-  const n = Math.min(a.length, b.length);
-  let i = 0;
-  while (i < n && a[i] === b[i]) i++;
-  return i;
-}
-
-interface EntryRecord extends RetentionCandidate {
-  entry: PromptCacheEntry;
-  bytes: number;
-  lastUsed: number;
-  lastUsedMs: number;
-  /** Ref-counted retain: the donor holds one share (entry.retain); every
-   *  clone handed out by take() holds another. */
-  share: { acquire(): () => void };
-  resident: boolean;
-  sessions: number;
-}
-
-export class PromptCache implements PrefixCache<Cache[], CheckpointAttachment[]> {
-  /** Exact tensor objects share this store's byte budget, eviction and writer.
-   * Empty token history denotes no conversation; the private namespace is the
-   * complete object identity. No fabricated language tokens enter the index. */
-  readonly #objects: import("./contracts/object-cache").ObjectCache<CheckpointAttachment[]> = {
-    take: key => this.#takeObject(`object:${key}`),
-    put: (key, value) => this.put([], [], `object:${key}`, undefined, value),
-  };
-
-  get objects(): import("./contracts/object-cache").ObjectCache<CheckpointAttachment[]> | undefined {
-    return this.maxBytes > 0 || this.#cold ? this.#objects : undefined;
-  }
-
-  async #takeObject(ns: string): Promise<{ value: CheckpointAttachment[]; dispose(): void } | null> {
-    let entry = this.findExact([], ns);
-    if (entry) this.objectHits++;
-    else this.objectMisses++;
-    if (!entry) {
-      const cold = this.#cold;
-      const hit = cold?.findExact?.([], ns);
-      if (!hit || !cold?.restoreObjectAsync) return null;
-      const epoch = this.#epoch;
-      const loaded = await cold.restoreObjectAsync(hit.handle);
-      if (!loaded) return null;
-      this.objectRestores++;
-      if (epoch !== this.#epoch) { this.#disposeEntry({ ...loaded, ns }, false); return null; }
-      // Another request may have published this immutable object during IO.
-      entry = this.findExact([], ns);
-      if (entry) this.#disposeEntry({ ...loaded, ns }, false);
-      else {
-        this.#publishingPrefetch = true;
-        try { this.put([], loaded.caches, ns, loaded.retain, loaded.attachments); }
-        catch (error) { this.#disposeEntry({ ...loaded, ns }, false); throw error; }
-        finally { this.#publishingPrefetch = false; }
-        entry = this.findExact([], ns);
-      }
-    }
-    if (!entry) return null;
-    const borrowed = this.#borrow(this.#records.get(entry)!, 0, ns);
-    try { this.evictToBudget(); this.reclaim(); }
-    catch (error) { this.#disposeEntry(borrowed, false); throw error; }
-    let released = false;
-    return { value: borrowed.attachments ?? [], dispose: () => {
-      if (released) return; released = true;
-      this.#disposeEntry(borrowed, false);
-    } };
-  }
-
-  readonly maxBytes: number;
-  #entries: EntryRecord[] = [];
-  #clock = 0;
-  #sessions = new Map<string, Map<string, EntryRecord>>();
-  #records = new WeakMap<PromptCacheEntry, EntryRecord>();
-  sessionHits = 0;
-  sessionMisses = 0;
-  prefixScans = 0;
-  objectHits = 0;
-  objectMisses = 0;
-  objectRestores = 0;
-
-  closeSession(sessionId: string): void {
-    for (const record of this.#sessions.get(sessionId)?.values() ?? []) record.sessions--;
-    this.#sessions.delete(sessionId);
-  }
-
-  #associate(sessionId: string | undefined, ns: string, record: EntryRecord): void {
-    if (!sessionId) return;
-    let entries = this.#sessions.get(sessionId);
-    if (!entries) this.#sessions.set(sessionId, entries = new Map());
-    const old = entries.get(ns);
-    if (old === record) return;
-    if (old) old.sessions--;
-    entries.set(ns, record); record.sessions++;
-  }
-
-  #sessionRecord(prompt: number[], ns: string, sessionId?: string): EntryRecord | undefined {
-    const record = sessionId ? this.#sessions.get(sessionId)?.get(ns) : undefined;
-    // A session labels affinity. Full-history callers may edit or branch;
-    // incompatible history uses ordinary lookup, never a request refusal.
-    return record && record.entry.tokens.length < prompt.length &&
-      commonPrefixLength(record.entry.tokens, prompt) === record.entry.tokens.length ? record : undefined;
-  }
-
-  #coldHit(prompt: number[], ns: string, sessionId?: string) {
-    const known = this.#sessionRecord(prompt, ns, sessionId);
-    return (known && this.#cold?.findExact?.(known.entry.tokens, ns)) || this.#cold?.find(prompt, ns) || null;
-  }
-
-  hits = 0;
-  misses = 0;
-  demotions = 0;
-  /** Cold-tier spill sinks (see SpillSink): spillOwned gets clones and
-   *  runs the write off-lock; spillSync gets the live entry inline.
-   *  Failures are the sink's problem; eviction proceeds regardless. */
-  readonly #spillOwned: ((entry: PromptCacheEntry) => void) | null;
-  readonly #spillSync: ((entry: PromptCacheEntry) => void) | null;
-  /** Optional cold tier: take() tiers over it, demoteIdle() spills into it. */
-  readonly #cold: ColdTier | null;
-  /** Observer after publication, independent of persistence policy. */
-  onPut: ((tokens: number[], ns: string) => void) | null = null;
-  /** The backend measures pressure; this store chooses which optional
-   * snapshots to release. No storage I/O runs at a reclaim boundary. */
-  pressure: { overBudget(): boolean; releasePending?(): void } | null = null;
-  /** Tier policy: a pending write is not a lower-tier copy. This controls
-   * eviction only; persistence completion does not itself remove an entry. */
-  canEvict: ((entry: PromptCacheEntry) => boolean) | null = null;
-  protected promoteRestores = false;
-  retention: RetentionPolicy = new LruRetention();
-  #prefetches = new Map<unknown, Promise<PromptCacheEntry | null>>();
-  #pins = new Map<PromptCacheEntry, number>();
-  #publishingPrefetch = false;
-  #epoch = 0;
-
-  /** Start cold reads during request preparation, coalescing identical IO.
-   * Active request interest keeps a just-restored donor resident until the
-   * caller has consumed it. Cancellation releases interest through the lease. */
-  async prefetch(prompt: number[], ns = "", sessionId?: string): Promise<() => void> {
-    const known = this.#sessionRecord(prompt, ns, sessionId);
-    if (known?.resident) return this.#pin(known.entry);
-    const cold = this.#cold;
-    const hit = cold?.restoreAsync ? this.#coldHit(prompt, ns, sessionId) : null;
-    if (!cold?.restoreAsync || !hit || (!known && hit.prefixLen <= this.peekPrefixLen(prompt, ns))) return () => {};
-    let pending = this.#prefetches.get(hit.handle);
-    if (!pending) {
-      const epoch = this.#epoch;
-      pending = (async () => {
-        const loaded = await cold.restoreAsync!(hit.handle);
-        if (!loaded) return null;
-        if (epoch !== this.#epoch) { this.#disposeEntry({ ...loaded, ns }, false); return null; }
-        this.#publishingPrefetch = true;
-        try { this.put(loaded.tokens, loaded.caches, ns, loaded.retain, loaded.attachments); }
-        catch (error) { this.#disposeEntry({ ...loaded, ns }, false); throw error; }
-        finally { this.#publishingPrefetch = false; }
-        return this.findExact(loaded.tokens, ns);
-      })();
-      this.#prefetches.set(hit.handle, pending);
-    }
-    let entry: PromptCacheEntry | null;
-    try { entry = await pending; }
-    finally { if (this.#prefetches.get(hit.handle) === pending) this.#prefetches.delete(hit.handle); }
-    if (!entry) return () => {};
-    this.#associate(sessionId, ns, this.#records.get(entry)!);
-    return this.#pin(entry);
-  }
-
-  #pin(entry: PromptCacheEntry): () => void {
-    this.#pins.set(entry, (this.#pins.get(entry) ?? 0) + 1);
-    let released = false;
-    return () => {
-      if (released) return; released = true;
-      const count = this.#pins.get(entry!) ?? 1;
-      if (count === 1) this.#pins.delete(entry!); else this.#pins.set(entry!, count - 1);
-      this.evictToBudget(); this.reclaim();
-    };
-  }
-
-  #victim(): EntryRecord | undefined {
-    const candidates = this.#entries.filter(record => !this.#pins.has(record.entry));
-    const unrelated = candidates.filter(record => record.sessions === 0);
-    return candidates.length ? this.retention.victim(unrelated.length ? unrelated : candidates) : undefined;
-  }
-
-  evictToBudget(): void {
-    if (this.#publishingPrefetch) return;
-    while (this.totalBytes > this.maxBytes) {
-      if (!this.#entries.length) break;
-      const oldest = this.#victim();
-      if (!oldest) break;
-      if (this.canEvict && !this.canEvict(oldest.entry)) break;
-      this.retention.evicted(oldest);
-      this.#entries = this.#entries.filter(r => r !== oldest);
-      try { this.#disposeEntry(oldest.entry, !this.canEvict); }
-      catch (error) { console.warn(`[prompt-cache] entry cleanup failed: ${error}`); }
-    }
-  }
-
-  reclaim(): void {
-    const pressure = this.pressure;
-    if (!pressure?.overBudget()) return;
-    pressure.releasePending?.();
-    while (this.#entries.length && pressure.overBudget()) {
-      const oldest = this.#victim();
-      if (!oldest) break;
-      if (this.canEvict && !this.canEvict(oldest.entry)) break;
-      this.retention.evicted(oldest);
-      this.#entries = this.#entries.filter(r => r !== oldest);
-      this.#disposeEntry(oldest.entry, false);
-      this.demotions++;
-    }
-  }
-
-  /** Zero-copy view cloner — injectable so model-free tests can stub it. */
-  readonly #clone: (caches: Cache[]) => Cache[];
-
-  constructor(
-    maxBytes: number,
-    // Bare function = legacy sync spill (tests, simple embedders).
-    spill: SpillSink | ((entry: PromptCacheEntry) => void) | null = null,
-    cold: ColdTier | null = null,
-    clone: (caches: Cache[]) => Cache[] = cloneKvCaches,
-    private readonly materialize: (caches: Cache[], attachments?: CheckpointAttachment[]) => void = materializeCheckpoint,
-  ) {
-    this.maxBytes = maxBytes;
-    this.#spillOwned = typeof spill === "function" ? null : spill?.spillOwned ?? null;
-    this.#spillSync = typeof spill === "function" ? spill : spill?.spillSync ?? null;
-    this.#cold = cold;
-    this.#clone = clone;
-  }
-
-  #disposeEntry(entry: PromptCacheEntry, spillFirst: boolean): void {
-    const record = this.#records.get(entry);
-    if (record) record.resident = false;
-    if (spillFirst) {
-      if (this.#spillOwned) {
-        // NON-BLOCKING spill: clone BEFORE disposing (clones are zero-copy
-        // views of the live arrays) and hand ownership to the sink, which
-        // flushes off the generation lock. The underlying GPU buffers stay
-        // alive until the sink disposes the clones — for demoteIdle that
-        // means the memory is freed one bounded write-chain hop later, not
-        // instantly; acceptable because the chain is serial and each write
-        // disposes its clones on settle (fulfill AND reject).
-        // Retain contract: entry.retain still runs below, possibly while
-        // the clones are in flight — safe because cold-restored entries
-        // OWN their bytes (streamed copy-restore, 2026-07-07): there is
-        // no backing mmap for retain to unmap, so nothing the clones
-        // could alias goes away.
-        let clones: Cache[] | null = null;
-        let attachments: CheckpointAttachment[] | undefined;
-        try {
-          clones = this.#clone(entry.caches);
-          attachments = cloneAttachments(entry.attachments);
-          this.#spillOwned({ tokens: [...entry.tokens], caches: clones, ns: entry.ns, attachments });
-        } catch {
-          // Failure degrades to no-spill; eviction/demotion never unwinds.
-          if (clones) for (const c of clones) try { c.dispose(); } catch {}
-          try { disposeAttachments(attachments); } catch {}
-        }
-      } else {
-        const spill =
-          this.#spillSync ??
-          (this.#cold
-            ? (e: PromptCacheEntry) => this.#cold!.store(e.tokens, e.caches, e.ns, e.attachments)
-            : null);
-        if (spill) try { spill(entry); } catch {}
-      }
-    }
-    disposeResources([...entry.caches, { dispose: () => disposeAttachments(entry.attachments) }]);
-    entry.retain?.(); // backing release requires successful cache disposal
-  }
-
-  get totalBytes(): number {
-    return this.#entries.reduce((a, e) => a + e.bytes, 0);
-  }
-
-  get size(): number {
-    return this.#entries.length;
-  }
-
-  /** Release LRU entries until the caller's live-memory check is satisfied.
-   * Unlike background eviction, pressure relief must finish the SSD write
-   * and release the buffers now, not enqueue clones that keep them alive.
-   * Call under the engine lock, after taking the request's reusable prefix. */
-  relievePressure(overBudget: () => boolean): number {
-    let count = 0;
-    while (this.#entries.length && overBudget()) {
-      const victim = this.#victim();
-      if (!victim) break;
-      const oldest = this.#entries.indexOf(victim);
-      const { entry } = victim;
-      // Keep the entry resident if persistence throws. A failed SSD write
-      // must not be presented as a successful demotion.
-      if (this.#cold) {
-        if (this.#cold.store(entry.tokens, entry.caches, entry.ns, entry.attachments) === false) break;
-      }
-      else this.#spillSync?.(entry);
-      this.retention.evicted(victim);
-      this.#entries.splice(oldest, 1);
-      this.#disposeEntry(entry, false);
-      this.demotions++;
-      count++;
-    }
-    return count;
-  }
-
-  /** Find the entry with the longest usable common prefix of `prompt`
-   *  across BOTH tiers. Every RAM entry lends zero-copy views and retains its
-   *  donor, including exact-boundary recurrent state and method attachments.
-   *  The caller owns the returned views —
-   *  dispose them or put() an extended entry, honoring `retain` either way.
-   *
-   *  Usable prefix = common prefix capped at prompt.length - 1 (at least
-   *  one token must be forwarded to produce logits). Entries longer than
-   *  the prefix need cache.trim(); ring caches lose trimability once
-   *  wrapped — those entries only match in full.
-   *
-   *  Tier order: the cold tier wins only with a STRICTLY longer usable
-   *  prefix (restore reads storage and materializes native arrays). A
-   *  cold entry with an untrimmable divergent tail (ring post-wrap, SSM)
-   *  is dropped and the RAM candidate (or a fresh prefill) serves.
-   *  hits/misses count RAM candidacy only (cold restores are counted by
-   *  the tier itself), preserving the /stats meaning. */
-  take(prompt: number[], ns = "", sessionId?: string): PromptCacheEntry | null {
-    const hit = this.#take(prompt, ns, sessionId);
-    // Keep the chosen state's owned views before releasing optional donors.
-    // This also covers a cold restore materializing its bytes at admission.
-    try { this.reclaim(); return hit; }
-    catch (error) {
-      if (hit) this.#disposeEntry(hit, false);
-      throw error;
-    }
-  }
-
-  #take(prompt: number[], ns: string, sessionId?: string): PromptCacheEntry | null {
-    const known = this.#sessionRecord(prompt, ns, sessionId);
-    if (known?.resident) {
-      this.sessionHits++; this.hits++;
-      return this.#borrow(known, known.entry.tokens.length, ns);
-    }
-    if (sessionId) this.sessionMisses++;
-    this.prefixScans++;
-    let bestIdx = -1;
-    let bestLen = 0;
-    for (let i = 0; i < this.#entries.length; i++) {
-      const e = this.#entries[i]!.entry;
-      if (e.ns !== ns) continue;
-      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p <= bestLen || p < minimumReusableOffset(e.caches)) continue;
-      const trimNeeded = e.tokens.length - p;
-      if (trimNeeded > 0 && (!!e.attachments?.length || !e.caches.every((c) => c.isTrimmable()))) continue;
-      bestLen = p;
-      bestIdx = i;
-    }
-    if (bestIdx === -1) this.misses++;
-    else this.hits++;
-
-    // Cold tier: strictly longer usable prefix ⇒ restore + trim; the RAM
-    // candidate stays put UNTOUCHED (the old serial-lane flow trimmed it
-    // before comparing — a needless entry degradation, gone now).
-    if (this.#cold) {
-      const hit = this.#coldHit(prompt, ns, sessionId);
-      if (hit && hit.prefixLen > bestLen) {
-        const loaded = this.#cold.restore(hit.handle);
-        if (loaded) {
-          const owner = ownResource(loaded, (entry) =>
-            this.#disposeEntry({ ...entry, ns }, false));
-          const result = withResource(owner, (entry) => {
-            const trimNeeded = entry.tokens.length - hit.prefixLen;
-            if (hit.prefixLen < minimumReusableOffset(entry.caches)) return null;
-            if (trimNeeded > 0 && (!!entry.attachments?.length || !entry.caches.every((cache) => cache.isTrimmable()))) {
-              console.warn(
-                `[prompt-cache] cold entry unusable: ${entry.tokens.length} tokens, ` +
-                `needs trim ${trimNeeded} but untrimmable (wrapped ring/SSM) — re-prefilling`,
-              );
-              return null;
-            }
-            if (trimNeeded > 0) for (const cache of entry.caches) cache.trim(trimNeeded);
-            const tokens = entry.tokens.slice(0, hit.prefixLen);
-            return { ...owner.transfer(), tokens, ns };
-          });
-          if (result) {
-            if (this.promoteRestores) {
-              // Cache the restored prefix immediately. A second request can
-              // reuse RAM even before this borrower finishes its generation.
-              const share = makeSharedRetain(result.retain);
-              result.retain = share.acquire();
-              const retained = share.acquire();
-              let clones: Cache[] = [];
-              let attachments: CheckpointAttachment[] | undefined;
-              try {
-                clones = this.#clone(result.caches);
-                attachments = cloneAttachments(result.attachments);
-                this.put([...result.tokens], clones, ns, retained, attachments, sessionId);
-              } catch (error) {
-                cleanupFailure(error, () => disposeResources([...clones,
-                  { dispose: () => disposeAttachments(attachments) },
-                  { dispose: retained }, { dispose: () => this.#disposeEntry(result, false) }]));
-              }
-            }
-            return result;
-          }
-        }
-      }
-    }
-
-    if (bestIdx === -1) return null;
-    // PREFIX SHARING (non-consuming serve): the caller gets zero-copy
-    // CLONES of the donor's caches, trimmed to the matched prefix; the
-    // donor stays in the cache untouched, ready to serve the next agent /
-    // session with the same prefix. Safe because mlx cache updates are
-    // FUNCTIONAL (updateAndFetch reassigns fresh arrays; buffer donation
-    // only fires at refcount 1, and the donor always holds a ref), so a
-    // clone extending itself can never mutate the donor's bytes. The old
-    // consume-and-trim semantics CANNIBALIZED donors: agent B borrowing
-    // agent A's 2k system prompt destroyed A's 10k entry to do it.
-    const rec = this.#entries[bestIdx]!;
-    this.#associate(sessionId, ns, rec);
-    return this.#borrow(rec, bestLen, ns);
-  }
-
-  #borrow(rec: EntryRecord, bestLen: number, ns: string): PromptCacheEntry {
-    rec.lastUsed = ++this.#clock;
-    rec.lastUsedMs = Date.now();
-    rec.uses++; rec.cost = bestLen; this.retention.accessed(rec);
-    // Lend immutable views for every backend state. RAM retention and SSD
-    // durability remain cache responsibilities while the caller extends its
-    // own views. This keeps synchronous persistence out of a RAM hit.
-    const clones = this.#clone(rec.entry.caches);
-    let attachments: CheckpointAttachment[] | undefined;
-    const trimNeeded = rec.entry.tokens.length - bestLen;
-    try {
-      attachments = cloneAttachments(rec.entry.attachments);
-      if (trimNeeded > 0) for (const c of clones) c.trim(trimNeeded);
-    } catch (error) {
-      cleanupFailure(error, () => disposeResources([...clones, { dispose: () => disposeAttachments(attachments) }]));
-    }
-    return {
-      tokens: rec.entry.tokens.slice(0, bestLen),
-      caches: clones,
-      attachments,
-      ns,
-      retain: rec.share.acquire(),
-    };
-  }
-
-  /** Longest usable prefix length available for `prompt` WITHOUT taking
-   *  (the SSD tier compares this against its own best before restoring). */
-  peekPrefixLen(prompt: number[], ns = ""): number {
-    let best = 0;
-    for (const { entry: e } of this.#entries) {
-      if (e.ns !== ns) continue;
-      const p = Math.min(commonPrefixLength(e.tokens, prompt), prompt.length - 1);
-      if (p <= best || p < minimumReusableOffset(e.caches)) continue;
-      const trimNeeded = e.tokens.length - p; // same usability rule as take()
-      if (trimNeeded > 0 && (!!e.attachments?.length || !e.caches.every((c) => c.isTrimmable()))) continue;
-      best = p;
-    }
-    return best;
-  }
-
-  /** Read-only lookup of an immutable, cache-owned checkpoint. */
-  findExact(tokens: number[], ns = ""): PromptCacheEntry | null {
-    for (const { entry: e } of this.#entries) {
-      if (e.ns !== ns || e.tokens.length !== tokens.length) continue;
-      if (commonPrefixLength(e.tokens, tokens) === tokens.length) return e;
-    }
-    return null;
-  }
-
-  /** Insert (or reinsert) an entry; evicts LRU entries over the byte cap
-   *  (spilling each to the cold tier first, when one is attached). If the
-   *  entry itself exceeds the cap it is spilled + disposed, not stored. */
-  put(tokens: number[], caches: Cache[], ns = "", retain?: () => void, attachments?: CheckpointAttachment[], sessionId?: string): void {
-    // All cache-provided validation runs before ownership is adopted.
-    const bytes = cacheBytes(caches) + attachmentBytes(attachments);
-    const trimmable = !attachments?.length && caches.every((cache) => cache.isTrimmable());
-    // Storage takes resolved snapshots, independently of their producer.
-    this.materialize(caches, attachments);
-    const discard = (entry: PromptCacheEntry, spill: boolean) => {
-      try { this.#disposeEntry(entry, spill); }
-      catch (error) {
-        // The new entry is already owned here. Throwing back to its previous
-        // owner would invite disposal of state still held by this store.
-        console.warn(`[prompt-cache] entry cleanup failed: ${error}`);
-      }
-    };
-    const entry: PromptCacheEntry = { tokens, caches, ns, retain, attachments };
-    if (bytes > this.maxBytes && !this.canEvict) {
-      discard(entry, true);
-      return;
-    }
-    // The donor's retain becomes ref-counted so take()'s clones can hold
-    // shares; the entry's own retain is the donor's share.
-    const share = makeSharedRetain(retain);
-    entry.retain = share.acquire();
-    const rec: EntryRecord = { entry, bytes, cost: tokens.length, uses: 1, priority: 0, lastUsed: ++this.#clock, lastUsedMs: Date.now(), share, resident: true, sessions: 0 };
-    this.retention.accessed(rec);
-    this.#entries.push(rec);
-    this.#records.set(entry, rec);
-    this.#associate(sessionId, ns, rec);
-    // Exact duplicates (equal tokens) are redundant REGARDLESS of
-    // trimmability — the new entry serves exactly the matches the old one
-    // did (mlx-lm's trie replaces in place). Without this, every gemma ctx
-    // repeat left another full-size wrapped-ring entry behind (measured
-    // ~0.4 GB retained per repeat on the 12B bench).
-    for (const old of [...this.#entries]) {
-      if (old === rec || old.entry.ns !== ns) continue;
-      if (old.entry.tokens.length !== tokens.length) continue;
-      if (commonPrefixLength(old.entry.tokens, tokens) !== tokens.length) continue;
-      // Equivalent publication replaces the cache object, not the checkpoint
-      // identity. Every session sharing it follows the replacement.
-      if (old.sessions) for (const sessions of this.#sessions.values()) {
-        if (sessions.get(ns) !== old) continue;
-        sessions.set(ns, rec); old.sessions--; rec.sessions++;
-      }
-      this.#entries = this.#entries.filter((r) => r !== old);
-      discard(old.entry, false);
-    }
-    // SUPERSEDE strict prefix-ancestors (the SSD store's rule, brought to
-    // RAM — prefix sharing would otherwise grow an entry per turn): same-ns
-    // entries whose tokens are a prefix of the new entry are redundant
-    // and disposed WITHOUT spill (the new entry serves their prefixes and
-    // its own write-behind persists them) — but ONLY when the new entry's
-    // caches are all trimmable. An untrimmable new entry (wrapped ring)
-    // can only serve exact-length matches, so shorter ancestors — e.g.
-    // the prompt-boundary snapshot that makes drifted multi-turn agents
-    // hit — must survive.
-    if (trimmable) {
-      for (const old of [...this.#entries]) {
-        if (old === rec || old.entry.ns !== ns) continue;
-        if (old.entry.tokens.length >= tokens.length || old.entry.tokens.length < minimumReusableOffset(caches)) continue;
-        if (commonPrefixLength(old.entry.tokens, tokens) !== old.entry.tokens.length) continue;
-        this.#entries = this.#entries.filter((r) => r !== old);
-        discard(old.entry, false);
-      }
-    }
-    this.evictToBudget();
-    try { this.onPut?.(tokens, ns); } catch {}
-  }
-
-  /** Evict entries older than idleMs. Tiered residency policy first ensures
-   * SSD coverage. No-op without a cold tier. Returns entries removed from RAM. */
-  demoteIdle(idleMs: number, now = Date.now()): number {
-    if (!this.#cold) return 0;
-    let n = 0;
-    for (const rec of [...this.#entries]) {
-      if (now - rec.lastUsedMs < idleMs || this.#pins.has(rec.entry)) continue;
-      if (this.canEvict && !this.canEvict(rec.entry)) continue;
-      this.#entries = this.#entries.filter((r) => r !== rec);
-      this.#disposeEntry(rec.entry, !this.canEvict); // persisted entries need no duplicate write
-      this.demotions++;
-      n++;
-    }
-    return n;
-  }
-
-  clear(): void {
-    this.#sessions.clear();
-    this.#epoch++;
-    const entries = this.#entries;
-    this.#entries = [];
-    disposeResources(entries.map(({ entry }) => ({ dispose: () => this.#disposeEntry(entry, false) })));
-  }
-}