mlx-bun 0.5.0 → 1.0.0

package/src/server/routes.ts

@@ -0,0 +1,205 @@
+import type { KvCacheService } from "@mlx-bun/inference/state";
+import { runtimeValue } from "@mlx-bun/inference/runtime/config";
+import { createPromptResponseTrace } from "@mlx-bun/inference/runtime/trace";
+import type { LoadedModelContext as ModelContext } from "../engine/model-host";
+import type { CompletionEngine } from "../engine/completion";
+import type { ModelBinding } from "../engine/model-binding";
+import type { GenerationGateway } from "../engine/generation-gateway";
+import type { PreparationExecutor } from "../engine/preparation";
+import { ChatRequest, TextCompletionRequest, type ChatRequestParams } from "./chat-request";
+import { ChatStage } from "./chat-stage";
+import { TextCompletionStage } from "./text-completion-stage";
+import { CompletionExecutor } from "./completion-executor";
+import { InferenceStage } from "./inference-request";
+import { createRequestPrep, type RequestPrepOptions } from "./request-prep";
+import type { PromptTokenHistory } from "./generated-token-history";
+import type { ModelPromptBuilder } from "./prompt-contracts";
+import { admit, errorResponse, respondJson, respondStream, type ErrorFormatter } from "./http";
+import { chatCompletionJson, chatCompletionStream, textCompletionJson, textCompletionStream } from "./openai-wire";
+import { createDiscoveryRoutes, type TranscriptionInfo } from "./discovery-routes";
+import { createDecisionRoutes } from "./decision-routes";
+import { anthropicToChatBody, chatJsonToAnthropic, createAnthropicStreamProtocol, type AnthropicRequest } from "./anthropic";
+
+import { ResponseStore, resolveResponsesConversation, responsesToChatBody, chatJsonToResponses,
+  createResponsesStreamProtocol, type ResponseHistory, type ResponsesRequest } from "./responses";
+
+const anthropicError: ErrorFormatter = (status, message, body) => Response.json({ type: "error",
+  error: { ...body, type: status >= 500 ? "api_error" : "invalid_request_error", message },
+}, { status });
+
+const responsesError: ErrorFormatter = (status, message, body) => Response.json({ error: {
+  ...body, message, type: status >= 500 ? "server_error" : "invalid_request_error", param: null, code: body.code ?? null,
+} }, { status });
+
+/** HTTP adapters borrow the engine; the application owns its lifetime. No
+ * socket, loaded model, cache, or scheduler is created by these routes. */
+export function createCompletionRoutes(engine: {
+  context: ModelContext;
+  completion: CompletionEngine;
+  preparation: PreparationExecutor;
+  binding: Pick<ModelBinding, "discovery" | "embed" | "decide"> & Partial<Pick<ModelBinding, "gateway">>;
+  gateway: Pick<GenerationGateway, "runExclusive"> & Partial<Pick<GenerationGateway, "runPreparation">>;
+}, options: RequestPrepOptions & {
+  /** The loaded model's cache service: the prefix peek and the encoder-output object cache. */
+  promptCache: Pick<KvCacheService, "peekPrefixLen"> & Partial<Pick<KvCacheService, "objects">>;
+  contextLimit: number | null;
+  defaultGeneratedTokens?: number;
+  defaultAdapter?: string;
+  tokenHistory?: PromptTokenHistory;
+  buildPrompt?: ModelPromptBuilder;
+  /** Process-local history by default; composition owns any replacement store. */
+  responseHistory?: ResponseHistory;
+  /** The Whisper companion `/v1/models` lists beside the chat model; absent means none is configured. */
+  transcription?: () => Promise<TranscriptionInfo | null>;
+}) {
+  const ctx = engine.context;
+  const decisions = createDecisionRoutes(engine, options);
+  const responseHistory = options.responseHistory ?? new ResponseStore();
+  const prep = createRequestPrep({ ctx, serverOptions: options,
+    defaultGeneratedTokens: options.defaultGeneratedTokens, tokenHistory: options.tokenHistory });
+  const chat = new ChatStage(ctx, prep, options.promptCache, options.contextLimit,
+    options.defaultAdapter, engine.preparation, options.buildPrompt);
+  const text = new TextCompletionStage(ctx, prep, options.contextLimit, options.defaultGeneratedTokens, options.defaultAdapter, engine.preparation);
+  const inference = new InferenceStage(new CompletionExecutor(engine.completion));
+  const discovery = createDiscoveryRoutes(ctx, engine.binding, Date.now(), options.transcription, undefined,
+    { temperature: options.defaultTemperature, topP: options.defaultTopP, topK: options.defaultTopK });
+  const applyCacheSession = (body: ChatRequestParams, request: Request, original: unknown = body) => {
+    const fields = original as { session_id?: unknown; prompt_cache_key?: unknown };
+    const session = typeof fields.session_id === "string" ? fields.session_id :
+      typeof fields.prompt_cache_key === "string" ? fields.prompt_cache_key :
+      request.headers.get("x-session-affinity") ?? request.headers.get("session_id") ?? undefined;
+    if (runtimeValue("MLX_BUN_SESSION_CACHE") !== "0") body.session_id = session;
+    else { delete body.session_id; delete body.prompt_cache_key; }
+  };
+  return {
+    responseStats: () => ({ entries: responseHistory.size, bytes: responseHistory.totalBytes,
+      max_bytes: responseHistory.maxBytes, ttl_ms: responseHistory.ttlMs }),
+    /** null means another app surface may handle this request. */
+    async handle(request: Request): Promise<Response | null> {
+      const url = new URL(request.url);
+      const format = url.pathname === "/v1/messages" ? anthropicError
+        : url.pathname === "/v1/responses" ? responsesError : undefined;
+      try {
+        const decision = await decisions.handle(request);
+        if (decision) return decision;
+        const discovered = await discovery.handle(url, request);
+        if (discovered) return discovered;
+        if (request.method !== "POST") return null;
+        if (url.pathname === "/v1/messages") {
+          let body: AnthropicRequest;
+          try {
+            body = await request.json();
+            if (!body || typeof body !== "object" || Array.isArray(body)) throw new Error("invalid JSON body");
+          } catch (error) {
+            return request.signal.aborted ? errorResponse(error, url.pathname, format, request.signal)
+              : anthropicError(400, "invalid JSON body", {});
+          }
+          let chatBody: ChatRequestParams;
+          try {
+            chatBody = anthropicToChatBody(body) as unknown as ChatRequestParams;
+            applyCacheSession(chatBody, request, body);
+          } catch (error) {
+            return anthropicError(400, error instanceof Error ? error.message : String(error), {});
+          }
+          const id = `chatcmpl-${crypto.randomUUID()}`;
+          const admitted = await admit(inference, () => chat.run(new ChatRequest(chatBody), id, request.signal),
+            undefined, "anthropic request", anthropicError, request.signal);
+          if ("response" in admitted) return admitted.response;
+          return body.stream
+            ? respondStream(inference, admitted.admitted, createAnthropicStreamProtocol(ctx.modelId), request.signal)
+            : respondJson(inference, admitted.admitted, result => chatJsonToAnthropic(chatCompletionJson(result,
+              { id, created: Math.floor(Date.now() / 1000), model: ctx.modelId }), ctx.modelId), request.signal, undefined, anthropicError);
+        }
+        if (url.pathname === "/v1/responses") {
+          let body: ResponsesRequest;
+          try {
+            body = await request.json();
+            if (!body || typeof body !== "object" || Array.isArray(body)) throw new Error("invalid JSON body");
+          } catch (error) {
+            return request.signal.aborted ? errorResponse(error, url.pathname, format, request.signal)
+              : responsesError(400, "invalid JSON body", {});
+          }
+          // Validate the conversation container before resolving a previous ID;
+          // only a missing/expired previous response is a 404.
+          if ((body.input != null && typeof body.input !== "string" && !Array.isArray(body.input)) ||
+              (body.previous_response_id != null && typeof body.previous_response_id !== "string"))
+            return responsesError(400, "input must be a string or array; previous_response_id must be a string", {});
+          let conversation: ReturnType<typeof resolveResponsesConversation>;
+          try { conversation = resolveResponsesConversation(body, responseHistory); }
+          catch (error) { return responsesError(404, error instanceof Error ? error.message : String(error), {}); }
+          body = conversation.body;
+          const remember = (final: Record<string, unknown>) => responseHistory.put(final.id as string, {
+            input: conversation.input, output: final.output as unknown[], instructions: conversation.instructions,
+          });
+          let chatBody: ChatRequestParams;
+          try {
+            chatBody = responsesToChatBody(body) as unknown as ChatRequestParams;
+            applyCacheSession(chatBody, request, body);
+          } catch (error) {
+            return responsesError(400, error instanceof Error ? error.message : String(error), {});
+          }
+          const id = `chatcmpl-${crypto.randomUUID()}`;
+          const admitted = await admit(inference, () => chat.run(new ChatRequest(chatBody), id, request.signal),
+            undefined, "responses request", responsesError, request.signal);
+          if ("response" in admitted) return admitted.response;
+          return body.stream
+            ? respondStream(inference, admitted.admitted,
+              createResponsesStreamProtocol(ctx.modelId, conversation.previousId, remember), request.signal)
+            : respondJson(inference, admitted.admitted, result => {
+              const response = chatJsonToResponses(chatCompletionJson(result,
+                { id, created: Math.floor(Date.now() / 1000), model: ctx.modelId }), ctx.modelId, conversation.previousId);
+              remember(response);
+              return response;
+            }, request.signal, undefined, responsesError);
+        }
+        if (url.pathname === "/v1/embeddings") {
+          const embed = engine.binding.embed?.bind(engine.binding);
+          if (!embed) return Response.json({ error: {
+            message: `served model "${ctx.modelId}" is not an embedding model; serve an embedding model (e.g. Qwen3-Embedding) to use /v1/embeddings`,
+            type: "invalid_request_error",
+          } }, { status: 400 });
+          let body: { input?: unknown; instruction?: unknown };
+          try { body = await request.json(); }
+          catch (error) {
+            if (request.signal.aborted) return errorResponse(error, url.pathname, undefined, request.signal);
+            return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+          }
+          const inputs = Array.isArray(body?.input) ? body.input : body?.input != null ? [body.input] : [];
+          if (!inputs.length || !inputs.every(input => typeof input === "string"))
+            return Response.json({ error: { message: "`input` must be a string or array of strings", type: "invalid_request_error" } }, { status: 400 });
+          const instruction = typeof body.instruction === "string" ? body.instruction : undefined;
+          const results = await engine.gateway.runExclusive(async () => embed(inputs, instruction), undefined, request.signal);
+          request.signal.throwIfAborted();
+          const total = results.reduce((sum, result) => sum + result.tokens, 0);
+          return Response.json({ object: "list", model: ctx.modelId,
+            data: results.map((result, index) => ({ object: "embedding", index, embedding: Array.from(result.vector) })),
+            usage: { prompt_tokens: total, total_tokens: total } });
+        }
+        const chatRoute = url.pathname === "/v1/chat/completions";
+        if (!chatRoute && url.pathname !== "/v1/completions") return null;
+        const id = `${chatRoute ? "chatcmpl" : "cmpl"}-${crypto.randomUUID()}`;
+        const trace = createPromptResponseTrace({ traceId: request.headers.get("x-mlx-bun-trace-id") ?? id, requestId: id, route: url.pathname });
+        let body: ChatRequestParams;
+        const closeParse = trace?.begin("request.body_parse");
+        try {
+          const parsed = await request.json();
+          if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) throw new Error("invalid JSON body");
+          body = parsed;
+        } catch {
+          trace?.finish(request.signal.aborted ? "abort" : "error", { stage: "body_parse" });
+          if (request.signal.aborted) return errorResponse(request.signal.reason, url.pathname, undefined, request.signal);
+          return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+        } finally { closeParse?.(); }
+        applyCacheSession(body, request);
+        const admitted = await admit(inference, () => chatRoute
+          ? chat.run(new ChatRequest(body), id, request.signal)
+          : text.run(new TextCompletionRequest(body), id, request.signal), trace, chatRoute ? "chat request" : "text completion", undefined, request.signal);
+        if ("response" in admitted) return admitted.response;
+        const meta = { id, created: Math.floor(Date.now() / 1000), model: ctx.modelId };
+        return body.stream
+          ? respondStream(inference, admitted.admitted, chatRoute ? chatCompletionStream(meta) : textCompletionStream(meta), request.signal, trace)
+          : respondJson(inference, admitted.admitted, result => chatRoute ? chatCompletionJson(result, meta) : textCompletionJson(result, meta), request.signal, trace);
+      } catch (error) { return errorResponse(error, url.pathname, format, request.signal); }
+    },
+  };
+}

package/src/server/start.ts

@@ -0,0 +1,72 @@
+import { chmodSync, rmSync } from "node:fs";
+import type { ModuleSocketData, ModuleSockets } from "@mlx-bun/app-services/portable";
+import type { createCompletionRoutes } from "./routes";
+
+// A path no mounted handler owns answers 404. That includes the worker's
+// private admin routes (`/admin/lease`, `/admin/drain`, `/admin/memory/complete`)
+// and `/engine` under `--in-process`, which exist only on the surface that
+// mounts them, and main's lab pages (/generate, /signal, /curves,
+// /curve-terrain, /dag), which are not product surface.
+
+/** Bind the app's HTTP/WebSocket surfaces. Ownership transfers on entry: a
+ * failed bind or close releases the caller's engine resources. The modules own
+ * their sockets (`sockets`) and close them when they stop, in `beforeDrain`.
+ * The web handler serves already-built assets and never loads a model. */
+export async function startServer(input: {
+  routes: Pick<ReturnType<typeof createCompletionRoutes>, "handle">;
+  web(request: Request): Response | null;
+  /** The WebSocket paths the installed modules declare; `createModuleSockets([])` for a host with none. */
+  sockets: ModuleSockets;
+  /** Stop background producers, cancel managed jobs and close the modules' sockets while the engine is alive,
+   * before waiting for HTTP/SSE responses to drain. */
+  beforeDrain?(): void | Promise<void>;
+  closeEngine(): Promise<void>;
+}, options: { port?: number;
+  /** Omitted: loopback (127.0.0.1). null: Bun's default, every interface. */
+  hostname?: string | null;
+  /** Internal (worker mode): bind this Unix socket path instead of TCP; `port`
+   * and `hostname` are then ignored. A stale file is replaced, the socket is
+   * narrowed to owner-only right after the bind, and close removes it. */
+  unix?: string } = {}) {
+  let closing: Promise<void> | undefined;
+  let stopped = false;
+  let server: ReturnType<typeof Bun.serve<ModuleSocketData>> | undefined;
+  const close = () => closing ??= (async () => {
+    stopped = true;
+    const errors: unknown[] = [];
+    try { await input.beforeDrain?.(); } catch (error) { errors.push(error); }
+    try { await server?.stop(false); } catch (error) { errors.push(error); }
+    finally { if (options.unix) rmSync(options.unix, { force: true }); }
+    try { await input.closeEngine(); } catch (error) { errors.push(error); }
+    if (errors.length) throw new AggregateError(errors, "server cleanup failed");
+  })();
+  try {
+    if (options.unix) rmSync(options.unix, { force: true });
+    const handlers = {
+      idleTimeout: 0,
+      websocket: input.sockets.websocket,
+      async fetch(request: Request, listener: ReturnType<typeof Bun.serve<ModuleSocketData>>) {
+        if (stopped) return Response.json({ error: { message: "server is shutting down" } }, { status: 503 });
+        const socket = input.sockets.upgrade(request, listener);
+        if (socket !== null) return socket;
+        const page = input.web(request);
+        if (page) return page;
+        const response = await input.routes.handle(request);
+        if (response) return response;
+        return Response.json({ error: { message: "Not found" } }, { status: 404 });
+      },
+    };
+    // bun-types declares idleTimeout for TCP listeners only; the runtime takes
+    // the same options for a Unix listener, where a job's idle lease connection
+    // must never time out either.
+    server = options.unix
+      ? Bun.serve<ModuleSocketData>({ unix: options.unix, ...handlers } as unknown as Parameters<typeof Bun.serve<ModuleSocketData>>[0])
+      : Bun.serve<ModuleSocketData>({ port: options.port ?? 8080,
+        ...(options.hostname === null ? {} : { hostname: options.hostname ?? "127.0.0.1" }), ...handlers });
+    if (options.unix) chmodSync(options.unix, 0o600);
+    return { server, close };
+  } catch (error) {
+    try { await close(); } catch (cleanup) { throw new AggregateError([error, cleanup], "server startup and cleanup failed"); }
+    throw error;
+  }
+}

package/src/server/status-routes.ts

@@ -0,0 +1,108 @@
+import { detectChip, fit, skuMatrix, thisMachine, type MachineSpec } from "@mlx-bun/inference/execution/fit";
+import type { KvCacheService } from "@mlx-bun/inference/state";
+import type { createCacheServices } from "../engine/cache-services";
+import type { GenerationGateway } from "../engine/generation-gateway";
+import type { LoadedModelContext } from "../engine/model-host";
+import type { ModelBinding } from "../engine/model-binding";
+
+type CacheServices = Awaited<ReturnType<typeof createCacheServices>>;
+
+/** Read-only HTTP views borrow counters and metadata. They neither open model
+ * storage nor acquire execution leases, and fit estimates never set admission.
+ * What the cache holds is the cache service's own report. */
+export function createStatusRoutes(input: {
+  context: Pick<LoadedModelContext, "modelId" | "memoryPlan"> & {
+    model: Pick<LoadedModelContext["model"], "config" | "weightsBytes">;
+  };
+  cache: Pick<KvCacheService, "report">;
+  resolvedKvScheme: Pick<CacheServices["resolvedKvScheme"], "fitOptions" | "describe">;
+  gateway: Pick<GenerationGateway, "activeRows" | "pendingRows" | "submittedRows">;
+  diagnostics: ModelBinding["diagnostics"];
+  responseStats(): { entries: number; bytes: number; max_bytes: number; ttl_ms: number };
+  /** The served artifact's registry sizes; a caller-built context may know neither (main's 0 and null). */
+  artifact: { expertsBytes: number; sizeBytes: number | null };
+  capacity: number;
+  contextLimit: number | null;
+  startedAt: number;
+  owner?: string;
+  ssdCacheDir?: string;
+  /** Main's explicit `--memory-budget`: the usable envelope for the estimates below. */
+  memoryBudgetBytes?: number;
+  /** Composition/test inputs; omitted values use the hub's machine detection. */
+  machine?: MachineSpec;
+  chip?: string | null;
+}) {
+  const { context: ctx, gateway } = input;
+  const kvScheme = input.resolvedKvScheme;
+  const machine = input.machine ?? thisMachine();
+  const chip = input.chip === undefined ? detectChip().name : input.chip;
+  const plan = ctx.memoryPlan;
+  const admission = plan ?? fit(ctx.model.config, ctx.model.weightsBytes, 1,
+    machine, undefined, 0, input.memoryBudgetBytes, kvScheme.fitOptions);
+
+  return { async handle(request: Request): Promise<Response | null> {
+    if (request.method !== "GET") return null;
+    const path = new URL(request.url).pathname;
+    if (path === "/stats") {
+      const held = input.cache.report(), pc = held.prefixes, disk = held.saved, pending = held.pending;
+      return Response.json({
+        server: { owner: input.owner ?? "embedded", model: ctx.modelId, started_at: input.startedAt },
+        prompt_cache: { entries: pc.entries, bytes: pc.bytes, max_bytes: pc.maxBytes,
+          hits: pc.hits, misses: pc.misses, session_hits: pc.sessionHits, session_misses: pc.sessionMisses,
+          prefix_scans: pc.prefixScans, object_hits: pc.objectHits, object_misses: pc.objectMisses, object_restores: pc.objectRestores },
+        ...(disk ? { ssd_cache: { dir: input.ssdCacheDir, entries: disk.entries, bytes: disk.bytes,
+          max_bytes: Number.isFinite(disk.maxBytes) ? disk.maxBytes : null,
+          restores: disk.restores, spills: disk.spills, restore_ms_last: Math.round(disk.restoreMsLast),
+          demotions: pc.demotions, pending_snapshots: pending.pendingSnapshots, pending_spills: pending.pendingSpills,
+          pending_spill_bytes: pending.pendingSpillBytes, dropped_spills: pending.droppedSpills, failed_spills: pending.failedSpills,
+          longest_durable_prefix_tokens: disk.longestDurablePrefixTokens } } : {}),
+        response_store: input.responseStats(),
+        kv_quant: kvScheme.describe(ctx.model.config),
+        admission: { max_safe_context: admission.maxSafeContext, enforced_context_tokens: input.contextLimit,
+          memory_budget_bytes: plan?.usableBytes ?? input.memoryBudgetBytes ?? null, usable_bytes: admission.usableBytes, weights_bytes: ctx.model.weightsBytes },
+        ...input.diagnostics(),
+        batch: { configured: input.capacity, mode: "batch", batched: true,
+          active_rows: gateway.activeRows, pending_rows: gateway.pendingRows, submitted_rows: gateway.submittedRows,
+          kv_bytes: held.kv.projectedBytes, kv_budget_bytes: held.kv.budgetBytes },
+      });
+    }
+    if (path !== "/fit") return null;
+    const machineWire = { chip, ram_bytes: machine.ramBytes, bandwidth_gbs: machine.bandwidthGBs };
+    if (plan) {
+      // A runtime that plans its own memory: the plan's own accounting, with
+      // its streamed tiers reported apart from resident weights.
+      const streamed = plan.streamedWeights ?? {};
+      const streamedBytes = Object.values(streamed).reduce((sum, bytes) => sum + bytes, 0);
+      return Response.json({ machine: machineWire, context_tokens: plan.contextTokens,
+        typical_context_tokens: plan.contextTokens, typical_decode_tps: null,
+        measured_decode_tps: null, measured_at: null,
+        report: { fits: true, weights_bytes: plan.weightsBytes - streamedBytes, kv_bytes: plan.kvBytes,
+          transient_bytes: plan.transientBytes + plan.reserveBytes,
+          total_bytes: plan.totalBytes, usable_bytes: plan.usableBytes,
+          max_safe_context: plan.contextTokens, predicted_decode_tps: null },
+        plan: { artifact_disk_bytes: input.artifact.sizeBytes, streamed_weights: streamed,
+          max_generation_tokens: plan.maxGenerationTokens ?? null },
+        sku_matrix_ctx: plan.contextTokens,
+        sku_matrix: [{ sku: chip, ram_gb: Math.round(machine.ramBytes / 2 ** 30), fits: true,
+          max_context: plan.contextTokens, decode_tps: null }],
+      });
+    }
+    const report = fit(ctx.model.config, ctx.model.weightsBytes, admission.maxSafeContext,
+      machine, undefined, input.artifact.expertsBytes, input.memoryBudgetBytes, kvScheme.fitOptions);
+    const typicalContext = Math.min(8192, admission.maxSafeContext);
+    return Response.json({ machine: machineWire, context_tokens: admission.maxSafeContext,
+      typical_context_tokens: typicalContext,
+      typical_decode_tps: fit(ctx.model.config, ctx.model.weightsBytes, typicalContext,
+        machine, undefined, input.artifact.expertsBytes, input.memoryBudgetBytes, kvScheme.fitOptions).predictedDecodeTps,
+      // The old EvalDB and machine-specific historical measurements are not
+      // live evidence for this server. Preserve the fields without inventing data.
+      measured_decode_tps: null, measured_at: null,
+      report: { fits: report.fits, weights_bytes: report.weightsBytes, kv_bytes: report.kvBytes,
+        transient_bytes: report.transientBytes, total_bytes: report.totalBytes,
+        usable_bytes: report.usableBytes, max_safe_context: report.maxSafeContext, predicted_decode_tps: report.predictedDecodeTps },
+      sku_matrix_ctx: 32768,
+      sku_matrix: skuMatrix(ctx.model.config, ctx.model.weightsBytes, 32768, input.artifact.expertsBytes)
+        .map(r => ({ sku: r.sku, ram_gb: r.ramGB, fits: r.fits, max_context: r.maxContext, decode_tps: r.decodeTps })),
+    });
+  } };
+}

package/src/{serve → server}/text-completion-stage.ts

@@ -4,7 +4,7 @@
 // special-token handling, exactly mlx-lm's `tokenizer.encode(request.prompt)`).
 // Same admission / adapter / lane path as chat; no tool router or thinking
 // splitter (raw text in, raw text out).
-import type { ServingContext } from "./model-host";
+import type { LoadedModelContext as ModelContext } from "../engine/model-host";
 import type { TextCompletionRequest } from "./chat-request";
 import type { InferenceRequest } from "./inference-request";
 import { RequestError } from "./pipeline";
@@ -14,17 +14,38 @@ import { StopMatcher, ThinkingTagSplitter, ToolAwareStream } from "./token-strea
 
 export class TextCompletionStage {
   constructor(
-    private readonly ctx: ServingContext,
+    private readonly ctx: ModelContext,
     private readonly prep: RequestPrep,
     private readonly contextLimit: number | null,
-    /** `--max-tokens` / the GLM plan's reservation. No invented fallback:
+    /** `--max-tokens` / a runtime plan's reservation. No invented fallback:
      *  unset runs to EOS or the admitted context (with no template an EOS
      *  may never come — that run is bounded by admission, the real limit). */
     private readonly defaultGeneratedTokens: number | undefined,
     private readonly defaultAdapter?: string,
+    private readonly preparation?: import("../engine/preparation").PreparationExecutor,
   ) {}
 
-  async run(request: TextCompletionRequest, requestId?: string): Promise<InferenceRequest> {
+  async run(request: TextCompletionRequest, requestId?: string, signal?: AbortSignal): Promise<InferenceRequest> {
+    signal?.throwIfAborted();
+    const body = request.params;
+    const constrained = body.response_format != null || !!body.guided_grammar ||
+      !!body.guided_regex || !!body.guided_choice?.length || body.structured_outputs != null;
+    let reservation = constrained ? await this.preparation?.reserve?.("constraint", signal) : undefined;
+    try {
+      const build = async () => {
+        signal?.throwIfAborted();
+        const result = await this.runPrepared(request, requestId);
+        if (signal?.aborted) { result.plan.ownership.dispose(); signal.throwIfAborted(); }
+        return result;
+      };
+      const result = constrained && this.preparation ? await this.preparation.run(build, signal) : await build();
+      if (reservation) result.plan.ownership.retain(reservation);
+      reservation = undefined;
+      return result;
+    } finally { reservation?.dispose(); }
+  }
+
+  private async runPrepared(request: TextCompletionRequest, requestId?: string): Promise<InferenceRequest> {
     const { ctx, prep } = this;
     const body = request.params;
     const id = requestId ?? `cmpl-${crypto.randomUUID()}`;

package/src/server/text-prompt.ts

@@ -0,0 +1,12 @@
+import type { ChatMessage, ChatTemplate, LoadedTokenizer } from "@mlx-bun/inference/input";
+
+/** Template rendering shared by app surfaces. CLI excludes tokenizer-added
+ * specials; HTTP preserves its existing encoding and duplicate-BOS correction. */
+export function textPrompt(template: ChatTemplate, tokenizer: LoadedTokenizer,
+  messages: ChatMessage[], options: Parameters<ChatTemplate["render"]>[1],
+  addSpecialTokens?: boolean): { rendered: string; ids: number[] } {
+  const rendered = template.render(messages, options);
+  const ids = addSpecialTokens === undefined ? tokenizer.encode(rendered) : tokenizer.encode(rendered, addSpecialTokens);
+  return { rendered, ids: addSpecialTokens === undefined && ids[0] === ids[1] && ids[0] === tokenizer.bosTokenId
+    ? ids.slice(1) : ids };
+}

package/src/{serve → server}/token-streams.ts

@@ -2,13 +2,10 @@
 // StreamDecoder (mlx-lm byte-parity incremental detokenizer), ToolAwareStream
 // (tool-call / reasoning-channel routing), StopMatcher (decoded-text stop
 // sequences), ThinkingTagSplitter (<think> markup → reasoning deltas).
-// Extracted from src/server.ts (repo-taming Phase 4).
-import type { ToolDefinition } from "../chat-template";
-import type { LoadedTokenizer } from "../tokenizer";
-import {
-  CHANNEL_END, CHANNEL_START, parseGeneratedToolCalls, parseToolCalls,
-  TOOL_CALL_END, TOOL_CALL_START,
-} from "../tool-call";
+import type { ToolDefinition } from "@mlx-bun/inference/input";
+import type { LoadedTokenizer } from "@mlx-bun/inference/input";
+import type { SentinelTokens } from "@mlx-bun/inference/input";
+import { parseGeneratedToolCalls, parseToolCalls } from "@mlx-bun/inference/input/tool-call";
 
 export interface OpenAIToolCall {
   id: string;
@@ -16,28 +13,26 @@ export interface OpenAIToolCall {
   function: { name: string; arguments: string };
 }
 
-export type ToolStreamMode = "gemma-sentinel" | "plain" | "buffered-text";
+export type ToolStreamMode = "sentinel-tokens" | "plain" | "buffered-text";
 
-/** Pick the stream router for a model family. The token-id sentinel router is
- *  Gemma-4-ONLY: ids 48/49 (<|tool_call>/<tool_call|>) and 100/101
- *  (<|channel>/<channel|>) are special tokens of that tokenizer family
- *  (src/tool-call.ts). On every other tokenizer — MiniCPM5, Qwen3/3.5, and
- *  the Tier-0 generics (llama, qwen2, phi3, …) — those ids are ordinary
- *  low-id vocab, so the sentinel router would silently swallow output into a
- *  phantom tool/reasoning segment. Everyone else parses tool calls from
- *  decoded text (buffered-text; parseGeneratedToolCalls covers the
- *  OpenAI-JSON <tool_call>, Qwen <function=…>, and MiniCPM5 <function name=…>
- *  shapes) when tools are present, and streams plain otherwise. Models whose
- *  markup isn't covered fail soft: the markup stays in content. */
-export function selectToolStreamMode(modelType: string, hasTools: boolean): ToolStreamMode {
-  if (modelType.startsWith("gemma4")) return "gemma-sentinel";
+/** Pick the stream router. A model whose profile declares sentinel tokens
+ *  (special tokens delimiting tool calls and the reasoning channel, resolved
+ *  against its tokenizer) is routed at the token level: on any other tokenizer
+ *  the same ids would be ordinary vocab, so a router must never guess them.
+ *  Everyone else parses tool calls from decoded text (buffered-text;
+ *  parseGeneratedToolCalls covers the OpenAI-JSON <tool_call>, Qwen
+ *  <function=…>, and MiniCPM5 <function name=…> shapes) when tools are present,
+ *  and streams plain otherwise. Models whose markup isn't covered fail soft:
+ *  the markup stays in content. */
+export function selectToolStreamMode(sentinels: SentinelTokens | null | undefined, hasTools: boolean): ToolStreamMode {
+  if (sentinels) return "sentinel-tokens";
   return hasTools ? "buffered-text" : "plain";
 }
 
-/** Routes generated tokens. Gemma uses family-specific sentinel token ids;
- *  MiniCPM5 and other text-template models use decoded-text parsing so
- *  ordinary tokenizer ids like "<" are never swallowed globally. Exported for
- *  unit tests (gemma-channel reasoning split). */
+/** Routes generated tokens. Models declaring sentinel tokens are routed by id;
+ *  text-template models use decoded-text parsing so ordinary tokenizer ids
+ *  like "<" are never swallowed globally. Exported for unit tests
+ *  (reasoning-channel split). */
 export class ToolAwareStream {
   readonly #decoder: StreamDecoder;
   #inTool = false;
@@ -51,9 +46,9 @@ export class ToolAwareStream {
   #textToolParseFailed = false;
   readonly toolSegments: number[][] = [];
 
-  /** Gemma reasoning-channel state (gemma-sentinel mode). The model wraps
-   *  chain-of-thought as `<|channel>thought\n…<channel|>` using special tokens
-   *  100/101 that the content decoder strips, so reasoning is captured here at
+  /** Reasoning-channel state (sentinel-tokens mode). The model wraps
+   *  chain-of-thought as `<|channel>thought\n…<channel|>` using the channel
+   *  sentinel tokens, which the content decoder strips, so reasoning is captured here at
    *  the token level. A SEPARATE decoder keeps the reasoning byte-stream's
    *  incremental state independent of content's. The `thought` channel-name
    *  word is stripped before the reasoning text (mlx-lm's think-start marker
@@ -74,11 +69,14 @@ export class ToolAwareStream {
     readonly mode: ToolStreamMode,
     readonly tools: ToolDefinition[] | null,
     /** Fired when parseGeneratedToolCalls rejects the emitted markup. The
-     *  fill table (src/fill/) subscribes: a request whose tool call the
+     *  fill table (packages/inference/src/generation/fill/) subscribes: a request whose tool call the
      *  parser refuses is a request whose template rendering disagreed with
      *  what the model emits, so its strict rows are disarmed. */
     private readonly onParseFailure?: () => void,
+    /** The resolved sentinel ids; required in sentinel-tokens mode. */
+    private readonly sentinels?: SentinelTokens | null,
   ) {
+    if (mode === "sentinel-tokens" && !sentinels) throw new Error("sentinel-tokens routing needs the resolved sentinel ids");
     this.#decoder = new StreamDecoder(tokenizer, mode !== "buffered-text");
     this.#channelDecoder = new StreamDecoder(tokenizer, true);
   }
@@ -101,7 +99,7 @@ export class ToolAwareStream {
     return rest;
   }
 
-  /** Drain reasoning captured since the last call (gemma-channel thinking). */
+  /** Drain reasoning captured since the last call (channel-token thinking). */
   takeReasoning(): string {
     const r = this.#reasoning;
     this.#reasoning = "";
@@ -139,7 +137,7 @@ export class ToolAwareStream {
 
   /** Returns the content text delta for this token ("" while capturing). */
   push(token: number): string {
-    if (this.mode !== "gemma-sentinel") {
+    if (this.mode !== "sentinel-tokens") {
       this.#text += this.#decoder.push(token);
       if (this.mode === "plain") {
         const out = this.#text.slice(this.#sent);
@@ -149,7 +147,7 @@ export class ToolAwareStream {
       return this.#textDelta();
     }
     if (this.#inTool) {
-      if (token === TOOL_CALL_END) {
+      if (token === this.sentinels!.toolCallEnd) {
         this.#inTool = false;
         this.toolSegments.push(this.#toolTokens);
         this.#toolTokens = [];
@@ -160,11 +158,11 @@ export class ToolAwareStream {
     }
     // Reasoning channel: tokens between <|channel> and <channel|> are thought,
     // captured to #reasoning (drained via takeReasoning), never content. An
-    // empty block (<|channel>thought\n<channel|>, emitted by larger Gemmas even
+    // empty block (<|channel>thought\n<channel|>, emitted by some checkpoints even
     // with thinking off) yields only the "\n" as reasoning (mlx-lm parity) and
     // leaks nothing into content.
     if (this.#inChannel) {
-      if (token === CHANNEL_END) {
+      if (token === this.sentinels!.channelEnd) {
         this.#inChannel = false;
         this.#reasoning += this.#feedChannel(this.#channelDecoder.flush());
       } else {
@@ -172,13 +170,13 @@ export class ToolAwareStream {
       }
       return "";
     }
-    if (token === CHANNEL_START) {
+    if (token === this.sentinels!.channelStart) {
       this.#inChannel = true;
       this.#channelNameDone = false;
       this.#channelNamePending = "";
       return "";
     }
-    if (token === TOOL_CALL_START) {
+    if (token === this.sentinels!.toolCallStart) {
       this.#inTool = true;
       return "";
     }
@@ -186,7 +184,7 @@ export class ToolAwareStream {
   }
 
   flush(): string {
-    if (this.mode !== "gemma-sentinel") {
+    if (this.mode !== "sentinel-tokens") {
       this.#text += this.#decoder.flush();
       if (this.mode === "buffered-text") {
         const calls = this.toolCalls();
@@ -221,7 +219,7 @@ export class ToolAwareStream {
   }
 
   toolCalls(): OpenAIToolCall[] {
-    if (this.mode !== "gemma-sentinel") {
+    if (this.mode !== "sentinel-tokens") {
       if (this.#textToolCalls) return this.#textToolCalls;
       try {
         this.#textToolCalls = parseGeneratedToolCalls(this.#text, this.tools ?? []).map((c) => ({