mlx-bun 0.4.0 → 1.0.0

package/src/{responses.ts → server/responses.ts}

@@ -1,5 +1,5 @@
 // OpenAI Responses API translation layer + response store for the
-// mlx-bun server (Phase 11).
+// mlx-bun server.
 //
 // Ported from: mlx-optiq optiq/responses_shim.py (MIT) —
 // responses_to_openai_body / openai_to_responses_response /
@@ -20,7 +20,7 @@
 import type {
   CompletionEvent,
   CompletionStreamProtocol,
-} from "./serve/completion-sink";
+} from "./completion-sink";
 
 // ---------------------------------------------------------------------
 // Response store (previous_response_id resumption)
@@ -815,6 +815,7 @@ export function createResponsesStreamProtocol(
 ): CompletionStreamProtocol {
   const translator = new ResponsesStreamTranslator(model, previousResponseId);
   const chunk = (delta: Record<string, unknown>) => ({ choices: [{ delta }] });
+  let failure: string | undefined;
   return {
     start: () => translator.addChunk(chunk({ role: "assistant", content: "" })),
     addEvents(events: CompletionEvent[]) {
@@ -831,14 +832,19 @@ export function createResponsesStreamProtocol(
       });
     },
     finish(reason, usage) {
-      const frames = [
-        ...translator.addChunk({ choices: [{ delta: {}, finish_reason: reason }], usage }),
-        ...translator.finalize(),
-      ];
+      const frames = translator.addChunk({ choices: [{ delta: {}, finish_reason: reason }], usage });
+      // The shared sink finishes after reporting an error. Give clients a
+      // terminal response without saving failed output as resumable history.
+      if (failure !== undefined) return [...frames, sse("response.failed", {
+        type: "response.failed", response: { ...translator.finalResponse(), status: "failed",
+          error: { code: "server_error", message: failure } },
+      })];
+      frames.push(...translator.finalize());
       onComplete(translator.finalResponse());
       return frames;
     },
     error(message) {
+      failure = message;
       return [sse("error", {
         type: "error",
         code: "server_error",
@@ -848,90 +854,3 @@ export function createResponsesStreamProtocol(
     },
   };
 }
-
-/** Wrap our OpenAI SSE byte stream as a Responses SSE byte stream.
- *  onComplete fires after response.completed with the final payload so
- *  the route can store it for previous_response_id resumption. */
-export function translateOpenAiSseToResponses(
-  upstream: ReadableStream<Uint8Array>,
-  model: string,
-  previousResponseId: string | null,
-  onComplete: (final: Record<string, unknown>) => void,
-): ReadableStream<Uint8Array> {
-  const translator = new ResponsesStreamTranslator(model, previousResponseId);
-  const enc = new TextEncoder();
-  const dec = new TextDecoder();
-  let buf = "";
-  let finalized = false;
-  let reader: ReadableStreamDefaultReader<Uint8Array> | null = null;
-  let cancelled = false;
-  return new ReadableStream<Uint8Array>({
-    start(controller) {
-      const emit = (frames: string[]) => {
-        for (const f of frames) controller.enqueue(enc.encode(f));
-      };
-      const finish = () => {
-        if (finalized) return;
-        finalized = true;
-        emit(translator.finalize());
-        onComplete(translator.finalResponse());
-      };
-      reader = upstream.getReader();
-      void (async () => {
-        try {
-          while (!cancelled) {
-            const { done, value } = await reader!.read();
-            if (done) break;
-            buf += dec.decode(value, { stream: true });
-            let sep: number;
-            while ((sep = buf.indexOf("\n\n")) !== -1) {
-              const frame = buf.slice(0, sep);
-              buf = buf.slice(sep + 2);
-              const data = frame
-                .split("\n")
-                .find((l) => l.startsWith("data: "))
-                ?.slice(6);
-              if (data == null) continue;
-              if (data === "[DONE]" || data === '"[DONE]"') {
-                finish();
-                continue;
-              }
-              const parsed = JSON.parse(data);
-              if (parsed?.error) {
-                emit([
-                  sse("error", {
-                    type: "error",
-                    code: "server_error",
-                    message: parsed.error.message ?? "generation failed",
-                    param: null,
-                  }),
-                ]);
-                continue;
-              }
-              emit(translator.addChunk(parsed));
-            }
-          }
-          if (!cancelled) finish();
-        } catch (e) {
-          if (!cancelled)
-            emit([
-              sse("error", {
-                type: "error",
-                code: "server_error",
-                message: (e as Error).message,
-                param: null,
-              }),
-            ]);
-        } finally {
-          reader?.releaseLock();
-          reader = null;
-          if (!cancelled) controller.close();
-        }
-      })();
-    },
-    async cancel(reason) {
-      cancelled = true;
-      await reader?.cancel(reason);
-    },
-  });
-}

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 };
+}