mlx-bun 0.4.0 → 1.0.0

package/src/cli/worker-entry.ts

@@ -0,0 +1,272 @@
+// The worker half of runtime isolation, reached only through the parent's
+// spawn helper (jobs/worker-process.ts): the launch record arrives as the
+// first stdin line, the ready line leaves on stdout, and the end of stdin
+// means the parent is gone. No flag selects it. Two launch forms, both
+// carrying the app's package version (a record with another one exits 2):
+// - model (the isolated server's worker, one per resident model): only the model-scoped host (serve-host.ts) over the
+//   parent's Unix socket, with the persistent services stubbed because the
+//   parent owns them, plus the memory task model the parent's synthesis calls
+//   through the admin surface (only the default worker is ever asked);
+// - app: the whole app, composed by `runServe` from serve arguments exactly as
+//   the CLI composes it, listening on the parent's socket instead of TCP.
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+import { createEventHub, createModuleSockets } from "@mlx-bun/app-services/portable";
+import { decodeLaunch, formatWorkerMessage, WORKER_PROTOCOL_VERSION } from "../jobs/worker-process";
+import type { WorkerMemory } from "../jobs/worker-memory";
+import { ResponseStore } from "../server/responses";
+import { createWorkerRoutes } from "../server/worker-routes";
+import type { CommandArgs } from "./args";
+// The app form's composition loads only when a launch asks for it.
+import type { ServeDependencies, ServedApp, SignalPort, startModelServer, startTranscriptionServer } from "./serve";
+import { startModelHost } from "./serve-host";
+import { createServedModelHost } from "./served-model-host";
+import type { ServeOptions } from "./serve-options";
+import type { AppState, ModelHostLink, RouteGroup } from "./serve-state";
+
+/** The model form: the socket to bind, the model record the parent selected,
+ * and the serve options with every query already resolved to a directory
+ * (draft, Whisper). The parent's port and hostname are ignored. */
+export interface ModelWorkerLaunch { kind?: "model"; version: string; socketPath: string; model: ModelRecord; options: ServeOptions }
+/** The app form: serve arguments as the CLI takes them (`["--model", model, ...]`),
+ * validated by `validateAppLaunchArgv` and composed by `runServe` on the socket. */
+export interface AppWorkerLaunch { kind: "app"; version: string; socketPath: string; argv: string[] }
+export type WorkerLaunch = ModelWorkerLaunch | AppWorkerLaunch;
+
+export function parseWorkerLaunch(text: string): WorkerLaunch {
+  let value: unknown;
+  try { value = decodeLaunch(text); } catch (error) { throw new Error(`worker launch is not JSON: ${error instanceof Error ? error.message : String(error)}`); }
+  const record = value as Record<string, unknown> | null;
+  if (!record || typeof record !== "object" || Array.isArray(record)) throw new Error("worker launch must be an object");
+  if (typeof record.socketPath !== "string" || !record.socketPath.trim()) throw new Error("worker launch needs socketPath");
+  if (record.version !== WORKER_PROTOCOL_VERSION)
+    throw new Error(`worker protocol version mismatch: the launch record is ${typeof record.version === "string" ? record.version : "unversioned"}, this worker is ${WORKER_PROTOCOL_VERSION}`);
+  const socketPath = record.socketPath, version = WORKER_PROTOCOL_VERSION;
+  if (record.kind === "app") {
+    if (!Array.isArray(record.argv) || !record.argv.every(item => typeof item === "string"))
+      throw new Error("worker launch needs argv, the serve arguments");
+    return { kind: "app", version, socketPath, argv: record.argv as string[] };
+  }
+  if (record.kind !== undefined && record.kind !== "model") throw new Error(`worker launch kind ${JSON.stringify(record.kind)} is unknown`);
+  const model = record.model as Partial<ModelRecord> | undefined;
+  if (!model || typeof model !== "object" || typeof model.repoId !== "string" || typeof model.path !== "string")
+    throw new Error("worker launch needs model.repoId and model.path");
+  const options = record.options as Partial<ServeOptions> | undefined;
+  if (!options || typeof options !== "object" || !options.cache || !options.request || typeof options.capacity !== "number")
+    throw new Error("worker launch needs resolved serve options");
+  return { ...(record.kind === "model" ? { kind: "model" as const } : {}), version, socketPath, model: model as ModelRecord, options: options as ServeOptions };
+}
+
+/** The persistent half a worker does not own: nothing served, no producers,
+ * and the host's link kept where the admin routes can reach it. */
+export function createWorkerState(options: ServeOptions, link: { current?: ModelHostLink }): AppState {
+  const none: RouteGroup = { handle: async () => null };
+  const events = createEventHub();
+  return {
+    web: () => null,
+    downloads: { active: [], snapshot: () => [], start() { throw new Error("a worker owns no downloads"); }, async close() {} },
+    // A worker publishes its own engine's events for the modules it hosts; it runs no jobs.
+    events,
+    modelHost: createServedModelHost({ link: () => link.current,
+      fetch: (request, supplied) => fetch(request, supplied.unix ? { unix: supplied.unix } as RequestInit : undefined) }),
+    responses: new ResponseStore(),
+    // Chat and its memory tools are the parent's: the worker never opens a vault, so the paths are placeholders.
+    memoryPaths: options.memoryPaths ?? { vault: "", skills: "" },
+    storagePaths: options.storagePaths ?? {},
+    sockets: createModuleSockets([]),
+    routes: { jobs: none, appModules: none, publishing: none },
+    attach(supplied) {
+      link.current = supplied;
+      return () => { if (link.current === supplied) link.current = undefined; };
+    },
+    async close() { events.close(); },
+  };
+}
+
+/** This process's MLX memory, read where the parent cannot (it has no native module); undefined without the native runtime. */
+async function processMemory(): Promise<(() => WorkerMemory | undefined) | undefined> {
+  const mlx = await import("@mlx-bun/mlx/ffi").catch(() => undefined);
+  if (!mlx) return undefined;
+  let workingSetBytes: number | undefined;
+  const workingSet = () => workingSetBytes ??= (() => { try { return mlx.maxRecommendedWorkingSetSize(); } catch { return 0; } })();
+  return () => {
+    try { return { activeBytes: mlx.activeMemory(), cacheBytes: mlx.cacheMemory(), peakBytes: mlx.peakMemory(), workingSetBytes: workingSet() }; }
+    catch { return undefined; }
+  };
+}
+
+/** What the app form composes with; `runServe`'s own defaults otherwise. */
+export interface AppWorkerDependencies {
+  resolve: ServeDependencies["resolve"];
+  start: typeof startModelServer;
+  startTranscription: typeof startTranscriptionServer;
+  log(message: string): void;
+}
+
+export interface WorkerEntryPorts {
+  /** The launch record is the first line; the stream's end means the parent left. */
+  stdin: ReadableStream<Uint8Array>;
+  write(line: string): void;
+  signals: SignalPort;
+  /** Internal (tests): stand-ins for the app form's selection, compositions,
+   * and log; a stand-in composition still receives the socket hooks. */
+  app?: Partial<AppWorkerDependencies>;
+}
+
+const defaults: WorkerEntryPorts = { stdin: Bun.stdin.stream(), write: line => { process.stdout.write(line + "\n"); }, signals: process };
+
+/** Runs until the parent signals or leaves; the exit code is the process's. */
+export async function runWorkerEntry(ports: WorkerEntryPorts = defaults): Promise<number> {
+  const reader = ports.stdin.getReader(), decoder = new TextDecoder();
+  let buffered = "";
+  const readLine = async (): Promise<string | null> => {
+    for (;;) {
+      const newline = buffered.indexOf("\n");
+      if (newline !== -1) { const line = buffered.slice(0, newline); buffered = buffered.slice(newline + 1); return line; }
+      const { done, value } = await reader.read();
+      if (done) { const rest = buffered; buffered = ""; return rest.trim() ? rest : null; }
+      buffered += decoder.decode(value, { stream: true });
+    }
+  };
+  let launch: WorkerLaunch, args: CommandArgs | undefined;
+  try {
+    const first = await readLine();
+    if (first === null) throw new Error("worker launch: stdin ended before the launch record");
+    launch = parseWorkerLaunch(first);
+    // The app form's arguments fail here like the CLI's, before anything composes.
+    if (launch.kind === "app") args = (await import("./serve")).validateAppLaunchArgv(launch.argv);
+  } catch (error) {
+    console.error(error instanceof Error ? error.message : String(error));
+    return 2;
+  }
+  const parentLeft = () => (async () => { while ((await readLine()) !== null) { /* nothing else is expected on stdin */ } })().catch(() => {});
+  if (launch.kind === "app") return runAppWorker(launch, args!, ports, parentLeft());
+  const link: { current?: ModelHostLink } = {};
+  // Main's memory task model, loaded by the first call and kept until this
+  // worker stops: the admin surface runs each call under the host's execution
+  // lease, and shutdown joins those calls, then closes the task model, both
+  // ahead of the host's engine.
+  const memory = (await import("./memory-engine")).createInProcessMemoryClient();
+  // The engine's events stay in this process's bus; the parent subscribes to them over the admin surface.
+  const state = createWorkerState(launch.options, link);
+  // The parent's models module reaches this worker's model through the host's own adapter operation.
+  const adapterOf = () => { const operation = link.current?.adapters; if (!operation) throw new Error("no model host is attached"); return operation; };
+  const memoryNow = await processMemory();
+  const admin = createWorkerRoutes({ modelId: launch.model.repoId, pid: process.pid, ...(memoryNow ? { memory: memoryNow } : {}),
+    acquireExecutionLease(signal) {
+      if (!link.current) return Promise.reject(new Error("no model host is attached"));
+      return link.current.acquireExecutionLease(signal);
+    }, memoryTaskModel: memory, events: state.events,
+    adapters: { list: async signal => adapterOf().list(signal), mount: async (id, path, signal) => adapterOf().mount(id, path, signal),
+      unmount: async (id, signal) => adapterOf().unmount(id, signal), merge: async (request, signal) => adapterOf().merge(request, signal) } });
+  let host: Awaited<ReturnType<typeof startModelHost>>;
+  try {
+    // One model per worker: the parent holds the residency of the workers, so this host never loads another.
+    host = await startModelHost(state, launch.model, launch.options,
+      { unix: launch.socketPath, oneModel: true, routes: model => admin.wrap(model),
+        beforeDrain: async () => { try { await admin.close(); } finally { await memory.close(); } } });
+  } catch (error) {
+    // Nothing was served, so the task model never loaded; closing only refuses later calls.
+    await memory.close();
+    console.error(`worker startup failed: ${error instanceof Error ? error.message : String(error)}`);
+    return 1;
+  }
+  ports.write(formatWorkerMessage({ type: "ready", socketPath: launch.socketPath, modelId: launch.model.repoId, pid: process.pid, version: WORKER_PROTOCOL_VERSION }));
+  // Stay up until the parent signals or its pipe ends. One close; a later
+  // signal waits on it, and the parent's grace period bounds the wait.
+  return new Promise<number>(resolve => {
+    let closing = false;
+    const stop = () => {
+      if (closing) return;
+      closing = true;
+      ports.signals.removeListener("SIGTERM", stop); ports.signals.removeListener("SIGINT", stop);
+      // A close whose saved state is not durable exits 3, so the parent that asked for it can say the flush was incomplete.
+      host.close().then(result => resolve(result.durable ? 0 : 3), error => { console.error(error instanceof Error ? error.message : String(error)); resolve(1); });
+    };
+    ports.signals.on("SIGTERM", stop); ports.signals.on("SIGINT", stop);
+    void parentLeft().finally(stop);
+  });
+}
+
+/** The app form's signal port: the process's SIGTERM and SIGINT, plus the
+ * parent leaving (the end of stdin) delivered as SIGTERM, so runServe's
+ * startup abort and its shutdown handlers see all three. */
+function parentSignals(base: SignalPort) {
+  const listeners = { SIGINT: new Set<() => void>(), SIGTERM: new Set<() => void>() };
+  let stopped = false;
+  const emit = (signal: "SIGINT" | "SIGTERM") => { stopped = true; for (const listener of [...listeners[signal]]) listener(); };
+  const onInt = () => emit("SIGINT"), onTerm = () => emit("SIGTERM");
+  base.on("SIGINT", onInt); base.on("SIGTERM", onTerm);
+  return {
+    get stopped() { return stopped; },
+    on(signal: "SIGINT" | "SIGTERM", listener: () => void) { listeners[signal].add(listener); },
+    removeListener(signal: "SIGINT" | "SIGTERM", listener: () => void) { listeners[signal].delete(listener); },
+    leave: () => emit("SIGTERM"),
+    dispose() { base.removeListener("SIGINT", onInt); base.removeListener("SIGTERM", onTerm); },
+  };
+}
+
+/** The app form: `runServe` composes the whole app from the launch's serve
+ * arguments, handing the compositions the socket hooks (the Unix socket, the
+ * admin surface ahead of the routes, the admin close ahead of the producers).
+ * A stop before ready aborts startup through runServe's startup signal, and
+ * nothing is announced; after ready, runServe's shutdown handlers close the
+ * app once under the CLI's deadline, and their exit code is the worker's. */
+async function runAppWorker(launch: AppWorkerLaunch, args: CommandArgs, ports: WorkerEntryPorts, parentLeft: Promise<void>): Promise<number> {
+  const { runServe, startModelServer, startTranscriptionServer, StartupCancelledError } = await import("./serve");
+  const signals = parentSignals(ports.signals);
+  void parentLeft.finally(signals.leave);
+  const exit = Promise.withResolvers<number>();
+  const supplied = ports.app ?? {};
+  const start = supplied.start ?? startModelServer, startTranscription = supplied.startTranscription ?? startTranscriptionServer;
+  const link: { current?: ModelHostLink } = {};
+  let modelId: string | undefined;
+  const memoryNow = await processMemory();
+  // One admin surface for the model the app serves; the transcription-only app has no execution lease.
+  const socket = (model: ModelRecord, lease: boolean) => {
+    modelId = model.repoId;
+    const admin = createWorkerRoutes({ modelId: model.repoId, pid: process.pid, ...(memoryNow ? { memory: memoryNow } : {}), ...(lease ? {
+      acquireExecutionLease(signal: AbortSignal) {
+        if (!link.current) return Promise.reject(new Error("no model host is attached"));
+        return link.current.acquireExecutionLease(signal);
+      } } : {}) });
+    return { unix: launch.socketPath, routes: (routes: RouteGroup) => admin.wrap(routes), beforeDrain: () => admin.close() };
+  };
+  try {
+    let app: ServedApp;
+    try {
+      app = await runServe(args, {
+        ...(supplied.resolve ? { resolve: supplied.resolve } : {}), ...(supplied.log ? { log: supplied.log } : {}),
+        // A worker serves: the Neural Engine install step belongs to the process that launched it (the CLI runs
+        // `mlx-bun compile-ane --quiet` before every model load), never to the serving process.
+        installAne: async () => {},
+        start: (model, options) => start(model, options, { ...socket(model, true), link }),
+        startTranscription: (model, options) => startTranscription(model, options, socket(model, false)),
+        // The app listens on the socket, so `--port` names no listener to check (the parent may serve it).
+        interactive: false, probe: async () => null, open() {}, signals, exit: code => exit.resolve(code),
+      });
+    } catch (error) {
+      const message = error instanceof Error ? error.message : String(error);
+      // Only runServe's own cancellation (a stop before any app existed) is clean.
+      // A load failure, or a failed close of an app a stop interrupted, fails the worker with its error.
+      if (error instanceof StartupCancelledError) { console.error(`worker startup cancelled: ${message}`); return 0; }
+      console.error(`worker startup failed: ${message}`);
+      return 1;
+    }
+    // A stop during the load closed the app inside runServe, successfully; nothing was announced.
+    if (app.aborted) return 0;
+    // A stop after startup reached the shutdown handlers, which are closing the app: no ready line.
+    if (!signals.stopped) {
+      try { ports.write(formatWorkerMessage({ type: "ready", socketPath: launch.socketPath, modelId: modelId!, pid: process.pid, version: WORKER_PROTOCOL_VERSION })); }
+      catch (error) {
+        console.error(`worker ready line failed: ${error instanceof Error ? error.message : String(error)}`);
+        try { await app.close(); } catch (failure) { console.error(failure instanceof Error ? failure.message : String(failure)); }
+        return 1;
+      }
+    }
+    return await exit.promise;
+  } finally { signals.dispose(); }
+}
+
+// The process owner exits after the host has released the model; native
+// handles the runtime still holds must not keep a stopped worker alive.
+if (import.meta.main) process.exit(await runWorkerEntry());

package/src/cli/worker-unit.ts

@@ -0,0 +1,179 @@
+// One model worker, as the isolated host holds it resident (residency/model-residency.ts).
+// A unit is a supervised worker process (jobs/worker-supervisor.ts) that serves one model on
+// its own Unix socket: the host leases it per request, forwards to it, pauses it for a
+// managed job, and closes it by draining the worker, letting it flush its saved state, and
+// waiting for it to exit. A crash respawns that worker with the same model within the
+// supervisor's budget, and it resumes from the saved state. Nothing here loads a model or
+// reaches the native library: the worker does, in its own process.
+import { join } from "node:path";
+import { fileURLToPath } from "node:url";
+import type { AppEvent, ModelOperation, ModelOperations, ModelRole } from "@mlx-bun/app-core";
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+import type { DisposableResource } from "@mlx-bun/inference/contracts/portable";
+import { heldBytes, parseMemoryHealth, parseMemoryLine, type WorkerMemory } from "../jobs/worker-memory";
+import { WORKER_PROTOCOL_VERSION } from "../jobs/worker-process";
+import { superviseWorker, type WorkerRestartBudget, type WorkerSupervisor } from "../jobs/worker-supervisor";
+import type { UnitClosed } from "../residency/model-residency";
+import { forwardToWorker } from "../server/proxy-routes";
+import type { RoutedUnit } from "../server/model-routes";
+import { shutdownTimeoutMs, type ServeOptions } from "./serve-options";
+import { workerAdapters } from "./worker-adapters";
+
+export interface WorkerUnit extends RoutedUnit {
+  readonly record: ModelRecord;
+  readonly role: ModelRole;
+  readonly supervisor: WorkerSupervisor;
+  /** Milliseconds from spawning the worker to its ready line (process start, native runtime, weights). */
+  readonly readyMs: number;
+  /** The MLX memory the worker last reported (active, cache, peak, the device working set); undefined until its first report and while it is down. */
+  measured(): WorkerMemory | undefined;
+}
+
+/** What the host supplies for every worker it spawns. */
+export interface WorkerUnitContext {
+  options: ServeOptions;
+  /** The model the server started with: its `--draft-*`, `--adapter` and `--mtp` belong to it alone. */
+  startup: ModelRecord;
+  /** A private directory (0700) the host removes; one socket per worker. */
+  socketDir: string;
+  publish(event: AppEvent): void;
+  entry?: string;
+  restarts?: Partial<WorkerRestartBudget>;
+  env?: Record<string, string | undefined>;
+  readyTimeoutMs?: number;
+  /** Overrides the CLI's shutdown budget: how long a worker may take to drain and to stop before it is killed. */
+  graceMs?: number;
+  notice(line: string): void;
+  log?(line: string): void;
+  error?(line: string): void;
+}
+
+let sockets = 0;
+
+/** The launch record of a chat worker: the resolved model and the options, one model only. */
+function chatLaunch(context: WorkerUnitContext, record: ModelRecord, socketPath: string) {
+  const { options } = context;
+  // The worker serves one model, in-process, with no Whisper of its own: Whisper is its own worker.
+  const own = record.repoId === context.startup.repoId ? options : { ...options, draft: undefined, adapterDir: undefined, mtp: undefined };
+  return { version: WORKER_PROTOCOL_VERSION, socketPath, model: record, options: { ...own, inProcess: true, isolate: undefined, whisper: undefined } };
+}
+
+/** Whisper runs the transcription-only server on its socket, retaining the configured weight residency policy. */
+function transcriptionLaunch(context: WorkerUnitContext, record: ModelRecord, socketPath: string) {
+  const whisper = context.options.whisper;
+  return { kind: "app" as const, version: WORKER_PROTOCOL_VERSION, socketPath, argv: ["--model", record.path,
+    ...(whisper?.resident ? ["--whisper-resident"] : []),
+    ...(whisper?.idleUnloadSec !== undefined ? ["--whisper-idle-unload", String(whisper.idleUnloadSec)] : [])] };
+}
+
+/** Publish what the worker's engine publishes on its own bus (the worker's load and unload are the host's to report), and hand
+ * its memory readings to `measure`. A worker that is gone measures nothing: its replacement will need what it did. */
+function relayEvents(supervisor: WorkerSupervisor, publish: (event: AppEvent) => void, measure: (memory: WorkerMemory | undefined) => void, signal: AbortSignal): void {
+  void (async () => {
+    while (!signal.aborted) {
+      try {
+        await supervisor.whenReady(signal);
+        const response = await supervisor.fetch("http://engine/admin/events", { signal });
+        if (!response.body) return;
+        const reader = response.body.getReader(), decoder = new TextDecoder();
+        let buffered = "";
+        for (;;) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buffered += decoder.decode(value, { stream: true });
+          for (let newline = buffered.indexOf("\n"); newline !== -1; newline = buffered.indexOf("\n")) {
+            const line = buffered.slice(0, newline).trim();
+            buffered = buffered.slice(newline + 1);
+            if (!line) continue;
+            try {
+              const parsed = JSON.parse(line) as AppEvent;
+              const memory = parseMemoryLine(parsed);
+              if (memory) measure(memory);
+              else if (parsed.type !== "model.load" && parsed.type !== "model.unload") publish(parsed);
+            } catch { /* a torn line is dropped */ }
+          }
+        }
+      } catch { /* the worker went away or is not ready: wait for its respawn */ }
+      if (!signal.aborted) measure(undefined);
+      if (signal.aborted || supervisor.state === "exhausted" || supervisor.state === "closed") return;
+      await new Promise(resolve => setTimeout(resolve, 200));
+    }
+  })();
+}
+
+/** Spawn a worker for `record` and resolve once it serves; a worker that fails to start rejects and leaves nothing behind. */
+export async function spawnWorkerUnit(context: WorkerUnitContext, record: ModelRecord, role: ModelRole, estimate: number): Promise<WorkerUnit> {
+  const socketPath = join(context.socketDir, `worker-${++sockets}.sock`);
+  const started = performance.now();
+  const budget = context.graceMs ?? shutdownTimeoutMs();
+  const supervisor = superviseWorker({
+    entry: context.entry ?? fileURLToPath(new URL("./worker-entry.ts", import.meta.url)), socketPath,
+    launch: role === "primary" ? chatLaunch(context, record, socketPath) : transcriptionLaunch(context, record, socketPath),
+    ...(context.restarts ? { restarts: context.restarts } : {}),
+    ...(context.env ? { env: context.env } : {}),
+    ...(context.readyTimeoutMs !== undefined ? { readyTimeoutMs: context.readyTimeoutMs } : {}),
+    // The worker flushes its model's saved state when it closes, and a kill would cut that: it gets the CLI's shutdown
+    // budget to drain and then to stop.
+    drainTimeoutMs: budget, graceMs: budget,
+    notice: context.notice,
+    ...(context.log ? { log: context.log } : {}),
+    ...(context.error ? { error: context.error } : {}),
+  });
+  try { await supervisor.ready; }
+  catch (error) { await supervisor.close().catch(() => {}); throw error; }
+  const readyMs = performance.now() - started;
+  context.notice(`${role === "primary" ? "engine" : "companion"} worker pid ${supervisor.pid} ready for ${record.repoId} in ${readyMs.toFixed(0)} ms (socket ${supervisor.socketPath})`);
+
+  // What the worker reports of the model it loaded: its weights, whether saved state was found for it, and (both roles) the
+  // MLX memory its process holds, which residency counts from now on in place of the estimate.
+  let weightsBytes = role === "companion" ? record.sizeBytes : 0, resumed = false, measured: WorkerMemory | undefined;
+  try {
+    const response = await supervisor.fetch("http://engine/health", { signal: AbortSignal.timeout(5_000) });
+    measured = parseMemoryHealth(await response.json());
+  } catch { /* the estimate stands until a report arrives */ }
+  if (role === "primary") {
+    try {
+      const response = await supervisor.fetch("http://engine/stats", { signal: AbortSignal.timeout(5_000) });
+      const stats = await response.json() as { admission?: { weights_bytes?: number }; ssd_cache?: { entries?: number } };
+      weightsBytes = stats.admission?.weights_bytes ?? 0;
+      resumed = (stats.ssd_cache?.entries ?? 0) > 0;
+    } catch { /* the estimate stands */ }
+  }
+  const stopEvents = new AbortController();
+  relayEvents(supervisor, context.publish, memory => { measured = memory; }, stopEvents.signal);
+  const operations: readonly ModelOperation[] = role === "primary" ? ["generate"] : ["transcribe"];
+  const currentMemory = () => supervisor.state === "ready" ? measured : undefined;
+  let closing: Promise<UnitClosed> | undefined;
+  return {
+    id: record.repoId, record, role, supervisor, readyMs, operations, resumed,
+    routes: { handle: request => forwardToWorker(supervisor, request) },
+    measured: currentMemory,
+    // The estimate stands only until the worker's first report; after it, the process's own active and cache memory, never
+    // below the weights it reported (MLX maps weights lazily, so a worker that has not run yet can read low).
+    bytes: () => {
+      const memory = currentMemory();
+      return memory ? Math.max(heldBytes(memory), weightsBytes) : Math.max(estimate, weightsBytes);
+    },
+    // A decision about who fits reads each worker's memory now rather than trusting the last pushed line.
+    refresh: async () => {
+      if (supervisor.state !== "ready") return;
+      try {
+        const response = await supervisor.fetch("http://engine/health", { signal: AbortSignal.timeout(2_000) });
+        const reading = parseMemoryHealth(await response.json());
+        if (reading) measured = reading;
+      } catch { /* the last reading stands */ }
+    },
+    memory: () => ({ weightsBytes, kvBytes: 0, prefixCacheBytes: 0 }),
+    operationsFor: () => (role === "primary" ? { generate: (request: Request) => forwardToWorker(supervisor, request), adapters: workerAdapters(supervisor) } : {}) as Partial<ModelOperations>,
+    broken: () => supervisor.state === "exhausted" || supervisor.state === "closed",
+    // A job needs the GPU to itself: the worker's own execution lease. Whisper's decode is brief and holds its own lock.
+    pause: signal => role === "primary" ? supervisor.acquireExecutionLease(signal ?? new AbortController().signal) : Promise.resolve({ dispose() {} } as DisposableResource),
+    close: () => closing ??= (async () => {
+      stopEvents.abort();
+      const drained = performance.now();
+      await supervisor.close();
+      // The worker exits 0 only when its saved state was durable; a kill after the budget, or an incomplete flush, is not.
+      return { flushed: supervisor.lastExit?.code === 0, drainMs: performance.now() - drained };
+    })(),
+  };
+}

package/src/engine/adapter-operation.ts

@@ -0,0 +1,26 @@
+// The `adapters` operation of a loaded model: mounting, unmounting and merging LoRA adapters on its own weights, each
+// under the engine's execution lock. This is where the numerical work happens (native MLX), so it runs in the process
+// that holds the model; a host that isolates its models reaches it through the worker's private route.
+import type { AdapterMergeStats, AdapterOperation, MountedAdapter } from "@mlx-bun/app-core";
+import type { AdapterInfo } from "@mlx-bun/inference/adapters";
+import type { GenerationGateway } from "./generation-gateway";
+import type { LoadedModelContext } from "./model-host";
+
+const mountedAdapter = (info: AdapterInfo): MountedAdapter => ({ id: info.id, path: info.path, rank: info.rank, scale: info.scale,
+  sizeBytes: info.sizeBytes, mountedLayers: info.mountedLayers, ramBytes: info.ramBytes });
+
+export function createAdapterOperation(context: Pick<LoadedModelContext, "adapters">, gateway: Pick<GenerationGateway, "runExclusive">,
+  /** Test seam: the merge, default the training library's. */
+  merge?: (sources: string[], output: string, scales?: number[]) => Promise<object>): AdapterOperation {
+  return {
+    async list() { return context.adapters.list().map(mountedAdapter); },
+    mount: (id, directory, signal) => gateway.runExclusive(async () => mountedAdapter(await context.adapters.mount(id, directory)), undefined, signal),
+    unmount: (id, signal) => gateway.runExclusive(async () => context.adapters.unmount(id), undefined, signal),
+    // The library releases every tensor it materialized before it resolves, so the lock is held until cleanup finished.
+    merge: (request, signal) => gateway.runExclusive(async () => {
+      const run = merge ?? (await import("@mlx-bun/training/merge")).mergeAdapters;
+      signal?.throwIfAborted();
+      return await run([...request.adapters], request.output, request.scales ? [...request.scales] : undefined) as AdapterMergeStats;
+    }, undefined, signal),
+  };
+}

package/src/engine/cache-services.ts

@@ -0,0 +1,114 @@
+import { readFileSync } from "node:fs";
+import { modelIdentity } from "@mlx-bun/hub/registry";
+import { kvSchemeOf } from "@mlx-bun/inference/state/kv-scheme";
+import type { CacheCodecProvider, KvCacheService, KvCacheServiceOptions, CacheServiceGraph } from "@mlx-bun/inference/state";
+import type { LoadedModelContext } from "./model-host";
+import type { ModelBinding } from "./model-binding";
+import { createKvBudget, type KvBudget } from "./kv-budget";
+
+/** Cache policy: the serve flags the composition's cache service is built
+ * with. The KV scheme and paging are the loaded composition's
+ * (`LoadedModelContext.composition`), resolved with the model; nothing here
+ * resolves them again. */
+export interface CacheServiceOptions {
+  promptCacheBytes?: number;
+  /** Aggregate projected KV bytes of admitted requests (`--kv-budget`); unset means no budget. */
+  kvBudgetBytes?: number;
+  /** The saved-state root; each model's store lives in its own directory under it. */
+  ssdCacheDir?: string;
+  /** Bytes the whole root may hold, across every model's directory; unset means unlimited. */
+  ssdCacheMaxBytes?: number;
+  /** One budget shared by every store in the process; without it this call makes one over `ssdCacheDir`. */
+  ssdBudget?: KvBudget;
+  ssdCacheVerify?: boolean;
+  ssdDemoteIdleSec?: number;
+  generationCheckpointTokens?: number;
+  /** An explicit allocator cap also limits optional RAM cache residency. */
+  allocatorLimitBytes?: number;
+}
+
+/** Construction ports let CPU tests and replacement runtimes observe and
+ * replace what the app supplies to the library's cache service. */
+export interface CacheServiceDependencies {
+  defaultStateCodecs: CacheCodecProvider;
+  createService(graph: CacheServiceGraph, options: KvCacheServiceOptions): KvCacheService;
+  costSizeRetention(): NonNullable<KvCacheServiceOptions["retention"]>;
+  configFingerprint(config: LoadedModelContext["model"]["config"]): string;
+  modelIdentity(modelDir: string): string;
+  activeMemory(): number;
+  maxWorkingSet(): number;
+  scheduleDemotion(run: () => void, intervalMs: number): () => void;
+}
+
+async function defaultDependencies(): Promise<CacheServiceDependencies> {
+  const [state, artifacts, memory] = await Promise.all([
+    import("@mlx-bun/inference/state"), import("@mlx-bun/inference/artifacts"), import("@mlx-bun/mlx/ffi"),
+  ]);
+  return {
+    defaultStateCodecs: state.legacyCacheCodecs,
+    createService: (graph, options) => state.createKvCacheService(graph, options),
+    costSizeRetention: () => new state.CostSizeRetention(), configFingerprint: artifacts.configFingerprint,
+    modelIdentity,
+    activeMemory: memory.activeMemory, maxWorkingSet: memory.maxRecommendedWorkingSetSize,
+    scheduleDemotion(run, interval) { const timer = setInterval(run, interval); timer.unref(); return () => clearInterval(timer); },
+  };
+}
+
+/** The loaded model's cache service (`state/cache-service`), built from the
+ * serve flags without starting execution: the scheduler, the server and the
+ * model host all reach request state, prefix reuse, saved state and its budget
+ * through it. The caller borrows it until execution has drained, then closes
+ * it before disposing the model. SSD is opt-in; plain KV and 8 GB RAM retain
+ * main's defaults. Native library imports are deferred until construction. */
+export async function createCacheServices(context: LoadedModelContext, binding: ModelBinding,
+  options: CacheServiceOptions = {}, supplied?: CacheServiceDependencies) {
+  if (options.generationCheckpointTokens !== undefined) {
+    if (!Number.isInteger(options.generationCheckpointTokens) || options.generationCheckpointTokens < 1)
+      throw new Error("generationCheckpointTokens must be a positive integer");
+    if (!options.ssdCacheDir) throw new Error("generation checkpoints require an SSD cache directory");
+  }
+  const cap = options.promptCacheBytes ?? 8e9;
+  if (options.ssdCacheDir && cap <= 0) throw new Error("SSD cache requires a nonzero RAM prompt cache");
+  if (options.ssdCacheDir && !binding.stateCompatibility?.length)
+    throw new Error("SSD cache requires the model binding's stateCompatibility identity");
+  // The composition's KV scheme: what the graph's caches compose, keys the saved
+  // state and describes `/stats`. A context an engine-owned implementation built
+  // composes none, and stores bf16.
+  const resolvedKvScheme = kvSchemeOf(context.composition?.kv ?? { kind: "bf16" });
+  const deps = supplied ?? await defaultDependencies();
+  const runtime = binding.gateway.runtime;
+  const stateCodecs = context.stateCodecs ?? deps.defaultStateCodecs;
+  let saved: KvCacheServiceOptions["saved"];
+  if (options.ssdCacheDir) {
+    const tokenizer = readFileSync(`${context.model.config.modelDir}/tokenizer.json`);
+    // The hub owns model identity. Snapshot revisions have distinct paths;
+    // ordinary restarts reuse the same directory without reading weight bytes.
+    const architecture = deps.configFingerprint(context.model.config);
+    const artifact = deps.modelIdentity(context.model.config.modelDir);
+    const rawQueueGb = Number(runtime.value("MLX_BUN_SSD_SPILL_QUEUE_GB"));
+    saved = { dir: options.ssdCacheDir,
+      identity: { modelId: context.modelId,
+        configFingerprint: `${architecture}-${resolvedKvScheme.cacheKey}-${Bun.hash(binding.stateCompatibility).toString(16)}-${Bun.hash(artifact).toString(16)}`,
+        tokenizerHash: Bun.hash(tokenizer).toString(16) },
+      budget: options.ssdBudget ?? createKvBudget(options.ssdCacheDir, options.ssdCacheMaxBytes ?? Infinity),
+      verify: options.ssdCacheVerify,
+      storage: { layout: runtime.value("MLX_BUN_SSD_LAYOUT") === "blocks" ? "blocks" : "whole",
+        segmented: runtime.value("MLX_BUN_SSD_SEGMENTED") !== "0" },
+      writeBehind: runtime.value("MLX_BUN_SSD_WRITEBEHIND") !== "0",
+      prefetch: runtime.value("MLX_BUN_SSD_PREFETCH") !== "0",
+      demoteIdleMs: (options.ssdDemoteIdleSec ?? 300) * 1000, schedule: deps.scheduleDemotion,
+      ...(options.generationCheckpointTokens ? { checkpointEveryTokens: options.generationCheckpointTokens,
+        checkpointQueueBytes: (Number.isFinite(rawQueueGb) && rawQueueGb >= 0 ? rawQueueGb : 2) * 1024 ** 3 } : {}) };
+  }
+  const workingSet = Math.min(deps.maxWorkingSet(), options.allocatorLimitBytes ?? Infinity);
+  const cache = deps.createService(binding.gateway.graph, {
+    promptCacheBytes: cap, codecs: stateCodecs, runtime,
+    ...(options.kvBudgetBytes !== undefined ? { kvBudgetBytes: options.kvBudgetBytes } : {}),
+    ...(runtime.value("MLX_BUN_CACHE_RETENTION") === "cost-size" ? { retention: deps.costSizeRetention() } : {}),
+    pressure: held => Math.max(deps.activeMemory(), context.model.weightsBytes + held) > workingSet * 0.85,
+    identity: { artifact: context.profile.artifact, implementation: context.profile.profile.execution,
+      stateAbi: "legacy-cache-array-v1", codecs: stateCodecs.id },
+    ...(saved ? { saved } : {}),
+  });
+  return { cache, resolvedKvScheme, adapterNamespace: (adapters: string[]) => context.adapters.cacheNamespace(adapters) };
+}

package/src/engine/completion.ts

@@ -0,0 +1,39 @@
+import type { GenerateOptions, GenerateStats } from "@mlx-bun/inference/generation";
+import type { ExecutionReason, ExecutionRequirements, ResolvedExecution, OnToken } from "@mlx-bun/inference/contracts/portable";
+import type { Vision } from "@mlx-bun/inference/contracts/mlx";
+import type { PromptResponseTrace } from "@mlx-bun/inference/runtime/trace";
+
+export type { OnToken } from "@mlx-bun/inference/contracts/portable";
+export type { Vision } from "@mlx-bun/inference/contracts/mlx";
+export interface RequestShape extends ExecutionRequirements {}
+export interface GenerationPlacement {
+  readonly shape: RequestShape;
+  readonly mechanism: "continuous";
+  readonly execution: ResolvedExecution;
+}
+
+const sharedExecutionExclusions = new Set<ExecutionReason>([
+  "media-batch-unsupported", "grammar-batch-unsupported", "method-batch-unsupported",
+]);
+
+/** A migrated request shape whose shared executor is not available yet. The
+ * message names the missing capabilities (`reasons`), never a model.
+ * Transport layers can distinguish this capability gap from an execution failure. */
+export class UnsupportedExecutionError extends Error {
+  readonly reasons: readonly ExecutionReason[];
+  constructor(readonly modelType: string, readonly method: ResolvedExecution["method"], reasons: readonly ExecutionReason[]) {
+    const exclusions = reasons.filter(reason => sharedExecutionExclusions.has(reason));
+    super(`method ${method} does not support shared execution${exclusions.length ? `: ${exclusions.join(", ")}` : ""}`);
+    this.name = "UnsupportedExecutionError";
+    this.reasons = Object.freeze(exclusions);
+  }
+}
+
+/** App-owned completion boundary. HTTP parsing and output formatting stay above it. */
+export interface CompletionEngine {
+  place(shape: RequestShape, options?: GenerateOptions): GenerationPlacement;
+  run(promptIds: number[], options: GenerateOptions & { stopSequences?: string[] },
+    onToken: OnToken, vision: Vision | undefined, shape: RequestShape,
+    placement: GenerationPlacement, signal?: AbortSignal,
+    trace?: PromptResponseTrace): Promise<GenerateStats>;
+}