mlx-bun 0.0.8 → 0.0.9

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (133) hide show
  1. package/README.md +112 -20
  2. package/docs/design/batching-v2-plan.md +515 -0
  3. package/docs/design/bucketing-stage.md +196 -0
  4. package/docs/design/compat-cli-surface-design.md +11 -11
  5. package/docs/design/docs-reorg-plan.md +1 -1
  6. package/docs/design/dreaming-nightly-pipeline.md +291 -0
  7. package/docs/design/dspark-speculative-decoding.md +106 -0
  8. package/docs/design/generic-model-support.md +381 -0
  9. package/docs/design/memory-inference-path.md +422 -0
  10. package/docs/design/memory-synthesis.md +287 -0
  11. package/docs/design/minicpm5-decode-megakernel.md +484 -0
  12. package/docs/design/mlx-lm-tool-parity-plan.md +535 -0
  13. package/docs/design/orpo-dynamic-lambda.md +41 -0
  14. package/docs/design/orpo-training.md +35 -5
  15. package/docs/design/parallel-slots.md +12 -1
  16. package/docs/design/parity-tier-dag.md +13 -5
  17. package/docs/design/spec-decode-larger-targets.md +229 -0
  18. package/docs/design/tauri-desktop-app.md +2 -2
  19. package/docs/design/the-dreaming-handoff.md +141 -0
  20. package/docs/design/the-dreaming-master-plan.md +446 -0
  21. package/docs/design/write-pipeline-entity.md +468 -0
  22. package/docs/investigations/curve-bisector-routing.md +1 -1
  23. package/docs/investigations/decode-roofline-lookagain.md +262 -0
  24. package/docs/investigations/dspark-handoff.md +58 -0
  25. package/docs/investigations/kernel-perf-review-2026-07.md +92 -0
  26. package/docs/planning/PRODUCT_ROADMAP.md +2 -2
  27. package/docs/planning/curve-sampler-research-plan.md +331 -0
  28. package/docs/planning/memory-docs-and-dag-plan.md +185 -0
  29. package/docs/planning/pi-integration-review.md +178 -0
  30. package/docs/planning/release-notes-v0.0.9.md +122 -0
  31. package/docs/planning/web-ui-pass-plan.md +213 -0
  32. package/docs/planning/website-readme-pass-plan.md +224 -0
  33. package/docs/reference/cli.md +255 -0
  34. package/docs/reference/embedding.md +7 -2
  35. package/docs/reference/library-api.md +22 -3
  36. package/docs/reference/memory.md +5 -5
  37. package/docs/reference/models.md +137 -0
  38. package/docs/reference/orpo-quickstart.md +6 -0
  39. package/docs/reference/server-api.md +187 -9
  40. package/docs/reference/server-config.md +40 -9
  41. package/docs/reference/training.md +25 -11
  42. package/package.json +1 -1
  43. package/src/anthropic.ts +22 -0
  44. package/src/assets/curve-designer.html +1 -1
  45. package/src/chat-template.ts +25 -7
  46. package/src/cli.ts +1029 -86
  47. package/src/config.ts +4 -1
  48. package/src/download.ts +99 -8
  49. package/src/embed.ts +77 -0
  50. package/src/eval/perplexity.ts +164 -0
  51. package/src/eval/runner.ts +15 -6
  52. package/src/eval/tasks/bfcl.ts +1 -1
  53. package/src/generate.ts +141 -18
  54. package/src/harness-pi.ts +43 -17
  55. package/src/index.ts +2 -0
  56. package/src/lora.ts +2 -2
  57. package/src/memory/article.ts +558 -0
  58. package/src/memory/chunk-validation.ts +150 -0
  59. package/src/memory/chunk.ts +305 -12
  60. package/src/memory/cluster.ts +344 -17
  61. package/src/memory/coverage.ts +134 -0
  62. package/src/memory/crosslink.ts +0 -0
  63. package/src/memory/db.ts +184 -45
  64. package/src/memory/entity.ts +203 -0
  65. package/src/memory/gate.ts +187 -0
  66. package/src/memory/history.ts +152 -0
  67. package/src/memory/ingest.ts +324 -11
  68. package/src/memory/model.ts +321 -0
  69. package/src/memory/normalize.ts +671 -0
  70. package/src/memory/parse.ts +22 -0
  71. package/src/memory/pipeline.ts +242 -30
  72. package/src/memory/prompts.ts +42 -0
  73. package/src/memory/query.ts +334 -0
  74. package/src/memory/reconcile.ts +734 -0
  75. package/src/memory/redirect.ts +108 -0
  76. package/src/memory/reindex.ts +325 -0
  77. package/src/memory/resolve.ts +307 -0
  78. package/src/memory/route.ts +272 -0
  79. package/src/memory/schedule.ts +6 -4
  80. package/src/memory/stages.ts +774 -0
  81. package/src/memory/synthesize.ts +1629 -20
  82. package/src/memory/tools.ts +446 -96
  83. package/src/memory/vault.ts +462 -35
  84. package/src/memory/wikify.ts +1057 -16
  85. package/src/mlx/ffi.ts +3 -0
  86. package/src/mlx/ops.ts +46 -0
  87. package/src/model/compiled-decode.ts +43 -8
  88. package/src/model/factory.ts +29 -8
  89. package/src/model/flash-attention.ts +4 -2
  90. package/src/model/fused-decode-kernel.ts +14 -7
  91. package/src/model/fused-mlp-kernel.ts +169 -0
  92. package/src/model/fused-swiglu-kernel.ts +240 -0
  93. package/src/model/gemma4-base.ts +62 -6
  94. package/src/model/gemma4.ts +33 -2
  95. package/src/model/minicpm5.ts +106 -8
  96. package/src/model/qwen3.ts +276 -0
  97. package/src/model/steel-linear-kernel.ts +96 -0
  98. package/src/model/support.ts +38 -8
  99. package/src/model/universal/archs.ts +332 -0
  100. package/src/model/universal/dense.ts +576 -0
  101. package/src/model/universal/modules.ts +163 -0
  102. package/src/model/universal/rope.ts +366 -0
  103. package/src/pi-provider.ts +1 -1
  104. package/src/pi-terminal.ts +29 -8
  105. package/src/pi-web.ts +19 -2
  106. package/src/quantize/config-writer.ts +1 -1
  107. package/src/registry.ts +215 -7
  108. package/src/responses.ts +23 -0
  109. package/src/sampler.ts +220 -9
  110. package/src/serve/batch-scheduler.ts +272 -80
  111. package/src/serve/generation-gateway.ts +86 -18
  112. package/src/server.ts +612 -49
  113. package/src/spec/drafter.ts +81 -49
  114. package/src/spec/dspark/data-dflash.ts +156 -0
  115. package/src/spec/dspark/data.ts +201 -0
  116. package/src/spec/dspark/generate-dflash.ts +193 -0
  117. package/src/spec/dspark/generate.ts +290 -0
  118. package/src/spec/dspark/loss.ts +213 -0
  119. package/src/spec/dspark/module-dflash.ts +361 -0
  120. package/src/spec/dspark/module.ts +540 -0
  121. package/src/spec/dspark/sample.ts +112 -0
  122. package/src/spec/generate.ts +40 -19
  123. package/src/status-page.html +1 -1
  124. package/src/tokenizer.ts +7 -0
  125. package/src/train/flash-cce.ts +70 -15
  126. package/src/train/fuse.ts +243 -0
  127. package/src/train/index.ts +3 -0
  128. package/src/train/job.ts +8 -0
  129. package/src/train/loss.ts +289 -30
  130. package/src/train/prefix-shared.ts +33 -5
  131. package/src/train/segmented.ts +67 -11
  132. package/src/train/trainer.ts +92 -28
  133. package/src/web/skills/memory/SKILL.md +40 -17
package/src/server.ts CHANGED
@@ -26,8 +26,8 @@ import { Gemma4Model } from "./model/gemma4";
26
26
  import { DiffusionGemmaModel } from "./model/diffusion-gemma";
27
27
  import { spliceImageTokens } from "./vision/diffusion-vision";
28
28
  import { createModel, type RuntimeModel } from "./model/factory";
29
- import { isMiniCPM5Config, isQwen35Config, isSupportedModelRecord } from "./model/support";
30
- import { generate, type GenerateOptions } from "./generate";
29
+ import { isMiniCPM5Config, isSupportedModelRecord } from "./model/support";
30
+ import { generate, type GenerateOptions, type TokenLogprobs } from "./generate";
31
31
  import type { HlgConfig } from "./sampler";
32
32
  import { isMonotone, CURVE_UMIN, type CurveParams } from "./curve-sampler";
33
33
  const CURVE_PAGE = curveDesignerHtml as unknown as string;
@@ -51,6 +51,7 @@ import {
51
51
  type ResponsesRequest,
52
52
  } from "./responses";
53
53
  import { AdapterManager, listAvailableAdapters } from "./lora";
54
+ import { embedMany, isEmbeddingModel } from "./embed";
54
55
  import { fit } from "./fit";
55
56
  import { setMemoryLimit } from "./mlx/ffi";
56
57
  import { VisionTower } from "./vision/embedder";
@@ -85,8 +86,10 @@ export interface ServerOptions {
85
86
  * it). Exposed at /stats.server so other mlx-bun processes can
86
87
  * warn before attaching to a server that may vanish. */
87
88
  owner?: "serve" | "pi-session" | "embedded";
88
- /** Interface to bind (Bun.serve hostname). Default: Bun's default
89
- * (all interfaces); pass "127.0.0.1" for loopback-only. */
89
+ /** Interface to bind (Bun.serve hostname). Unset ⇒ Bun's default
90
+ * (all interfaces) — embedded/library use. The CLI always passes one:
91
+ * "127.0.0.1" (loopback, mlx_lm.server parity) unless --host says
92
+ * otherwise. */
90
93
  hostname?: string;
91
94
  /** Server-wide default for the chat template's `enable_thinking`
92
95
  * variable (MiniCPM5/CPM and other hybrid-reasoning models). A
@@ -101,6 +104,10 @@ export interface ServerOptions {
101
104
  defaultTemperature?: number;
102
105
  defaultTopP?: number;
103
106
  defaultTopK?: number;
107
+ /** Completion cap when the request omits max_tokens (`--max-tokens`).
108
+ * mlx_lm.server's flag; its default there is 512 — ours stays 65,536 so
109
+ * thinking traces never truncate. `--max-tokens 512` = mlx-lm behavior. */
110
+ defaultMaxTokens?: number;
104
111
  /** Max concurrent requests batched through the mlx-lm-parity engine
105
112
  * (`--batch N`). Default 1 = today's serialized single-queue path. >1 opts
106
113
  * the WHOLE server into the batched engine (continuous batching, B floats
@@ -111,6 +118,11 @@ export interface ServerOptions {
111
118
  /** HLG tone-curve sampling default (set via --hlg-sampling on + sub-knobs).
112
119
  * A per-request `hlg` object overrides it field-by-field. Off when unset. */
113
120
  hlg?: HlgConfig;
121
+ /** Adapter id mounted at startup via `serve --adapter <dir>` (mlx-lm's
122
+ * `--adapter-path`). Used when a request sends no `adapter` field; a
123
+ * request's explicit `adapter` (including "none") always wins. Hot-swap
124
+ * via /v1/adapters is unchanged. */
125
+ defaultAdapter?: string;
114
126
  }
115
127
 
116
128
  export interface ServerContext {
@@ -262,6 +274,34 @@ interface ChatRequest {
262
274
  top_k?: number;
263
275
  seed?: number;
264
276
  repetition_penalty?: number;
277
+ /** mlx-lm extension: recent-token window for repetition_penalty
278
+ * (default 20; 0 = whole history, Python `[-0:]` semantics). */
279
+ repetition_context_size?: number;
280
+ /** min-p sampling (mlx_lm.server's `min_p`): keep tokens whose probability
281
+ * is ≥ min_p · p(top token). 0 = off. */
282
+ min_p?: number;
283
+ /** XTC sampling (mlx_lm.server names): with probability `xtc_probability`
284
+ * per step, remove every token above `xtc_threshold` except the least
285
+ * likely of them. EOS + the newline token are always exempt (the server
286
+ * injects them as xtc special tokens, matching mlx_lm.server). */
287
+ xtc_probability?: number;
288
+ xtc_threshold?: number;
289
+ /** OpenAI logit_bias: {tokenId: additive bias}. JSON object keys arrive as
290
+ * strings; coerced to int keys / float values like mlx-lm (400 on failure). */
291
+ logit_bias?: Record<string, number>;
292
+ /** OpenAI presence/frequency penalties + mlx-lm's context-size extensions
293
+ * (window of recent tokens the penalty looks at; default 20). */
294
+ presence_penalty?: number;
295
+ presence_context_size?: number;
296
+ frequency_penalty?: number;
297
+ frequency_context_size?: number;
298
+ /** mlx_lm.server logprobs: `logprobs` is a BOOL (even on /v1/completions —
299
+ * not OpenAI's legacy int), `top_logprobs` an int in [0, 11] or the -1
300
+ * "unset" sentinel (server.py validates exactly that; OpenAI's cap is 20,
301
+ * mlx-lm's is 11 — we copy the reference). Non-stream responses carry
302
+ * mlx-lm's logprobs block; stream chunks never do (reference behavior). */
303
+ logprobs?: boolean;
304
+ top_logprobs?: number;
265
305
  /** OpenAI stop sequences: plain string or array (spec allows up to 4).
266
306
  * Matched on DECODED text, not token ids — see StopMatcher. */
267
307
  stop?: string | string[];
@@ -290,6 +330,15 @@ interface ChatRequest {
290
330
  };
291
331
  }
292
332
 
333
+ /** POST /v1/completions body (mlx_lm.server's raw text completion — no chat
334
+ * template). Sampling/penalty/stop fields are the same names as ChatRequest;
335
+ * `prompt` replaces `messages`. mlx_lm.server accepts only a STRING prompt
336
+ * (it calls `tokenizer.encode(request.prompt)` directly) — token-array
337
+ * prompts are rejected there too, so we match. No `echo` (mlx-lm has none). */
338
+ type TextCompletionRequest = Omit<ChatRequest, "messages" | "tools" | "tool_choice"> & {
339
+ prompt?: unknown;
340
+ };
341
+
293
342
  /** Per-field default HLG knobs when enabling without specifying them. */
294
343
  const HLG_DEFAULTS = { width: 4, shoulder: 4, toe: 6, pivotOffset: 6 } as const;
295
344
 
@@ -319,7 +368,23 @@ interface OpenAIToolCall {
319
368
  function: { name: string; arguments: string };
320
369
  }
321
370
 
322
- type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
371
+ export type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
372
+
373
+ /** Pick the stream router for a model family. The token-id sentinel router is
374
+ * Gemma-4-ONLY: ids 48/49 (<|tool_call>/<tool_call|>) and 100/101
375
+ * (<|channel>/<channel|>) are special tokens of that tokenizer family
376
+ * (src/tool-call.ts). On every other tokenizer — MiniCPM5, Qwen3/3.5, and
377
+ * the Tier-0 generics (llama, qwen2, phi3, …) — those ids are ordinary
378
+ * low-id vocab, so the sentinel router would silently swallow output into a
379
+ * phantom tool/reasoning segment. Everyone else parses tool calls from
380
+ * decoded text (buffered-text; parseGeneratedToolCalls covers the
381
+ * OpenAI-JSON <tool_call>, Qwen <function=…>, and MiniCPM5 <function name=…>
382
+ * shapes) when tools are present, and streams plain otherwise. Models whose
383
+ * markup isn't covered fail soft: the markup stays in content. */
384
+ export function selectToolStreamMode(modelType: string, hasTools: boolean): ToolStreamMode {
385
+ if (modelType.startsWith("gemma4")) return "gemma-sentinel";
386
+ return hasTools ? "buffered-text" : "plain";
387
+ }
323
388
 
324
389
  /** Routes generated tokens. Gemma uses family-specific sentinel token ids;
325
390
  * MiniCPM5 and other text-template models use decoded-text parsing so
@@ -729,6 +794,129 @@ export function normalizeMessages(messages: ChatMessage[]): ChatMessage[] {
729
794
  });
730
795
  }
731
796
 
797
+ /** Coerce a wire `logit_bias` ({"tokenId": bias}) to numeric keys/values —
798
+ * mlx_lm.server's `{int(k): float(v)}` coercion; throws its exact error
799
+ * message on anything non-numeric (surfaced as a 400). JSON object keys
800
+ * always arrive as strings, hence the coercion. */
801
+ export function parseLogitBias(
802
+ raw: Record<string, number> | undefined | null,
803
+ ): Record<number, number> | undefined {
804
+ if (raw == null) return undefined;
805
+ if (typeof raw !== "object" || Array.isArray(raw))
806
+ throw new Error("logit_bias must be a dict of int to float");
807
+ const out: Record<number, number> = {};
808
+ for (const [k, v] of Object.entries(raw)) {
809
+ const id = Number(k);
810
+ const bias = Number(v);
811
+ if (!Number.isInteger(id) || !Number.isFinite(bias))
812
+ throw new Error("logit_bias must be a dict of int to float");
813
+ out[id] = bias;
814
+ }
815
+ return Object.keys(out).length ? out : undefined;
816
+ }
817
+
818
+ /** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
819
+ * request-unique: under `--batch N` the batch lane serves ONLY default-seed
820
+ * requests (explicit-seed requests route serial, see GenerationGateway.willBatch),
821
+ * so two identical prompts arriving in the same millisecond would share a seed
822
+ * and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
823
+ * byte-identical completions, silently collapsing best-of-N diversity. Mix a
824
+ * per-process Weyl counter (golden-ratio increment, period 2^32) into the
825
+ * timestamp so every call yields a distinct uint32 within any given ms.
826
+ * Determinism contract unchanged: reproducibility is only promised for an
827
+ * EXPLICIT request seed (`req.seed ?? nextDefaultSeed()` — explicit wins,
828
+ * byte-identical to before); a default seed is fresh entropy per request. */
829
+ let seedWeyl = 0;
830
+ export function nextDefaultSeed(): number {
831
+ seedWeyl = (seedWeyl + 0x9e3779b9) >>> 0;
832
+ return ((Date.now() & 0xffffffff) ^ seedWeyl) >>> 0;
833
+ }
834
+
835
+ /** mlx_lm.server's logprobs request validation, copied exactly (server.py
836
+ * APIHandler.validate_model_parameters: `_validate("logprobs", bool)` and
837
+ * `_validate("top_logprobs", int, min_val=0, max_val=11, whitelist=[-1])`
838
+ * with defaults logprobs=False / top_logprobs=-1). Returns the reference's
839
+ * error message (→ 400) or null when valid. Note mlx-lm caps top_logprobs
840
+ * at 11, not OpenAI's 20 — we mirror the reference. */
841
+ export function validateLogprobsParams(body: {
842
+ logprobs?: unknown;
843
+ top_logprobs?: unknown;
844
+ }): string | null {
845
+ const lp = body.logprobs ?? false;
846
+ if (typeof lp !== "boolean") return "logprobs must be of type bool";
847
+ const tl = body.top_logprobs ?? -1;
848
+ if (typeof tl !== "number" || !Number.isInteger(tl))
849
+ return "top_logprobs must be of type int";
850
+ if (tl === -1) return null; // the "unset" whitelist sentinel
851
+ if (tl < 0) return "top_logprobs must be at least 0";
852
+ if (tl > 11) return "top_logprobs must be at most 11";
853
+ return null;
854
+ }
855
+
856
+ /** Collects per-token logprob info and shapes mlx_lm.server's response block
857
+ * (server.py generate_response L1317-1327). NOT OpenAI's shape — entries
858
+ * carry token *ids* (and raw vocab token strings in the top-k form), and the
859
+ * SAME block is attached under choices[0].logprobs for chat AND text
860
+ * completions:
861
+ * - top_logprobs > 0 → {content: [{id, token, logprob,
862
+ * top_logprobs: [{id, token, logprob}, …]}, …]} — each entry is the
863
+ * top-1 candidate merged with its own top-k list (`dict(i[0],
864
+ * top_logprobs=i)`); mlx-lm leaves argpartition order unspecified, we
865
+ * sort descending, so the entry is deterministically the argmax.
866
+ * - logprobs=true (and top_logprobs ≤ 0) → {content: [{id, logprob}, …]}
867
+ * with the SAMPLED token's logprob.
868
+ * When both are set, only the top_logprobs form is emitted (the reference's
869
+ * if/elif). Stream chunks never carry logprobs — mlx-lm's streaming
870
+ * generate_response calls pass no token_logprobs/top_tokens. */
871
+ class LogprobsCollector {
872
+ readonly #tokens: number[] = [];
873
+ readonly #tokenLogprobs: number[] = [];
874
+ readonly #topTokens: { id: number; token: string; logprob: number }[][] = [];
875
+
876
+ constructor(
877
+ private readonly wantLogprobs: boolean,
878
+ private readonly topK: number,
879
+ private readonly idToToken: (id: number) => string,
880
+ ) {}
881
+
882
+ get active(): boolean {
883
+ return this.wantLogprobs || this.topK > 0;
884
+ }
885
+
886
+ push(token: number, info?: TokenLogprobs): void {
887
+ if (!this.active) return;
888
+ this.#tokens.push(token);
889
+ if (this.wantLogprobs) this.#tokenLogprobs.push(info?.logprob ?? NaN);
890
+ if (this.topK > 0)
891
+ this.#topTokens.push(
892
+ (info?.top ?? []).map((t) => ({
893
+ id: t.id,
894
+ token: this.idToToken(t.id),
895
+ logprob: t.logprob,
896
+ })),
897
+ );
898
+ }
899
+
900
+ /** choices[0].logprobs value, or null when nothing to attach (mirrors the
901
+ * reference: the key is omitted entirely for zero collected tokens). */
902
+ payload(): { content: Record<string, unknown>[] } | null {
903
+ if (this.#topTokens.length)
904
+ return {
905
+ content: this.#topTokens.map((t) =>
906
+ t.length ? { ...t[0]!, top_logprobs: t } : {},
907
+ ),
908
+ };
909
+ if (this.#tokenLogprobs.length)
910
+ return {
911
+ content: this.#tokens.map((id, i) => ({
912
+ id,
913
+ logprob: this.#tokenLogprobs[i]!,
914
+ })),
915
+ };
916
+ return null;
917
+ }
918
+ }
919
+
732
920
  /** Incremental detokenizer: emits the longest stable decoded prefix. */
733
921
  class StreamDecoder {
734
922
  #ids: number[] = [];
@@ -788,13 +976,6 @@ export function createServer(
788
976
  // lane (see GenerationGateway.willBatch).
789
977
  const batch = Math.max(1, Math.floor(serverOptions.batch ?? 1));
790
978
 
791
- let queue: Promise<unknown> = Promise.resolve();
792
- const enqueue = <T>(fn: () => Promise<T>): Promise<T> => {
793
- const run = queue.then(fn, fn);
794
- queue = run.catch(() => {});
795
- return run;
796
- };
797
-
798
979
  const promptCache = new PromptCache(serverOptions.promptCacheBytes ?? 2e9);
799
980
  // Responses-API store for previous_response_id resumption (Phase 11):
800
981
  // TTL + byte-capped LRU, port of optiq/response_store.py. Pairs with
@@ -802,8 +983,11 @@ export function createServer(
802
983
  // so its KV prefill is already cached.
803
984
  const responseStore = new ResponseStore();
804
985
 
805
- /** Run one generation with prompt-cache reuse. Must be called inside
806
- * the queue. onToken returning `false` halts generation early (stop
986
+ /** Run one generation with prompt-cache reuse. Must be called under the
987
+ * gateway's exclusive lock — the gateway invokes it on the serial lane;
988
+ * the curve /generate endpoint wraps it in gateway.runExclusive. (The old
989
+ * `enqueue` promise queue is gone; the gateway's AsyncMutex is THE lock.)
990
+ * onToken returning `false` halts generation early (stop
807
991
  * sequence fired); the cache snapshot stays valid and is still kept.
808
992
  * Vision requests bypass the prompt cache: image tokens are
809
993
  * identical placeholder ids, so prefix matching across different
@@ -811,7 +995,7 @@ export function createServer(
811
995
  const runGeneration = async (
812
996
  promptIds: number[],
813
997
  options: GenerateOptions,
814
- onToken: (token: number) => void | boolean | Promise<void | boolean>,
998
+ onToken: (token: number, logprobs?: TokenLogprobs) => void | boolean | Promise<void | boolean>,
815
999
  vision?: { embeddings: import("./mlx/array").MlxArray; imageMask: import("./mlx/array").MlxArray },
816
1000
  ) => {
817
1001
  // Cache entries are adapter-specific: KV computed under one adapter
@@ -826,7 +1010,7 @@ export function createServer(
826
1010
  ...(vision ? { promptEmbeddings: vision.embeddings, imageMask: vision.imageMask } : {}),
827
1011
  });
828
1012
  for await (const t of gen) {
829
- if ((await onToken(t.token)) === false) break;
1013
+ if ((await onToken(t.token, t.logprobs)) === false) break;
830
1014
  }
831
1015
  const s = gen.stats!; // set on completion AND on early break
832
1016
  if (vision) {
@@ -914,6 +1098,17 @@ export function createServer(
914
1098
  `Omit --kv-quant to batch in bf16. (docs/design/parallel-slots.md)`,
915
1099
  );
916
1100
 
1101
+ // XTC never removes EOS or the newline token: mlx_lm.server passes
1102
+ // [tokenizer.eos_token_id, tokenizer.encode("\n")] as xtc_special_tokens.
1103
+ // We pass the flat equivalent — ALL configured EOS ids plus the bare newline
1104
+ // id(s), encoded WITHOUT special tokens (mlx-lm's encode("\n") can drag a
1105
+ // BOS along on some tokenizers; protecting BOS from XTC is a no-op, so the
1106
+ // flat/no-BOS form is behaviorally identical and cleaner).
1107
+ const xtcSpecialTokens = [
1108
+ ...ctx.model.config.eosTokenIds,
1109
+ ...ctx.tokenizer.encode("\n", false),
1110
+ ];
1111
+
917
1112
  // Effective enable_thinking for a request, with the same precedence the chat
918
1113
  // template uses (extracted so sampling and template rendering can't disagree):
919
1114
  // explicit chat_template_kwargs.enable_thinking → reasoning_effort ("none" =
@@ -947,12 +1142,25 @@ export function createServer(
947
1142
  // truncates the visible reply. Default very generously (the model's
948
1143
  // context is far larger); only an explicit max_tokens narrows it. The
949
1144
  // model still stops at its eos_token well before this in normal replies.
950
- maxTokens: req.max_completion_tokens ?? req.max_tokens ?? 65_536,
1145
+ maxTokens: req.max_completion_tokens ?? req.max_tokens ?? serverOptions.defaultMaxTokens ?? 65_536,
951
1146
  temperature: req.temperature ?? serverOptions.defaultTemperature ?? defaultTemp,
952
1147
  topP: req.top_p ?? serverOptions.defaultTopP ?? ctx.genDefaults.topP ?? 0,
953
1148
  topK: req.top_k ?? serverOptions.defaultTopK ?? ctx.genDefaults.topK ?? 0,
954
- seed: req.seed ?? (Date.now() & 0xffffffff),
1149
+ seed: req.seed ?? nextDefaultSeed(),
955
1150
  repetitionPenalty: req.repetition_penalty ?? ctx.genDefaults.repetitionPenalty,
1151
+ // mlx_lm.server sampling extensions (defaults mirror server.py: min_p /
1152
+ // xtc off, context windows 20). parseLogitBias throws on non-numeric
1153
+ // input — callers surface it as a 400 (mlx-lm's coercion error).
1154
+ minP: req.min_p ?? 0,
1155
+ xtcProbability: req.xtc_probability ?? 0,
1156
+ xtcThreshold: req.xtc_threshold ?? 0,
1157
+ ...(req.xtc_probability ? { xtcSpecialTokens } : {}),
1158
+ logitBias: parseLogitBias(req.logit_bias),
1159
+ repetitionContextSize: req.repetition_context_size ?? 20,
1160
+ presencePenalty: req.presence_penalty ?? 0,
1161
+ presenceContextSize: req.presence_context_size ?? 20,
1162
+ frequencyPenalty: req.frequency_penalty ?? 0,
1163
+ frequencyContextSize: req.frequency_context_size ?? 20,
956
1164
  hlg: resolveHlg(req.hlg, serverOptions.hlg),
957
1165
  // generate() yields token ids; stop sequences match on decoded text
958
1166
  // (they can span token boundaries), so the StopMatcher sits at the
@@ -980,15 +1188,11 @@ export function createServer(
980
1188
  return { ids: trimmed, startInThinking: promptEndsInOpenThink(rendered) };
981
1189
  };
982
1190
 
983
- const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode => {
984
- // MiniCPM5 and Qwen3.5 emit tool calls as DECODED TEXT (Qwen:
985
- // <tool_call><function=name><parameter=…>; both handled by
986
- // parseGeneratedToolCalls), not Gemma's token sentinels — so they take the
987
- // buffered-text path. Gemma keeps the token-level sentinel path.
988
- if (isMiniCPM5Config(ctx.model.config) || isQwen35Config(ctx.model.config))
989
- return tools?.length ? "buffered-text" : "plain";
990
- return "gemma-sentinel";
991
- };
1191
+ const toolStreamMode = (tools: ToolDefinition[] | null): ToolStreamMode =>
1192
+ // Gemma-4 keeps the token-level sentinel path; every other family
1193
+ // (MiniCPM5, Qwen3/3.5, Tier-0 generics) emits tool calls as DECODED TEXT
1194
+ // — see selectToolStreamMode for why sentinel ids must stay family-gated.
1195
+ selectToolStreamMode(ctx.model.config.modelType, !!tools?.length);
992
1196
 
993
1197
  const toolRouter = (tools: ToolDefinition[] | null): ToolAwareStream =>
994
1198
  new ToolAwareStream(ctx.tokenizer, toolStreamMode(tools), tools);
@@ -1067,9 +1271,14 @@ export function createServer(
1067
1271
  // surface on their own. (scan() is INSERT-OR-REPLACE + prunes deleted,
1068
1272
  // so it's idempotent and cheap for a local cache.)
1069
1273
  await reg.scan();
1274
+ const { visionCapable } = await import("./registry");
1275
+ const { supportTier } = await import("./model/support");
1070
1276
  const rows = [];
1071
- for (const m of reg.list()) {
1072
- const supported = isSupportedModelRecord(m.modelType, m.repoId);
1277
+ // listCanonical: one row per repo (refs/main) — duplicate snapshots
1278
+ // from upstream re-pushes stay visible only in `ls --all-revisions`.
1279
+ for (const m of reg.listCanonical()) {
1280
+ const tier = supportTier(m.modelType, m.repoId);
1281
+ const supported = tier !== null;
1073
1282
  let assessment = null;
1074
1283
  try {
1075
1284
  const config = await loadModelConfig(m.path);
@@ -1083,7 +1292,7 @@ export function createServer(
1083
1292
  rows.push({
1084
1293
  repo_id: m.repoId, model_type: m.modelType,
1085
1294
  size_bytes: m.sizeBytes, quant_bits: m.quantBits,
1086
- vision: m.hasVisionSidecar, supported,
1295
+ vision: visionCapable(m), supported, support_tier: tier,
1087
1296
  serving: m.repoId === ctx.modelId,
1088
1297
  assessment,
1089
1298
  });
@@ -1098,6 +1307,39 @@ export function createServer(
1098
1307
  return Response.json({ downloads: downloadsSnapshot() });
1099
1308
  }
1100
1309
 
1310
+ // Memory synthesis progress (P8-T5). The same DAG the nightly launchd job
1311
+ // runs (`mlx-bun memory synthesize`), streamed as Server-Sent Events so the
1312
+ // status page can show live stage/log/done progress. `?dry=1` plans the DAG
1313
+ // without any model call or vault write — the safe wiring-verification path
1314
+ // (the FULL-corpus run is USER-ACTION, P6-T5). GET so EventSource can drive it.
1315
+ if (url.pathname === "/v1/memory/synthesize" && request.method === "GET") {
1316
+ const dryRun = url.searchParams.get("dry") === "1";
1317
+ const { runSynthesis } = await import("./memory/pipeline");
1318
+ const enc = new TextEncoder();
1319
+ const stream = new ReadableStream({
1320
+ async start(controller) {
1321
+ const send = (e: unknown) =>
1322
+ controller.enqueue(enc.encode(`data: ${JSON.stringify(e)}\n\n`));
1323
+ try {
1324
+ const summary = await runSynthesis({ dryRun }, (ev) => send(ev));
1325
+ send({ type: "summary", ...summary });
1326
+ controller.enqueue(enc.encode("data: [DONE]\n\n"));
1327
+ } catch (e) {
1328
+ send({ type: "error", message: (e as Error).message });
1329
+ } finally {
1330
+ controller.close();
1331
+ }
1332
+ },
1333
+ });
1334
+ return new Response(stream, {
1335
+ headers: {
1336
+ "content-type": "text/event-stream",
1337
+ "cache-control": "no-cache",
1338
+ connection: "keep-alive",
1339
+ },
1340
+ });
1341
+ }
1342
+
1101
1343
  if (url.pathname === "/fit" && request.method === "GET") {
1102
1344
  // Fit assessment for the status page: this-machine report at the
1103
1345
  // admission ceiling + the Apple SKU matrix at a fixed 32k.
@@ -1160,8 +1402,9 @@ export function createServer(
1160
1402
  return Response.json({
1161
1403
  name: "mlx-bun", version: pkgVersion, model: ctx.modelId,
1162
1404
  endpoints: [
1163
- "POST /v1/chat/completions", "POST /v1/messages", "POST /v1/responses",
1164
- "GET /v1/models", "GET/POST/DELETE /v1/adapters",
1405
+ "POST /v1/chat/completions", "POST /v1/completions", "POST /v1/messages",
1406
+ "POST /v1/responses", "POST /v1/embeddings", "GET /v1/models",
1407
+ "GET/POST/DELETE /v1/adapters", "GET /health",
1165
1408
  "GET /stats", "GET /fit", "GET /library", "GET /downloads",
1166
1409
  ],
1167
1410
  });
@@ -1229,16 +1472,95 @@ export function createServer(
1229
1472
  });
1230
1473
  }
1231
1474
 
1232
- if (url.pathname === "/v1/models" && request.method === "GET") {
1475
+ // mlx_lm.server parity: GET /health → the exact body it writes
1476
+ // ('{"status": "ok"}', note the space) so byte-for-byte health checks pass.
1477
+ if (url.pathname === "/health" && request.method === "GET") {
1478
+ return new Response('{"status": "ok"}', {
1479
+ headers: { "content-type": "application/json" },
1480
+ });
1481
+ }
1482
+
1483
+ // GET /v1/models — the served model FIRST (with our capability extras),
1484
+ // then every other servable model the registry knows (mlx-lm scans the
1485
+ // HF cache here; our registry is that scan, filtered to supported
1486
+ // architectures). GET /v1/models/<id> filters to that id — same list
1487
+ // shape, matching mlx_lm.server's handle_models_request.
1488
+ if (
1489
+ (url.pathname === "/v1/models" || url.pathname.startsWith("/v1/models/")) &&
1490
+ request.method === "GET"
1491
+ ) {
1492
+ const filterId = url.pathname.length > "/v1/models/".length - 1
1493
+ ? decodeURIComponent(url.pathname.slice("/v1/models/".length))
1494
+ : null;
1495
+ const created = Math.floor(startedAt / 1000);
1496
+ const data: Array<Record<string, unknown>> = [{
1497
+ id: ctx.modelId, object: "model", created, owned_by: "mlx-bun",
1498
+ context_window: ctx.model.config.text.maxPositionEmbeddings,
1499
+ // Capability flags for clients (CLI/external pi) that build a
1500
+ // provider from discovery — `reasoning` gates the thinking toggle,
1501
+ // `vision` the image input declaration (tower loaded or lazily
1502
+ // loadable; same signal the web embed uses).
1503
+ reasoning: ctx.template.supportsThinking,
1504
+ vision: !!(ctx.vision || ctx.loadVision),
1505
+ }];
1506
+ try {
1507
+ const { Registry } = await import("./registry");
1508
+ const reg = new Registry();
1509
+ try {
1510
+ if (reg.list().length === 0) await reg.scan();
1511
+ // Canonical rows only — a repo with N snapshots is ONE model id.
1512
+ for (const m of reg.listCanonical()) {
1513
+ if (m.repoId === ctx.modelId) continue;
1514
+ if (!isSupportedModelRecord(m.modelType, m.repoId)) continue;
1515
+ data.push({ id: m.repoId, object: "model", created });
1516
+ }
1517
+ } finally {
1518
+ reg.close();
1519
+ }
1520
+ } catch { /* registry unavailable → served model only */ }
1521
+ return Response.json({
1522
+ object: "list",
1523
+ data: filterId ? data.filter((m) => m.id === filterId) : data,
1524
+ });
1525
+ }
1526
+
1527
+ // OpenAI embeddings API. Works when the SERVED model is an embedding model
1528
+ // (plain Qwen3 / Qwen3-Embedding) — consistent with the single-model server
1529
+ // design: `mlx-bun serve <embedding-model>` to use this. Optional non-standard
1530
+ // `instruction` applies Qwen3-Embedding's query format. Embedding is a pure
1531
+ // forward (no decode loop / gateway), so it runs inline.
1532
+ if (url.pathname === "/v1/embeddings" && request.method === "POST") {
1533
+ if (!isEmbeddingModel(ctx.model))
1534
+ return Response.json({
1535
+ error: {
1536
+ message: `served model "${ctx.modelId}" is not an embedding model; ` +
1537
+ `serve an embedding model (e.g. Qwen3-Embedding) to use /v1/embeddings`,
1538
+ type: "invalid_request_error",
1539
+ },
1540
+ }, { status: 400 });
1541
+ let body: { input?: string | string[]; instruction?: string };
1542
+ try {
1543
+ body = (await request.json()) as typeof body;
1544
+ } catch {
1545
+ return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
1546
+ }
1547
+ const inputs = Array.isArray(body.input) ? body.input : body.input != null ? [body.input] : [];
1548
+ if (inputs.length === 0 || !inputs.every((s) => typeof s === "string"))
1549
+ return Response.json({
1550
+ error: { message: "`input` must be a string or array of strings", type: "invalid_request_error" },
1551
+ }, { status: 400 });
1552
+ const instruction = typeof body.instruction === "string" ? body.instruction : undefined;
1553
+ const results = embedMany(ctx.model, ctx.tokenizer, inputs, instruction);
1554
+ let totalTokens = 0;
1555
+ const data = results.map((r, index) => {
1556
+ totalTokens += r.tokens;
1557
+ return { object: "embedding", index, embedding: Array.from(r.vector) };
1558
+ });
1233
1559
  return Response.json({
1234
1560
  object: "list",
1235
- data: [{
1236
- id: ctx.modelId, object: "model", created: 0, owned_by: "mlx-bun",
1237
- context_window: ctx.model.config.text.maxPositionEmbeddings,
1238
- // Capability flags for clients (CLI/external pi) that build a
1239
- // provider from discovery — `reasoning` gates the thinking toggle.
1240
- reasoning: ctx.template.supportsThinking,
1241
- }],
1561
+ data,
1562
+ model: ctx.modelId,
1563
+ usage: { prompt_tokens: totalTokens, total_tokens: totalTokens },
1242
1564
  });
1243
1565
  }
1244
1566
 
@@ -1287,7 +1609,10 @@ export function createServer(
1287
1609
  if (!body.id || !body.path)
1288
1610
  return Response.json({ error: { message: "id and path required" } }, { status: 400 });
1289
1611
  try {
1290
- const info = await enqueue(() => ctx.adapters.mount(body.id!, body.path!));
1612
+ // Under the gateway lock: mount/unmount mutate the shared adapter
1613
+ // registry (unmount disposes arrays a running generation could still
1614
+ // hold) — one mutual-exclusion domain with generation (D3).
1615
+ const info = await gateway.runExclusive(() => ctx.adapters.mount(body.id!, body.path!));
1291
1616
  return Response.json({
1292
1617
  id: info.id, mounted_layers: info.mountedLayers,
1293
1618
  rank: info.rank, scale: info.scale,
@@ -1298,7 +1623,7 @@ export function createServer(
1298
1623
  }
1299
1624
  if (url.pathname.startsWith("/v1/adapters/") && request.method === "DELETE") {
1300
1625
  const id = decodeURIComponent(url.pathname.slice("/v1/adapters/".length));
1301
- const removed = await enqueue(async () => ctx.adapters.unmount(id));
1626
+ const removed = await gateway.runExclusive(async () => ctx.adapters.unmount(id));
1302
1627
  return removed > 0
1303
1628
  ? Response.json({ id, removed_layers: removed })
1304
1629
  : Response.json({ error: { message: `adapter ${id} not mounted` } }, { status: 404 });
@@ -1317,7 +1642,10 @@ export function createServer(
1317
1642
  if (sids[0] === sids[1] && sids[0] === ctx.tokenizer.bosTokenId) sids = sids.slice(1);
1318
1643
  try {
1319
1644
  const NB = 80;
1320
- const result = await enqueue(async () => {
1645
+ // gateway.runExclusive: this raw forward must not run concurrently
1646
+ // with batched decode steps or a serial generation (GPU + shared
1647
+ // model state are single-owner — D3, one lock).
1648
+ const result = await gateway.runExclusive(async () => {
1321
1649
  const cache = ctx.model.makeCache();
1322
1650
  try {
1323
1651
  const logits = ctx.model.forward(sids, cache); // [1, L, V]
@@ -1372,7 +1700,9 @@ export function createServer(
1372
1700
  const genOpts: GenerateOptions = useCurve
1373
1701
  ? { curve, seed: baseSeed + i, maxTokens, ...kvScheme }
1374
1702
  : { temperature: recipe.temperature, topP: recipe.topP, topK: recipe.topK, seed: baseSeed + i, maxTokens, ...kvScheme };
1375
- await enqueue(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
1703
+ // runExclusive: runGeneration touches the GPU + prompt cache, so
1704
+ // it needs the gateway lock (this endpoint bypasses gateway.run).
1705
+ await gateway.runExclusive(() => runGeneration(ids, genOpts, (t) => { toks.push(t); }));
1376
1706
  const text = ctx.tokenizer.decode(toks, true).trim();
1377
1707
  samples.push({ text, junk: curveJunk(text) });
1378
1708
  }
@@ -1390,6 +1720,10 @@ export function createServer(
1390
1720
  const handleChat = async (body: ChatRequest): Promise<Response> => {
1391
1721
  if (!Array.isArray(body.messages) || body.messages.length === 0)
1392
1722
  return Response.json({ error: { message: "messages required" } }, { status: 400 });
1723
+ // mlx-lm validates logprobs params up front (ValueError → 400)
1724
+ const lpParamError = validateLogprobsParams(body);
1725
+ if (lpParamError)
1726
+ return Response.json({ error: { message: lpParamError } }, { status: 400 });
1393
1727
 
1394
1728
  const id = `chatcmpl-${crypto.randomUUID()}`;
1395
1729
  const created = Math.floor(Date.now() / 1000);
@@ -1459,7 +1793,16 @@ export function createServer(
1459
1793
  { status: 400 },
1460
1794
  );
1461
1795
  }
1462
- const options = toOptions(body);
1796
+ let options: ReturnType<typeof toOptions>;
1797
+ try {
1798
+ options = toOptions(body);
1799
+ } catch (e) {
1800
+ // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
1801
+ vision?.embeddings.dispose();
1802
+ vision?.imageMask.dispose();
1803
+ diffusionPixels?.dispose();
1804
+ return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1805
+ }
1463
1806
  if (diffusionPixels) options.visionPixels = diffusionPixels;
1464
1807
  // Admission: reject what cannot finish within the memory budget
1465
1808
  // (the GPU OOM it would otherwise hit is uncatchable and kills
@@ -1485,18 +1828,41 @@ export function createServer(
1485
1828
  );
1486
1829
  }
1487
1830
  try {
1488
- const adapterIds = ctx.adapters.resolveSpec(body.adapter);
1831
+ // A request's explicit `adapter` (incl. "none") wins over the
1832
+ // startup default from `serve --adapter <dir>`.
1833
+ const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
1489
1834
  if (adapterIds.length) options.adapters = adapterIds;
1490
1835
  } catch (e) {
1491
1836
  return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
1492
1837
  }
1493
1838
 
1494
- // What lane this request takes (vision / adapters / repetition penalty /
1839
+ // logprobs capture: non-stream only — stream chunks never carry
1840
+ // logprobs (mirroring mlx_lm.server, whose streaming generate_response
1841
+ // calls pass no token_logprobs/top_tokens), so streaming requests skip
1842
+ // the capture cost entirely and stay batchable.
1843
+ const wantLogprobs = body.logprobs === true;
1844
+ const topLogprobs =
1845
+ typeof body.top_logprobs === "number" && body.top_logprobs > 0
1846
+ ? body.top_logprobs : 0;
1847
+ const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
1848
+ if (captureLogprobs) {
1849
+ options.logprobs = wantLogprobs;
1850
+ options.topLogprobs = topLogprobs;
1851
+ }
1852
+
1853
+ // What lane this request takes (vision / adapters / logits processors /
1495
1854
  // a user-fixed seed → serial; everything else batches when --batch N).
1496
1855
  const shape = {
1497
1856
  hasVision: !!vision,
1498
1857
  hasAdapters: !!options.adapters?.length,
1499
1858
  hasRepetitionPenalty: !!options.repetitionPenalty,
1859
+ // min_p / XTC / logit_bias / presence+frequency penalties: serial-only
1860
+ // in v1 (see GenerationGateway.RequestShape.hasLogitsExtras).
1861
+ hasLogitsExtras: !!(
1862
+ options.minP || options.xtcProbability || options.logitBias ||
1863
+ options.presencePenalty || options.frequencyPenalty
1864
+ ),
1865
+ wantsLogprobs: captureLogprobs,
1500
1866
  userSeed: body.seed !== undefined,
1501
1867
  kvQuant: !!(options.kvConfig?.length || options.kvBits),
1502
1868
  };
@@ -1611,9 +1977,15 @@ export function createServer(
1611
1977
  const router = toolRouter(tools);
1612
1978
  const stopper = new StopMatcher(options.stopSequences);
1613
1979
  const thinking = new ThinkingTagSplitter(ctx.template.thinkingFormat === "think-tag", startInThinking);
1980
+ // mlx-lm collects logprobs across EVERY generated token (reasoning
1981
+ // and tool tokens included), not just visible content — same here.
1982
+ const lpc = captureLogprobs
1983
+ ? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
1984
+ : null;
1614
1985
  let content = "";
1615
1986
  let reasoning = "";
1616
- const s = await gateway.run(promptIds, options, (token) => {
1987
+ const s = await gateway.run(promptIds, options, (token, lpInfo) => {
1988
+ lpc?.push(token, lpInfo);
1617
1989
  const rawContent = router.push(token);
1618
1990
  reasoning += router.takeReasoning(); // gemma-channel thinking
1619
1991
  const parts = thinking.push(stopper.push(rawContent));
@@ -1640,6 +2012,7 @@ export function createServer(
1640
2012
  ? "tool_calls"
1641
2013
  : stopper.stopped ? "stop"
1642
2014
  : s.generatedTokens >= (options.maxTokens ?? 1024) ? "length" : "stop";
2015
+ const logprobsBlock = lpc?.payload() ?? null;
1643
2016
  return Response.json({
1644
2017
  id, object: "chat.completion", created, model: ctx.modelId,
1645
2018
  choices: [{
@@ -1650,6 +2023,7 @@ export function createServer(
1650
2023
  ...(reasoning ? { reasoning } : {}),
1651
2024
  ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
1652
2025
  },
2026
+ ...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
1653
2027
  finish_reason: finish,
1654
2028
  }],
1655
2029
  usage: {
@@ -1675,6 +2049,195 @@ export function createServer(
1675
2049
  return handleChat(body);
1676
2050
  }
1677
2051
 
2052
+ // Raw text completion (mlx_lm.server's /v1/completions, request_type
2053
+ // "text"): NO chat template — the prompt string is tokenized directly
2054
+ // (tokenizer.encode with the tokenizer's own special-token handling,
2055
+ // exactly mlx-lm's `tokenizer.encode(request.prompt)`). Rides the same
2056
+ // GenerationGateway + admission + adapter path as chat; no tool router
2057
+ // or thinking splitter (raw text in, raw text out).
2058
+ if (url.pathname === "/v1/completions" && request.method === "POST") {
2059
+ let body: TextCompletionRequest;
2060
+ try {
2061
+ body = (await request.json()) as TextCompletionRequest;
2062
+ } catch {
2063
+ return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
2064
+ }
2065
+ if (typeof body.prompt !== "string" || body.prompt.length === 0)
2066
+ return Response.json(
2067
+ {
2068
+ error: {
2069
+ message: "prompt (a non-empty string) is required " +
2070
+ "(token-array prompts are not accepted, matching mlx_lm.server)",
2071
+ },
2072
+ },
2073
+ { status: 400 },
2074
+ );
2075
+ // mlx-lm validates logprobs params up front (ValueError → 400)
2076
+ const lpParamError = validateLogprobsParams(body);
2077
+ if (lpParamError)
2078
+ return Response.json({ error: { message: lpParamError } }, { status: 400 });
2079
+ const id = `cmpl-${crypto.randomUUID()}`;
2080
+ const created = Math.floor(Date.now() / 1000);
2081
+ const promptIds = ctx.tokenizer.encode(body.prompt);
2082
+ let options: ReturnType<typeof toOptions>;
2083
+ try {
2084
+ options = toOptions(body as unknown as ChatRequest);
2085
+ } catch (e) {
2086
+ // bad logit_bias (non-numeric keys/values) — mlx-lm's coercion error
2087
+ return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2088
+ }
2089
+ // mlx_lm.server's default max_tokens is 512 (its --max-tokens CLI
2090
+ // default). The chat lane's very generous default is wrong for raw
2091
+ // completion: with no template an EOS may never come.
2092
+ options.maxTokens = body.max_completion_tokens ?? body.max_tokens ?? serverOptions.defaultMaxTokens ?? 512;
2093
+ const requiredCtx = promptIds.length + options.maxTokens;
2094
+ if (requiredCtx > admission.maxSafeContext)
2095
+ return Response.json(
2096
+ {
2097
+ error: {
2098
+ message:
2099
+ `request needs ${requiredCtx} tokens of context ` +
2100
+ `(prompt ${promptIds.length} + max_tokens ${options.maxTokens}) but the ` +
2101
+ `memory budget caps safe context at ${admission.maxSafeContext} — ` +
2102
+ `shorten the prompt or lower max_tokens`,
2103
+ type: "memory_admission",
2104
+ code: "context_over_budget",
2105
+ },
2106
+ },
2107
+ { status: 400 },
2108
+ );
2109
+ try {
2110
+ const adapterIds = ctx.adapters.resolveSpec(body.adapter ?? serverOptions.defaultAdapter);
2111
+ if (adapterIds.length) options.adapters = adapterIds;
2112
+ } catch (e) {
2113
+ return Response.json({ error: { message: (e as Error).message } }, { status: 400 });
2114
+ }
2115
+ // logprobs capture — same mlx-lm block as chat (generate_response is
2116
+ // shared in the reference too); non-stream only, stream chunks never
2117
+ // carry logprobs.
2118
+ const wantLogprobs = body.logprobs === true;
2119
+ const topLogprobs =
2120
+ typeof body.top_logprobs === "number" && body.top_logprobs > 0
2121
+ ? body.top_logprobs : 0;
2122
+ const captureLogprobs = !body.stream && (wantLogprobs || topLogprobs > 0);
2123
+ if (captureLogprobs) {
2124
+ options.logprobs = wantLogprobs;
2125
+ options.topLogprobs = topLogprobs;
2126
+ }
2127
+ const shape = {
2128
+ hasVision: false,
2129
+ hasAdapters: !!options.adapters?.length,
2130
+ hasRepetitionPenalty: !!options.repetitionPenalty,
2131
+ hasLogitsExtras: !!(
2132
+ options.minP || options.xtcProbability || options.logitBias ||
2133
+ options.presencePenalty || options.frequencyPenalty
2134
+ ),
2135
+ wantsLogprobs: captureLogprobs,
2136
+ userSeed: body.seed !== undefined,
2137
+ kvQuant: !!(options.kvConfig?.length || options.kvBits),
2138
+ };
2139
+ const finishReason = (stopped: boolean, generated: number): "stop" | "length" =>
2140
+ stopped ? "stop" : generated >= (options.maxTokens ?? 512) ? "length" : "stop";
2141
+
2142
+ if (body.stream) {
2143
+ const stream = new ReadableStream<Uint8Array>({
2144
+ async start(controller) {
2145
+ const enc = new TextEncoder();
2146
+ const send = (obj: unknown) =>
2147
+ controller.enqueue(enc.encode(`data: ${JSON.stringify(obj)}\n\n`));
2148
+ const chunk = (text: string, finish: string | null) => ({
2149
+ id, object: "text_completion", created, model: ctx.modelId,
2150
+ choices: [{ index: 0, text, finish_reason: finish }],
2151
+ });
2152
+ try {
2153
+ const decoder = new StreamDecoder(ctx.tokenizer);
2154
+ const stopper = new StopMatcher(options.stopSequences);
2155
+ // ≥25 ms macrotask hop, same reason as the chat lane: keep the
2156
+ // serial decode loop from starving the socket (SSE bursts).
2157
+ let lastFlush = performance.now();
2158
+ const s = await gateway.run(promptIds, options, (token) => {
2159
+ const text = stopper.push(decoder.push(token));
2160
+ if (text) send(chunk(text, null));
2161
+ if (stopper.stopped) return false; // halt generation
2162
+ if (text) {
2163
+ const now = performance.now();
2164
+ if (now - lastFlush >= 25) {
2165
+ lastFlush = now;
2166
+ return new Promise<void>((r) => setImmediate(r));
2167
+ }
2168
+ }
2169
+ }, undefined, shape);
2170
+ if (!stopper.stopped) {
2171
+ let tail = stopper.push(decoder.flush());
2172
+ if (!stopper.stopped) tail += stopper.flush();
2173
+ if (tail) send(chunk(tail, null));
2174
+ }
2175
+ // final chunk: finish_reason + usage (mlx-lm gates usage behind
2176
+ // stream_options.include_usage; we always attach it, matching
2177
+ // our chat lane — an additive superset OpenAI clients ignore)
2178
+ send({
2179
+ ...chunk("", finishReason(stopper.stopped, s.generatedTokens)),
2180
+ usage: {
2181
+ prompt_tokens: s.promptTokens,
2182
+ completion_tokens: s.generatedTokens,
2183
+ total_tokens: s.promptTokens + s.generatedTokens,
2184
+ prompt_tokens_details: { cached_tokens: s.cachedTokens },
2185
+ },
2186
+ });
2187
+ controller.enqueue(enc.encode("data: [DONE]\n\n"));
2188
+ } catch (e) {
2189
+ send({ error: { message: (e as Error).message } });
2190
+ } finally {
2191
+ controller.close();
2192
+ }
2193
+ },
2194
+ });
2195
+ return new Response(stream, {
2196
+ headers: {
2197
+ "content-type": "text/event-stream",
2198
+ "cache-control": "no-cache",
2199
+ connection: "keep-alive",
2200
+ },
2201
+ });
2202
+ }
2203
+
2204
+ try {
2205
+ const decoder = new StreamDecoder(ctx.tokenizer);
2206
+ const stopper = new StopMatcher(options.stopSequences);
2207
+ const lpc = captureLogprobs
2208
+ ? new LogprobsCollector(wantLogprobs, topLogprobs, (tid) => ctx.tokenizer.idToToken(tid))
2209
+ : null;
2210
+ let text = "";
2211
+ const s = await gateway.run(promptIds, options, (token, lpInfo) => {
2212
+ lpc?.push(token, lpInfo);
2213
+ text += stopper.push(decoder.push(token));
2214
+ if (stopper.stopped) return false; // halt generation
2215
+ }, undefined, shape);
2216
+ if (!stopper.stopped) {
2217
+ let tail = stopper.push(decoder.flush());
2218
+ if (!stopper.stopped) tail += stopper.flush();
2219
+ text += tail;
2220
+ }
2221
+ const logprobsBlock = lpc?.payload() ?? null;
2222
+ return Response.json({
2223
+ id, object: "text_completion", created, model: ctx.modelId,
2224
+ choices: [{
2225
+ index: 0, text,
2226
+ ...(logprobsBlock ? { logprobs: logprobsBlock } : {}),
2227
+ finish_reason: finishReason(stopper.stopped, s.generatedTokens),
2228
+ }],
2229
+ usage: {
2230
+ prompt_tokens: s.promptTokens,
2231
+ completion_tokens: s.generatedTokens,
2232
+ total_tokens: s.promptTokens + s.generatedTokens,
2233
+ prompt_tokens_details: { cached_tokens: s.cachedTokens },
2234
+ },
2235
+ });
2236
+ } catch (e) {
2237
+ return Response.json({ error: { message: (e as Error).message } }, { status: 500 });
2238
+ }
2239
+ }
2240
+
1678
2241
  // Anthropic Messages API (Phase 11) — on by default, mirroring
1679
2242
  // optiq serve (--anthropic defaults True; the drop-in claim
1680
2243
  // depends on it). Oracle: optiq/anthropic_shim.py, ported in
@@ -1932,7 +2495,7 @@ export function createServer(
1932
2495
  const snapMatch = body.model_id.match(/(models--[^/]+)\/snapshots\//);
1933
2496
  let org = "local", name: string;
1934
2497
  if (snapMatch) {
1935
- const parts = snapMatch[1].split("--"); // ["models", org, ...name]
2498
+ const parts = snapMatch[1]!.split("--"); // ["models", org, ...name]
1936
2499
  org = parts[1] ?? "local";
1937
2500
  name = parts.slice(2).join("--");
1938
2501
  } else if (body.model_id.includes("/") && !body.model_id.startsWith("/") && !body.model_id.startsWith("~")) {