mlx-bun 0.4.0 → 1.0.0

package/src/residency/resident-estimate.ts

@@ -0,0 +1,50 @@
+// What a model holds once it is resident, for residency decisions: its weights
+// plus the KV cache and prefill working set of a typical serving context, from
+// the same fit model `/fit` reports. Loading replaces the estimate with what the
+// loaded model reports (its real weights bytes and the state it carries).
+import { totalmem } from "node:os";
+import type { ModelConfig } from "@mlx-bun/inference/artifacts/config";
+import { fit, WIRED_FRACTION, type MachineSpec } from "@mlx-bun/inference/execution/fit";
+import type { KvSchemeOptions } from "@mlx-bun/inference/state/kv-scheme";
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+
+/** The context the estimate reserves KV for; longer contexts grow into the budget's headroom or evict others. */
+export const SERVING_CONTEXT_TOKENS = 8192;
+
+/** The share of the device working set all resident models may use together by default; the rest is headroom for
+ * KV growth, prefill transients and the system. */
+export const DEFAULT_BUDGET_FRACTION = 0.7;
+
+/** The default residency budget, one rule for both compositions: a share of the device working set (Metal's recommended
+ * maximum, which `cache-services` also plans against), or of the RAM the GPU may wire (`fit`'s rule) when the device does
+ * not say. The in-process host reads the working set itself; an isolated host is told it by its first worker. */
+export function defaultBudgetBytes(workingSetBytes?: number): number {
+  const workingSet = workingSetBytes !== undefined && workingSetBytes > 0 ? workingSetBytes : totalmem() * WIRED_FRACTION;
+  return Math.floor(workingSet * DEFAULT_BUDGET_FRACTION);
+}
+
+/** Bytes a model's KV and prefill working set take at the serving context. */
+export function servingReserveBytes(config: ModelConfig, weightsBytes: number,
+  options: { expertsBytes?: number; kvScheme?: KvSchemeOptions; machine?: MachineSpec } = {}): number {
+  const context = Math.max(1, Math.min(SERVING_CONTEXT_TOKENS, config.text.maxPositionEmbeddings));
+  const report = fit(config, weightsBytes, context, options.machine, undefined, options.expertsBytes ?? 0, undefined, options.kvScheme);
+  const reserve = report.kvBytes + report.transientBytes;
+  // A config the fit model cannot size adds nothing: residency then counts the weights alone.
+  return Number.isFinite(reserve) && reserve > 0 ? reserve : 0;
+}
+
+/** What loading `record` would take: its weights and a typical context's KV and working set. A runtime that plans its
+ * own memory is served alone, so it needs the whole budget. Reads the model's config only; loads nothing. */
+export async function estimateRecordBytes(record: ModelRecord, input: {
+  budgetBytes: number; machine?: MachineSpec;
+  cache: { kvQuant?: "off" | "config" | number; turboQuant?: import("@mlx-bun/inference/artifacts/config").TurboQuantScheme; quantizedKvStart?: number };
+}): Promise<number> {
+  const [{ loadModelConfig }, { resolveModelProfile }, { plansMemory }, { resolveKvScheme }] = await Promise.all([
+    import("@mlx-bun/inference/artifacts/config"), import("@mlx-bun/inference/models/profile"), import("@mlx-bun/inference/models/memory-plan"),
+    import("@mlx-bun/inference/state/kv-scheme")]);
+  const config = await loadModelConfig(record.path);
+  if (plansMemory(resolveModelProfile(config))) return input.budgetBytes;
+  const kvScheme = resolveKvScheme({ override: input.cache.kvQuant, turboQuant: input.cache.turboQuant,
+    quantizedKvStart: input.cache.quantizedKvStart, config: config.kvQuant }).fitOptions;
+  return record.sizeBytes + servingReserveBytes(config, record.sizeBytes, { expertsBytes: record.expertsBytes, kvScheme, ...(input.machine ? { machine: input.machine } : {}) });
+}

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

@@ -8,7 +8,7 @@
 //
 // Documented upgrades over the oracle (possible because our
 // chat-completions layer is natively capable where mlx-lm's gemma path
-// is not — same divergence class as the Phase 4 bidirectional-mask fix):
+// is not — same divergence class as the bidirectional-mask fix):
 //   - tools: Anthropic `tools`/`tool_use`/`tool_result` map onto our REAL
 //     OpenAI tools surface (token-level gemma tool calling) instead of
 //     the oracle's Qwen-style <tool_call> inline-text hack ("out of
@@ -22,11 +22,11 @@
 //   - prior-turn `thinking` blocks are dropped on re-ingest (Anthropic
 //     semantics) instead of the oracle's json.dumps fallback.
 
-import type { ChatMessage, ToolDefinition } from "./chat-template";
+import type { ChatMessage, ToolDefinition } from "@mlx-bun/inference/input";
 import type {
   CompletionEvent,
   CompletionStreamProtocol,
-} from "./serve/completion-sink";
+} from "./completion-sink";
 
 // ---------------------------------------------------------------------
 // Request translation: Anthropic → OpenAI (our ChatRequest shape)
@@ -545,81 +545,3 @@ export function createAnthropicStreamProtocol(model: string): CompletionStreamPr
     },
   };
 }
-
-/** Wrap our OpenAI SSE byte stream as an Anthropic SSE byte stream.
- *  Frames translate as they arrive (per-chunk flush preserved — the
- *  upstream already rate-limits socket flushes). */
-export function translateOpenAiSse(
-  upstream: ReadableStream<Uint8Array>,
-  model: string,
-): ReadableStream<Uint8Array> {
-  const translator = new AnthropicStreamTranslator(model);
-  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));
-      };
-      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]"') {
-                if (!finalized) {
-                  finalized = true;
-                  emit(translator.finalize());
-                }
-                continue;
-              }
-              const parsed = JSON.parse(data);
-              if (parsed?.error) {
-                emit([
-                  sse("error", {
-                    type: "error",
-                    error: { type: "api_error", message: parsed.error.message ?? "generation failed" },
-                  }),
-                ]);
-                continue;
-              }
-              emit(translator.addChunk(parsed));
-            }
-          }
-          if (!cancelled && !finalized) emit(translator.finalize());
-        } catch (e) {
-          if (!cancelled)
-            emit([
-              sse("error", {
-                type: "error",
-                error: { type: "api_error", message: (e as Error).message },
-              }),
-            ]);
-        } finally {
-          reader?.releaseLock();
-          reader = null;
-          if (!cancelled) controller.close();
-        }
-      })();
-    },
-    async cancel(reason) {
-      cancelled = true;
-      await reader?.cancel(reason);
-    },
-  });
-}

package/src/server/cache-routes.ts

@@ -0,0 +1,36 @@
+import type { KvCacheService } from "@mlx-bun/inference/state";
+import { errorResponse } from "./http";
+
+export type CacheRouteServices = Pick<KvCacheService, "closeSession" | "flush" | "report">;
+
+/** Administrative HTTP policy over the loaded model's cache service, preserving
+ * main's `/admin/cache/*` wire contract. Both routes borrow the service: no
+ * execution lease is taken, a session close never waits for a flush in
+ * progress, and a close neither deletes checkpoints nor rejects an unknown
+ * session. The flush answer adds the service's own report of saved state. */
+export function createCacheRoutes(cache: CacheRouteServices) {
+  return { async handle(request: Request): Promise<Response | null> {
+    const route = `${request.method} ${new URL(request.url).pathname}`;
+    if (route !== "POST /admin/cache/session/close" && route !== "POST /admin/cache/flush") return null;
+    try {
+      if (route === "POST /admin/cache/session/close") {
+        let body: unknown;
+        try {
+          const text = await request.text();
+          body = text.trim() ? JSON.parse(text) : {};
+        } catch (error) {
+          if (request.signal.aborted) return errorResponse(error, route, undefined, request.signal);
+          return Response.json({ error: { message: "invalid JSON body" } }, { status: 400 });
+        }
+        const id = body && typeof body === "object" && !Array.isArray(body) && "session_id" in body ? body.session_id : undefined;
+        if (typeof id === "string") cache.closeSession(id);
+        return Response.json({ closed: typeof id === "string" });
+      }
+      const result = await cache.flush();
+      const saved = cache.report().saved;
+      return Response.json({ ...result, entries: saved?.entries ?? 0,
+        longest_durable_prefix_tokens: saved?.longestDurablePrefixTokens ?? 0 },
+        { status: result.durable ? 200 : 503 });
+    } catch (error) { return errorResponse(error, route, undefined, request.signal); }
+  } };
+}

package/src/{serve → server}/chat-request.ts

@@ -1,12 +1,14 @@
 // Wire-level chat/completion request shapes and the pure per-request
 // helpers: message normalization, sampling-field coercion, grammar-degrade
 // prompt injection, and the mlx_lm.server-parity validators. No model, no
-// server state. Extracted from src/server.ts (repo-taming Phase 4).
-import type { ChatMessage, ToolDefinition } from "../chat-template";
-import type { HlgConfig } from "../sampler";
+// server state.
+import type { ChatMessage, ToolDefinition } from "@mlx-bun/inference/input";
+import type { HlgConfig } from "@mlx-bun/inference/sampling";
 import { RequestError } from "./pipeline";
 
 export interface ChatRequestParams {
+  session_id?: string;
+  prompt_cache_key?: string;
   messages: ChatMessage[];
   stream?: boolean;
   max_tokens?: number;
@@ -54,7 +56,9 @@ export interface ChatRequestParams {
    *  adds preserve_thinking (keep think blocks from history; template default
    *  true — better prompt-cache reuse in agent loops). */
   chat_template_kwargs?: {
-    enable_thinking?: boolean;
+    /** Absent inherits server/model defaults; a boolean overrides them.
+     * null leaves thinking to the chat template itself, bypassing both defaults. */
+    enable_thinking?: boolean | null;
     preserve_thinking?: boolean;
     [key: string]: unknown;
   };
@@ -68,7 +72,7 @@ export interface ChatRequestParams {
   /** Mounted LoRA adapter selection: "id", "a+b" (stacked), or "none". */
   adapter?: string;
   /** HLG tone-curve sampling override (per request). Snake_case wire fields,
-   *  merged over the server's --hlg-sampling config. docs/archive/hlg-sampling.md. */
+   *  merged over the server's --hlg-sampling config. `02d723a:docs/archive/hlg-sampling.md`. */
   hlg?: {
     enabled?: boolean;
     width?: number;
@@ -78,7 +82,7 @@ export interface ChatRequestParams {
   };
   /** OpenAI structured output: {type:"json_object"} | {type:"json_schema",
    *  json_schema:{name,schema,strict?}} | {type:"text"}. Enforced at the
-   *  sampler via xgrammar token-bitmasks (src/grammar.ts). L2-class (oMLX
+   *  sampler via xgrammar token-bitmasks (packages/inference/src/sampling/grammar.ts). L2-class (oMLX
    *  oracle). On compile failure, degrades to a system-prompt injection +
    *  Warning header (oMLX parity), never 500. */
   response_format?: unknown;
@@ -230,10 +234,8 @@ export function parseLogitBias(
 }
 
 /** Default seed for a request that didn't pin one. `Date.now()` alone is NOT
- *  request-unique: under `--batch N` the batch lane serves ONLY default-seed
- *  requests (explicit-seed requests use the serial mechanism, see
- *  GenerationGateway.place),
- *  so two identical prompts arriving in the same millisecond would share a seed
+ *  request-unique: concurrent requests share the continuous batch, so two
+ *  identical prompts arriving in the same millisecond would share a seed
  *  and — with per-row RNG keyed as stepKey(seed, generatedCount) — produce
  *  byte-identical completions, silently collapsing best-of-N diversity. Mix a
  *  per-process Weyl counter (golden-ratio increment, period 2^32) into the

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

@@ -4,16 +4,16 @@
 // controller, and the token→text pipeline the model needs. Every chat-shaped
 // surface (OpenAI chat, Anthropic Messages, Responses) runs this stage; the
 // wire format is not its business.
-import type { PromptCache } from "../prompt-cache";
+import type { KvCacheService } from "@mlx-bun/inference/state";
 import {
   applyGrammarDegrade,
   type ChatRequest,
   type ChatRequestParams,
 } from "./chat-request";
 import type { InferenceRequest } from "./inference-request";
-import type { ServingContext } from "./model-host";
-import type { BuiltPrompt, ModelPromptBuilder } from "./model-binding";
-import { modelPromptBuilder } from "../backends/mlx/model-serving";
+import type { LoadedModelContext as ModelContext } from "../engine/model-host";
+import type { BuiltPrompt, ModelPromptBuilder, PromptNativeWork } from "./prompt-contracts";
+import { modelPromptBuilder } from "./prompt-builder";
 import { RequestError } from "./pipeline";
 import { RequestOwnership } from "./request-plan";
 import type { RequestPrep } from "./request-prep";
@@ -21,15 +21,15 @@ import { StopMatcher, ThinkingTagSplitter } from "./token-streams";
 
 export class ChatStage {
   constructor(
-    private readonly ctx: ServingContext,
+    private readonly ctx: ModelContext,
     private readonly prep: RequestPrep,
-    private readonly promptCache: Pick<PromptCache, "peekPrefixLen">,
-    /** Admission ceiling (fit() / the GLM memory plan), resolved once. */
+    private readonly promptCache: Pick<KvCacheService, "peekPrefixLen"> & Partial<Pick<KvCacheService, "objects">>,
+    /** Admission ceiling (fit() / a runtime memory plan), resolved once. */
     private readonly contextLimit: number | null,
     /** `serve --adapter <dir>` startup default; a request's explicit
      *  `adapter` (incl. "none") wins. */
     private readonly defaultAdapter?: string,
-    private readonly preparation?: import("./preparation").PreparationExecutor,
+    private readonly preparation?: import("../engine/preparation").PreparationExecutor,
     private readonly buildPrompt: ModelPromptBuilder = modelPromptBuilder(ctx),
   ) {}
 
@@ -40,9 +40,14 @@ export class ChatStage {
         message.content.some((part) => part.type !== "text"));
     const native = media || body.response_format != null || !!body.guided_grammar ||
       !!body.guided_regex || !!body.guided_choice?.length || body.structured_outputs != null;
+    const nativeWork: PromptNativeWork = async (work) => {
+      signal?.throwIfAborted();
+      const entered = async () => { signal?.throwIfAborted(); return work(); };
+      return this.preparation ? this.preparation.run(entered, signal) : entered();
+    };
     const build = async () => {
       signal?.throwIfAborted();
-      const result = await this.runPrepared(request, requestId);
+      const result = await this.runPrepared(request, requestId, nativeWork);
       if (signal?.aborted) {
         result.plan.ownership.dispose();
         signal.throwIfAborted();
@@ -52,14 +57,15 @@ export class ChatStage {
     if (!native || !this.preparation) return build();
     let reservation = await this.preparation.reserve?.(media ? "media" : "constraint", signal);
     try {
-      const result = await this.preparation.run(build, signal);
+      const result = await build();
       if (reservation) result.plan.ownership.retain(reservation);
       reservation = undefined;
       return result;
     } finally { reservation?.dispose(); }
   }
 
-  private async runPrepared(request: ChatRequest, requestId?: string): Promise<InferenceRequest> {
+  private async runPrepared(request: ChatRequest, requestId: string | undefined,
+    nativeWork: PromptNativeWork): Promise<InferenceRequest> {
     const { ctx, prep } = this;
     let body = request.params;
     const id = requestId ?? `chatcmpl-${crypto.randomUUID()}`;
@@ -72,18 +78,18 @@ export class ChatStage {
       throw new RequestError(status, message);
     };
 
-    // Grammar-constrained decoding (src/grammar.ts). Compile BEFORE prompt
+    // Grammar-constrained decoding (packages/inference/src/sampling/grammar.ts). Compile BEFORE prompt
     // rendering: on the degrade path (compile failed but a constraint was
     // requested) inject a system message instructing valid JSON so the
     // model still best-efforts schema-conformant output (oMLX parity —
     // _compile_grammar_for_request returning None + build_json_system_prompt).
-    let grammarCtrl: import("../grammar").GrammarController | null = null;
+    let grammarCtrl: import("@mlx-bun/inference/sampling/grammar").GrammarController | null = null;
     const grammarReq =
       body.response_format != null || !!body.guided_grammar ||
       !!body.guided_regex || !!body.guided_choice?.length ||
       body.structured_outputs != null;
     if (grammarReq) {
-      const g = await prep.compileGrammarForRequest(body);
+      const g = await nativeWork(() => prep.compileGrammarForRequest(body));
       grammarCtrl = ownership.own(g.controller);
       if (!g.controller && g.degradeHint) {
         const degraded = applyGrammarDegrade(body, g.degradeHint);
@@ -94,7 +100,7 @@ export class ChatStage {
 
     let built: BuiltPrompt;
     try {
-      built = await this.buildPrompt(body, tools, ownership, prep);
+      built = await this.buildPrompt(body, tools, ownership, prep, nativeWork, this.promptCache.objects);
     } catch (e) {
       if (e instanceof RequestError) return reject(e.status, e.message);
       return reject(400, `prompt build failed: ${(e as Error).message}`);
@@ -131,13 +137,11 @@ export class ChatStage {
     // Token fast-forwarding (K3, MLX_BUN_FILL=strict): the determined-span
     // table for this request's tools. Refused for the shapes only this stage
     // can see — a compiled grammar (forced tokens are its job), media prompts
-    // (embeddings prefill / mRoPE), a mounted draft model (the spec loop is a
-    // different executor and never reads options.fill), and quantized-KV
-    // schemes (post-conversion multi-token append is unvalidated). The
+    // (embeddings prefill / mRoPE). Provider support
+    // is declared by the engine's method binding. The
     // body-level refusals live in prep.fillPlanFor.
-    let fillSession: import("../fill/fill-session").FillSession | null = null;
-    if (!grammarCtrl && !vision && !diffusionPixels && !ctx.draft &&
-        !options.kvBits && !options.kvConfig?.length && !options.turboQuant) {
+    let fillSession: import("@mlx-bun/inference/generation/fill").FillSession | null = null;
+    if (!grammarCtrl && !vision && !diffusionPixels) {
       fillSession = prep.fillPlanFor(body, tools, promptIds);
       if (fillSession) options.fill = fillSession;
     }
@@ -171,6 +175,7 @@ export class ChatStage {
         topLogprobs,
         adapterIds,
         hasVision: !!vision,
+        hasPreparedPrefixIdentity: !!vision?.prefixIdentity,
         userSeed: body.seed !== undefined,
         hasGrammar: !!grammarCtrl,
         hasDraft: !!ctx.draft,
@@ -183,7 +188,7 @@ export class ChatStage {
         router: prep.toolRouter(tools, fillSession ? () => fillSession!.noteParseFailure() : undefined),
         stopper: new StopMatcher(options.stopSequences),
         thinking: new ThinkingTagSplitter(
-          ctx.template.thinkingFormat === "think-tag",
+          ctx.template?.thinkingFormat === "think-tag",
           startInThinking,
         ),
         collectToolCalls: true,

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

@@ -1,5 +1,31 @@
-import type { CompletionClient } from "./contracts/completion";
-import type { EngineHost } from "./contracts/host";
+// The public `mlx-bun/client` entry (main's src/client.ts with its host and
+// completion contracts): an HTTP completion client over a borrowed transport,
+// and an adapter that gives an in-process request handler the same host shape.
+// It imports nothing at runtime, so a process without the MLX library can use it.
+import type { Cancellation } from "@mlx-bun/inference/contracts/portable";
+
+/** Transport-specific request/response values stay outside portable policy. */
+export interface EngineHost<Request, Response> {
+  readonly ready: Promise<void>;
+  forward(request: Request): Promise<Response>;
+  /** Stop accepting work and wait until the owned worker has exited. */
+  close(): Promise<void>;
+}
+
+/** Application-facing calls. The implementation owns tokenization, numerical
+ * method and transport; callers provide request data and consume results. */
+export interface CompletionClient<Request, Result> {
+  complete(request: Request): Promise<Result>;
+}
+
+export interface BatchCompletionClient<Request, Result> extends CompletionClient<Request, Result> {
+  /** Results preserve input order; the implementation chooses execution groups. */
+  completeBatch(requests: readonly Request[]): Promise<Result[]>;
+}
+
+export interface TaskClient<Request, Progress, Result> {
+  run(request: Request, report: (progress: Progress) => void, cancellation?: Cancellation): Promise<Result>;
+}
 
 export interface CompletionCall {
   readonly body: Readonly<Record<string, unknown>>;

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

@@ -2,14 +2,13 @@ import type {
   GenerateOptions,
   GenerateStats,
   TokenLogprobs,
-} from "../generate";
+} from "@mlx-bun/inference/generation";
 import type {
-  GenerationMechanism,
   GenerationPlacement,
   OnToken,
   RequestShape,
   Vision,
-} from "./generation-gateway";
+} from "../engine/completion";
 import {
   type CompletionEvent,
   type CompletionFinishReason,
@@ -19,29 +18,19 @@ import {
   CompletionSink,
   type CompletionToolCall,
 } from "./completion-sink";
-import { recordLane, type Lane } from "./lane-registry";
+type Lane = "batched";
+type GenerationMechanism = "continuous";
 import {
   planRequest,
   type RequestPlan,
   type RequestPlanInput,
   type RequestRejection,
 } from "./request-plan";
-import type { PromptResponseTrace } from "./prompt-response-trace";
-import { cleanupFailure } from "../engine/resources";
+import type { PromptResponseTrace } from "@mlx-bun/inference/runtime/trace";
+import { cleanupFailure } from "@mlx-bun/inference/runtime/resources";
 
-export interface CompletionEngine {
-  place(shape: RequestShape, options?: GenerateOptions): GenerationPlacement;
-  run(
-    promptIds: number[],
-    options: GenerateOptions & { stopSequences?: string[] },
-    onToken: OnToken,
-    vision: Vision | undefined,
-    shape: RequestShape,
-    placement: GenerationPlacement,
-    signal?: AbortSignal,
-    trace?: PromptResponseTrace,
-  ): Promise<GenerateStats>;
-}
+import type { CompletionEngine } from "../engine/completion";
+export type { CompletionEngine } from "../engine/completion";
 
 export interface CompletionPreparation {
   requestId: string;
@@ -119,7 +108,7 @@ export interface CompletionUsage {
   completionTokens: number;
   totalTokens: number;
   speculation?: GenerateStats["spec"];
-  /** Token fast-forwarding telemetry (serial lane, MLX_BUN_FILL=strict). */
+  /** Token fast-forwarding telemetry. */
   fill?: GenerateStats["fill"];
 }
 
@@ -191,20 +180,6 @@ class LogprobsCollector {
   }
 }
 
-function initialLane(mechanism: GenerationMechanism, shape: RequestShape): Lane {
-  if (mechanism === "continuous") return "batched";
-  return shape.hasDraft ? "serial+spec" : "serial";
-}
-
-function finalLane(
-  mechanism: GenerationMechanism,
-  lane: Lane,
-  stats: GenerateStats,
-): Lane {
-  if (mechanism === "continuous") return "batched";
-  return stats.spec ? "serial+spec" : "serial";
-}
-
 function finishReason(
   result: CompletionResult,
   stats: GenerateStats,
@@ -230,8 +205,28 @@ function combineControl(
 /** Owns one completion attempt from admission through terminal accounting.
  * Model math and scheduling policy remain in the supplied CompletionEngine. */
 export class CompletionExecutor {
+  readonly #placements = new WeakMap<PreparedCompletion, GenerationPlacement>();
   constructor(private readonly engine: CompletionEngine) {}
 
+  /** Resolve capability gaps before a transport commits streaming headers.
+   * A rejected placement releases the prepared request's owned resources. */
+  place(prepared: PreparedCompletion): GenerationPlacement {
+    const input = preparedData.get(prepared);
+    if (!input) throw new Error("prepared completion has already been executed");
+    const existing = this.#placements.get(prepared);
+    if (existing) return existing;
+    try {
+      const placement = this.engine.place(input.plan.shape, input.plan.options);
+      if (placement.shape !== input.plan.shape)
+        throw new Error("generation placement does not belong to this request shape");
+      this.#placements.set(prepared, placement);
+      return placement;
+    } catch (error) {
+      preparedData.delete(prepared);
+      return cleanupFailure(error, () => input.plan.dispose());
+    }
+  }
+
   async execute(
     prepared: PreparedCompletion,
     control: CompletionControl = {},
@@ -254,15 +249,13 @@ export class CompletionExecutor {
     try {
       if (usageProgress) control.onUsageProgress!(usageProgress);
       const closePlacement = control.trace?.begin("completion.placement");
-      const enginePlacement = this.engine.place(planned.shape, planned.options);
+      const enginePlacement = this.#placements.get(prepared) ?? this.engine.place(planned.shape, planned.options);
+      this.#placements.delete(prepared);
       closePlacement?.();
       if (enginePlacement.shape !== planned.shape)
         throw new Error("generation placement does not belong to this request shape");
       const mechanism = enginePlacement.mechanism;
-      const lane = enginePlacement.execution
-        ? (mechanism === "continuous" ? "batched" : enginePlacement.execution.method === "speculative" ? "serial+spec" : "serial")
-        : initialLane(mechanism, planned.shape);
-      recordLane(input.requestId, lane);
+      const lane: Lane = "batched";
       const placement = { mechanism, lane, shape: planned.shape };
       input.onPlacement?.(placement);
       const flowControl = input.createFlowControl?.(placement);
@@ -280,9 +273,17 @@ export class CompletionExecutor {
       control.signal?.throwIfAborted();
       planned.transferOwnership();
 
+      // Bounded diagnostic of hidden markers versus visible events. No token
+      // readbacks are added: these IDs already crossed the sampling interface.
+      let traceTokens = control.trace ? 8 : 0;
       const consumeToken: OnToken = (token, tokenLogprobs) => {
         logprobs?.push(token, tokenLogprobs);
-        const pushed = sink.push(token);
+        const route = traceTokens ? { token, index: 8 - traceTokens, events: 0 } : undefined;
+        const closeRoute = route ? control.trace!.begin("response.token_route", route) : undefined;
+        let pushed: ReturnType<CompletionSink["push"]>;
+        try { pushed = sink.push(token); if (route) route.events = pushed.events.length; }
+        finally { closeRoute?.(); }
+        if (traceTokens) traceTokens = pushed.events.length ? 0 : traceTokens - 1;
         const observed = pushed.events.length
           ? control.onEvents?.(pushed.events)
           : undefined;
@@ -310,8 +311,7 @@ export class CompletionExecutor {
       const result = sink.finish();
       if (result.events.length) await control.onEvents?.(result.events);
 
-      const resolvedLane = finalLane(mechanism, lane, stats);
-      if (resolvedLane !== lane) recordLane(input.requestId, resolvedLane);
+      const resolvedLane = lane;
       if (usageProgress) {
         usageProgress.promptTokens = stats.promptTokens;
         usageProgress.cachedTokens = stats.cachedTokens;

package/src/server/decision-body.ts

@@ -0,0 +1,55 @@
+import { gunzipSync, zstdDecompressSync } from "node:zlib";
+import { RequestError } from "./pipeline";
+
+const MAX_BYTES = 32 * 1024 * 1024;
+/** Preserve compressed and multipart bodies through model routing. The same
+ * parser finds the routing id and reads the final decision handler's input. */
+export async function readDecisionBody(request: Request): Promise<unknown> {
+  const contentType = request.headers.get("content-type") ?? "application/json";
+  const encoding = request.headers.get("content-encoding")?.trim().toLowerCase();
+  const reader = request.body?.getReader(), chunks: Uint8Array[] = [];
+  let length = 0;
+  if (reader) try {
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      length += value.length;
+      if (length > MAX_BYTES) { await reader.cancel(); throw new RequestError(413, "decision body exceeds 32 MiB"); }
+      chunks.push(value);
+    }
+  } finally { reader.releaseLock(); }
+  let bytes: Uint8Array = Buffer.concat(chunks, length);
+  try {
+    if (encoding === "gzip" || encoding === "x-gzip") bytes = gunzipSync(bytes, { maxOutputLength: MAX_BYTES });
+    else if (encoding === "zstd") bytes = zstdDecompressSync(bytes, { maxOutputLength: MAX_BYTES });
+    else if (encoding && encoding !== "identity") throw new RequestError(415, `unsupported Content-Encoding: ${encoding}`);
+  } catch (error) {
+    if (error instanceof RequestError) throw error;
+    if ((error as { code?: string }).code === "ERR_BUFFER_TOO_LARGE") throw new RequestError(413, "decompressed decision body exceeds 32 MiB");
+    throw new RequestError(400, "malformed compressed decision body");
+  }
+  try {
+    if (contentType.toLowerCase().startsWith("multipart/form-data")) {
+      const form = await new Response(new Uint8Array(bytes), { headers: { "content-type": contentType } }).formData();
+      const payload = form.get("request");
+      if (!payload) throw new Error("multipart body needs a request part");
+      const body = JSON.parse(typeof payload === "string" ? payload : await payload.text());
+      const used = new Set<string>(["request"]);
+      for (const asset of Object.values(body.assets ?? {}) as Record<string, unknown>[]) {
+        if (typeof asset?.part !== "string") continue;
+        if (used.has(asset.part)) throw new Error("duplicate image part");
+        used.add(asset.part);
+        const image = form.get(asset.part);
+        if (!(image instanceof File)) throw new Error("missing image part");
+        asset.format = asset.format ?? (image.type === "image/png" ? "png" : "jpg");
+        asset.data_base64 = Buffer.from(await image.arrayBuffer()).toString("base64");
+        delete asset.part;
+      }
+      if ([...form.keys()].some(key => !used.has(key)) || [...form.keys()].length !== used.size) throw new Error("unused or duplicate multipart part");
+      return body;
+    }
+    const json = new TextDecoder("utf-8", { fatal: true }).decode(bytes);
+    if (bytes[0] === 0xef && bytes[1] === 0xbb && bytes[2] === 0xbf) throw new Error("JSON BOM is unsupported");
+    return JSON.parse(json);
+  } catch (error) { throw new RequestError(400, error instanceof Error ? error.message : "invalid decision JSON body"); }
+}