mlx-bun 0.4.0 → 1.0.0

package/src/cli/inference.ts

@@ -0,0 +1,214 @@
+import { existsSync } from "node:fs";
+import { basename, join, resolve } from "node:path";
+import type { Registry } from "@mlx-bun/hub/registry";
+import type { KvQuantOverride } from "@mlx-bun/inference/state/kv-scheme";
+import { runtimeValue } from "@mlx-bun/inference/runtime/config";
+import { parseTurboQuantScheme, type TurboQuantScheme } from "@mlx-bun/inference/artifacts/config";
+import type { GenerateOptions } from "@mlx-bun/inference/generation";
+import type { LoadContextOptions, LoadedModelContext } 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 { planRequest, RequestOwnership } from "../server/request-plan";
+import { textPrompt } from "../server/text-prompt";
+import type { CommandArgs } from "./args";
+import { numericalPolicy } from "./numerical-policy";
+import { isEmbeddingModelType } from "@mlx-bun/inference/models/support";
+import { openRegistry } from "../storage/paths";
+
+type InferenceCommand = "generate" | "embed";
+interface SelectedModel { path: string; repoId: string; }
+type ModelRegistry = Pick<Registry, "resolve" | "list" | "scan" | "close">;
+
+/** One-shot selection has no starter or background downloads. Generate keeps
+ * main's direct registry resolution; embed alone scans an empty cache. */
+export async function resolveInferenceModel(command: InferenceCommand, query: string,
+  registry: () => ModelRegistry = () => openRegistry()): Promise<SelectedModel> {
+  if (query && existsSync(join(query, "config.json"))) {
+    const path = resolve(query), hf = /models--([^/]+)--([^/]+)\/snapshots\//.exec(path);
+    return { path, repoId: hf ? `${hf[1]}/${hf[2]}` : basename(path) };
+  }
+  const reg = registry();
+  try {
+    if (command === "generate") return reg.resolve(query);
+    if (reg.list().length === 0) await reg.scan();
+    if (query) return reg.resolve(query);
+    const model = reg.list().find(model => isEmbeddingModelType(model.modelType));
+    if (!model) throw new Error("no embedding model downloaded — try: mlx-bun get mlx-community/Qwen3-Embedding-4B-4bit-DWQ");
+    return model;
+  } finally { reg.close(); }
+}
+
+export interface OneShotEngine {
+  context: LoadedModelContext;
+  completion: CompletionEngine;
+  binding: Pick<ModelBinding, "embed">;
+  gateway: Pick<GenerationGateway, "runExclusive">;
+  close(): Promise<void>;
+}
+export interface InferenceDependencies {
+  resolve: typeof resolveInferenceModel;
+  /** `composition`: the command's KV flags. Adapters mount through the loaded context. */
+  load(model: SelectedModel, maxTokens?: number, composition?: Pick<LoadContextOptions, "kv">): Promise<LoadedModelContext>;
+  /** Takes ownership of the context even if construction rejects. */
+  engine(context: LoadedModelContext): Promise<OneShotEngine>;
+  write(text: string): void;
+  stdin(signal?: AbortSignal): Promise<string>;
+  stdinIsTTY(): boolean;
+}
+const defaults: InferenceDependencies = {
+  resolve: resolveInferenceModel,
+  async load(model, maxTokens, composition) {
+    // The startup step: compile any Neural Engine program the system cache lacks before the model loads.
+    await (await import("./ane-install")).ensureAneInstall();
+    const { loadContext } = await import("../engine/model-host");
+    // One-shot commands run one row (the engine below has capacity 1).
+    return loadContext(model.path, model.repoId, { requireChatTemplate: false,
+      runtime: { nativeDraft: false, maxGenerationTokens: maxTokens }, ...composition, maxRows: 1 });
+  },
+  async engine(context) {
+    const { createAppEngine } = await import("../engine");
+    return createAppEngine(context, { capacity: 1 });
+  },
+  write: text => { process.stdout.write(text); },
+  async stdin(signal) {
+    const reader = Bun.stdin.stream().getReader(), decoder = new TextDecoder();
+    let text = "";
+    const abort = () => { void reader.cancel(signal?.reason).catch(() => {}); };
+    signal?.addEventListener("abort", abort, { once: true });
+    try {
+      signal?.throwIfAborted();
+      for (;;) {
+        const next = await reader.read(); signal?.throwIfAborted();
+        if (next.done) return text + decoder.decode();
+        text += decoder.decode(next.value, { stream: true });
+      }
+    } finally { signal?.removeEventListener("abort", abort); reader.releaseLock(); }
+  },
+  stdinIsTTY: () => !!process.stdin.isTTY,
+};
+
+function option(args: CommandArgs, name: string): string | undefined {
+  const value = args.values[name]; return typeof value === "string" ? value : undefined;
+}
+function numeric(args: CommandArgs, name: string, input: { integer?: boolean; minimum?: number; maximum?: number } = {}): number | undefined {
+  const raw = option(args, name);
+  if (raw === undefined) return undefined;
+  const value = Number(raw);
+  if (!raw.trim() || !Number.isFinite(value) || (input.integer && !Number.isSafeInteger(value)) ||
+      value < (input.minimum ?? -Infinity) || value > (input.maximum ?? Infinity)) throw new Error(`invalid --${name}: ${raw}`);
+  return value;
+}
+
+/** App CLI policy: raw generated text, greedy/full-precision defaults, no
+ * model-author server sampling defaults or thinking/tool output filtering. */
+export function generateOptions(args: CommandArgs): { prompt: string; raw: boolean; options: GenerateOptions;
+  kvQuant: KvQuantOverride; turboQuant?: TurboQuantScheme; fusedSdpa: boolean; system?: string; adapter?: string; quantizedKvStart?: number } {
+  const prompt = option(args, "prompt") ?? args.positionals[1];
+  if (!prompt) throw new Error('usage: mlx-bun generate [query] --prompt "…" [--raw] [--max-tokens N]');
+  const { kv, fusedSdpa } = numericalPolicy(args);
+  // Main's TurboQuant spec, turbo (k8v3) or turbo:k<bits>v<bits>, is its own scheme beside affine KV.
+  const turboQuant = kv === undefined ? null : parseTurboQuantScheme(kv);
+  if (kv !== undefined && !turboQuant && !["off", "config", "4", "8"].includes(kv))
+    throw new Error("--kv-quant must be off, config, 4, 8, or turbo[:k<bits>v<bits>]");
+  const adapter = option(args, "adapter") ?? option(args, "adapter-path");
+  if (adapter !== undefined && !adapter.trim()) throw new Error("--adapter expects a directory");
+  const system = option(args, "system-prompt");
+  const quantizedKvStart = numeric(args, "quantized-kv-start", { integer: true, minimum: 0 });
+  return { prompt, raw: args.values.raw === true, fusedSdpa,
+    ...(quantizedKvStart !== undefined ? { quantizedKvStart } : {}),
+    ...(system !== undefined ? { system } : {}), ...(adapter !== undefined ? { adapter } : {}),
+    kvQuant: turboQuant || kv === undefined ? undefined : kv === "4" || kv === "8" ? Number(kv) : kv as KvQuantOverride,
+    ...(turboQuant ? { turboQuant } : {}),
+    options: {
+      maxTokens: numeric(args, "max-tokens", { integer: true, minimum: 1 }) ?? 256,
+      temperature: numeric(args, "temperature", { minimum: 0 }) ?? numeric(args, "temp", { minimum: 0 }) ?? 0,
+      topP: numeric(args, "top-p", { minimum: 0, maximum: 1 }) ?? 0,
+      topK: numeric(args, "top-k", { integer: true, minimum: 0 }) ?? 0,
+      ...(option(args, "min-p") !== undefined ? { minP: numeric(args, "min-p", { minimum: 0, maximum: 1 })! } : {}),
+      ...(option(args, "min-tokens-to-keep") !== undefined ? { minTokensToKeep: numeric(args, "min-tokens-to-keep", { integer: true, minimum: 1 })! } : {}),
+      ...(option(args, "xtc-probability") !== undefined ? { xtcProbability: numeric(args, "xtc-probability", { minimum: 0, maximum: 1 })! } : {}),
+      ...(option(args, "xtc-threshold") !== undefined ? { xtcThreshold: numeric(args, "xtc-threshold", { minimum: 0, maximum: 0.5 })! } : {}),
+      ...(option(args, "seed") !== undefined ? { seed: numeric(args, "seed", { integer: true })! } : {}),
+    },
+  };
+}
+
+export async function runInference(command: InferenceCommand, args: CommandArgs,
+  supplied: Partial<InferenceDependencies> = {}, signal?: AbortSignal): Promise<void> {
+  const deps = { ...defaults, ...supplied };
+  const generation = command === "generate" ? generateOptions(args) : undefined;
+  let texts: string[] = [];
+  if (command === "embed") {
+    const one = option(args, "text") ?? args.positionals[1];
+    texts = one !== undefined ? [one] : !deps.stdinIsTTY()
+      ? (await deps.stdin(signal)).split("\n").map(line => line.trim()).filter(Boolean) : [];
+    if (!texts.length) throw new Error('usage: mlx-bun embed [query] --text "…" (or pipe text, one per line)');
+  }
+  // `--prompt -` reads the prompt from stdin, as mlx_lm.generate does.
+  if (generation?.prompt === "-") generation.prompt = await deps.stdin(signal);
+  signal?.throwIfAborted();
+  const model = await deps.resolve(command, args.positionals[0] ?? option(args, "query") ?? "");
+  signal?.throwIfAborted();
+  // The command's numerical policy is the loaded composition's: its KV flags and
+  // the affine read kernels (`--fused-sdpa`), resolved once with the model.
+  const kv = { ...(generation?.turboQuant ? { turboQuant: generation.turboQuant } : {}), override: generation?.kvQuant ??
+    (generation && runtimeValue("MLX_BUN_EVAL_KV_QUANT") === "1" ? "config" as const : "off" as const),
+    ...(generation?.quantizedKvStart !== undefined ? { quantizedKvStart: generation.quantizedKvStart }
+      : generation?.kvQuant === undefined && runtimeValue("MLX_BUN_EVAL_KV_QUANT") === "1" ? { quantizedKvStart: 0 } : {}),
+    ...(generation ? { fusedSdpa: generation.fusedSdpa } : {}) };
+  const context = await deps.load(model, generation?.options.maxTokens, { kv });
+  let close: (() => void | Promise<void>) | undefined = () => context.dispose();
+  let failure: unknown, failed = false;
+  try {
+    signal?.throwIfAborted();
+    close = undefined; // engine construction owns failure cleanup from here
+    const engine = await deps.engine(context);
+    close = () => engine.close();
+    signal?.throwIfAborted();
+    if (generation) {
+      const ids = !generation.raw && context.template
+        ? textPrompt(context.template, context.tokenizer, [...generation.system !== undefined ? [{ role: "system" as const, content: generation.system }] : [],
+          { role: "user" as const, content: generation.prompt }],
+          { addGenerationPrompt: true, enableThinking: runtimeValue("MLX_BUN_EVAL_THINK") === "1" }, false).ids
+        : context.tokenizer.encode(generation.prompt, true);
+      // XTC never removes EOS or the newline token (as the server's request policy).
+      const xtc = generation.options.xtcProbability ? { xtcSpecialTokens: [...context.model.config.eosTokenIds, ...context.tokenizer.encode("\n", false)] } : {};
+      const adapterIds: string[] = [];
+      if (generation.adapter) {
+        const id = basename(resolve(generation.adapter));
+        await engine.gateway.runExclusive(async () => { await context.adapters.mount(id, generation.adapter!); }, undefined, signal);
+        adapterIds.push(id);
+      }
+      const plan = planRequest({ promptIds: ids, options: { ...generation.options, ...xtc, stopSequences: [] },
+        requestedMaxTokens: generation.options.maxTokens!, contextLimit: context.memoryPlan?.contextTokens ?? null,
+        stream: false, wantLogprobs: false, topLogprobs: 0, adapterIds, hasVision: false,
+        userSeed: generation.options.seed !== undefined, hasGrammar: false, hasDraft: false, ownership: new RequestOwnership() });
+      if (!plan.ok) { plan.dispose(); throw new Error(plan.error.message); }
+      const tokens: number[] = [];
+      try {
+        const placement = engine.completion.place(plan.shape, plan.options);
+        plan.transferOwnership();
+        await engine.completion.run(plan.promptIds, plan.options, token => { tokens.push(token); }, undefined,
+          plan.shape, placement, signal);
+        signal?.throwIfAborted();
+        const text = context.tokenizer.decode(tokens, true);
+        deps.write(text.endsWith("\n") ? text : text + "\n");
+      } finally { plan.dispose(); }
+    } else {
+      const embed = engine.binding.embed?.bind(engine.binding);
+      if (!embed) throw new Error(`"${model.repoId}" is not an embedding model (its graph declares no pooled embeddings; serve an embedding model, e.g. Qwen3-Embedding).`);
+      const results = await engine.gateway.runExclusive(async () => embed(texts, option(args, "instruct")), undefined, signal);
+      signal?.throwIfAborted();
+      if (args.values.json === true) {
+        const total = results.reduce((sum, result) => sum + result.tokens, 0);
+        deps.write(JSON.stringify({ object: "list", data: results.map((result, index) => ({ object: "embedding", index,
+          embedding: Array.from(result.vector) })), model: model.repoId, usage: { prompt_tokens: total, total_tokens: total } }) + "\n");
+      } else for (const result of results) deps.write(JSON.stringify(Array.from(result.vector)) + "\n");
+    }
+  } catch (error) { failed = true; failure = error; throw error; }
+  finally {
+    try { await close?.(); }
+    catch (error) { if (failed) throw new AggregateError([failure, error], "inference and cleanup failed"); throw error; }
+  }
+}

package/src/cli/job-entry.ts

@@ -0,0 +1,97 @@
+import { jobError, JobStore, nowIso } from "../jobs/db";
+import { makeEmit } from "../jobs/events";
+import { runtimeFlag, runtimeValue } from "@mlx-bun/inference/runtime/config";
+import type { JobRunner } from "../jobs/protocol";
+
+/** A job's runner and what to release after it: the installed module that registered the kind. */
+async function resolveRunner(kind: string): Promise<{ runner: JobRunner; release?(): Promise<void> }> {
+  const [{ installedModules }, { loadModules }, { createHostServices }] = await Promise.all([
+    import("../modules"), import("@mlx-bun/app-host"), import("@mlx-bun/app-services") ]);
+  const [owner] = await installedModules(manifest => manifest.jobs?.some(job => job.kind === kind) ?? false);
+  if (!owner) throw new Error(`no runner registered for kind "${kind}"`);
+  // The child activates only the module that owns the kind, over the services a job needs; it submits no jobs of its own.
+  const host = createHostServices({ log() {} });
+  const loaded = await loadModules([owner], { services: { catalog: host.bindings.catalog, storage: host.bindings.storage,
+    jobs: () => ({ submit: () => Promise.reject(new Error("a job child submits no jobs")), get: async () => undefined, list: async () => [],
+      cancel: async () => {}, events: () => (async function* () {})() }) } });
+  return { runner: loaded.jobs.get(kind)!.runner as unknown as JobRunner, release: async () => { try { await loaded.stop(); } finally { await host.whisper.close(); } } };
+}
+
+/** The job host spawns this child as the leader of its own process group
+ * (jobs/runner.ts), which the terminal's signals no longer reach, and holds a
+ * pipe on its stdin for the child's life. The pipe's end means the host is
+ * gone: the child then stops itself and everything it started, as the host's
+ * own shutdown would have (SIGTERM to the group). */
+export function stopWithParent(stdin: ReadableStream<Uint8Array>): void {
+  void (async () => {
+    const reader = stdin.getReader();
+    try { while (!(await reader.read()).done) { /* nothing is sent on this pipe */ } }
+    catch { /* a broken pipe is the host's end too */ }
+    try { process.kill(-process.pid, "SIGTERM"); }
+    catch { process.kill(process.pid, "SIGTERM"); } // not a group leader
+  })();
+}
+
+export async function runJobEntry(jobId = process.argv[2]): Promise<number> {
+  if (!jobId) {
+    console.error("usage: bun job-entry.ts <jobId>");
+    return 2;
+  }
+
+  if (runtimeFlag("MLX_BUN_JOB_PARENT_PIPE", false)) stopWithParent(Bun.stdin.stream());
+
+  // Own fresh connection — overrides come from env so a spawned child finds
+  // the same DB/logs the parent used (tests set these to a tmp dir).
+  const store = new JobStore(
+    runtimeValue("MLX_BUN_JOBS_DB") || undefined,
+    runtimeValue("MLX_BUN_JOBS_DIR") || undefined,
+  );
+
+  const row = store.get(jobId);
+  if (!row) {
+    console.error(`job ${jobId} not found`);
+    store.close();
+    return 2;
+  }
+
+  const emit = makeEmit(store, jobId, row.log_path);
+  store.setStatus(jobId, "running");
+  emit({ type: "started", ts: Date.now() });
+
+  let config: Record<string, unknown>;
+  try {
+    config = JSON.parse(row.config_json) as Record<string, unknown>;
+  } catch (e) {
+    const error = `bad config_json: ${e instanceof Error ? e.message : String(e)}`;
+    store.setStatus(jobId, "failed", { error, endedAt: nowIso() });
+    emit({ type: "failed", error, ts: Date.now() });
+    store.close();
+    return 1;
+  }
+
+  let release: (() => Promise<void>) | undefined;
+  try {
+    const resolved = await resolveRunner(row.kind);
+    release = resolved.release;
+    const result = await resolved.runner(emit, config);
+    const out = result?.outputPath;
+    if (out) store.setOutputPath(jobId, out);
+    store.setProgress(jobId, 1);
+    store.setStatus(jobId, "done", { endedAt: nowIso() });
+    emit({ type: "done", ts: Date.now(), output_dir: out ?? row.output_path ?? undefined });
+    store.close();
+    await release?.().catch(() => {});
+    return 0;
+  } catch (e) {
+    const error = jobError(e);
+    store.setStatus(jobId, "failed", { error, endedAt: nowIso() });
+    emit({ type: "failed", error, ts: Date.now() });
+    store.close();
+    await release?.().catch(() => {});
+    return 1;
+  }
+}
+
+// The process owner exits after durable terminal publication and store cleanup;
+// producer-owned lingering handles must not retain the parent execution lease.
+if (import.meta.main) process.exit(await runJobEntry());

package/src/cli/library-host.ts

@@ -0,0 +1,117 @@
+// `openIsolatedHost` behind the public `mlx-bun/engine` entry (engine-entry.ts),
+// imported only when a host is opened: one full-app worker (the `__worker` app
+// launch form, worker-entry.ts) on a private Unix socket, supervised with the
+// CLI's restart budget (jobs/worker-supervisor.ts) and shutdown deadline, and
+// adapted to the `EngineHost` contract main's EngineChild implemented. The
+// worker runs `runServe` from `["--model", model, ...arguments]`, so the host
+// forwards to the whole app: completions, discovery, jobs, sessions, memory.
+import { mkdtempSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { fileURLToPath } from "node:url";
+import { WORKER_PROTOCOL_VERSION } from "../jobs/worker-process";
+import { superviseWorker, type WorkerSupervisor, type WorkerSupervisorOptions } from "../jobs/worker-supervisor";
+import type { EngineHost } from "../server/client";
+import { stripHopByHop } from "../server/proxy-routes";
+import { shutdownTimeoutMs, validateAppLaunchArgv } from "./serve";
+import type { AppWorkerLaunch } from "./worker-entry";
+
+/** `openIsolatedHost`'s options; engine-entry.ts documents them. */
+export interface IsolatedHostOptions { arguments?: readonly string[]; command?: readonly string[]; readyTimeoutMs?: number }
+
+/** Internal (tests), never a public option: stand-ins for this module's path
+ * (the compiled-consumer check), process creation, and the restart budget. */
+export interface LibraryHostSeams {
+  modulePath?: string;
+  spawn?: WorkerSupervisorOptions["spawn"];
+  restarts?: WorkerSupervisorOptions["restarts"];
+}
+
+/** Main's rule: a source or package consumer runs this package's CLI source
+ * under its own Bun; a compiled consumer's executable is not the mlx-bun CLI,
+ * so it must name one. */
+function workerCommand(command: readonly string[] | undefined, modulePath: string): readonly string[] {
+  if (!command && modulePath.includes("$bunfs"))
+    throw new Error("a compiled library consumer must supply the mlx-bun executable as command");
+  const resolved = command ?? [process.execPath, fileURLToPath(new URL("./main.ts", import.meta.url))];
+  if (!resolved.length) throw new Error("engine command must not be empty");
+  return resolved;
+}
+
+/** Refuses what the CLI would refuse before anything is spawned, starts the
+ * worker, and resolves once it serves. A worker that fails before ready is
+ * stopped and its socket directory removed; the rejection is its reason. */
+export async function openLibraryHost(model: string, options: IsolatedHostOptions = {}, seams: LibraryHostSeams = {}): Promise<EngineHost<Request, Response>> {
+  if (typeof model !== "string" || !model.trim())
+    throw new Error("openIsolatedHost needs a model: an empty query selects one automatically and may download the starter model");
+  const argv = ["--model", model, ...(options.arguments ?? [])];
+  validateAppLaunchArgv(argv);
+  const command = workerCommand(options.command, seams.modulePath ?? fileURLToPath(import.meta.url));
+  // A private 0700 directory, as in the isolated server: a short path within the
+  // 104-byte socket limit, removed after the worker has exited.
+  const socketDir = mkdtempSync(join(tmpdir(), "mlx-worker-"));
+  const removeSocketDir = () => rmSync(socketDir, { recursive: true, force: true });
+  const socketPath = join(socketDir, "engine.sock");
+  const launch: AppWorkerLaunch = { kind: "app", version: WORKER_PROTOCOL_VERSION, socketPath, argv };
+  // The worker's app closes under the CLI's deadline; the parent drains and
+  // then waits that long after SIGTERM before SIGKILL.
+  const budget = shutdownTimeoutMs();
+  let supervisor: WorkerSupervisor;
+  try {
+    supervisor = superviseWorker({ command, socketPath, launch, drainTimeoutMs: budget, graceMs: budget,
+      ...(options.readyTimeoutMs !== undefined ? { readyTimeoutMs: options.readyTimeoutMs } : {}),
+      // The worker's own output reaches the consumer's console unprefixed, as main's inherited stdio did.
+      log: line => console.log(line), error: line => console.error(line),
+      ...(seams.spawn ? { spawn: seams.spawn } : {}), ...(seams.restarts ? { restarts: seams.restarts } : {}) });
+  } catch (error) { removeSocketDir(); throw error; }
+  const host = supervisedHost(supervisor, removeSocketDir);
+  try { await supervisor.ready; }
+  catch (error) { await host.close(); throw error; }
+  return host;
+}
+
+/** Main's EngineChild forwarding over the supervisor: wait (abortably) for a
+ * serving worker, strip hop-by-hop headers both ways, stream both bodies,
+ * retry a GET or HEAD once after 250 ms, never replay anything else. */
+function supervisedHost(supervisor: WorkerSupervisor, removeSocketDir: () => void): EngineHost<Request, Response> {
+  let closing: Promise<void> | undefined;
+  const closed = () => new Error("engine host is closed");
+  // After close the supervisor's refusals ("shutting down", an aborted fetch)
+  // mean the host is closed; a caller's own abort keeps its reason.
+  const rethrow = (error: unknown, signal?: AbortSignal): never => { throw closing && !signal?.aborted ? closed() : error; };
+  const forwardOnce = async (request: Request): Promise<Response> => {
+    if (closing) throw closed();
+    request.signal.throwIfAborted();
+    // Startup or a respawn: nothing is sent until a worker serves.
+    try { await supervisor.whenReady(request.signal); } catch (error) { rethrow(error, request.signal); }
+    if (closing) throw closed();
+    request.signal.throwIfAborted();
+    let upstream: Response;
+    try {
+      upstream = await supervisor.fetch(request.url, { method: request.method, headers: stripHopByHop(request.headers),
+        body: request.body, signal: request.signal, redirect: "manual", duplex: "half" });
+    } catch (error) { return rethrow(error, request.signal); }
+    return new Response(upstream.body, { status: upstream.status, headers: stripHopByHop(upstream.headers) });
+  };
+  return {
+    /** Follows the current worker: pending during a respawn, rejected once the restart budget is spent or the host is closed. */
+    get ready() {
+      const ready = closing ? Promise.reject(closed()) : supervisor.whenReady().catch(error => rethrow(error));
+      ready.catch(() => {}); // observed by whoever reads it
+      return ready;
+    },
+    async forward(request) {
+      try { return await forwardOnce(request); }
+      catch (error) {
+        // A bodyless request that raced a worker exit waits out the respawn once.
+        if ((request.method !== "GET" && request.method !== "HEAD") || request.signal.aborted || closing) throw error;
+        await new Promise(resolve => setTimeout(resolve, 250));
+        return await forwardOnce(request);
+      }
+    },
+    /** Drain, SIGTERM, SIGKILL after the grace; the socket directory goes once the worker has exited. */
+    close() {
+      return closing ??= (async () => { try { await supervisor.close(); } finally { removeSocketDir(); } })();
+    },
+  };
+}

package/src/cli/main.ts

@@ -0,0 +1,57 @@
+#!/usr/bin/env bun
+import "../jobs/executable";
+import { commandInvocation, help, isCommand, parseCommand } from "./args";
+import { invokedAlias, translateAlias, type AliasResult } from "./mlx-lm-aliases";
+import { installedVerbs, runInstalledVerb } from "./module-verbs";
+import { renderHelp } from "./terminal";
+import pkg from "../../package.json" with { type: "json" };
+
+try {
+  // Started as `mlx-bun.<cmd>` (an mlx-lm alias): the verb and its arguments come from the translation.
+  const alias = invokedAlias();
+  const aliased: AliasResult | undefined = alias ? translateAlias(alias, process.argv.slice(2)) : undefined;
+  const { command, args } = aliased && "command" in aliased ? { command: aliased.command as string, args: [] as string[] }
+    : commandInvocation(process.argv.slice(2));
+  const parsedArgs = aliased && "parsed" in aliased ? aliased.parsed : undefined;
+  if (aliased && "help" in aliased) {
+    console.log(renderHelp(aliased.help));
+  } else if (command === "__job") {
+    process.exit(await (await import("./job-entry")).runJobEntry(args[0]));
+  } else if (command === "__worker") {
+    // Private: the isolation worker a parent spawns (jobs/worker-process.ts);
+    // its launch record arrives on stdin. No user flag selects this mode.
+    process.exit(await (await import("./worker-entry")).runWorkerEntry());
+  } else if (["--version", "-v", "version"].includes(command ?? "")) {
+    console.log(`mlx-bun ${pkg.version}`);
+  } else if (command === "help" || command === "--help" || command === "-h") {
+    console.log(renderHelp(help(args[0])));
+  } else if (installedVerbs().has(command)) {
+    // A verb an installed module declares (`src/modules.ts`); an alias (`mlx-bun.convert`) arrives already translated to the verb's options.
+    if (args.includes("--help") || args.includes("-h")) console.log(renderHelp(help(command)));
+    else process.exitCode = await runInstalledVerb(command, parsedArgs ?? args);
+  } else if (!isCommand(command)) {
+    throw new Error(`Unknown command: ${command}. Use mlx-bun --help.`);
+  } else if (args.includes("--help") || args.includes("-h")) {
+    console.log(renderHelp(help(command)));
+  } else {
+    const parsed = parsedArgs ?? parseCommand(command, args);
+    if (command === "serve") {
+      const { runServe } = await import("./serve");
+      await runServe(parsed);
+    } else if (command === "generate" || command === "embed") {
+      // SIGINT/SIGTERM abort the one-shot work; the verb rejects with the reason and exits 1.
+      const cancellation = new AbortController();
+      const stop = () => cancellation.abort(new Error("inference cancelled"));
+      process.on("SIGINT", stop); process.on("SIGTERM", stop);
+      try { await (await import("./inference")).runInference(command, parsed, {}, cancellation.signal); }
+      finally { process.off("SIGINT", stop); process.off("SIGTERM", stop); }
+    } else if (command === "perplexity") {
+      await (await import("./perplexity")).runPerplexity(parsed);
+    } else if (command === "compile-ane") {
+      process.exitCode = await (await import("./ane-install")).runCompileAne(parsed);
+    }
+  }
+} catch (error) {
+  console.error(error instanceof Error ? error.message : String(error));
+  process.exitCode = 1;
+}

package/src/cli/memory-engine.ts

@@ -0,0 +1,163 @@
+// The `memory` verb's in-process task model: main's `memory synthesize` (and the
+// nightly job) loaded its own Gemma-4 e4b on first use instead of needing a
+// serving mlx-bun. Nothing loads until a stage actually asks for a completion,
+// and every call runs through the app engine's continuous gateway, where one
+// call is a batch of one; there is no second server. Policy stays here and in memory/model:
+// the task model, the chunk adapter only for the chunk stage, greedy argmax
+// with the model's EOS and no stop strings, full-precision KV, and the raw
+// decoded text main returned.
+
+import type { LoadContextOptions, LoadedModelContext } from "../engine/model-host";
+import type { CompletionEngine } from "../engine/completion";
+import type { GenerationGateway } from "../engine/generation-gateway";
+import {
+  MEMORY_TASK_MODEL, adapterDirFor, locateTaskModel, memoryBatchSize, memoryPromptIds,
+  type MemoryCompletionClient, type MemoryCompletionRequest,
+} from "../modules";
+import { planRequest, RequestOwnership } from "../server/request-plan";
+
+/** The chunk stage's trained adapter, mounted once under this id when present. */
+export const CHUNK_ADAPTER = "memory-chunk";
+
+export interface MemoryEngine {
+  context: LoadedModelContext;
+  completion: CompletionEngine;
+  gateway: Pick<GenerationGateway, "runExclusive">;
+  close(): Promise<void>;
+}
+
+export interface MemoryEngineDependencies {
+  /** The task model's snapshot directory; never downloads. */
+  locate(repoId: string): Promise<string>;
+  /** `composition`: the batch cap. The chunk adapter mounts through the loaded context. */
+  load(path: string, repoId: string, composition?: Pick<LoadContextOptions, "maxRows">): Promise<LoadedModelContext>;
+  /** Takes ownership of the context even if construction rejects. */
+  engine(context: LoadedModelContext, capacity: number): Promise<MemoryEngine>;
+  adapterDir(stage: string): string | undefined;
+}
+
+/** The shipped wiring: the cached task model, the app engine, main's adapter directory. */
+export const defaultMemoryEngineDependencies: MemoryEngineDependencies = {
+  locate: repoId => locateTaskModel(repoId),
+  async load(path, repoId, composition) {
+    const { loadContext } = await import("../engine/model-host");
+    return loadContext(path, repoId, composition);
+  },
+  async engine(context, capacity) {
+    const { createAppEngine } = await import("../engine");
+    return createAppEngine(context, { capacity });
+  },
+  adapterDir: adapterDirFor,
+};
+
+interface MemoryRuntime { engine: MemoryEngine; chunkAdapter: boolean }
+
+/** The task model's client, and a view of it whose every completion ends when a signal aborts. */
+export interface InProcessMemoryClient {
+  client: MemoryCompletionClient;
+  /** `snapshot`: the task model directory the caller selected. A view's call
+   * that starts the load loads exactly it instead of locating one; once
+   * loaded, the task model stays what it is. */
+  clientFor(signal: AbortSignal, snapshot?: string): MemoryCompletionClient;
+  close(): Promise<void>;
+}
+
+/** A lazily loaded memory client. The first completion loads the task model
+ *  once (every caller shares that one initialization); a failed load stays
+ *  failed for this client, as in main. `close()` is safe before, during and
+ *  after initialization, including after it failed. */
+export function createInProcessMemoryClient(deps: MemoryEngineDependencies = defaultMemoryEngineDependencies): InProcessMemoryClient {
+  let closed = false, init: Promise<MemoryRuntime> | undefined, closing: Promise<void> | undefined;
+  const closedError = () => new Error("memory: the task model client is closed");
+  const start = async (snapshot?: string): Promise<MemoryRuntime> => {
+    const path = snapshot ?? await deps.locate(MEMORY_TASK_MODEL);
+    if (closed) throw closedError();
+    const dir = deps.adapterDir("chunk");
+    const context = await deps.load(path, MEMORY_TASK_MODEL, { maxRows: memoryBatchSize() });
+    if (closed) { context.dispose(); throw closedError(); }
+    // Main decoded with a full-precision cache, whatever the artifact's kv_config says: the load names no KV flags.
+    const engine = await deps.engine(context, memoryBatchSize());
+    try {
+      if (closed) throw closedError();
+      if (dir) await engine.gateway.runExclusive(async () => { await engine.context.adapters.mount(CHUNK_ADAPTER, dir); });
+      if (closed) throw closedError();
+      return { engine, chunkAdapter: dir !== undefined };
+    } catch (error) { await engine.close(); throw error; }
+  };
+  const runtime = (snapshot?: string) => closed ? Promise.reject(closedError()) : (init ??= start(snapshot));
+
+  const run = async (rt: MemoryRuntime, request: MemoryCompletionRequest, signal?: AbortSignal): Promise<string> => {
+    const { context, completion } = rt.engine;
+    if (!context.template) throw new Error(`memory: ${MEMORY_TASK_MODEL} has no chat template`);
+    const tokenizer = context.tokenizer;
+    const promptIds = memoryPromptIds(request.stage, request.input,
+      { encode: text => tokenizer.encode(text), bosTokenId: tokenizer.bosTokenId ?? -1 }, context.template);
+    const plan = planRequest({ promptIds,
+      options: { maxTokens: request.maxTokens, temperature: 0, stopSequences: [] },
+      requestedMaxTokens: request.maxTokens, contextLimit: context.memoryPlan?.contextTokens ?? null,
+      stream: false, wantLogprobs: false, topLogprobs: 0,
+      adapterIds: request.stage === "chunk" && rt.chunkAdapter ? [CHUNK_ADAPTER] : [],
+      hasVision: false, userSeed: false, hasGrammar: false, hasDraft: false, ownership: new RequestOwnership() });
+    if (!plan.ok) { plan.dispose(); throw new Error(plan.error.message); }
+    const tokens: number[] = [];
+    try {
+      const placement = completion.place(plan.shape, plan.options);
+      plan.transferOwnership();
+      await completion.run(plan.promptIds, plan.options, token => { tokens.push(token); }, undefined, plan.shape, placement, signal);
+      signal?.throwIfAborted();
+      return tokenizer.decode(tokens, true);
+    } finally { plan.dispose(); }
+  };
+
+  const clientFor = (outer?: AbortSignal, snapshot?: string): MemoryCompletionClient => ({
+    async complete(request) {
+      // Cancelled work never starts the shared load; a load another caller started continues.
+      outer?.throwIfAborted();
+      const rt = await runtime(snapshot);
+      outer?.throwIfAborted();
+      return run(rt, request, outer);
+    },
+    async completeBatch(requests) {
+      if (!requests.length) return [];
+      outer?.throwIfAborted();
+      const rt = await runtime(snapshot);
+      outer?.throwIfAborted();
+      const first = requests[0]!;
+      // Main's rule: a mixed-stage or mixed-budget batch runs its rows one after another.
+      if (memoryBatchSize() <= 1 || requests.length === 1 ||
+        !requests.every(request => request.stage === first.stage && request.maxTokens === first.maxTokens)) {
+        const results: string[] = [];
+        for (const request of requests) results.push(await run(rt, request, outer));
+        return results;
+      }
+      // Bounded producers, as main's scheduler and the loopback client: at most
+      // the batch width of rows in flight, results in input order. One row's
+      // failure cancels its siblings and stops new rows; every started row
+      // settles before the batch rejects.
+      const abort = new AbortController(), signal = outer ? AbortSignal.any([outer, abort.signal]) : abort.signal;
+      const results = new Array<string>(requests.length);
+      let next = 0, failure: { error: unknown } | undefined;
+      const producer = async () => {
+        while (!failure && next < requests.length) {
+          const index = next++;
+          try { results[index] = await run(rt, requests[index]!, signal); }
+          catch (error) { failure ??= { error }; abort.abort(error); }
+        }
+      };
+      await Promise.all(Array.from({ length: Math.min(memoryBatchSize(), requests.length) }, producer));
+      if (failure) throw failure.error;
+      return results;
+    },
+  });
+  return {
+    client: clientFor(),
+    clientFor,
+    close() {
+      closed = true;
+      return closing ??= (async () => {
+        const rt = await init?.catch(() => undefined);
+        await rt?.engine.close();
+      })();
+    },
+  };
+}