mlx-bun 0.2.0 → 0.4.0

package/src/serve/chat-request.ts

@@ -0,0 +1,391 @@
+// Wire-level chat/completion request shapes and the pure per-request
+// helpers: message normalization, sampling-field coercion, grammar-degrade
+// prompt injection, and the mlx_lm.server-parity validators. No model, no
+// server state. Extracted from src/server.ts (repo-taming Phase 4).
+import type { ChatMessage, ToolDefinition } from "../chat-template";
+import type { HlgConfig } from "../sampler";
+import { RequestError } from "./pipeline";
+
+export interface ChatRequestParams {
+  messages: ChatMessage[];
+  stream?: boolean;
+  max_tokens?: number;
+  max_completion_tokens?: number;
+  temperature?: number;
+  top_p?: number;
+  top_k?: number;
+  seed?: number;
+  repetition_penalty?: number;
+  /** mlx-lm extension: recent-token window for repetition_penalty
+   *  (default 20; 0 = whole history, Python `[-0:]` semantics). */
+  repetition_context_size?: number;
+  /** min-p sampling (mlx_lm.server's `min_p`): keep tokens whose probability
+   *  is ≥ min_p · p(top token). 0 = off. */
+  min_p?: number;
+  /** XTC sampling (mlx_lm.server names): with probability `xtc_probability`
+   *  per step, remove every token above `xtc_threshold` except the least
+   *  likely of them. EOS + the newline token are always exempt (the server
+   *  injects them as xtc special tokens, matching mlx_lm.server). */
+  xtc_probability?: number;
+  xtc_threshold?: number;
+  /** OpenAI logit_bias: {tokenId: additive bias}. JSON object keys arrive as
+   *  strings; coerced to int keys / float values like mlx-lm (400 on failure). */
+  logit_bias?: Record<string, number>;
+  /** OpenAI presence/frequency penalties + mlx-lm's context-size extensions
+   *  (window of recent tokens the penalty looks at; default 20). */
+  presence_penalty?: number;
+  presence_context_size?: number;
+  frequency_penalty?: number;
+  frequency_context_size?: number;
+  /** mlx_lm.server logprobs: `logprobs` is a BOOL (even on /v1/completions —
+   *  not OpenAI's legacy int), `top_logprobs` an int in [0, 11] or the -1
+   *  "unset" sentinel (server.py validates exactly that; OpenAI's cap is 20,
+   *  mlx-lm's is 11 — we copy the reference). Non-stream responses carry
+   *  mlx-lm's logprobs block; stream chunks never do (reference behavior). */
+  logprobs?: boolean;
+  top_logprobs?: number;
+  /** OpenAI stop sequences: plain string or array (spec allows up to 4).
+   *  Matched on DECODED text, not token ids — see StopMatcher. */
+  stop?: string | string[];
+  tools?: ToolDefinition[];
+  tool_choice?: "auto" | "none" | { type: string; function?: { name: string } };
+  /** Forwarded to HF chat templates, matching optiq serve. MiniCPM5 uses
+   *  enable_thinking to select direct answers vs the <think> channel; Qwen3.8
+   *  adds preserve_thinking (keep think blocks from history; template default
+   *  true — better prompt-cache reuse in agent loops). */
+  chat_template_kwargs?: {
+    enable_thinking?: boolean;
+    preserve_thinking?: boolean;
+    [key: string]: unknown;
+  };
+  /** OpenAI reasoning control. For models with a switchable <think> channel
+   *  (Qwen3.5/MiniCPM5) it gates enable_thinking: "none" → off, any level → on.
+   *  This is what Pi sends when the provider advertises reasoning. For
+   *  templates that read a reasoning-depth variable (Qwen3.8: xhigh/medium/
+   *  low), the level ALSO maps into the template via qwenReasoningEffort —
+   *  "xhigh" accepted as a first-class value (Qwen3.8's default/deepest). */
+  reasoning_effort?: "none" | "minimal" | "low" | "medium" | "high" | "xhigh";
+  /** Mounted LoRA adapter selection: "id", "a+b" (stacked), or "none". */
+  adapter?: string;
+  /** HLG tone-curve sampling override (per request). Snake_case wire fields,
+   *  merged over the server's --hlg-sampling config. docs/archive/hlg-sampling.md. */
+  hlg?: {
+    enabled?: boolean;
+    width?: number;
+    shoulder?: number;
+    toe?: number;
+    pivot_offset?: number;
+  };
+  /** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
+   *  json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
+   *  sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
+   *  oracle). On compile failure, degrades to a system-prompt injection +
+   *  Warning header (oMLX parity), never 500. */
+  response_format?: unknown;
+  /** vLLM/oMLX grammar aliases (all compiled via xgrammar): raw EBNF/LARK
+   *  grammar, regex, enum choice, and a bare JSON-schema object. Precedence
+   *  (guided_grammar > json_schema > json_object > structured_outputs >
+   *  guided_regex > guided_choice) mirrors oMLX _effective_guided_grammar. */
+  guided_grammar?: string;
+  guided_regex?: string;
+  guided_choice?: string[];
+  structured_outputs?: unknown;
+}
+
+/** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
+ *  template). Sampling/penalty/stop fields are the same names as ChatRequestParams;
+ *  `prompt` replaces `messages`. mlx_lm.server accepts only a STRING prompt
+ *  (it calls `tokenizer.encode(request.prompt)` directly) — token-array
+ *  prompts are rejected there too, so we match. No `echo` (mlx-lm has none). */
+export type TextCompletionParams = Omit<ChatRequestParams, "messages" | "tools" | "tool_choice"> & {
+  prompt?: unknown;
+};
+
+
+/** Per-field default HLG knobs when enabling without specifying them. */
+export const HLG_DEFAULTS = { width: 4, shoulder: 4, toe: 6, pivotOffset: 6 } as const;
+
+/** Resolve the effective HLG config: a per-request `hlg` object overrides the
+ *  server's --hlg-sampling default field-by-field. Returns undefined (HLG off)
+ *  unless enabled by the request or the server. */
+export function resolveHlg(
+  reqHlg: ChatRequestParams["hlg"],
+  serverHlg: HlgConfig | undefined,
+): HlgConfig | undefined {
+  const enabled = reqHlg?.enabled ?? serverHlg?.enabled ?? false;
+  if (!enabled) return undefined;
+  const base = serverHlg ?? HLG_DEFAULTS;
+  return {
+    enabled: true,
+    width: reqHlg?.width ?? base.width,
+    shoulder: reqHlg?.shoulder ?? base.shoulder,
+    toe: reqHlg?.toe ?? base.toe,
+    pivotOffset: reqHlg?.pivot_offset ?? base.pivotOffset,
+    pivot: "top",
+  };
+}
+
+/** True when a rendered prompt ends INSIDE an unclosed `<think>` block — the
+ *  generation prompt primed reasoning (Qwen3.5 / MiniCPM5 with thinking on),
+ *  so the model continues the chain-of-thought and emits only the closing
+ *  `</think>`. Seeds ThinkingTagSplitter so reasoning is split out correctly.
+ *  Thinking-off primes a CLOSED empty block (`<think>\n\n</think>`), and
+ *  no-thinking templates have no `<think>` at all — both return false. */
+export function promptEndsInOpenThink(rendered: string): boolean {
+  const open = rendered.lastIndexOf("<think>");
+  return open !== -1 && open > rendered.lastIndexOf("</think>");
+}
+
+/** Concatenate the text of an OpenAI content-part array, ignoring non-text
+ *  parts. Tolerant of the part shapes clients actually send: `{type:"text",
+ *  text}` (OpenAI/pi), `{type:"input_text", text}` (Responses-style), or any
+ *  part carrying a string `text`. */
+export function contentPartsToText(parts: Array<Record<string, unknown>>): string {
+  return parts
+    .map((p) => (p && typeof p.text === "string" ? p.text : ""))
+    .join("");
+}
+
+/** True if any content part is media — image OR audio — so the multimodal
+ *  path must keep the array form for extractImages/extractAudio. */
+export function hasMediaPart(parts: Array<Record<string, unknown>>): boolean {
+  return parts.some(
+    (p) =>
+      p &&
+      (p.type === "image" || p.type === "image_url" ||
+        p.type === "audio" || p.type === "input_audio" || p.type === "audio_url" ||
+        p.type === "video" || p.type === "video_url"),
+  );
+}
+
+/** OpenAI sends assistant tool_call arguments as JSON strings; the
+ *  template renders the object form natively — normalize before render.
+ *
+ *  Two more wire-format → template-format fixes, both "match the format the
+ *  model's chat template expects":
+ *
+ *  1. Map the OpenAI reasoning-model "developer" role to "system": pi-ai (and
+ *     OpenAI's own SDKs) rename the system prompt to `developer` whenever a
+ *     model advertises reasoning, but our chat templates only know
+ *     system/user/assistant/tool and raise "Unexpected message role." This is
+ *     why Qwen3.5/MiniCPM5 chat got no messages while Gemma (non-reasoning, so
+ *     pi keeps `system`) worked. `developer` IS the system prompt, so the remap
+ *     is semantics-preserving.
+ *
+ *  2. Flatten text-only content-part arrays to a plain string. pi (and any
+ *     OpenAI multimodal client) sends user content as `[{type:"text",text}]`,
+ *     but non-vision chat templates expect `content` to be a STRING and render
+ *     nothing for an array — so the user's turn silently vanishes and the model
+ *     replies "I don't see any message." Arrays that carry a media part
+ *     (image or audio) are left intact for the multimodal path
+ *     (extractImages/extractAudio); only text-only arrays are collapsed
+ *     here, which is a no-op for the vision templates too. */
+export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
+  return messages.map((raw) => {
+    let m = raw.role === "developer" ? { ...raw, role: "system" } : raw;
+    if (Array.isArray(m.content) && !hasMediaPart(m.content)) {
+      m = { ...m, content: contentPartsToText(m.content) };
+    }
+    // Pi and mlx-lm return `reasoning`; Qwen templates read
+    // `reasoning_content`. Keep the original field and immutable history,
+    // and let an explicitly supplied canonical string (including "") win.
+    if (typeof m.reasoning_content !== "string" && typeof m.reasoning === "string")
+      m = { ...m, reasoning_content: m.reasoning };
+    if (!m.tool_calls) return m;
+    return {
+      ...m,
+      tool_calls: m.tool_calls.map((tc) => ({
+        ...tc,
+        function: {
+          ...tc.function,
+          arguments:
+            typeof tc.function.arguments === "string"
+              ? (JSON.parse(tc.function.arguments) as Record<string, unknown>)
+              : tc.function.arguments,
+        },
+      })),
+    };
+  });
+}
+
+/** Coerce a wire `logit_bias` ({"tokenId": bias}) to numeric keys/values —
+ *  mlx_lm.server's `{int(k): float(v)}` coercion; throws its exact error
+ *  message on anything non-numeric (surfaced as a 400). JSON object keys
+ *  always arrive as strings, hence the coercion. */
+export function parseLogitBias(
+  raw: Record<string, number> | undefined | null,
+): Record<number, number> | undefined {
+  if (raw == null) return undefined;
+  if (typeof raw !== "object" || Array.isArray(raw))
+    throw new Error("logit_bias must be a dict of int to float");
+  const out: Record<number, number> = {};
+  for (const [k, v] of Object.entries(raw)) {
+    const id = Number(k);
+    const bias = Number(v);
+    if (!Number.isInteger(id) || !Number.isFinite(bias))
+      throw new Error("logit_bias must be a dict of int to float");
+    out[id] = bias;
+  }
+  return Object.keys(out).length ? out : undefined;
+}
+
+/** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
+ *  request-unique: under `--batch N` the batch lane serves ONLY default-seed
+ *  requests (explicit-seed requests use the serial mechanism, see
+ *  GenerationGateway.place),
+ *  so two identical prompts arriving in the same millisecond would share a seed
+ *  and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
+ *  byte-identical completions, silently collapsing best-of-N diversity. Mix a
+ *  per-process Weyl counter (golden-ratio increment, period 2^32) into the
+ *  timestamp so every call yields a distinct uint32 within any given ms.
+ *  Determinism contract unchanged: reproducibility is only promised for an
+ *  EXPLICIT request seed (`req.seed ?? nextDefaultSeed()` — explicit wins,
+ *  byte-identical to before); a default seed is fresh entropy per request. */
+let seedWeyl = 0;
+export function nextDefaultSeed(): number {
+  seedWeyl = (seedWeyl + 0x9e3779b9) >>> 0;
+  return ((Date.now() & 0xffffffff) ^ seedWeyl) >>> 0;
+}
+
+/** mlx_lm.server's logprobs request validation, copied exactly (server.py
+ *  APIHandler.validate_model_parameters: `_validate("logprobs", bool)` and
+ *  `_validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])`
+ *  with defaults logprobs=False / top_logprobs=-1). Returns the reference's
+ *  error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
+ *  at 11, not OpenAI's 20 — we mirror the reference. */
+/** Build a degrade-path system-prompt instruction for JSON output, mirroring
+ *  oMLX's api.tool_calling.build_json_system_prompt (used when xgrammar
+ *  compile fails — the response_format degrades to prompt injection rather
+ *  than a 500, oMLX parity). Returns null for {type:"text"} / unset. */
+export function degradeJsonSystemPrompt(body: ChatRequestParams): string | null {
+  if (body.guided_grammar) {
+    return "You must respond with text matching this grammar:\n\n" + body.guided_grammar;
+  }
+  const rf = body.response_format as { type?: string; json_schema?: { name?: string; description?: string; schema?: unknown } } | undefined;
+  if (rf && typeof rf === "object" && rf.type === "json_object") {
+    return "You must respond with valid JSON only. " +
+      "Do not include any explanation or text outside the JSON object.";
+  }
+  const schemaSpec = rf && typeof rf === "object" && rf.type === "json_schema"
+    ? rf.json_schema
+    : body.structured_outputs && typeof body.structured_outputs === "object"
+      ? { schema: body.structured_outputs }
+      : null;
+  if (schemaSpec) {
+    const spec = schemaSpec;
+    const schema = spec.schema ?? body.structured_outputs ?? {};
+    const name = spec.name ?? "response";
+    const description = spec.description ?? "";
+    let prompt = `You must respond with valid JSON matching the '${name}' schema.`;
+    if (description) prompt += ` ${description}`;
+    prompt += `\n\nJSON Schema:\n\`\`\`json\n${JSON.stringify(schema, null, 2)}\n\`\`\`\n\n` +
+      "Respond with only the JSON object, no additional text or explanation.";
+    return prompt;
+  }
+  if (body.guided_regex) {
+    return "You must respond with text matching this regular expression:\n\n" +
+      body.guided_regex;
+  }
+  if (body.guided_choice?.length) {
+    return "You must respond with exactly one of these choices:\n\n" +
+      body.guided_choice.map((choice) => JSON.stringify(choice)).join("\n");
+  }
+  return null;
+}
+
+export function applyGrammarDegrade(
+  body: ChatRequestParams,
+  degradeHint: string,
+): { body: ChatRequestParams; warning: string } {
+  const content = degradeJsonSystemPrompt(body) ??
+    "Follow the requested output constraint exactly.";
+  return {
+    warning:
+      `grammar not enforced: ${degradeHint} - falling back to prompt injection`,
+    body: {
+      ...body,
+      messages: [{ role: "system", content }, ...body.messages],
+    },
+  };
+}
+
+export function validateLogprobsParams(body: {
+  logprobs?: unknown;
+  top_logprobs?: unknown;
+}): string | null {
+  const lp = body.logprobs ?? false;
+  if (typeof lp !== "boolean") return "logprobs must be of type bool";
+  const tl = body.top_logprobs ?? -1;
+  if (typeof tl !== "number" || !Number.isInteger(tl))
+    return "top_logprobs must be of type int";
+  if (tl === -1) return null; // the "unset" whitelist sentinel
+  if (tl < 0) return "top_logprobs must be at least 0";
+  if (tl > 11) return "top_logprobs must be at most 11";
+  return null;
+}
+
+const REASONING_EFFORT_LEVELS =
+  ["none", "minimal", "low", "medium", "high", "xhigh"] as const;
+
+/** Validate `reasoning_effort` against the accepted level names. An invalid
+ *  string must NOT pass through: resolveEnableThinking treats any defined
+ *  value ≠ "none" as thinking-ON, so a typo like "hihg" silently flipped the
+ *  thinking channel (and its hotter temperature) on thinking-off models and
+ *  burned a full default-depth reasoning budget (2026-08-18 review). Same
+ *  contract style as validateLogprobsParams above. */
+export function validateReasoningEffort(body: {
+  reasoning_effort?: unknown;
+}): string | null {
+  const v = body.reasoning_effort;
+  if (v === undefined || v === null) return null;
+  if (typeof v !== "string" || !(REASONING_EFFORT_LEVELS as readonly string[]).includes(v))
+    return "reasoning_effort must be one of " +
+      "'none', 'minimal', 'low', 'medium', 'high', 'xhigh'";
+  return null;
+}
+
+// ---------------------------------------------------------------------------
+// The request objects a caller constructs and hands to the pipeline. The
+// constructor is the validation boundary: structural checks that need no
+// model (mlx_lm.server parity messages) throw RequestError(400) here, so a
+// stage never sees a malformed request.
+
+export class ChatRequest {
+  readonly params: ChatRequestParams;
+
+  constructor(params: ChatRequestParams) {
+    if (!Array.isArray(params.messages) || params.messages.length === 0)
+      throw new RequestError(400, "messages required");
+    const lpParamError = validateLogprobsParams(params);
+    if (lpParamError) throw new RequestError(400, lpParamError);
+    const effortError = validateReasoningEffort(params);
+    if (effortError) throw new RequestError(400, effortError);
+    this.params = params;
+  }
+
+  get stream(): boolean {
+    return this.params.stream === true;
+  }
+}
+
+/** POST /v1/completions — raw text, no chat template (mlx_lm.server's
+ *  request_type "text"). Only a non-empty STRING prompt is accepted. */
+export class TextCompletionRequest {
+  readonly params: TextCompletionParams & { prompt: string };
+
+  constructor(params: TextCompletionParams) {
+    if (typeof params.prompt !== "string" || params.prompt.length === 0)
+      throw new RequestError(
+        400,
+        "prompt (a non-empty string) is required " +
+          "(token-array prompts are not accepted, matching mlx_lm.server)",
+      );
+    const lpParamError = validateLogprobsParams(params);
+    if (lpParamError) throw new RequestError(400, lpParamError);
+    this.params = params as TextCompletionParams & { prompt: string };
+  }
+
+  get stream(): boolean {
+    return this.params.stream === true;
+  }
+}

package/src/serve/chat-stage.ts

@@ -0,0 +1,195 @@
+// ChatStage: ChatRequest in, InferenceRequest out. Turns a chat conversation
+// (messages, tools, media parts, sampling fields) into the prompt token ids,
+// media embeddings, resolved generation options, adapter selection, grammar
+// controller, and the token→text pipeline the model needs. Every chat-shaped
+// surface (OpenAI chat, Anthropic Messages, Responses) runs this stage; the
+// wire format is not its business.
+import type { PromptCache } from "../prompt-cache";
+import {
+  applyGrammarDegrade,
+  type ChatRequest,
+  type ChatRequestParams,
+} from "./chat-request";
+import type { InferenceRequest } from "./inference-request";
+import type { ServingContext } from "./model-host";
+import type { BuiltPrompt, ModelPromptBuilder } from "./model-binding";
+import { modelPromptBuilder } from "../backends/mlx/model-serving";
+import { RequestError } from "./pipeline";
+import { RequestOwnership } from "./request-plan";
+import type { RequestPrep } from "./request-prep";
+import { StopMatcher, ThinkingTagSplitter } from "./token-streams";
+
+export class ChatStage {
+  constructor(
+    private readonly ctx: ServingContext,
+    private readonly prep: RequestPrep,
+    private readonly promptCache: Pick<PromptCache, "peekPrefixLen">,
+    /** Admission ceiling (fit() / the GLM memory plan), resolved once. */
+    private readonly contextLimit: number | null,
+    /** `serve --adapter <dir>` startup default; a request's explicit
+     *  `adapter` (incl. "none") wins. */
+    private readonly defaultAdapter?: string,
+    private readonly preparation?: import("./preparation").PreparationExecutor,
+    private readonly buildPrompt: ModelPromptBuilder = modelPromptBuilder(ctx),
+  ) {}
+
+  async run(request: ChatRequest, requestId?: string, signal?: AbortSignal): Promise<InferenceRequest> {
+    signal?.throwIfAborted();
+    const body = request.params;
+    const media = body.messages.some((message) => Array.isArray(message.content) &&
+        message.content.some((part) => part.type !== "text"));
+    const native = media || body.response_format != null || !!body.guided_grammar ||
+      !!body.guided_regex || !!body.guided_choice?.length || body.structured_outputs != null;
+    const build = async () => {
+      signal?.throwIfAborted();
+      const result = await this.runPrepared(request, requestId);
+      if (signal?.aborted) {
+        result.plan.ownership.dispose();
+        signal.throwIfAborted();
+      }
+      return result;
+    };
+    if (!native || !this.preparation) return build();
+    let reservation = await this.preparation.reserve?.(media ? "media" : "constraint", signal);
+    try {
+      const result = await this.preparation.run(build, signal);
+      if (reservation) result.plan.ownership.retain(reservation);
+      reservation = undefined;
+      return result;
+    } finally { reservation?.dispose(); }
+  }
+
+  private async runPrepared(request: ChatRequest, requestId?: string): Promise<InferenceRequest> {
+    const { ctx, prep } = this;
+    let body = request.params;
+    const id = requestId ?? `chatcmpl-${crypto.randomUUID()}`;
+    const tools =
+      body.tool_choice === "none" ? null : (body.tools?.length ? body.tools : null);
+    const ownership = new RequestOwnership();
+    const warnings: string[] = [];
+    const reject = (status: number, message: string): never => {
+      ownership.dispose();
+      throw new RequestError(status, message);
+    };
+
+    // Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
+    // rendering: on the degrade path (compile failed but a constraint was
+    // requested) inject a system message instructing valid JSON so the
+    // model still best-efforts schema-conformant output (oMLX parity —
+    // _compile_grammar_for_request returning None + build_json_system_prompt).
+    let grammarCtrl: import("../grammar").GrammarController | null = null;
+    const grammarReq =
+      body.response_format != null || !!body.guided_grammar ||
+      !!body.guided_regex || !!body.guided_choice?.length ||
+      body.structured_outputs != null;
+    if (grammarReq) {
+      const g = await prep.compileGrammarForRequest(body);
+      grammarCtrl = ownership.own(g.controller);
+      if (!g.controller && g.degradeHint) {
+        const degraded = applyGrammarDegrade(body, g.degradeHint);
+        warnings.push(degraded.warning);
+        body = degraded.body;
+      }
+    }
+
+    let built: BuiltPrompt;
+    try {
+      built = await this.buildPrompt(body, tools, ownership, prep);
+    } catch (e) {
+      if (e instanceof RequestError) return reject(e.status, e.message);
+      return reject(400, `prompt build failed: ${(e as Error).message}`);
+    }
+    const { promptIds, vision, startInThinking, diffusionPixels } = built;
+    ownership.own(vision?.embeddings);
+    ownership.own(vision?.imageMask);
+    ownership.own(vision?.multimodalMask);
+    ownership.own(diffusionPixels);
+
+    let options: ReturnType<RequestPrep["toOptions"]>;
+    try {
+      options = prep.toOptions(body);
+      // Stable cache boundary for the prompt-boundary snapshot (chat
+      // prompts end in a generation primer the next turn won't contain).
+      // The probe costs a second full render+encode (~150 ms at 16k), so
+      // run it ONLY when a snapshot can actually be taken: the cache
+      // must not already hold the strict prefix (warm repeats peek at
+      // len-1 → skip; the snapshot gate needs boundary > cached prefix).
+      if (built.probeStableLen) {
+        const cachePeek = this.promptCache.peekPrefixLen(
+          promptIds, options.adapters?.length ? ctx.adapters.cacheNamespace(options.adapters) : "");
+        if (cachePeek < promptIds.length - 1)
+          options.snapshotAt = prep.stableLenFor(body, tools, promptIds);
+      }
+    } catch (e) {
+      // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
+      return reject(400, (e as Error).message);
+    }
+    if (diffusionPixels) options.visionPixels = diffusionPixels;
+    // Attach the compiled grammar controller (null when no constraint /
+    // degrade — generate() runs the unmasked fast pipelined loop).
+    if (grammarCtrl) options.grammar = grammarCtrl;
+    // Token fast-forwarding (K3, MLX_BUN_FILL=strict): the determined-span
+    // table for this request's tools. Refused for the shapes only this stage
+    // can see — a compiled grammar (forced tokens are its job), media prompts
+    // (embeddings prefill / mRoPE), a mounted draft model (the spec loop is a
+    // different executor and never reads options.fill), and quantized-KV
+    // schemes (post-conversion multi-token append is unvalidated). The
+    // body-level refusals live in prep.fillPlanFor.
+    let fillSession: import("../fill/fill-session").FillSession | null = null;
+    if (!grammarCtrl && !vision && !diffusionPixels && !ctx.draft &&
+        !options.kvBits && !options.kvConfig?.length && !options.turboQuant) {
+      fillSession = prep.fillPlanFor(body, tools, promptIds);
+      if (fillSession) options.fill = fillSession;
+    }
+    // Unset = no request-level ceiling; admission clamps to the admitted
+    // context (the program is the cap).
+    const requestedMaxTokens = options.maxTokens ?? Infinity;
+    let adapterIds: string[];
+    try {
+      // A request's explicit `adapter` (incl. "none") wins over the
+      // startup default from `serve --adapter <dir>`.
+      adapterIds = ctx.adapters.resolveSpec(body.adapter ?? this.defaultAdapter);
+    } catch (e) {
+      return reject(400, (e as Error).message);
+    }
+
+    const wantLogprobs = body.logprobs === true;
+    const topLogprobs =
+      typeof body.top_logprobs === "number" && body.top_logprobs > 0
+        ? body.top_logprobs : 0;
+    return {
+      requestId: id,
+      stream: request.stream,
+      warnings,
+      plan: {
+        promptIds,
+        options,
+        requestedMaxTokens,
+        contextLimit: this.contextLimit,
+        stream: request.stream,
+        wantLogprobs,
+        topLogprobs,
+        adapterIds,
+        hasVision: !!vision,
+        userSeed: body.seed !== undefined,
+        hasGrammar: !!grammarCtrl,
+        hasDraft: !!ctx.draft,
+        ownership,
+      },
+      ...(vision ? { vision } : {}),
+      pipeline: {
+        // The router's parse-failure hook disarms this request's strict fill
+        // rows (and records `usage.fill.parseFallback`).
+        router: prep.toolRouter(tools, fillSession ? () => fillSession!.noteParseFailure() : undefined),
+        stopper: new StopMatcher(options.stopSequences),
+        thinking: new ThinkingTagSplitter(
+          ctx.template.thinkingFormat === "think-tag",
+          startInThinking,
+        ),
+        collectToolCalls: true,
+      },
+      idToToken: (tokenId) => ctx.tokenizer.idToToken(tokenId),
+    };
+  }
+
+}

package/src/serve/checkpoint-identity.ts

@@ -0,0 +1,64 @@
+import type { GenerateOptions } from "../generate";
+import type { ResolvedExecution } from "../contracts/execution";
+import { createHash } from "node:crypto";
+
+/** Sort object fields only; array order is meaningful for adapter composition,
+ * stop handling and quantization recipes. */
+function canonical(value: unknown): unknown {
+  if (Array.isArray(value)) return value.map(canonical);
+  if (value && typeof value === "object")
+    return Object.fromEntries(Object.entries(value).sort(([a], [b]) => a.localeCompare(b))
+      .map(([key, item]) => [key, canonical(item)]));
+  return value;
+}
+
+/** Stable identity for an exactly resumable generation. The prompt already
+ * includes rendered messages and tool schemas; the remaining fields cover
+ * every sampling policy that can change the continuation. */
+export function generationCheckpointKey(
+  promptIds: number[], options: GenerateOptions & { stopSequences?: readonly string[] }, cacheNs = "",
+  execution?: ResolvedExecution, bindingIdentity?: unknown,
+): string {
+  const policy = {
+    version: 4,
+    execution: execution && { method: execution.method, mechanism: execution.mechanism,
+      pagedKv: execution.pagedKv, fill: execution.fill,
+      compiledDecode: execution.compiledDecode, grammarJump: execution.grammarJump },
+    bindingIdentity,
+    stopSequences: options.stopSequences ?? [],
+    cacheNs,
+    promptIds,
+    maxTokens: options.maxTokens,
+    eosTokenIds: options.eosTokenIds,
+    temperature: options.temperature,
+    topP: options.topP,
+    topK: options.topK,
+    minP: options.minP,
+    minTokensToKeep: options.minTokensToKeep,
+    xtcProbability: options.xtcProbability,
+    xtcThreshold: options.xtcThreshold,
+    xtcSpecialTokens: options.xtcSpecialTokens,
+    // A missing client seed means "fresh randomness" for a new request, but
+    // an identical RETRY must find the durable checkpoint and restore the
+    // original random stream. Keep the generated seed out of request identity;
+    // it travels in checkpoint metadata instead. Explicit seeds remain part of
+    // the policy, so changing one cannot select an incompatible checkpoint.
+    seedMode: options.seedWasExplicit ? "explicit" : "server-default",
+    seed: options.seedWasExplicit ? options.seed : undefined,
+    hlg: options.hlg,
+    curve: options.curve,
+    logitBias: options.logitBias,
+    repetitionPenalty: options.repetitionPenalty,
+    repetitionContextSize: options.repetitionContextSize,
+    presencePenalty: options.presencePenalty,
+    presenceContextSize: options.presenceContextSize,
+    frequencyPenalty: options.frequencyPenalty,
+    frequencyContextSize: options.frequencyContextSize,
+    kvBits: options.kvBits,
+    kvGroupSize: options.kvGroupSize,
+    quantizedKvStart: options.quantizedKvStart,
+    kvConfig: options.kvConfig,
+    turboQuant: options.turboQuant,
+  };
+  return createHash("sha256").update(JSON.stringify(canonical(policy))).digest("hex");
+}