@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/prefix-cache.ts

@@ -0,0 +1,685 @@
+import { type CheckpointAttachment } from "../contracts/mlx/checkpoint";
+import { materializeCheckpoint, attachmentBytes, cloneAttachments, disposeAttachments } from "./checkpoint";
+import { leaseCacheState } from "./leases";
+import { minimumReusableOffset } from "./views";
+import { withResource, cleanupFailure, disposeResources, ownResource } from "../runtime/resources";
+// Byte-capped LRU prompt cache — the RAM tier of the Layer-0 KV store.
+//
+// The mlx-lm lesson: 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 cache service the scheduler acquires
+// through, direct generation) 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 "../contracts/mlx/cache";
+import type { PrefixCache, PrefixCacheHit } from "../contracts/portable/prefix-cache";
+import { LruRetention, type RetentionPolicy, type RetentionCandidate } from "./retention-policy";
+import { cloneKvCaches } from "./persistence";
+
+/** 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
+ *  (apps/mlx-bun/src/engine/cache-services.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/portable/object-cache").ObjectCache<CheckpointAttachment[]> = {
+    take: key => this.#takeObject(`object:${key}`),
+    put: (key, value) => this.put([], [], `object:${key}`, undefined, value),
+  };
+
+  get objects(): import("../contracts/portable/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) })));
+  }
+}