mlx-bun 0.0.9 → 0.0.11

package/src/tokenizer.ts

@@ -17,6 +17,270 @@ export interface LoadedTokenizer {
   idToToken(id: number): string;
   readonly bosTokenId: number | null;
   readonly eosTokenId: number | null;
+  /** Absolute path to tokenizer.json (the model snapshot dir's copy). Set by
+   *  loadTokenizer; read by src/grammar.ts to build the xgrammar TokenizerInfo
+   *  (vocab extraction + vocab-type detection for constrained decoding). */
+  readonly tokenizerJsonPath?: string;
+  /** config.json vocab_size, may exceed the tokenizer vocab length when the
+   *  embedding is padded to a power-of-two / 256. Read by src/grammar.ts to
+   *  size the logit bitmask against the model's logit width, not the
+   *  tokenizer's. Set by the server when it loads config; not set by the
+   *  bare loadTokenizer() helper (callers that need it pass modelDir + config). */
+  vocabSize?: number;
+  /** mlx-lm streaming-detokenizer parity: true for ByteLevel-decoder (BPE)
+   *  tokenizers, where mlx-lm's BPEStreamingDetokenizer drops ONE leading " "
+   *  at the very start of a generated sequence (tokenizer_utils.py
+   *  `_maybe_trim_space`: `elif not self.text: return current_text[1:]`).
+   *  Our decode() delegates to @huggingface/tokenizers, whose ByteLevel
+   *  decoder KEEPS that space — the serving StreamDecoder consults this flag
+   *  to match mlx-lm's rendered bytes (e.g. MiniCPM5 /v1/completions:
+   *  "2, 3, 5…" not " 2, 3, 5…"). SPM tokenizers need nothing: with a Strip
+   *  decoder the strip already happens inside decode(); without one
+   *  (gemma-4) mlx-lm uses trim_space=False and keeps the space too. */
+  readonly trimsLeadingSpace?: boolean;
+  /** mlx-lm streaming-detokenizer parity (companion to trimsLeadingSpace):
+   *  for ByteLevel-decoder (BPE) tokenizers, the vocab id of the bare-space
+   *  token "Ġ" — the ONLY token matching mlx-lm 0.31.3
+   *  BPEStreamingDetokenizer.add_token's hold condition
+   *  `len(v) == 1 and self._byte_decoder.get(v[0]) == 32`
+   *  (tokenizer_utils.py:206-218, "For single spaces wait until the next
+   *  token"). mlx_lm.server NEVER calls detokenizer.finalize() (grep: zero
+   *  hits in server.py 0.31.3), so a generation ENDING on this token silently
+   *  drops its space from the served bytes. The serving StreamDecoder
+   *  consults this id to mirror both behaviors (MiniCPM5: id 242). */
+  readonly bareSpaceTokenId?: number;
+  /** Encode-path counters (memo hits / incremental splices / fallbacks /
+   *  full encodes) — observability + test hook, no behavioral role. */
+  readonly encodeStats?: EncodeStats;
+}
+
+// Exact-input encode memo. Our pure-JS encode costs ~90 ms per 9.6k tokens
+// (vs Rust tokenizers ~15 ms) and the serving hot paths re-encode BYTE-
+// IDENTICAL text constantly: warm/ctx repeats re-render the same
+// conversation, and regenerate flows resend it. A string-keyed LRU makes
+// those free and is exact by construction (no BPE boundary reasoning).
+// Bounded: keys are capped by count and total chars, so at most a few MB.
+// The append-only multi-turn path (turn N+1's rendered prompt = turn N's
+// plus a suffix) goes through IncrementalEncoder below; the native port is
+// the endgame.
+const ENCODE_MEMO_MAX_ENTRIES = 16;
+const ENCODE_MEMO_MAX_CHARS = 1 << 21; // ~2M chars ≈ 4 MB of keys
+class EncodeMemo {
+  #map = new Map<string, number[]>(); // Map preserves insertion order → LRU
+  #chars = 0;
+  get(key: string): number[] | undefined {
+    const hit = this.#map.get(key);
+    if (!hit) return undefined;
+    this.#map.delete(key); // bump recency
+    this.#map.set(key, hit);
+    return hit;
+  }
+  set(key: string, ids: number[]): void {
+    if (this.#map.has(key)) { this.#chars -= key.length; this.#map.delete(key); }
+    this.#map.set(key, ids);
+    this.#chars += key.length;
+    while (this.#map.size > ENCODE_MEMO_MAX_ENTRIES || this.#chars > ENCODE_MEMO_MAX_CHARS) {
+      const oldest = this.#map.keys().next().value as string | undefined;
+      if (oldest === undefined) break;
+      this.#chars -= oldest.length;
+      this.#map.delete(oldest);
+    }
+  }
+  /** Longest-common-string-prefix scan over memo entries with the same
+   *  addSpecialTokens flag char ("S"/"s", the key's first char). ≤16 entries
+   *  × O(commonLen) char compares — sub-ms at memo scale. `>=` keeps the most
+   *  RECENT entry on common-length ties (Map iterates insertion order; recent
+   *  entries are re-inserted last by get()). */
+  findBestPrefix(
+    flagChar: string,
+    text: string,
+  ): { text: string; ids: number[]; common: number } | null {
+    let best: { text: string; ids: number[]; common: number } | null = null;
+    for (const [key, ids] of this.#map) {
+      if (key[0] !== flagChar) continue;
+      const ptext = key.slice(2);
+      const n = Math.min(ptext.length, text.length);
+      let i = 0;
+      while (i < n && ptext.charCodeAt(i) === text.charCodeAt(i)) i++;
+      if (best === null || i >= best.common) best = { text: ptext, ids, common: i };
+    }
+    return best && best.common > 0 ? best : null;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Incremental encode for append-only conversations.
+//
+// The installed @huggingface/tokenizers (0.1.3) Encoding exposes NO per-token
+// offsets (types/static/types.d.ts: `Encoding = { ids, tokens,
+// attention_mask }`), so char positions of token boundaries are recovered by
+// decoding token TAILS and checking them against the known source text:
+// decode(ids[T..N]) must be a literal suffix of the cached text (both target
+// decoders round-trip: cpm5's ByteLevel is pure byte-concat, e4b's
+// Sequence[Replace ▁→" ", ByteFallback, Fuse] inverts its normalizer; a tail
+// cut mid-UTF-8-char yields U+FFFD and simply fails endsWith → step back).
+// clean_up_tokenization_spaces must be forced FALSE here — the lib defaults
+// it to config??true (dist decode()), which mangles offsets.
+//
+// Exactness contract: the fast path returns cached[0..T] + encode(tail) only
+// after a seam VERIFICATION — re-encode from an anchor K tokens before the
+// cut and require the first K+min(K,|tail|) tokens to reproduce
+// cached[T-K..T] + tail[0..K] exactly. Any mismatch (BPE merges crossing the
+// cut, decoder round-trip failure, mid-char cut) falls back to a full
+// encode, so every returned result is either a verified splice or a
+// from-scratch encode. SLACK backs the cut off the divergence point for
+// merge locality; K≈32 tokens (~130 chars) exceeds any pre-token span the
+// target tokenizers produce.
+const INC_MIN_TEXT = 4096; // same gate as the memo: short prompts encode <1 ms
+const INC_MIN_COMMON = 1024; // shared-prefix chars below this: splice not worth it
+const INC_SLACK = 64; // chars backed off from the divergence point
+const INC_VERIFY_TOKENS = 32; // K: seam window re-encoded and compared
+const INC_CUT_ATTEMPTS = 8;
+
+export interface EncodeStats {
+  memoHits: number;
+  incremental: number;
+  incrementalFallbacks: number; // splice attempted, verification/cut refused
+  full: number;
+}
+
+/** The subset of @huggingface/tokenizers' Tokenizer we need (encode/decode). */
+type RawTokenizer = Pick<Tokenizer, "encode" | "decode">;
+
+export class IncrementalEncoder {
+  readonly stats: EncodeStats = { memoHits: 0, incremental: 0, incrementalFallbacks: 0, full: 0 };
+  #tok: RawTokenizer;
+  #memo = new EncodeMemo();
+  /** Whether the post_processor only PREPENDS specials (BOS). An APPENDED
+   *  special (EOS suffix template) would sit past the spliced tail where the
+   *  seam verification cannot see it, so such tokenizers get no incremental
+   *  path for addSpecialTokens=true. Both targets are prefix-only: cpm5
+   *  TemplateProcessing single=[<s>, A], e4b single=[A]. */
+  #prefixOnlySpecials: boolean;
+
+  constructor(tok: RawTokenizer, prefixOnlySpecials: boolean) {
+    this.#tok = tok;
+    this.#prefixOnlySpecials = prefixOnlySpecials;
+  }
+
+  /** Test hook: inject a memo entry as if `ids` were the encode of `text`. */
+  seed(text: string, ids: number[], addSpecialTokens: boolean): void {
+    this.#memo.set(`${addSpecialTokens ? "S" : "s"}|${text}`, ids.slice());
+  }
+
+  #rawEncode(text: string, special: boolean): number[] {
+    return this.#tok.encode(text, { add_special_tokens: special }).ids.map(Number);
+  }
+
+  /** Raw tail decode for offset math: specials render as their literal vocab
+   *  string (chat-template text contains them literally), cleanup forced off. */
+  #decodeTail(ids: number[]): string {
+    return this.#tok.decode(ids, {
+      skip_special_tokens: false,
+      clean_up_tokenization_spaces: false,
+    });
+  }
+
+  encode(text: string, special: boolean): number[] {
+    if (text.length < INC_MIN_TEXT) return this.#rawEncode(text, special);
+    const key = `${special ? "S" : "s"}|${text}`;
+    const hit = this.#memo.get(key);
+    if (hit) {
+      this.stats.memoHits++;
+      return hit.slice();
+    }
+    if (!special || this.#prefixOnlySpecials) {
+      const best = this.#memo.findBestPrefix(special ? "S" : "s", text);
+      if (best && best.common >= INC_MIN_COMMON) {
+        const inc = this.#tryIncremental(best.text, best.ids, best.common, text);
+        if (inc) {
+          this.stats.incremental++;
+          this.#memo.set(key, inc);
+          return inc.slice();
+        }
+        this.stats.incrementalFallbacks++;
+      }
+    }
+    const ids = this.#rawEncode(text, special);
+    this.stats.full++;
+    this.#memo.set(key, ids);
+    return ids.slice();
+  }
+
+  /** Splice cached prefix tokens with a fresh tail encode, or null to force
+   *  the full-encode fallback. `common` = shared-prefix chars between ptext
+   *  (the memoized source text of pids) and text. Cost is O(tail + seam), not
+   *  O(text): two tail decodes + two suffix encodes. */
+  #tryIncremental(
+    ptext: string,
+    pids: number[],
+    common: number,
+    text: string,
+  ): number[] | null {
+    const target = common - INC_SLACK; // keep cached tokens ending at/before here
+    const N = pids.length;
+    const K = INC_VERIFY_TOKENS;
+    if (target <= 0 || N <= K + 2) return null;
+    // Estimate the cut token index from mean chars/token, then correct by
+    // decoding the actual tail. charStart(T) = ptext.length - |decode(T..N)|,
+    // valid only when the decoded tail is a literal suffix of ptext.
+    const avg = Math.max(ptext.length / N, 0.25);
+    let T = Math.min(
+      N - 1,
+      Math.max(K + 1, N - Math.ceil((ptext.length - target) / avg) - 4),
+    );
+    let cutChar = -1;
+    for (let attempt = 0; attempt < INC_CUT_ATTEMPTS; attempt++) {
+      if (T <= K) return null;
+      const d = this.#decodeTail(pids.slice(T));
+      if (!ptext.endsWith(d)) {
+        T = Math.max(K, T - 8); // mid-char cut (byte-fallback/ByteLevel) → back off
+        continue;
+      }
+      const start = ptext.length - d.length;
+      if (start > target) {
+        T = Math.max(K, T - Math.ceil((start - target) / avg) - 2);
+        continue;
+      }
+      cutChar = start;
+      break;
+    }
+    if (cutChar < 0) return null;
+    // Anchor K tokens before the cut; its char position via the same tail trick.
+    const anchorTail = this.#decodeTail(pids.slice(T - K));
+    if (!ptext.endsWith(anchorTail)) return null;
+    const anchorChar = ptext.length - anchorTail.length;
+    if (anchorChar >= cutChar) return null; // zero-width seam window: refuse
+    const tailIds = this.#rawEncode(text.slice(cutChar), false);
+    // Seam verification: a fresh encode from the anchor must reproduce the
+    // cached tokens across the cut AND the spliced tail's start. When the
+    // tail is shorter than K the verify encode covers the ENTIRE remainder,
+    // so require full equality (no trailing extras).
+    const verifyIds = this.#rawEncode(text.slice(anchorChar), false);
+    const expectedLen = K + Math.min(K, tailIds.length);
+    if (tailIds.length <= K ? verifyIds.length !== expectedLen : verifyIds.length < expectedLen)
+      return null;
+    for (let i = 0; i < K; i++) if (verifyIds[i] !== pids[T - K + i]) return null;
+    for (let i = K; i < expectedLen; i++) if (verifyIds[i] !== tailIds[i - K]) return null;
+    return pids.slice(0, T).concat(tailIds);
+  }
+}
+
+/** True when the tokenizer.json post_processor adds specials only BEFORE the
+ *  sequence (or adds none): null, ByteLevel (offset-only, adds no ids), or a
+ *  TemplateProcessing whose `single` template ends with the Sequence slot.
+ *  Unknown processor types conservatively return false. */
+export function specialsArePrefixOnly(tokenizerJson: unknown): boolean {
+  const pp = (tokenizerJson as { post_processor?: { type?: string; single?: unknown[] } })
+    .post_processor;
+  if (pp == null) return true;
+  if (pp.type === "ByteLevel") return true;
+  if (pp.type === "TemplateProcessing" && Array.isArray(pp.single)) {
+    const single = pp.single as Record<string, unknown>[];
+    const seqIdx = single.findIndex((item) => "Sequence" in item);
+    return seqIdx >= 0 && seqIdx === single.length - 1;
+  }
+  return false;
 }
 
 export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer> {
@@ -25,6 +289,7 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
     Bun.file(`${modelDir}/tokenizer_config.json`).json(),
   ]);
   const tok = new Tokenizer(tokenizerJson, tokenizerConfig);
+  const enc = new IncrementalEncoder(tok, specialsArePrefixOnly(tokenizerJson));
 
   const idOf = (key: string): number | null => {
     const t = tokenizerConfig[key];
@@ -34,9 +299,23 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
     return enc.ids.length === 1 ? Number(enc.ids[0]) : null;
   };
 
+  // mirror mlx-lm's `_is_bpe_decoder` (tokenizer_utils.py): top-level
+  // decoder type "ByteLevel" selects the BPE streaming detokenizer.
+  const isByteLevel =
+    (tokenizerJson as { decoder?: { type?: string } }).decoder?.type === "ByteLevel";
+  // The bare-space token: in the GPT-2 ByteLevel byte-to-unicode map, byte 32
+  // (space) renders as "Ġ" (U+0120) — the only single-char vocab entry whose
+  // byte decodes to 32, i.e. the only token mlx-lm's BPEStreamingDetokenizer
+  // withholds in _unflushed. Look it up straight in the BPE vocab.
+  const bareSpaceTokenId = isByteLevel
+    ? (tokenizerJson as { model?: { vocab?: Record<string, number> } }).model
+        ?.vocab?.["Ġ"]
+    : undefined;
+
   return {
-    encode: (text, addSpecialTokens = true) =>
-      tok.encode(text, { add_special_tokens: addSpecialTokens }).ids.map(Number),
+    // Memo + incremental-splice path live in IncrementalEncoder; short
+    // prompts (<4096 chars, <1 ms) bypass both inside it.
+    encode: (text, addSpecialTokens = true) => enc.encode(text, addSpecialTokens),
     decode: (ids, skipSpecialTokens = false) =>
       // python's decode([]) === ""; the JS lib throws on empty input
       ids.length === 0 ? "" : tok.decode(ids, { skip_special_tokens: skipSpecialTokens }),
@@ -45,5 +324,12 @@ export async function loadTokenizer(modelDir: string): Promise<LoadedTokenizer>
         .id_to_token(id) ?? tok.decode([id], { skip_special_tokens: false }),
     bosTokenId: idOf("bos_token"),
     eosTokenId: idOf("eos_token"),
+    tokenizerJsonPath: `${modelDir}/tokenizer.json`,
+    // ByteLevel/BPE = the only detokenizer class with the sequence-start
+    // space-trim and the bare-space hold-back our decode() lacks (see the
+    // interface docs above).
+    trimsLeadingSpace: isByteLevel,
+    bareSpaceTokenId: typeof bareSpaceTokenId === "number" ? bareSpaceTokenId : undefined,
+    encodeStats: enc.stats,
   };
 }

package/src/tool-approvals.ts

@@ -0,0 +1,131 @@
+// Durable "always allow this tool" persistence for the web chat's
+// tool-approval gate (web-chat-redesign.md §5.4/§6.5/§9 Phase 2, risk #6).
+//
+// Convention mirrors src/hf-push.ts's ~/.mlx-bun/hf.json: a single plain
+// JSON file, mode 0600, no encryption-at-rest (the threat model is a single
+// local user on their own machine). Versioned shape from day one so a
+// future field addition can migrate instead of silently losing the file —
+// risk #6 explicitly calls out the Claude-Desktop-resets-on-update failure
+// mode as the thing NOT to repeat here.
+//
+// Keyed by tool NAME (stable identity — "bash", "edit", "write", or a
+// custom tool's name), not by call args: "always allow bash" means every
+// future bash call skips the card, matching every competitor's own
+// granularity for this feature (per-tool, not per-command).
+
+import { mkdirSync, readFileSync, writeFileSync, chmodSync } from "node:fs";
+import { join } from "node:path";
+import { homedir } from "node:os";
+
+/** Home dir, honoring $HOME at call time — matches hf-push.ts's home(). */
+function home(): string {
+  return process.env.HOME || homedir();
+}
+
+function configPath(): string {
+  return join(home(), ".mlx-bun", "tool-approvals.json");
+}
+
+export const TOOL_APPROVALS_VERSION = 1 as const;
+
+/** On-disk shape. `allows` maps tool name -> true (present+true = always
+ *  allow; absent = ask every time, the default). A plain boolean map (not a
+ *  richer per-tool record) is deliberately the whole shape for v1 — nothing
+ *  here needs migrating yet, but the version field exists so a future
+ *  richer shape (e.g. per-tool scoping, expiry) has somewhere to land. */
+export interface ToolApprovalsFileV1 {
+  version: 1;
+  allows: Record<string, true>;
+}
+
+/** Anything this module can load and return today. Widen this union (not
+ *  ToolApprovalsFileV1) when a v2 shape ships, and teach loadToolApprovals'
+ *  migration step to upgrade older files in place. */
+export type ToolApprovalsFile = ToolApprovalsFileV1;
+
+const EMPTY: ToolApprovalsFileV1 = { version: 1, allows: {} };
+
+/**
+ * Load the durable approvals file, migrating forward if needed.
+ *
+ * Never throws: a missing file, unreadable file, or corrupt JSON all
+ * degrade to the empty default (ask every time) rather than crashing the
+ * approval gate — a config-file bug must never make mutating tools MORE
+ * permissive by accident, and "ask every time" is the safe failure mode.
+ */
+export function loadToolApprovals(): ToolApprovalsFile {
+  let raw: string;
+  try {
+    raw = readFileSync(configPath(), "utf8");
+  } catch {
+    return { ...EMPTY, allows: {} };
+  }
+  let parsed: unknown;
+  try {
+    parsed = JSON.parse(raw);
+  } catch {
+    return { ...EMPTY, allows: {} };
+  }
+  return migrate(parsed);
+}
+
+/**
+ * Migrate an arbitrary parsed JSON value to the current ToolApprovalsFile
+ * shape. Unknown/missing `version` is treated as "pre-versioning" (there
+ * isn't one yet — this IS v1 — but the shape is defensive against a
+ * hand-edited or truncated file rather than assuming well-formedness).
+ * Unknown future versions (a downgrade scenario) fall back to empty rather
+ * than guessing at an unfamiliar shape.
+ */
+function migrate(parsed: unknown): ToolApprovalsFile {
+  if (!parsed || typeof parsed !== "object") return { ...EMPTY, allows: {} };
+  const obj = parsed as Record<string, unknown>;
+  const version = obj.version;
+  if (version === 1 || version === undefined) {
+    const allowsRaw = obj.allows;
+    const allows: Record<string, true> = {};
+    if (allowsRaw && typeof allowsRaw === "object") {
+      for (const [k, v] of Object.entries(allowsRaw as Record<string, unknown>)) {
+        if (v === true) allows[k] = true;
+      }
+    }
+    return { version: 1, allows };
+  }
+  // Unrecognized/future version: don't guess at its shape. Ask every time
+  // is always the safe degradation.
+  return { ...EMPTY, allows: {} };
+}
+
+function saveToolApprovals(file: ToolApprovalsFile): void {
+  const dir = join(home(), ".mlx-bun");
+  mkdirSync(dir, { recursive: true });
+  const path = configPath();
+  writeFileSync(path, JSON.stringify(file, null, 2) + "\n", { mode: 0o600 });
+  chmodSync(path, 0o600); // writeFileSync's mode is masked by umask on create
+}
+
+/** Whether `toolName` is currently in the durable always-allow set. */
+export function isToolAlwaysAllowed(toolName: string): boolean {
+  return loadToolApprovals().allows[toolName] === true;
+}
+
+/** Persist "always allow this tool" for `toolName`. Idempotent. */
+export function setToolAlwaysAllowed(toolName: string): ToolApprovalsFile {
+  const file = loadToolApprovals();
+  file.allows[toolName] = true;
+  saveToolApprovals(file);
+  return file;
+}
+
+/** Revoke a previously-granted always-allow (settings UI "forget" action). */
+export function revokeToolAlwaysAllowed(toolName: string): ToolApprovalsFile {
+  const file = loadToolApprovals();
+  delete file.allows[toolName];
+  saveToolApprovals(file);
+  return file;
+}
+
+/** Full always-allow set, for a settings-panel listing. */
+export function listAlwaysAllowedTools(): string[] {
+  return Object.keys(loadToolApprovals().allows).sort();
+}