mlx-bun 0.4.0 → 1.0.0

package/src/cli/server-entry.ts

@@ -0,0 +1,148 @@
+// The public `mlx-bun/server` entry: main's in-process `createServer` over a
+// caller-loaded model context, with main's `loadContext`. The server is the app
+// `mlx-bun serve` composes (serve.ts startApp: the persistent state, then
+// serve-host's startContextHost), so it serves the same routes, browser app,
+// Pi chat, jobs, and hub. Loading this module loads no native library and
+// starts nothing; `createServer` imports the composition when it is called.
+import type { DurabilityFlushResult } from "@mlx-bun/inference/state";
+import { releaseContext, type ContextOwnership, type LoadedModelContext } from "../engine/model-host";
+import type { ContextHost, ContextHostOptions } from "./serve-host";
+import type { ServeOptions } from "./serve-options";
+
+export { loadContext } from "../engine/model-host";
+export type { AdapterService, ContextOwnership, DraftKind, GenSamplingDefaults, LoadContextOptions, LoadedModelContext,
+  ModelContext, ModelHostSource, ServedModelInfo } from "../engine/model-host";
+export type { ModelBinding } from "../engine/model-binding";
+export type { BuiltPrompt, ModelPromptBuilder, PromptNativeWork } from "../server/prompt-contracts";
+export type { DurabilityFlushResult, DurabilitySnapshotStats } from "@mlx-bun/inference/state";
+
+/** Serving policy, `mlx-bun serve`'s options by their programmatic names. */
+export interface ServerOptions extends Partial<Pick<ServeOptions, "capacity" | "defaultGeneratedTokens" | "kvBudgetBytes" |
+  "memoryBudgetBytes" | "readOnly" | "cache" | "request" | "memoryPaths" | "chatPaths" | "storagePaths">>,
+  Omit<ContextHost, "ownership" | "owner"> {
+  /** The interface to bind. Omitted: Bun's default, every interface (main's
+   * library default); `mlx-bun serve` binds 127.0.0.1. */
+  hostname?: string;
+  /** The enforced context ceiling. Omitted: `MLX_BUN_RD_CONTEXT_LIMIT` when set, else none. */
+  contextLimit?: number | null;
+  /** The speech-to-text companion's checkpoint directory and residency; omitted,
+   * the first downloaded Whisper checkpoint is found on the first audio request. */
+  whisper?: Pick<NonNullable<ServeOptions["whisper"]>, "modelDir" | "modelId" | "idleUnloadSec" | "resident">;
+  /** Default "borrowed" (main's): the server never disposes the context. */
+  ownership?: ContextOwnership;
+}
+
+export interface ServerShutdownResult {
+  /** The drain finished: the listener stopped and every resource the server created was released. */
+  stopped: boolean;
+  /** The deadline passed first. The drain continues; nothing was released early. */
+  timedOut: boolean;
+  /** The final cache flush's result. After a timeout, the pending counters with `durable: false`. */
+  durability: DurabilityFlushResult;
+}
+
+export interface RunningServer {
+  /** The bound TCP port. */
+  readonly port: number;
+  /** Flush prompt-cache persistence while serving (main's `flushServerCacheDurability`). */
+  flush(): Promise<DurabilityFlushResult>;
+  /** Stop admission and drain (main's `shutdownServer`); see createServer. */
+  close(options?: { timeoutMs?: number }): Promise<ServerShutdownResult>;
+}
+
+/**
+ * Serve a loaded model context in this process: the full `mlx-bun serve` app
+ * (OpenAI, Anthropic, and Responses APIs, `/stats`, adapters, cache
+ * administration, the browser app, Pi web chat on `/ws/chat`, jobs, hub,
+ * memory, sessions, and publishing), over one shared continuous engine.
+ * Resolves once the listener is bound; `port` 0 binds an ephemeral port.
+ *
+ * The context comes from `loadContext` or a caller's own implementation.
+ * Without `options.binding` the built-in model classes supply the numerics; a
+ * supplied `binding` and `buildPrompt` reach the engine and routes as given,
+ * with no registry lookup, model class, or reload. Chat prompts render the
+ * context's template, and a context without one is refused, unless
+ * `buildPrompt` builds them (it then returns `probeStableLen: false`).
+ * `artifact` only describes the weights for `/fit`, `/stats`, and hub GC
+ * protection.
+ *
+ * The server applies no runtime switches and activates no expert offload; the
+ * process's runtime configuration and whatever loaded the context own those.
+ * It sets the MLX allocator limit from the context's runtime memory plan or
+ * `memoryBudgetBytes` and restores the previous limit after close. That limit
+ * is process-wide: servers running at the same time in one process share it,
+ * the last applied wins, and each close restores the value it found. Sequential
+ * servers over one context are supported. No signal handler, browser, or
+ * download is started; user stores default to HOME unless `memoryPaths`,
+ * `chatPaths`, or `storagePaths` name others.
+ *
+ * Ownership: by default the context is borrowed. Close, and any failed start
+ * (options, state, binding, caches, engine, or bind), release everything the
+ * server created and never the context, which stays usable (for another
+ * server) until the caller disposes it after close has resolved with
+ * `stopped: true`. With `ownership: "owned"` the server disposes the context
+ * exactly once, after execution drains or when startup fails, even when other
+ * cleanup fails.
+ *
+ * `close({ timeoutMs = 120_000 })` stops admission (new requests get 503),
+ * stops jobs and downloads, closes chat sessions, waits for responses in
+ * flight, then releases the engine, flushes and clears the caches, releases an
+ * owned context, and restores the allocator limit. One drain runs however often
+ * it is called; a failure rejects every call. It resolves `{ stopped: true,
+ * timedOut: false, durability }` with the final flush's result, which is
+ * `durable: false` when persistence is incomplete. If the deadline passes
+ * first it resolves `{ stopped: false, timedOut: true }` with the pending
+ * counters and `durable: false`, and the drain continues: weights are never
+ * released under running execution. Call close again (`timeoutMs: Infinity`
+ * waits without a deadline) before disposing a borrowed context.
+ */
+export async function createServer(context: LoadedModelContext, port = 0, options: ServerOptions = {}): Promise<RunningServer> {
+  const ownership = options.ownership ?? "borrowed";
+  // An owned context is released here if startup fails before the host takes
+  // it; the host owns its release, and every later cleanup, from its call on.
+  let transferred = false;
+  let app;
+  try {
+    const [{ startApp, profileContextLimit }, { startContextHost }, { createInProcessMemoryClient }] = await Promise.all([
+      import("./serve"), import("./serve-host"), import("./memory-engine"),
+    ]);
+    const serving: ContextHostOptions = { port, hostname: options.hostname ?? null, capacity: options.capacity ?? 8,
+      contextLimit: options.contextLimit === undefined ? profileContextLimit() : options.contextLimit,
+      cache: options.cache ?? {}, request: options.request ?? {},
+      defaultGeneratedTokens: options.defaultGeneratedTokens, kvBudgetBytes: options.kvBudgetBytes,
+      memoryBudgetBytes: options.memoryBudgetBytes, whisper: options.whisper };
+    const served: ContextHost = { ownership, binding: options.binding, buildPrompt: options.buildPrompt,
+      artifact: options.artifact, defaultAdapter: options.defaultAdapter };
+    // Memory synthesis gets main's task model, as in `mlx-bun serve`; nothing
+    // loads before a run. The host's drain joins the state's close (jobs,
+    // downloads, synthesis) while the engine and model are still alive.
+    app = await startApp({ port, memoryPaths: options.memoryPaths, chatPaths: options.chatPaths, readOnly: options.readOnly ?? false,
+      memoryTaskModel: () => createInProcessMemoryClient() }, options.storagePaths ?? {}, state => {
+      transferred = true;
+      return startContextHost(state, context, serving, served, { beforeDrain: () => state.close() });
+    });
+  } catch (error) {
+    if (!transferred) {
+      try { releaseContext(context, ownership); }
+      catch (failure) { throw new AggregateError([error, failure], "startup and cleanup failed"); }
+    }
+    throw error;
+  }
+  const { host } = app;
+  let drain: Promise<ServerShutdownResult> | undefined;
+  return {
+    port: host.port,
+    flush: () => host.flush(),
+    async close({ timeoutMs = 120_000 } = {}) {
+      const started = performance.now();
+      const work = drain ??= app.close().then(durability => ({ stopped: true, timedOut: false, durability }));
+      if (!Number.isFinite(timeoutMs)) return work;
+      let timer: ReturnType<typeof setTimeout> | undefined;
+      const deadline = new Promise<ServerShutdownResult>(resolve => {
+        timer = setTimeout(() => resolve({ stopped: false, timedOut: true, durability: { ...host.stats(), durable: false,
+          flushedSnapshots: 0, missingSnapshots: 0, elapsedMs: performance.now() - started } }), Math.max(0, timeoutMs));
+      });
+      try { return await Promise.race([work, deadline]); } finally { clearTimeout(timer); }
+    },
+  };
+}

package/src/cli/serving-unit.ts

@@ -0,0 +1,177 @@
+// One loaded model, served: its binding, caches, engine, and the routes bound
+// to them. A unit is what the model host holds resident (residency/model-residency.ts);
+// closing it drains its work, flushes its saved state durably, and releases the
+// context by the ownership rule. Persistent services and the listener belong to
+// the host that holds units (serve-host.ts), never to a unit.
+import { releaseContext, requireChatTemplate, type ContextOwnership, type LoadedModelContext } from "../engine/model-host";
+import type { ModelBinding } from "../engine/model-binding";
+import type { KvBudget } from "../engine/kv-budget";
+import type { UnitClosed, ResidentUnit } from "../residency/model-residency";
+import { servingReserveBytes } from "../residency/resident-estimate";
+import { fit } from "@mlx-bun/inference/execution/fit";
+import type { DurabilityFlushResult, DurabilitySnapshotStats } from "@mlx-bun/inference/state";
+import type { EventBus } from "@mlx-bun/app-core";
+import type { DisposableResource } from "@mlx-bun/inference/contracts/portable";
+import { createEngineTelemetry } from "../engine/telemetry";
+import type { TranscriptionInfo } from "../server/discovery-routes";
+import type { ModelPromptBuilder } from "../server/prompt-contracts";
+import type { ResponseHistory } from "../server/responses";
+import { resolveServingLimits, type ServeOptions } from "./serve-options";
+import type { RouteGroup } from "./serve-state";
+
+/** The serving policy one host applies to every unit: serve's options minus the loader's and the CLI process's own. */
+export type ContextHostOptions = Pick<ServeOptions, "port" | "capacity" | "contextLimit" | "defaultGeneratedTokens" |
+  "kvBudgetBytes" | "memoryBudgetBytes" | "modelBudgetBytes" | "whisper" | "cache" | "request"> & {
+  /** The interface to bind; null binds Bun's default, every interface. */
+  hostname: string | null;
+};
+
+/** What a unit serves besides its context. Nothing here is loaded or looked up. */
+export interface ContextHost {
+  /** Who releases the context; see ContextOwnership. */
+  ownership: ContextOwnership;
+  /** Replacement numerics for the context; default binds the built-in model classes. */
+  binding?: ModelBinding;
+  /** Replacement chat prompt construction; default renders the context's
+   * template. With one, a context without a template serves; the builder then
+   * returns `probeStableLen: false`. */
+  buildPrompt?: ModelPromptBuilder;
+  /** Descriptive artifact metadata for `/fit`, `/stats`, and hub GC protection;
+   * each field defaults to main's unknown value (path: the config's model directory). */
+  artifact?: { path?: string; sizeBytes?: number; expertsBytes?: number };
+  /** An adapter id already mounted on the context, used when a request names none. */
+  defaultAdapter?: string;
+  /** `/stats`' `server.owner`; default "embedded". */
+  owner?: string;
+  /** Milliseconds the loader took to load the model, published with its `model.load` event; absent for a context this host did not load. */
+  loadMs?: number;
+}
+
+/** Shared by every unit of one host. */
+export interface UnitShared {
+  readonly responses: ResponseHistory;
+  /** The `events` bus: each unit's engine publishes its request timings and samples through it. */
+  readonly events: Pick<EventBus, "publish">;
+  readonly toolApprovalsFile?: string;
+  readonly transcription: () => Promise<TranscriptionInfo | null>;
+  /** One saved-state budget across every unit's store. */
+  readonly ssdBudget?: KvBudget;
+  /** The process allocator limit the host applied; bounds each unit's optional RAM caches. */
+  readonly allocatorLimitBytes?: number;
+  /** Each close's durability result, for the host's own final report. */
+  readonly onClosed?: (result: DurabilityFlushResult) => void;
+}
+
+export interface ServingUnit extends ResidentUnit {
+  readonly context: LoadedModelContext;
+  readonly artifactPath: string;
+  /** Everything model-scoped: status, cache administration and the wire routes. */
+  readonly routes: RouteGroup;
+  /** Stop background work (idle demotion) ahead of a drain. */
+  stopBackground(): void;
+  flush(): Promise<DurabilityFlushResult>;
+  stats(): DurabilitySnapshotStats;
+  /** The engine's execution gateway, for callers that run work under its lock. */
+  readonly gateway: { acquireExecutionLease(signal?: AbortSignal): Promise<DisposableResource>; runExclusive<T>(fn: () => Promise<T>, trace?: undefined, signal?: AbortSignal): Promise<T> };
+}
+
+/** The context is released by `input.ownership` on close and on a failed
+ * start alike; everything else the unit created is released either way. A
+ * context without a chat template is refused unless `input.buildPrompt`
+ * builds its prompts. */
+export async function createServingUnit(context: LoadedModelContext, options: ContextHostOptions, input: ContextHost,
+  shared: UnitShared): Promise<ServingUnit> {
+  let cleanup: (() => void | Promise<unknown>) | undefined = () => releaseContext(context, input.ownership);
+  try {
+    const [{ modelServingBinding, createCacheServices, createAppEngine },
+      { createCompletionRoutes }, { GeneratedTokenHistory }, { createStatusRoutes }, { createCacheRoutes }, { createAdapterOperation }] = await Promise.all([
+      import("../engine"), import("../server/routes"), import("../server/generated-token-history"), import("../server/status-routes"),
+      import("../server/cache-routes"), import("../engine/adapter-operation"),
+    ]);
+    // The default prompt path renders the context's template; a supplied builder replaces it.
+    if (!input.buildPrompt) requireChatTemplate(context);
+    const binding = await modelServingBinding(context, input.binding);
+    // The loaded model's cache service, built with its budget (`--kv-budget`, `--ssd-cache-max`).
+    const caches = await createCacheServices(context, binding, { ...options.cache,
+      ...(options.kvBudgetBytes !== undefined ? { kvBudgetBytes: options.kvBudgetBytes } : {}),
+      ...(shared.allocatorLimitBytes ? { allocatorLimitBytes: shared.allocatorLimitBytes } : {}),
+      ...(shared.ssdBudget ? { ssdBudget: shared.ssdBudget } : {}) });
+    const cache = caches.cache;
+    // The persistence result of the final flush is the close's evidence.
+    let durability: DurabilityFlushResult | undefined;
+    const closeCaches = async () => {
+      const result = durability = await cache.close();
+      if (!result.durable) console.warn(`[server] cache flush incomplete: ${result.pendingSnapshots} snapshots, ${result.pendingSpills} spills, ${result.failedSpills} failed`);
+    };
+    cleanup = async () => { try { await closeCaches(); } finally { releaseContext(context, input.ownership); } };
+    // createAppEngine takes ownership even when its constructor rejects.
+    cleanup = undefined;
+    const engine = await createAppEngine(context, { capacity: options.capacity, binding, ownership: input.ownership,
+      gateway: { cache, adapterNamespace: caches.adapterNamespace },
+      beforeModelDispose: closeCaches,
+    });
+    cleanup = () => engine.close();
+    // The engine's events (request timings, batch and cache samples) come from one adapter around the completion engine.
+    const telemetry = createEngineTelemetry({ events: shared.events, model: context.modelId, capacity: options.capacity,
+      weightsBytes: context.model.weightsBytes, gateway: engine.gateway, cache });
+    engine.completion = telemetry.observe(engine.completion);
+    telemetry.sample();
+    const tokenHistory = new GeneratedTokenHistory(context.tokenizer);
+    for (const tokens of cache.savedPrefixes()) tokenHistory.remember(tokens);
+    cache.onPut = tokens => tokenHistory.remember(tokens);
+    const admission = context.memoryPlan ?? fit(context.model.config, context.model.weightsBytes, 1,
+      undefined, undefined, 0, options.memoryBudgetBytes, caches.resolvedKvScheme.fitOptions);
+    const limits = resolveServingLimits(options, context.memoryPlan, admission);
+    const completions = createCompletionRoutes(engine, { ...options.request, promptCache: cache,
+      ...limits, tokenHistory, responseHistory: shared.responses,
+      transcription: shared.transcription,
+      ...(input.buildPrompt ? { buildPrompt: input.buildPrompt } : {}),
+      ...(input.defaultAdapter ? { defaultAdapter: input.defaultAdapter } : {}) });
+    const status = createStatusRoutes({ ...(input.owner ? { owner: input.owner } : {}), context, cache,
+      resolvedKvScheme: caches.resolvedKvScheme, gateway: engine.gateway,
+      diagnostics: () => binding.diagnostics(), responseStats: completions.responseStats,
+      artifact: { expertsBytes: input.artifact?.expertsBytes ?? 0, sizeBytes: input.artifact?.sizeBytes ?? null },
+      capacity: options.capacity, contextLimit: limits.contextLimit, startedAt: Date.now(),
+      ssdCacheDir: options.cache.ssdCacheDir, memoryBudgetBytes: options.memoryBudgetBytes });
+    const cacheAdmin = createCacheRoutes(cache);
+    const adapters = createAdapterOperation(context, engine.gateway);
+    const routes: RouteGroup = { handle: async request => await status.handle(request) ?? await cacheAdmin.handle(request) ?? await completions.handle(request) };
+    const artifactPath = input.artifact?.path ?? context.model.config.modelDir;
+    // Weights plus the KV and working set of a typical context, plus the RAM prefix cache it holds now.
+    const reserve = context.memoryPlan ? 0 : servingReserveBytes(context.model.config, context.model.weightsBytes,
+      { expertsBytes: input.artifact?.expertsBytes, kvScheme: caches.resolvedKvScheme.fitOptions });
+    let closing: Promise<UnitClosed> | undefined;
+    const unit: ServingUnit = {
+      id: context.modelId, context, artifactPath, routes, gateway: engine.gateway,
+      // Saved state was found for this model: a request can resume from it.
+      resumed: (cache.report().saved?.entries ?? 0) > 0,
+      operations: ["generate", ...(binding.discovery.embeddings ? ["embed" as const] : [])],
+      bytes: () => context.memoryPlan ? context.memoryPlan.totalBytes : context.model.weightsBytes + reserve + cache.report().prefixes.bytes,
+      memory: () => {
+        const held = cache.report();
+        return { weightsBytes: context.model.weightsBytes, kvBytes: held.kv.projectedBytes, prefixCacheBytes: held.prefixes.bytes };
+      },
+      operationsFor: () => ({
+        generate: async request => await completions.handle(request) ?? Response.json({ error: { message: "Not found" } }, { status: 404 }),
+        adapters,
+      }),
+      pause: signal => engine.gateway.acquireExecutionLease(signal),
+      stopBackground: () => { telemetry.stop(); cache.stopIdleDemotion(); },
+      flush: () => cache.flush(),
+      stats: () => cache.stats(),
+      // Closes the engine (admission stops, work in flight drains), then flushes and clears the caches, then releases the context.
+      close: () => closing ??= (async () => {
+        const started = performance.now();
+        telemetry.stop();
+        await engine.close();
+        return { flushed: durability?.durable ?? true, drainMs: performance.now() - started - (durability?.elapsedMs ?? 0),
+          ...(durability ? { flushMs: durability.elapsedMs } : {}) };
+      })().finally(() => { if (durability) shared.onClosed?.(durability); }),
+    };
+    return unit;
+  } catch (error) {
+    try { await cleanup?.(); }
+    catch (failure) { throw new AggregateError([error, failure], "startup and cleanup failed"); }
+    throw error;
+  }
+}

package/src/{tui.ts → cli/terminal.ts}

@@ -40,6 +40,7 @@ const WORDMARK = `
 ██║ ╚═╝ ██║███████╗██╔╝ ██╗      ██████╔╝╚██████╔╝██║ ╚████║
 ╚═╝     ╚═╝╚══════╝╚═╝  ╚═╝      ╚═════╝  ╚═════╝ ╚═╝  ╚═══╝`.slice(1);
 
+/** Main's startup wordmark; one plain line when stdout is not a TTY. */
 export function banner(version: string, tagline = "local AI · apple silicon · one binary"): void {
   if (!tty()) {
     console.log(`mlx-bun ${version} — ${tagline}`);
@@ -86,17 +87,21 @@ export function step(text: string): Step {
   };
 }
 
-/** Rounded box with a gradient border. Lines may contain ANSI codes. */
-export function box(lines: string[], { pad = 1 }: { pad?: number } = {}): void {
+/** Rounded box with a gradient border, one output line per entry. Lines may contain ANSI codes. */
+export function boxLines(lines: string[], { pad = 1 }: { pad?: number } = {}): string[] {
   const visible = (s: string) => [...s.replace(/\x1b\[[0-9;]*m/g, "")].length;
   const width = Math.max(...lines.map(visible)) + pad * 2;
   const hue = (s: string) => (tty() ? fg(rampAt(0.15)) + s + RESET : s);
-  console.log(hue(`  ╭${"─".repeat(width)}╮`));
-  for (const l of lines) {
-    const fill = " ".repeat(width - visible(l) - pad * 2);
-    console.log(`${hue("  │")}${" ".repeat(pad)}${l}${fill}${" ".repeat(pad)}${hue("│")}`);
-  }
-  console.log(hue(`  ╰${"─".repeat(width)}╯`));
+  return [
+    hue(`  ╭${"─".repeat(width)}╮`),
+    ...lines.map((l) => `${hue("  │")}${" ".repeat(pad)}${l}${" ".repeat(width - visible(l) - pad * 2)}${" ".repeat(pad)}${hue("│")}`),
+    hue(`  ╰${"─".repeat(width)}╯`),
+  ];
+}
+
+/** Rounded box with a gradient border. Lines may contain ANSI codes. */
+export function box(lines: string[], options: { pad?: number } = {}): void {
+  for (const line of boxLines(lines, options)) console.log(line);
 }
 
 export interface Column {

package/src/cli/worker-adapters.ts

@@ -0,0 +1,34 @@
+// A worker's `adapters` operation as the isolated parent reaches it: each call is one POST on the worker's private socket
+// (`/admin/adapters`), which runs on the worker's own model under its engine lock (server/worker-routes.ts). Nothing here
+// touches native code.
+import type { AdapterMergeStats, AdapterOperation, MountedAdapter } from "@mlx-bun/app-core";
+import { ModelHostFailure } from "@mlx-bun/app-services/portable";
+import type { WorkerAdapterCall } from "../server/worker-routes";
+
+export function workerAdapters(worker: { fetch(url: string, init?: RequestInit): Promise<Response> }): AdapterOperation {
+  async function call<T>(body: WorkerAdapterCall, signal: AbortSignal | undefined): Promise<T> {
+    let response: Response;
+    try {
+      response = await worker.fetch("http://engine/admin/adapters", { method: "POST", headers: { "content-type": "application/json" },
+        body: JSON.stringify(body), ...(signal ? { signal } : {}) });
+    } catch (error) {
+      // A worker that is down, restarting or gone under the call is the model host being unavailable, not a refusal of the adapter.
+      signal?.throwIfAborted();
+      throw new ModelHostFailure("closed", error instanceof Error ? error.message : String(error), { cause: error });
+    }
+    const text = await response.text();
+    let parsed: { error?: { message?: string } } & Record<string, unknown> = {};
+    try { parsed = JSON.parse(text) as typeof parsed; } catch { /* answered with something else */ }
+    if (!response.ok) throw new Error(parsed.error?.message ?? `the model worker answered ${response.status}`);
+    return parsed as T;
+  }
+  return {
+    async list(signal) { return (await call<{ adapters: MountedAdapter[] }>({ op: "list" }, signal)).adapters; },
+    async mount(id, path, signal) { return (await call<{ adapter: MountedAdapter }>({ op: "mount", id, path }, signal)).adapter; },
+    async unmount(id, signal) { return (await call<{ removed: number }>({ op: "unmount", id }, signal)).removed; },
+    async merge(request, signal) {
+      return (await call<{ stats: AdapterMergeStats }>({ op: "merge", adapters: [request.adapters[0], request.adapters[1]], output: request.output,
+        ...(request.scales ? { scales: [...request.scales] } : {}) }, signal)).stats;
+    },
+  };
+}