mlx-bun 0.4.0 → 1.0.0

package/src/{serve → server}/http.ts

@@ -1,4 +1,5 @@
-import { AdmissionRejected } from "../engine/admission";
+import { AdmissionRejected } from "@mlx-bun/inference/execution/admission";
+import { UnsupportedExecutionError } from "../engine/completion";
 // The HTTP end of the pipeline: build + admit a request (errors → a JSON
 // error response in the surface's own shape), then run it and write the
 // result as JSON, or as protocol frames while events arrive. One writer for
@@ -11,7 +12,7 @@ import type {
 } from "./inference-request";
 import { openAiSummaryUsage, openAiUsage } from "./openai-wire";
 import { RequestError } from "./pipeline";
-import type { PromptResponseTrace } from "./prompt-response-trace";
+import type { PromptResponseTrace } from "@mlx-bun/inference/runtime/trace";
 
 export type ErrorFormatter = (status: number, message: string, body: Record<string, unknown>) => Response;
 
@@ -19,13 +20,21 @@ export type ErrorFormatter = (status: number, message: string, body: Record<stri
 export const openAiError: ErrorFormatter = (status, _message, body) =>
   Response.json({ error: body }, { status });
 
-/** RequestError → the surface's error response; anything else is a 500 with
- *  its stack logged (a 500 with no server-side trace is undebuggable). */
-export function errorResponse(e: unknown, context: string, format: ErrorFormatter = openAiError): Response {
+function errorMessage(error: unknown): string {
+  try { return error instanceof Error ? error.message : String(error); }
+  catch { return "Unknown error"; }
+}
+
+/** Request cancellation and capability gaps are expected responses, not 500s. */
+export function errorResponse(e: unknown, context: string, format: ErrorFormatter = openAiError, signal?: AbortSignal): Response {
+  if (signal?.aborted) return format(499, "request cancelled", { message: "request cancelled", type: "request_cancelled", code: "request_cancelled" });
+  if (e instanceof UnsupportedExecutionError) return format(501, e.message, {
+    message: e.message, type: "not_implemented", code: "unsupported_execution", reasons: e.reasons,
+  });
   if (e instanceof AdmissionRejected) return format(429, e.message, { message: e.message, type: "resource_admission", code: "queue_full" });
   if (e instanceof RequestError) return format(e.status, e.message, e.body);
-  console.error(`[serve] 500 on ${context}:\n${(e as Error).stack ?? e}`);
-  const message = (e as Error).message;
+  const message = errorMessage(e);
+  console.error(`[serve] 500 on ${context}:\n${e instanceof Error ? e.stack ?? message : message}`);
   return format(500, message, { message });
 }
 
@@ -38,6 +47,7 @@ export async function admit(
   trace: PromptResponseTrace | undefined,
   context: string,
   format: ErrorFormatter = openAiError,
+  signal?: AbortSignal,
 ): Promise<{ admitted: AdmittedRequest } | { response: Response }> {
   const closePrepare = trace?.begin("request.prompt_prepare");
   try {
@@ -46,8 +56,8 @@ export async function admit(
     return { admitted };
   } catch (e) {
     closePrepare?.();
-    trace?.finish("error", { stage: "prepare_completion" });
-    return { response: errorResponse(e, context, format) };
+    trace?.finish(signal?.aborted ? "abort" : "error", { stage: "prepare_completion" });
+    return { response: errorResponse(e, context, format, signal) };
   }
 }
 
@@ -65,13 +75,14 @@ export async function respondJson(
 ): Promise<Response> {
   try {
     const result = await stage.run(admitted, { signal, ...(trace ? { trace } : {}) });
+    signal.throwIfAborted();
     trace?.mark("response.final_write");
     const response = Response.json(json(result), { headers: warningHeaders(admitted) });
     trace?.finish("success");
     return response;
   } catch (e) {
     trace?.finish(signal.aborted ? "abort" : "error");
-    return errorResponse(e, `request ${admitted.requestId}`, format);
+    return errorResponse(e, `request ${admitted.requestId}`, format, signal);
   }
 }
 
@@ -107,8 +118,8 @@ export function respondStream(
         // ignored by clients and do not alter the completion event stream.
         const heartbeat = setInterval(() => emit([": keep-alive\n\n"]), heartbeatMs);
         try {
-          // The gateway owns lane selection + GPU exclusivity; this body
-          // runs per-request (concurrently in batched mode, each writing
+          // The gateway owns execution placement + GPU exclusivity; this body
+          // runs per-request (concurrently, each writing
           // its own stream — the per-row fan-out).
           emit(protocol.start());
           const result = await stage.run(admitted, {
@@ -123,13 +134,14 @@ export function respondStream(
             },
             onUsageProgress: (usage) => { latestUsage = usage; },
           });
+          generationSignal.throwIfAborted();
           emit(protocol.finish(result.finishReason, openAiSummaryUsage(result)));
           trace?.mark("response.final_write");
         } catch (e) {
           outcome = generationSignal.aborted ? "abort" : "error";
           if (!generationSignal.aborted) {
             emit([
-              ...protocol.error((e as Error).message),
+              ...protocol.error(errorMessage(e)),
               ...protocol.finish("stop", latestUsage ? openAiUsage(latestUsage) : {}),
             ]);
           }

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

@@ -2,7 +2,7 @@
 // resolved generation settings, ready for a model. InferenceStage first
 // ADMITS it (memory ceiling → reject or clamp; the resolved settings become
 // an opaque, single-use plan) and then RUNS it — the gateway routes it to
-// the lane it belongs to (serial / batch / spec) and the executor returns
+// continuous scheduling with the selected generation method and the executor returns
 // the terminal InferenceResult. Every completion surface — chat, raw text,
 // Anthropic, Responses — produces one of these and reads one of those.
 import {
@@ -14,7 +14,7 @@ import {
   type PreparedCompletion,
 } from "./completion-executor";
 import { createTimedFlowControl } from "./completion-sink";
-import { runtimeValue } from "../runtime-config";
+import { runtimeValue } from "@mlx-bun/inference/runtime/config";
 import { RequestError, type RunControl } from "./pipeline";
 
 export interface InferenceRequest
@@ -37,7 +37,7 @@ export interface AdmittedRequest {
 export type InferenceResult = CompletionSummary;
 
 export class InferenceStage {
-  constructor(private readonly executor: Pick<CompletionExecutor, "execute">) {}
+  constructor(private readonly executor: Pick<CompletionExecutor, "place" | "execute">) {}
 
   /** Admission: throws RequestError(400, memory_admission) when the prompt
    *  leaves no room to generate inside the safe context; otherwise clamps
@@ -49,8 +49,7 @@ export class InferenceStage {
         ...preparation,
         ...(stream
           ? {
-              createFlowControl: ({ mechanism }: { mechanism: string }) =>
-                createTimedFlowControl(mechanism === "serial"),
+              createFlowControl: () => createTimedFlowControl(false),
             }
           : {}),
         onPlacement: ({ mechanism, shape }) => {
@@ -61,6 +60,7 @@ export class InferenceStage {
             );
         },
       });
+      this.executor.place(prepared);
       return { requestId: request.requestId, warnings, prepared };
     } catch (error) {
       if (!(error instanceof CompletionRejected)) throw error;

package/src/server/job-routes.ts

@@ -0,0 +1,20 @@
+import type { JobStore } from "../jobs/db";
+import { streamJobResponse } from "../jobs/sse";
+
+export function createJobRoutes(host: { ensureStore(): JobStore; signal: AbortSignal }) {
+  return { async handle(request: Request): Promise<Response | null> {
+    const url = new URL(request.url);
+    if (request.method !== "GET") return null;
+    if (url.pathname === "/api/jobs") {
+      const limit = Number(url.searchParams.get("limit") ?? "50");
+      const kind = url.searchParams.get("kind") ?? undefined;
+      return Response.json({ ok: true, jobs: host.ensureStore().recent(limit, kind) });
+    }
+    const match = url.pathname.match(/^\/api\/jobs\/([^/]+?)(\/stream)?$/);
+    if (!match) return null;
+    const store = host.ensureStore(), jobId = match[1]!;
+    if (match[2]) return streamJobResponse(store, jobId, AbortSignal.any([host.signal, request.signal]));
+    const job = store.get(jobId);
+    return job ? Response.json({ ok: true, job }) : Response.json({ ok: false, error: "job not found" }, { status: 404 });
+  } };
+}

package/src/server/media-prompt.ts

@@ -0,0 +1,63 @@
+import { normalizeMessages, type ChatRequestParams } from "./chat-request";
+import { requireChatTemplate, type LoadedModelContext as ModelContext } from "../engine/model-host";
+import { mediaRefusalMessage } from "../engine/media-preparation";
+import { RequestError } from "./pipeline";
+import type { RequestOwnership } from "./request-plan";
+import type { RequestPrep } from "./request-prep";
+import type { BuiltPrompt, PromptNativeWork } from "./prompt-contracts";
+
+export async function buildModelPrompt(
+    ctx: ModelContext,
+    prep: RequestPrep,
+    body: ChatRequestParams,
+    tools: ChatRequestParams["tools"] | null,
+    _ownership: RequestOwnership,
+    nativeWork: PromptNativeWork = (work) => work(),
+    objects?: import("@mlx-bun/inference/contracts/portable").ObjectCache<import("@mlx-bun/inference/state").CheckpointAttachment[]>,
+  ): Promise<BuiltPrompt> {
+    requireChatTemplate(ctx);
+    const toolList = tools ?? null;
+    const partOf = (types: string[]) => body.messages.some(
+      (m) => Array.isArray(m.content) &&
+        m.content.some((p: any) => types.includes(p.type)),
+    );
+    const hasImages = partOf(["image_url", "image"]);
+    // Same shapes extractAudio accepts: OpenAI-canonical input_audio plus
+    // optiq's audio / audio_url aliases (`02d723a:docs/design/generic-model-support.md` §6.6).
+    const hasAudio = partOf(["input_audio", "audio", "audio_url"]);
+    // Video content parts (video_url / video with base64 data) — accepted only
+    // by graphs whose media route declares video (decoded to sampled frames via
+    // the AVFoundation sidecar, packages/inference/src/input/vision/video-frames.ts).
+    const hasVideos = partOf(["video_url", "video"]);
+
+    // Video never composes with audio — one early guard so no downstream
+    // branch can silently drop a video part.
+    if (hasVideos && (hasAudio || !ctx.media.video))
+      throw new RequestError(400, mediaRefusalMessage(hasAudio ? "video-with-audio" : "video-unsupported", ctx.modelId));
+    if (hasImages || hasAudio || hasVideos) {
+      const prepared = await ctx.media.prepare({
+        messages: normalizeMessages(body.messages),
+        parts: { images: hasImages, audio: hasAudio, videos: hasVideos },
+        tools: toolList,
+        templateOptions: () => prep.templateOptionsFor(body, toolList),
+        nativeWork,
+        objects,
+      });
+      if ("rejected" in prepared) throw new RequestError(400, mediaRefusalMessage(prepared.rejected, ctx.modelId));
+      return {
+        promptIds: prepared.promptIds,
+        vision: prepared.vision,
+        startInThinking: false,
+        probeStableLen: false,
+        diffusionPixels: prepared.pixels ?? null,
+      };
+    }
+    const text = prep.promptIdsFor(body, toolList);
+    return {
+      promptIds: text.ids,
+      vision: undefined,
+      startInThinking: text.startInThinking,
+      probeStableLen: true,
+      diffusionPixels: null,
+    };
+  }

package/src/server/memory-completion-client.ts

@@ -0,0 +1,46 @@
+import type { TaskCompletionClient as MemoryCompletionClient } from "@mlx-bun/app-core";
+import type { WorkerMemoryCall } from "./worker-routes";
+export interface MemoryTarget { worker: { fetch(url: string, init?: RequestInit): Promise<Response> }; snapshot: string }
+/** The isolated parent's client for one synthesis run: each `complete` or
+ * `completeBatch` is one POST of the rows and the selected snapshot to the
+ * private `/admin/memory/complete` of the worker `select` returns (the
+ * model worker), answered with the ordered raw outputs. A failed
+ * selection fails the call as it is. The worker's task model runs the call
+ * under the worker's own execution lease; none is taken here (one
+ * would wait on that worker). The run's signal aborts the request, and the
+ * worker then aborts and joins every row. A call is never retried: a worker
+ * that stops mid-call fails it rather than replay a POST. */
+export function createWorkerMemoryClient(select: (signal: AbortSignal) => Promise<MemoryTarget>, signal: AbortSignal): MemoryCompletionClient {
+  const send = async (call: Omit<WorkerMemoryCall, "snapshot">): Promise<string[]> => {
+    signal.throwIfAborted();
+    const { worker, snapshot } = await select(signal);
+    const body: WorkerMemoryCall = { call: call.call, snapshot, requests: call.requests };
+    let response: Response;
+    try {
+      response = await worker.fetch("http://engine/admin/memory/complete", { method: "POST",
+        headers: { "content-type": "application/json" }, body: JSON.stringify(body), signal });
+    } catch (error) {
+      if (signal.aborted) throw signal.reason;
+      throw new Error(`memory: the model worker did not complete the task model ${body.call} (${error instanceof Error ? error.message : String(error)}); it is not retried`);
+    }
+    if (!response.ok) {
+      const detail = await response.text().catch(() => "");
+      let message = detail;
+      try {
+        const parsed = JSON.parse(detail) as { error?: { message?: unknown } } | null;
+        if (typeof parsed?.error?.message === "string") message = parsed.error.message;
+      } catch { /* not JSON: the text itself */ }
+      throw new Error(`memory: the model worker's task model ${body.call} failed (${response.status}): ${message.slice(0, 400)}`);
+    }
+    const { outputs } = await response.json() as { outputs?: unknown };
+    signal.throwIfAborted();
+    if (!Array.isArray(outputs) || outputs.length !== body.requests.length || !outputs.every(output => typeof output === "string"))
+      throw new Error(`memory: the model worker's task model ${body.call} answered no ordered outputs`);
+    return outputs as string[];
+  };
+  return {
+    complete: async request => (await send({ call: "complete", requests: [request] }))[0]!,
+    // An empty batch is answered here, as the task model answers it: without a call.
+    completeBatch: async requests => requests.length ? send({ call: "completeBatch", requests: [...requests] }) : [],
+  };
+}

package/src/server/model-routes.ts

@@ -0,0 +1,192 @@
+// Routes model-scoped requests to the model that should answer them. The
+// inference endpoints carry a `model` id in the JSON body: an exact local id
+// (the ones `/v1/models` lists) leases that model from the model host, loading
+// it (and draining another) when it is not resident; anything else (no id, an
+// alias such as Pi's `local`, a name another server would know) is answered by
+// the current model. Every other model-scoped path (`/stats`, `/fit`, cache
+// administration) belongs to the current model too. The response
+// body holds the lease until it ends, so a model is never released under a
+// stream. `/v1/models` lists every local model with which are
+// resident and which is current; `/stats` adds the host's residency.
+import type { ResidencyHost, ResidentUnit } from "../residency/model-residency";
+import { ResidencyError } from "../residency/model-residency";
+import { errorResponse } from "./http";
+import { RequestError } from "./pipeline";
+import { readDecisionBody } from "./decision-body";
+
+/** Inference endpoints whose JSON body carries the routing `model` field. */
+export const MODEL_ROUTED = new Set(["/v1/chat/completions", "/v1/completions", "/v1/messages", "/v1/responses", "/v1/embeddings", "/v1/decisions", "/v1/systemone"]);
+
+const decoder = new TextDecoder();
+/** The routing key from a buffered JSON body. A body that is not a JSON object with a string `model` names no model; the current model answers it and its own parser rejects a malformed body. */
+export function modelField(body: Uint8Array | string): string | undefined {
+  try {
+    const value = JSON.parse(typeof body === "string" ? body : decoder.decode(body)) as { model?: unknown } | null;
+    return value && typeof value === "object" && typeof value.model === "string" ? value.model : undefined;
+  } catch { return undefined; }
+}
+
+/** A unit whose model-scoped routes this router can delegate to. */
+export interface RoutedUnit extends ResidentUnit {
+  readonly routes: { handle(request: Request): Promise<Response | null> };
+}
+
+export interface ModelRoutesOptions<U extends RoutedUnit> {
+  host: Pick<ResidencyHost<U>, "acquire" | "peek" | "resident" | "policy">;
+  /** The model that answers requests naming no local model. */
+  current(): string;
+  /** Whether `id` is an exact id this host can serve (the ids `/v1/models` lists), resident or not. */
+  serves(id: string): Promise<boolean>;
+  /** The companion model (Whisper) that answers this path, when it is one of its own; the router leases it like any model. */
+  companion?(pathname: string): Promise<string | undefined>;
+  /** The configured transcription companion, including before its worker starts. */
+  transcription?(): Promise<string | undefined>;
+}
+
+/** A response that keeps its lease until the body has ended, been cancelled, or failed. */
+function holding(response: Response, release: () => void): Response {
+  if (!response.body) { release(); return response; }
+  const reader = response.body.getReader();
+  let done = false;
+  const finish = () => { if (!done) { done = true; release(); } };
+  const body = new ReadableStream<Uint8Array>({
+    async pull(controller) {
+      try {
+        const chunk = await reader.read();
+        if (chunk.done) { finish(); controller.close(); } else controller.enqueue(chunk.value);
+      } catch (error) { finish(); controller.error(error); }
+    },
+    cancel(reason) { finish(); return reader.cancel(reason).catch(() => {}); },
+  });
+  return new Response(body, { status: response.status, statusText: response.statusText, headers: response.headers });
+}
+
+/** The typed answer for a model that could not be leased. */
+function unavailable(error: unknown, context: string, signal: AbortSignal): Response {
+  if (error instanceof ResidencyError) {
+    const status = error.code === "load-failed" ? 502 : error.code === "closed" ? 503 : 400;
+    const type = error.code === "load-failed" ? "model_load_failed" : error.code === "closed" ? "server_closing" : "invalid_request_error";
+    return errorResponse(new RequestError(status, error.message, { message: error.message, type, code: error.code }), context, undefined, signal);
+  }
+  return errorResponse(error, context, undefined, signal);
+}
+
+export function createModelRoutes<U extends RoutedUnit>(options: ModelRoutesOptions<U>) {
+  const { host } = options;
+  /** The unit that answers model-scoped requests that name no model, without loading anything. */
+  const active = (): U | undefined => {
+    const current = host.peek(options.current());
+    if (current) return current;
+    const latest = host.resident().filter(model => model.state === "ready" && model.role === "primary").sort((a, b) => b.lastUsedAt - a.lastUsedAt)[0];
+    return latest ? host.peek(latest.id) : undefined;
+  };
+  const listing = async (request: Request, path: string): Promise<Response | null> => {
+    const unit = active();
+    return unit ? await unit.routes.handle(new Request(new URL(path, request.url), { headers: request.headers, signal: request.signal })) : null;
+  };
+  const residentIds = () => new Set(host.resident().filter(model => model.state === "ready").map(model => model.id));
+
+  /** Lease `id`, run `serve` on its unit, and hold the lease until the response has been delivered. */
+  async function leased(id: string, request: Request, serve: (unit: U, lease: { release(): void }) => Promise<Response | null>): Promise<Response | null> {
+    const path = new URL(request.url).pathname;
+    let lease;
+    try { lease = await host.acquire(id, { signal: request.signal }); }
+    catch (error) { return unavailable(error, path, request.signal); }
+    try {
+      const response = await serve(lease.unit, lease);
+      return response ? holding(response, () => lease.release()) : (lease.release(), null);
+    } catch (error) { lease.release(); return errorResponse(error, path, undefined, request.signal); }
+  }
+
+  return {
+    async handle(request: Request): Promise<Response | null> {
+      const url = new URL(request.url);
+      const { pathname } = url;
+      if (request.method === "POST" && MODEL_ROUTED.has(pathname)) {
+        // The body is read once here and handed on as bytes; the model's own parser answers a malformed one.
+        const headers = new Headers(request.headers);
+        const body = new Uint8Array(await request.arrayBuffer());
+        let named = modelField(body);
+        if (pathname === "/v1/systemone" || pathname === "/v1/decisions") {
+          try {
+            const parsed = await readDecisionBody(new Request(request.url, { method: "POST", headers, body }));
+            if (parsed && typeof parsed === "object" && typeof (parsed as { model?: unknown }).model === "string")
+              named = (parsed as { model: string }).model;
+          } catch { /* The model's parser supplies the protocol's error. */ }
+        }
+        const id = named !== undefined && await options.serves(named) ? named : options.current();
+        return leased(id, request, unit => unit.operationsFor().generate!(new Request(request.url, {
+          method: request.method, headers, body, signal: request.signal })));
+      }
+      const listsModels = pathname === "/v1/models" || pathname.startsWith("/v1/models/");
+      if (request.method === "GET" && listsModels) {
+        // `/v1/models/{id}` lists that one model; `/v1/models` lists them all.
+        const filter = decodeURIComponent(pathname.slice("/v1/models/".length)) || null;
+        const own = filter === null ? undefined : options.host.peek(filter);
+        if (own && !own.operations.includes("transcribe")) return own.routes.handle(request);
+        const base = await listing(request, "/v1/models");
+        if (!base) return null;
+        // A unit that cannot answer (a worker that is down) is the answer, not something to decorate.
+        if (!base.ok) return base;
+        const list = await base.json() as { object: string; data: Record<string, unknown>[] };
+        // Chat workers have no companion of their own. Discovery belongs to this host's
+        // resolved companion, not whichever Whisper snapshot a worker discovers locally.
+        if (options.transcription) {
+          const id = await options.transcription();
+          list.data = list.data.filter(entry => entry.transcription !== true && entry.id !== id);
+          for (const entry of list.data) {
+            if (entry.id === options.current() && entry.capabilities)
+              entry.capabilities = { ...(entry.capabilities as Record<string, unknown>), transcription: id !== undefined };
+          }
+          if (id !== undefined) list.data.push({ id, object: "model", owned_by: "mlx-bun", transcription: true,
+            resident: false, capabilities: { transcription: true, translation: true, chat_completions: false } });
+        }
+        const resident = residentIds(), current = options.current();
+        // Every resident model answers for itself; the rest are listed as the registry knows them.
+        const full = await Promise.all([...resident].map(async id => {
+          const unit = host.peek(id);
+          const companion = unit?.operations.includes("transcribe");
+          const path = companion ? "/v1/models" : `/v1/models/${encodeURIComponent(id)}`;
+          const entry = unit ? await unit.routes.handle(new Request(new URL(path, request.url))) : null;
+          const rows = entry?.ok ? ((await entry.json()) as { data?: Record<string, unknown>[] }).data : undefined;
+          const row = companion ? rows?.find(row => row.transcription === true) : rows?.[0];
+          return [id, row ? { ...row, id } : undefined] as const;
+        }));
+        const bySelf = new Map(full.flatMap(([id, entry]) => entry ? [[id, entry] as const] : []));
+        const data: Record<string, unknown>[] = list.data.map(entry => {
+          const id = entry.id as string, row = bySelf.get(id) ?? entry;
+          // A companion's own row says whether it is loaded; the host knows when it holds it resident.
+          return entry.transcription === true ? row : { ...row, resident: resident.has(id), current: id === current };
+        });
+        const selected = filter === null ? data : data.filter(entry => entry.id === filter);
+        return Response.json({ ...list, data: selected,
+          // Jev clients use `models`, OpenAI clients use `data`; one discovery
+          // response can carry both views without changing either contract.
+          models: selected.filter(entry => entry.transcription !== true).map(entry => ({ name: entry.id,
+            description: "Local mlx-bun inference model", release_date: "" })) });
+      }
+      if (request.method === "GET" && pathname === "/stats") {
+        const unit = active();
+        const base = unit ? await unit.routes.handle(request) : null;
+        if (!base) return null;
+        if (!base.ok) return base;
+        const stats = await base.json() as Record<string, unknown>;
+        const models = host.resident();
+        return Response.json({ ...stats, models: { current: options.current(), budget_bytes: host.policy.budgetBytes,
+          resident_bytes: models.reduce((sum, model) => sum + model.bytes, 0),
+          resident: models.map(model => ({ id: model.id, state: model.state, bytes: model.bytes, pinned: model.pinned, leases: model.leases,
+            last_used_at: model.lastUsedAt })) } });
+      }
+      if (request.method === "GET" && pathname === "/health")
+        return await active()?.routes.handle(request) ?? Response.json({ status: "ok" });
+      const companion = await options.companion?.(pathname);
+      if (companion !== undefined) return leased(companion, request, served => served.routes.handle(request));
+      // Everything else model-scoped is the current model's, when one is resident: a read answers as it is, and
+      // a change (flushing a cache) holds the model resident until it is done.
+      const unit = active();
+      if (!unit) return null;
+      if (["GET", "HEAD"].includes(request.method)) return unit.routes.handle(request);
+      return leased(unit.id, request, served => served.routes.handle(request));
+    },
+  };
+}

package/src/{serve → server}/openai-wire.ts

@@ -1,6 +1,7 @@
 // OpenAI wire formats: InferenceResult → chat.completion / text_completion
 // JSON, and the SSE chunk protocols for both. Anthropic and Responses have
-// their own (src/anthropic.ts, src/responses.ts); all three implement the
+// their own (apps/mlx-bun/src/server/anthropic.ts, apps/mlx-bun/src/server/responses.ts);
+// all three implement the
 // same CompletionStreamProtocol so the HTTP writer (http.ts) is shared.
 import type { CompletionSummary, CompletionUsage } from "./completion-executor";
 import type { CompletionEvent, CompletionStreamProtocol } from "./completion-sink";
@@ -22,7 +23,7 @@ export const openAiUsage = (usage: CompletionUsage) => ({
   ...(usage.fill ? { fill: usage.fill } : {}),
 });
 
-/** Terminal usage: the wire block plus mlx-bun's `lane` (batch|serial|…). */
+/** Terminal usage: the wire block plus mlx-bun's `lane` (always `batched`: one request is a batch of one). */
 export const openAiSummaryUsage = (summary: CompletionSummary) => ({
   ...openAiUsage(summary.usage),
   lane: summary.lane,

package/src/{serve → server}/pipeline.ts

@@ -8,7 +8,7 @@
 // A stage knows only the interface it reads and the one it writes.
 import type { CompletionEvent } from "./completion-sink";
 import type { CompletionUsage } from "./completion-executor";
-import type { PromptResponseTrace } from "./prompt-response-trace";
+import type { PromptResponseTrace } from "@mlx-bun/inference/runtime/trace";
 
 /** A request a stage refuses. `status` is the HTTP status the wire adapter
  *  answers with; `body` the protocol-neutral error object (adapters reshape

package/src/server/prompt-builder.ts

@@ -0,0 +1,13 @@
+import type { LoadedModelContext as ModelContext } from "../engine/model-host";
+import type { ModelPromptBuilder } from "./prompt-contracts";
+export function modelPromptBuilder(context: ModelContext): ModelPromptBuilder {
+  return async (body, tools, ownership, prep, nativeWork, objects) => {
+    const media = body.messages.some(message => Array.isArray(message.content) && message.content.some(part => part.type !== "text"));
+    if (media) {
+      const { buildModelPrompt } = await import("./media-prompt");
+      return buildModelPrompt(context, prep, body, tools, ownership, nativeWork, objects);
+    }
+    const text = prep.promptIdsFor(body, tools ?? null);
+    return { promptIds: text.ids, vision: undefined, startInThinking: text.startInThinking, probeStableLen: true, diffusionPixels: null };
+  };
+}

package/src/server/prompt-contracts.ts

@@ -0,0 +1,17 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { ObjectCache } from "@mlx-bun/inference/contracts/portable";
+import type { CheckpointAttachment } from "@mlx-bun/inference/state";
+import type { Vision } from "../engine/completion";
+import type { ChatRequestParams } from "./chat-request";
+import type { RequestOwnership } from "./request-plan";
+import type { RequestPrep } from "./request-prep";
+export interface BuiltPrompt {
+  promptIds: number[]; vision: Vision | undefined; startInThinking: boolean;
+  probeStableLen: boolean; diffusionPixels: MlxArray | null;
+}
+export type PromptNativeWork = <T>(work: () => Promise<T>) => Promise<T>;
+export interface ModelPromptBuilder {
+  (body: ChatRequestParams, tools: ChatRequestParams["tools"] | null,
+    ownership: RequestOwnership, prep: RequestPrep, nativeWork?: PromptNativeWork,
+    objects?: ObjectCache<CheckpointAttachment[]>): Promise<BuiltPrompt>;
+}