mlx-bun 0.0.13 → 0.3.0

package/src/ssd-cache.ts

@@ -1,4 +1,4 @@
-// SSD cold tier for the prompt/KV cache (docs/design/ssd-kv-cold-tier.md).
+// SSD cold tier for the prompt/KV cache (docs/design/kv-cache.md).
 //
 // Files ARE the database: no sidecar index. Layout
 //   <dir>/<configFingerprint>/<nsHash>/<uuid>.mlxkv
@@ -22,6 +22,7 @@ import { randomUUID } from "node:crypto";
 import { commonPrefixLength } from "./prompt-cache";
 import {
   saveKvCache, saveKvCacheAsync, loadKvCache, readKvHeader,
+  cacheHeadersTrimmable, legacyCacheCodecs, type CacheCodecProvider,
   type KvSaveMeta, type KvLoadExpect, type LoadedKvCache,
 } from "./kv-store";
 import type { Cache } from "./model/gemma4-base";
@@ -38,22 +39,11 @@ export interface SsdIndexEntry {
    *  over a shorter usable one (the 2026-07-06 restart-0 defect: the
    *  [prompt+gen] file always won, always got rejected after restore). */
   trimmable: boolean;
+  generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>;
 }
 
-/** Trimmability from a file header alone (mirror of the live caches'
- *  isTrimmable(): kv/qkv/turboquant always trim; rotating kinds only pre-wrap
- *  (offset < maxSize); GLM compressed kinds always trim; ssm never). Keep in
- *  sync with gemma4-base/qwen3-delta/glm52-cache. */
-const headerTrimmable = (caches: { kind: string; offset: number; maxSize?: number }[]): boolean =>
-  caches.every((c) =>
-    c.kind === "kv" || c.kind === "qkv" || c.kind === "turboquant" ||
-      c.kind === "mla" || c.kind === "mla-dsa" || c.kind === "mtp-mla"
-      ? true
-      : c.kind === "rotating" || c.kind === "rotating-qkv"
-        ? c.offset < (c.maxSize ?? 0)
-        : false);
-
 export interface SsdStoreOptions {
+  codecs?: CacheCodecProvider;
   dir: string;
   maxBytes: number;
   /** Server identity — written into every file, enforced on scan/load. */
@@ -69,6 +59,7 @@ const nsHash = (ns: string): string =>
 
 export class SsdCacheStore {
   readonly #opts: SsdStoreOptions;
+  readonly #codecs: CacheCodecProvider;
   readonly #root: string; // <dir>/<configFingerprint>
   #index: SsdIndexEntry[] = [];
   #warnedWriteFailure = false;
@@ -76,6 +67,7 @@ export class SsdCacheStore {
 
   constructor(opts: SsdStoreOptions) {
     this.#opts = opts;
+    this.#codecs = opts.codecs ?? legacyCacheCodecs;
     this.#root = join(opts.dir, opts.configFingerprint);
   }
 
@@ -91,6 +83,21 @@ export class SsdCacheStore {
     return this.#opts.maxBytes;
   }
 
+  /** Longest token prefix whose atomic cache file is present in the index. */
+  get longestDurablePrefixTokens(): number {
+    return this.#index.reduce((best, entry) => Math.max(best, entry.tokens.length), 0);
+  }
+
+  /** True when the SSD index can seed every token in this prefix. A longer
+   * trimmable descendant is equivalent to an exact boundary snapshot. */
+  hasDurablePrefix(tokens: number[], ns = ""): boolean {
+    return this.#index.some((entry) => {
+      if (entry.ns !== ns || entry.tokens.length < tokens.length) return false;
+      if (commonPrefixLength(entry.tokens, tokens) !== tokens.length) return false;
+      return entry.tokens.length === tokens.length || entry.trimmable;
+    });
+  }
+
   /** Startup recovery: rebuild the in-memory index from disk. Corrupt
    *  headers and metadata mismatches are unlinked; `.tmp` orphans reaped.
    *  Only OUR fingerprint dir is touched. Returns entries indexed. */
@@ -118,10 +125,14 @@ export class SsdCacheStore {
             h.tokenizerHash !== this.#opts.tokenizerHash
           )
             throw new Error("metadata mismatch");
+          if ((h.codecProvider ?? legacyCacheCodecs.id) !== this.#codecs.id) continue;
           const st = statSync(path);
           this.#index.push({
             path, ns: h.ns ?? "", tokens: h.tokens, bytes: st.size, mtimeMs: st.mtimeMs,
-            trimmable: headerTrimmable(h.caches),
+            trimmable: cacheHeadersTrimmable(h.caches, this.#codecs),
+            ...(h.generationCheckpoint
+              ? { generationCheckpoint: h.generationCheckpoint }
+              : {}),
           });
         } catch {
           try { rmSync(path, { force: true }); } catch {}
@@ -168,7 +179,7 @@ export class SsdCacheStore {
         ns: entry.ns,
         verify: this.#opts.verify,
       };
-      const loaded = loadKvCache(entry.path, model, expect);
+      const loaded = loadKvCache(entry.path, model, expect, this.#codecs);
       const now = new Date();
       try { utimesSync(entry.path, now, now); entry.mtimeMs = now.getTime(); } catch {}
       this.stats.restores++;
@@ -195,7 +206,7 @@ export class SsdCacheStore {
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      saveKvCache(path, tokens, caches, this.#meta(ns));
+      saveKvCache(path, tokens, caches, this.#meta(ns), this.#codecs);
       return this.#indexStored(path, tokens, caches, ns);
     } catch (err) {
       return this.#storeFailed(path, err);
@@ -207,21 +218,75 @@ export class SsdCacheStore {
    *  yields the event loop between tensors so serving interleaves.
    *  Caller passes zero-copy CLONES it owns (a consistent snapshot no
    *  matter what the live entry does meanwhile) and disposes them after.
-   *  `waitTurn` gates every per-tensor step (see saveKvCacheAsync) — the
-   *  server passes the gateway's onIdle so the flush only progresses while
-   *  the engine is idle, never mid-decode. */
-  async storeAsync(tokens: number[], caches: Cache[], ns = "", waitTurn?: () => Promise<void>): Promise<boolean> {
+   *  `runStep` owns every per-tensor step (see saveKvCacheAsync) — the
+   *  server passes the gateway's exclusive runner so a request cannot start
+   *  between an idle check and the blocking MLX readback. */
+  async storeAsync(
+    tokens: number[], caches: Cache[], ns = "",
+    runStep?: <T>(step: () => T) => Promise<T>,
+  ): Promise<boolean> {
     const dir = join(this.#root, nsHash(ns));
     const path = join(dir, `${randomUUID()}.mlxkv`);
     try {
       mkdirSync(dir, { recursive: true });
-      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), waitTurn);
+      await saveKvCacheAsync(path, tokens, caches, this.#meta(ns), runStep, this.#codecs);
       return this.#indexStored(path, tokens, caches, ns);
     } catch (err) {
       return this.#storeFailed(path, err);
     }
   }
 
+  /** Longest durable in-flight generation for an identical request. The
+   *  checkpoint namespace is separate from the ordinary prompt-cache tier,
+   *  so incomplete assistant output can never seed an unrelated prefill. */
+  findGenerationCheckpoint(
+    prompt: number[], key: string, cacheNs = "",
+  ): SsdIndexEntry | null {
+    let best: SsdIndexEntry | null = null;
+    for (const entry of this.#index) {
+      const checkpoint = entry.generationCheckpoint;
+      if (!checkpoint || checkpoint.key !== key || checkpoint.cacheNs !== cacheNs)
+        continue;
+      if (checkpoint.originalPromptTokens !== prompt.length) continue;
+      if (entry.tokens.length <= prompt.length) continue;
+      if (commonPrefixLength(entry.tokens, prompt) !== prompt.length) continue;
+      if (!best || entry.tokens.length > best.tokens.length) best = entry;
+    }
+    return best;
+  }
+
+  /** Persist one in-flight generation atomically. A completed newer
+   *  checkpoint supersedes older checkpoints for the same request only after
+   *  its rename succeeds, so a crash during a write leaves the prior point. */
+  async storeGenerationCheckpoint(
+    tokens: number[], caches: Cache[], checkpoint: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+  ): Promise<boolean> {
+    const ns = `__generation_checkpoint__:${checkpoint.key}`;
+    const dir = join(this.#root, nsHash(ns));
+    const path = join(dir, `${randomUUID()}.mlxkv`);
+    try {
+      mkdirSync(dir, { recursive: true });
+      await saveKvCacheAsync(path, tokens, caches, {
+        ...this.#meta(ns), generationCheckpoint: checkpoint,
+      }, undefined, this.#codecs);
+      const stored = this.#indexStored(path, tokens, caches, ns, checkpoint);
+      if (stored) {
+        for (const entry of [...this.#index]) {
+          if (entry.path === path) continue;
+          if (entry.generationCheckpoint?.key === checkpoint.key) this.remove(entry.path);
+        }
+      }
+      return stored;
+    } catch (err) {
+      return this.#storeFailed(path, err);
+    }
+  }
+
+  removeGenerationCheckpoints(key: string): void {
+    for (const entry of [...this.#index])
+      if (entry.generationCheckpoint?.key === key) this.remove(entry.path);
+  }
+
   #meta(ns: string): KvSaveMeta {
     return {
       modelId: this.#opts.modelId,
@@ -231,7 +296,10 @@ export class SsdCacheStore {
     };
   }
 
-  #indexStored(path: string, tokens: number[], caches: Cache[], ns: string): boolean {
+  #indexStored(
+    path: string, tokens: number[], caches: Cache[], ns: string,
+    generationCheckpoint?: NonNullable<KvSaveMeta["generationCheckpoint"]>,
+  ): boolean {
     const st = statSync(path);
     if (st.size > this.#opts.maxBytes) {
       rmSync(path, { force: true });
@@ -257,7 +325,10 @@ export class SsdCacheStore {
         if (commonPrefixLength(e.tokens, tokens) === e.tokens.length) this.remove(e.path);
       }
     }
-    this.#index.push({ path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable });
+    this.#index.push({
+      path, ns, tokens, bytes: st.size, mtimeMs: Date.now(), trimmable,
+      ...(generationCheckpoint ? { generationCheckpoint } : {}),
+    });
     this.stats.spills++;
     this.evictToCap();
     return true;

package/src/ssd-durability.ts

@@ -0,0 +1,207 @@
+import type { Cache } from "./model/gemma4";
+import type { PromptCacheEntry } from "./prompt-cache";
+import type { SpillItem, SpillQueue } from "./kv-store";
+
+export interface DurabilityGateway {
+  readonly busy: boolean;
+  runExclusive<T>(fn: () => Promise<T>): Promise<T>;
+}
+
+export interface DurabilityPromptCache {
+  findExact(tokens: number[], ns?: string): PromptCacheEntry | null;
+}
+
+export interface DurabilitySnapshotStats {
+  pendingSnapshots: number;
+  pendingSpills: number;
+  pendingSpillBytes: number;
+  droppedSpills: number;
+  failedSpills: number;
+}
+
+export interface DurabilityFlushResult extends DurabilitySnapshotStats {
+  durable: boolean;
+  flushedSnapshots: number;
+  missingSnapshots: number;
+  elapsedMs: number;
+}
+
+interface DirtySnapshot {
+  key: string;
+  tokens: number[];
+  ns: string;
+}
+
+type StoreOutcome = "stored" | "missing" | "failed";
+
+/**
+ * Turns debounced prompt-cache puts into a durability boundary.
+ *
+ * A dirty record stays present until SpillQueue reports that its atomic store
+ * completed. Queue drops and write failures therefore remain retryable during
+ * an explicit flush or graceful shutdown.
+ */
+export class SsdDurabilityCoordinator {
+  readonly #dirty = new Map<string, DirtySnapshot>();
+  readonly #timers = new Map<string, ReturnType<typeof setTimeout>>();
+  readonly #attempts = new Map<string, Promise<StoreOutcome>>();
+  #flush: Promise<DurabilityFlushResult> | null = null;
+
+  constructor(
+    readonly gateway: DurabilityGateway,
+    readonly promptCache: DurabilityPromptCache,
+    readonly spillQueue: SpillQueue,
+    readonly cloneCaches: (caches: Cache[]) => Cache[],
+    readonly isAlreadyDurable: (tokens: number[], ns: string) => boolean = () => false,
+    readonly debounceMs = 1_000,
+    readonly busyRetryMs = 5_000,
+  ) {}
+
+  schedule(tokens: number[], ns = ""): void {
+    if (tokens.length === 0) return;
+    // PromptCache can hold unrelated entries with the same namespace and
+    // length. Include the tokens so one conversation cannot cancel another
+    // conversation's pending durability record.
+    const key = `${ns.length}:${ns}:${tokens.join(",")}`;
+    const rec: DirtySnapshot = { key, tokens: [...tokens], ns };
+    this.#dirty.set(key, rec);
+    this.#arm(key, this.debounceMs);
+  }
+
+  get stats(): DurabilitySnapshotStats {
+    return {
+      pendingSnapshots: this.#dirty.size,
+      pendingSpills: this.spillQueue.pendingCount,
+      pendingSpillBytes: this.spillQueue.pendingBytes,
+      droppedSpills: this.spillQueue.droppedCount,
+      failedSpills: this.spillQueue.failedCount,
+    };
+  }
+
+  /** Force every dirty RAM snapshot through the queue and wait for storage. */
+  flush(): Promise<DurabilityFlushResult> {
+    if (this.#flush) return this.#flush;
+    const started = performance.now();
+    this.#flush = this.#flushInner(started).finally(() => { this.#flush = null; });
+    return this.#flush;
+  }
+
+  #arm(key: string, delayMs: number): void {
+    const old = this.#timers.get(key);
+    if (old) clearTimeout(old);
+    const timer = setTimeout(() => {
+      this.#timers.delete(key);
+      void this.#attempt(key, false);
+    }, delayMs);
+    timer.unref?.();
+    this.#timers.set(key, timer);
+  }
+
+  async #attempt(key: string, force: boolean): Promise<StoreOutcome> {
+    const running = this.#attempts.get(key);
+    if (running) {
+      const outcome = await running;
+      if (!force || !this.#dirty.has(key)) return outcome;
+    }
+
+    const rec = this.#dirty.get(key);
+    if (!rec) return "stored";
+    if (!force && this.gateway.busy) {
+      this.#arm(key, this.busyRetryMs);
+      return "failed";
+    }
+
+    const task = this.#store(rec);
+    this.#attempts.set(key, task);
+    try {
+      const outcome = await task;
+      if (outcome === "stored" && this.#dirty.get(key) === rec)
+        this.#dirty.delete(key);
+      if (outcome !== "stored" && !force && this.#dirty.get(key) === rec)
+        this.#arm(key, this.busyRetryMs);
+      return outcome;
+    } finally {
+      if (this.#attempts.get(key) === task) this.#attempts.delete(key);
+    }
+  }
+
+  async #store(rec: DirtySnapshot): Promise<StoreOutcome> {
+    let snap: SpillItem | null = null;
+    try {
+      snap = await this.gateway.runExclusive(async () => {
+        if (this.isAlreadyDurable(rec.tokens, rec.ns)) return null;
+        const entry = this.promptCache.findExact(rec.tokens, rec.ns);
+        if (!entry) return null;
+        return {
+          tokens: [...entry.tokens],
+          caches: this.cloneCaches(entry.caches),
+          ns: rec.ns,
+        };
+      });
+    } catch {
+      return "failed";
+    }
+    if (!snap)
+      return this.isAlreadyDurable(rec.tokens, rec.ns) ? "stored" : "missing";
+    return await this.spillQueue.enqueue(snap) ? "stored" : "failed";
+  }
+
+  async #flushInner(started: number): Promise<DurabilityFlushResult> {
+    for (const timer of this.#timers.values()) clearTimeout(timer);
+    this.#timers.clear();
+
+    let flushedSnapshots = 0;
+    const missing = new Map<string, DirtySnapshot>();
+    const droppedBefore = this.spillQueue.droppedCount;
+    const failedBefore = this.spillQueue.failedCount;
+
+    if (this.#attempts.size > 0)
+      await Promise.allSettled([...this.#attempts.values()]);
+    await this.spillQueue.drain();
+
+    // Flush one entry at a time. This prevents the queue cap from dropping a
+    // boundary snapshot while a large final snapshot is already in flight.
+    // Attempt each record version once. Missing state stays dirty, so a
+    // second flush cannot report success without a real durable snapshot.
+    const attempted = new Set<DirtySnapshot>();
+    while (this.#dirty.size > 0) {
+      const records = [...this.#dirty.entries()].filter(([, rec]) => !attempted.has(rec));
+      if (records.length === 0) break;
+      for (const [key, rec] of records) {
+        if (this.#dirty.get(key) !== rec) continue;
+        attempted.add(rec);
+        const outcome = await this.#attempt(key, true);
+        await this.spillQueue.drain();
+        if (outcome === "stored") {
+          flushedSnapshots++;
+        } else if (outcome === "missing") {
+          missing.set(key, rec);
+        }
+      }
+    }
+
+    // A later trimmable snapshot may now cover a superseded RAM ancestor.
+    // Reconcile against committed SSD coverage after all writes settle;
+    // never clear a missing prefix merely because another write succeeded.
+    for (const [key, rec] of missing) {
+      if (!this.isAlreadyDurable(rec.tokens, rec.ns)) continue;
+      if (this.#dirty.get(key) === rec) this.#dirty.delete(key);
+      missing.delete(key);
+      flushedSnapshots++;
+    }
+
+    const stats = this.stats;
+    return {
+      ...stats,
+      durable:
+        stats.pendingSnapshots === 0 &&
+        stats.pendingSpills === 0 &&
+        stats.droppedSpills === droppedBefore &&
+        stats.failedSpills === failedBefore &&
+        missing.size === 0,
+      flushedSnapshots,
+      missingSnapshots: missing.size,
+      elapsedMs: performance.now() - started,
+    };
+  }
+}

package/src/tool-call.ts

@@ -161,7 +161,7 @@ export function parseToolCalls(text: string): ParsedToolCall[] {
 
 const TOOL_CALL_BLOCK_RE = /<tool_call>\s*([\s\S]*?)\s*<\/tool_call>/gi;
 const XML_FUNCTION_EQUALS_RE = /^\s*<function=([^>\s]+)>\s*([\s\S]*?)\s*<\/function>\s*$/i;
-const XML_PARAMETER_EQUALS_RE = /<parameter=([^>\s]+)>\s*([\s\S]*?)\s*<\/parameter>/gi;
+const XML_PARAMETER_EQUALS_RE = /<parameter=([^>\s]+)>([\s\S]*?)<\/parameter>/gi;
 const GLM52_ARG_KEY_RE = /<arg_key>/i;
 const GLM52_ARG_PAIR_RE = /<arg_key>\s*([\s\S]*?)\s*<\/arg_key>\s*<arg_value>\s*([\s\S]*?)\s*<\/arg_value>/gi;
 const XML_FUNCTION_ATTR_RE = /<function\s+name=["']([^"']+)["']\s*>\s*([\s\S]*?)\s*<\/function>/gi;
@@ -208,11 +208,9 @@ function toolParameterSchema(
   return null;
 }
 
-/** Oracle `_decode_tool_parameter_value`: string-typed params stay raw
- *  text (a path like "2025" must not become a number); everything else
- *  is JSON-decoded when possible. */
-function decodeToolValue(value: string, schema: Record<string, unknown> | null): unknown {
-  const text = decodeXml(value);
+/** Decode a value after its caller removes protocol-specific delimiters.
+ *  String-typed params stay raw; everything else is JSON-decoded when possible. */
+function decodeToolValue(text: string, schema: Record<string, unknown> | null): unknown {
   if (!text) return "";
   const types = schemaTypeNames(schema);
   if (types.size && [...types].every((t) => t === "string" || t === "null")) return text;
@@ -546,23 +544,59 @@ function parseJsonToolCall(block: string): ParsedToolCall | null {
   return repairs.length ? { ...call, repaired: true, repairs } : call;
 }
 
-function parseXmlEqualsToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
+function parseXmlEqualsToolCall(block: string, tools: ToolSpec[], strict = false): ParsedToolCall | null {
   const match = XML_FUNCTION_EQUALS_RE.exec(block);
   XML_FUNCTION_EQUALS_RE.lastIndex = 0;
   if (!match) return null;
+  if (strict && (match.index !== 0 || match[0].length !== block.length)) return null;
   const name = match[1]!.trim();
   const body = match[2]!;
   const args: Record<string, unknown> = {};
+  let cursor = 0;
   for (const param of body.matchAll(XML_PARAMETER_EQUALS_RE)) {
+    if (strict && body.slice(cursor, param.index).trim()) return null;
     const key = param[1]!.trim();
     if (!key) throw new Error(`tool '${name}' contains an empty parameter name`);
-    args[key] = decodeToolValue(param[2]!, toolParameterSchema(tools, name, key));
+    if (strict && Object.hasOwn(args, key)) return null;
+    // Qwen's format carries raw text, not XML entity-encoded text. Match
+    // mlx_lm.tool_parsers.qwen3_coder: remove one framing newline at each
+    // end, preserving indentation, literal entities and JSON string escapes.
+    const value = param[2]!.replace(/^\n|\n$/g, "");
+    args[key] = decodeToolValue(value, toolParameterSchema(tools, name, key));
+    cursor = param.index + param[0].length;
   }
   XML_PARAMETER_EQUALS_RE.lastIndex = 0;
+  if (strict && body.slice(cursor).trim()) return null;
   return { name, arguments: args };
 }
 
-function parseGlm52ToolCall(block: string, tools: ToolSpec[]): ParsedToolCall | null {
+/** One complete envelope, with full consumption and no syntax repair. Used
+ * to establish structural fill boundaries, never to verify model predictions. */
+export function parseStrictToolCall(text: string, tools: ToolSpec[]): ParsedToolCall | null {
+  const open = "<tool_call>", close = "</tool_call>";
+  if (!text.startsWith(open) || !text.endsWith(close)) return null;
+  const body = text.slice(open.length, -close.length).trim();
+  if (body.includes(open) || body.includes(close)) return null;
+  try {
+    let call: ParsedToolCall | null;
+    if (body.startsWith("{")) {
+      // JSON.parse must succeed before the existing payload interpretation.
+      JSON.parse(body);
+      call = parseJsonToolCall(body);
+    } else {
+      call = parseXmlEqualsToolCall(body, tools, true) ?? parseGlm52ToolCall(body, tools, true);
+    }
+    if (!call || call.repaired || !tools.some(tool => toolSpecName(tool) === call.name)) return null;
+    return call;
+  } catch { return null; }
+}
+
+/** A textual argument delimiter is complete only after its closing bracket. */
+export function hasCompleteToolValueDelimiter(text: string): boolean {
+  return /^(?:<\/parameter>|<\/arg_value>)/.test(text.trimStart());
+}
+
+function parseGlm52ToolCall(block: string, tools: ToolSpec[], strict = false): ParsedToolCall | null {
   const firstArg = block.search(GLM52_ARG_KEY_RE);
   const name = (firstArg < 0 ? block : block.slice(0, firstArg)).trim();
   if (!name || !/^[^\s<>]+$/.test(name)) return null;
@@ -576,7 +610,8 @@ function parseGlm52ToolCall(block: string, tools: ToolSpec[]): ParsedToolCall |
     if (body.slice(cursor, pair.index).trim()) return null;
     const key = pair[1]!.trim();
     if (!key) throw new Error(`tool '${name}' contains an empty argument name`);
-    args[key] = decodeToolValue(pair[2]!, toolParameterSchema(tools, name, key));
+    if (strict && Object.hasOwn(args, key)) return null;
+    args[key] = decodeToolValue(decodeXml(pair[2]!), toolParameterSchema(tools, name, key));
     cursor = pair.index + pair[0].length;
     matched = true;
   }
@@ -594,7 +629,7 @@ function parseXmlAttrToolCalls(text: string, tools: ToolSpec[]): ParsedToolCall[
     for (const param of body.matchAll(XML_PARAM_ATTR_RE)) {
       const key = param[1]!.trim();
       if (!key) throw new Error(`tool '${name}' contains an empty parameter name`);
-      args[key] = decodeToolValue(param[2]!, toolParameterSchema(tools, name, key));
+      args[key] = decodeToolValue(decodeXml(param[2]!), toolParameterSchema(tools, name, key));
     }
     XML_PARAM_ATTR_RE.lastIndex = 0;
     calls.push({ name, arguments: args });
@@ -616,9 +651,11 @@ export function parseGeneratedToolCalls(text: string, tools: ToolSpec[]): Parsed
   const calls: ParsedToolCall[] = [];
   for (const block of text.matchAll(TOOL_CALL_BLOCK_RE)) {
     const body = block[1]!.trim();
-    const parsed = parseJsonToolCall(body) ??
-      parseXmlEqualsToolCall(body, tools) ??
-      parseGlm52ToolCall(body, tools);
+    // A valid Qwen parameter may itself contain JSON. Dispatch by the outer
+    // format before JSON repair can mistake that value for the tool envelope.
+    const parsed = /^<function=/i.test(body)
+      ? parseXmlEqualsToolCall(body, tools)
+      : parseJsonToolCall(body) ?? parseGlm52ToolCall(body, tools);
     if (!parsed) throw new Error("unsupported tool_call payload format");
     calls.push(parsed);
   }

package/src/train/flash-cce.ts

@@ -37,6 +37,7 @@ import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { STEEL_QMM_HEADER } from "./steel-qmm-header";
+import { runtimeValue } from "../runtime-config";
 
 /** The quantized head pieces the kernel needs (packed weight + per-group affine
  *  scale/bias + bits/group_size). Mirrors HeadQuant in loss.ts. */
@@ -392,7 +393,7 @@ function fwdSgKernel(): MetalKernel {
       inputNames: ["h", "wp", "scales", "biases", "targets", "shape", "capv"],
       outputNames: ["blockMax_out", "blockSum_out", "blockTgt_out"],
       source: FWD_SG_SOURCE,
-      header: `#include <metal_simdgroup_matrix>\n#define TG ${TG}\n#define BB ${SG_BB}\n#define BV ${SG_BV}\n#define BD ${SG_BD}\n#define BDX ${SG_BD + 4}\n#define BVX ${SG_BV + 4}\n#define CCE_SG_SKIPDQ ${process.env.MLX_BUN_CCE_SG_SKIPDQ === "1" ? 1 : 0}\n`,
+      header: `#include <metal_simdgroup_matrix>\n#define TG ${TG}\n#define BB ${SG_BB}\n#define BV ${SG_BV}\n#define BD ${SG_BD}\n#define BDX ${SG_BD + 4}\n#define BVX ${SG_BV + 4}\n#define CCE_SG_SKIPDQ ${runtimeValue("MLX_BUN_CCE_SG_SKIPDQ") === "1" ? 1 : 0}\n`,
       ensureRowContiguous: true,
     });
   }
@@ -404,7 +405,7 @@ function fwdSgKernel(): MetalKernel {
 // compute-bound regime. Removing the staging redundancy (one thread per packed
 // word; scale/bias once per group) cut it 1629→848 ms, beating the scalar
 // kernel's 1097 ms. Fall back to scalar (MLX_BUN_CCE_SCALAR=1) when H % SG_BD != 0.
-const USE_SIMDGROUP_FWD = process.env.MLX_BUN_CCE_SCALAR !== "1" && process.env.MLX_BUN_CCE_LANE !== "1";
+const USE_SIMDGROUP_FWD = runtimeValue("MLX_BUN_CCE_SCALAR") !== "1" && runtimeValue("MLX_BUN_CCE_LANE") !== "1";
 
 // Backward: dh[t,d] = Σ_v coeff_v · dequant(W_v)[d], coeff_v = g_t·(onehot_v −
 // softmax_v)·sech²_v, softmax_v = exp(softcap(logit_v) − lse_t). Two-phase per
@@ -549,7 +550,7 @@ function bwdKernel(daccN: number): MetalKernel {
       inputNames: ["h", "wp", "scales", "biases", "targets", "lse", "gv", "shape", "capv"],
       outputNames: ["dh"],
       source: BWD_SOURCE,
-      header: `#define TG ${TG}\n#define BLOCK_B ${BLOCK_B_BWD}\n#define DACC_N ${daccN}\n#define CCE_BWD_SKIP_P2DQ ${process.env.MLX_BUN_CCE_BWD_SKIP_P2DQ === "1" ? 1 : 0}\n#define CCE_BWD_SKIP_P2 ${process.env.MLX_BUN_CCE_BWD_SKIP_P2 === "1" ? 1 : 0}\n`,
+      header: `#define TG ${TG}\n#define BLOCK_B ${BLOCK_B_BWD}\n#define DACC_N ${daccN}\n#define CCE_BWD_SKIP_P2DQ ${runtimeValue("MLX_BUN_CCE_BWD_SKIP_P2DQ") === "1" ? 1 : 0}\n#define CCE_BWD_SKIP_P2 ${runtimeValue("MLX_BUN_CCE_BWD_SKIP_P2") === "1" ? 1 : 0}\n`,
       ensureRowContiguous: true,
       atomicOutputs: true, // dh is summed across the NBLK vocab-block programs
     });
@@ -727,7 +728,7 @@ const BWD_SG_SOURCE = String.raw`
 // (flipped 2026-07-02 — the kernel-perf-review backlog #1 flip). The old
 // 0.66%→2.7% dh-accuracy scare was on RANDOM targets (flat softmax, the filter's
 // worst case — see PLAN.md "filter-on-real-data"); the real-data gate
-// (scripts/experiments/flash-cce-filter-realdata.ts) measured, on REAL chunk-ORPO
+// (flash-cce-filter-realdata.ts (deleted 2026-08-23; git history)) measured, on REAL chunk-ORPO
 // data (M1 Max, 2026-07-02): CPM5 eps 1e-5 = 0.343% dh for 1.41×; e4b = 0.158% dh
 // for 1.70× — and the teacher-forced full-logits fidelity shows the filter's added
 // error is far below the flash-vs-full fp-reassociation scatter that already
@@ -735,7 +736,7 @@ const BWD_SG_SOURCE = String.raw`
 // (The ~3.5× figure floating around older notes was the retired SG kernel, not
 // steel.) MLX_BUN_CCE_BWD_FILTER_EPS=0 restores exact gradients — eps "0"
 // compiles the filter OUT (#if CCE_BWD_FILTER 0), the identical pre-flip kernel.
-const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "1e-5";
+const BWD_FILTER_EPS = runtimeValue("MLX_BUN_CCE_BWD_FILTER_EPS") ?? "1e-5";
 // blockMax vocab-block early-exit eps (skip whole cold (token-block,vocab-block) programs
 // — attacks BOTH phases). DEFAULT ON at 1e-5 (flipped 2026-07-02). It is ~LOSSLESS when
 // it fires (only skips blocks whose max prob < eps AND which lack the target → ~0
@@ -744,7 +745,7 @@ const BWD_FILTER_EPS = process.env.MLX_BUN_CCE_BWD_FILTER_EPS ?? "1e-5";
 // at ≤0.004% dh; COMBINED with the coeff filter at 1e-5/1e-5 the backward is 1.71×
 // (CPM5) / 3.16× (e4b) vs exact. Fires only when the caller supplies blockMax (the
 // production head does); MLX_BUN_CCE_BWD_BLOCK_EPS=0 disables (compiled out).
-const BWD_BLOCK_EPS = process.env.MLX_BUN_CCE_BWD_BLOCK_EPS ?? "1e-5";
+const BWD_BLOCK_EPS = runtimeValue("MLX_BUN_CCE_BWD_BLOCK_EPS") ?? "1e-5";
 const _bwdSgKernels = new Map<string, MetalKernel>();
 function bwdSgKernel(dimTiles: number, filterEps: string, blockEps: string): MetalKernel {
   const filterOn = Number.parseFloat(filterEps) > 0 ? 1 : 0;
@@ -780,7 +781,7 @@ function bwdSgKernel(dimTiles: number, filterEps: string, blockEps: string): Met
 //     accumulated into D's frags LANE-LOCALLY (frag += frag — no get_coord math, the
 //     per-lane element mapping is identical between temp and D). dh atomic-added
 //     across the NBLK vocab-block programs. Validated dh ~1e-4 rel vs autograd.
-//     (Design + parity: scripts/experiments/steel-bwd-htile-test.ts.)
+//     (Design + parity: steel-bwd-htile-test.ts (deleted 2026-08-23; git history).)
 //   PHASE-2 DEQUANT: the Wd[32,HTw] staging uses MLX's vectorized+fused
 //     QuantizedBlockLoader (the SAME loader the forward uses), NOT a manual scalar
 //     dequant. The wrinkle: phase-2's GEMM contracts over vocab (K=32) and outputs H
@@ -1061,9 +1062,9 @@ export function flashCceForward(
   // tiling is clean (H,V,blockV all % 32). MLX_BUN_CCE_NOSTEEL=1 falls back to the
   // old kernels (simdgroup / scalar / lane).
   const blockV0 = Math.ceil(V / NBLK);
-  const useSteel = process.env.MLX_BUN_CCE_NOSTEEL !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0;
+  const useSteel = runtimeValue("MLX_BUN_CCE_NOSTEEL") !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0;
   const useSimd = !useSteel && USE_SIMDGROUP_FWD && H % SG_BD === 0;
-  if (!useSteel && process.env.MLX_BUN_CCE_NOSTEEL !== "1") {
+  if (!useSteel && runtimeValue("MLX_BUN_CCE_NOSTEEL") !== "1") {
     // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
     const fails = [
       H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,
@@ -1072,7 +1073,7 @@ export function flashCceForward(
     ].filter(Boolean).join(", ");
     warnSteelFallback("forward", fails, useSimd ? "simdgroup kernel" : "scalar/lane kernel");
   }
-  const useLane = !useSteel && !useSimd && process.env.MLX_BUN_CCE_LANE === "1" && WPR % 32 === 0;
+  const useLane = !useSteel && !useSimd && runtimeValue("MLX_BUN_CCE_LANE") === "1" && WPR % 32 === 0;
   const kernel = useSteel ? fwdSteelKernel(GS, head.bits) : useSimd ? fwdSgKernel() : useLane ? fwdLaneKernel() : fwdKernel();
   const tokTile = useSteel ? STEEL_BM : useSimd ? SG_BB : BLOCK_B;
   // steel + simdgroup read h as [M, H] (M on axis 0); lane + scalar read h TRANSPOSED
@@ -1178,11 +1179,11 @@ export function flashCceBackward(
   const blockV0 = Math.ceil(V / NBLK);
   // HTw == group_size for the phase-2 QuantizedBlockLoader (one group per H-tile); needs
   // H % GS == 0 (always, GR*GS=H) and GS % 32 == 0 (TNv = GS/32 frags per H-tile integral).
-  const useSteel = process.env.MLX_BUN_CCE_BWD_NOSTEEL !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0 && H % GS === 0 && GS % 32 === 0;
+  const useSteel = runtimeValue("MLX_BUN_CCE_BWD_NOSTEEL") !== "1" && H % 32 === 0 && V % 32 === 0 && blockV0 % 32 === 0 && H % GS === 0 && GS % 32 === 0;
   // simdgroup_matrix phase-2 GEMM is the next tier (5512→3913 ms, 1.41×) when the H
   // tiling is clean; MLX_BUN_CCE_BWD_SCALAR=1 forces the float4 scalar kernel.
-  const useSg = !useSteel && process.env.MLX_BUN_CCE_BWD_SCALAR !== "1" && H % (SGN * 8) === 0 && H % SG_BD === 0;
-  if (!useSteel && process.env.MLX_BUN_CCE_BWD_NOSTEEL !== "1") {
+  const useSg = !useSteel && runtimeValue("MLX_BUN_CCE_BWD_SCALAR") !== "1" && H % (SGN * 8) === 0 && H % SG_BD === 0;
+  if (!useSteel && runtimeValue("MLX_BUN_CCE_BWD_NOSTEEL") !== "1") {
     // NOTE blockV0 = ceil(V/NBLK) can violate %32 even when V%32==0 (e.g. V=50304)
     const fails = [
       H % 32 !== 0 ? `H%32!=0 (H=${H})` : null,