mlx-bun 0.4.0 → 1.0.0

package/src/modules.ts

@@ -0,0 +1,58 @@
+import type { MemoryModuleOptions } from "@mlx-bun/module-memory";
+import type { AppModule } from "@mlx-bun/app-core";
+import { manifest as benchmarks } from "@mlx-bun/module-benchmarks/manifest";
+import { manifest as chat } from "@mlx-bun/module-chat/manifest";
+import { manifest as datasets } from "@mlx-bun/module-datasets/manifest";
+import { manifest as memory } from "@mlx-bun/module-memory/manifest";
+import { manifest as metrics } from "@mlx-bun/module-metrics/manifest";
+import { manifest as models } from "@mlx-bun/module-models/manifest";
+import { manifest as quantize } from "@mlx-bun/module-quantize/manifest";
+import { manifest as train } from "@mlx-bun/module-train/manifest";
+import { manifest as transcription } from "@mlx-bun/module-transcription/manifest";
+
+// The one file that names module packages: import each `@mlx-bun/module-<id>`
+// here and list it, and add it to package.json's dependencies. Order is
+// activation order. The manifests are plain data, so the command list, `--help`
+// and argument parsing read them without loading a module; a module's code
+// loads only when a host activates it.
+export const manifests: readonly Omit<AppModule, "activate">[] = [transcription, datasets, metrics, quantize, benchmarks, train, models, chat, memory];
+
+/** Where a module activates: `state` is the app's persistent services (job runners, sockets, storage, the served
+ * model's wire: modules that require `jobs` or declare sockets), `model` the model host's (Whisper leases, the
+ * catalog). A CLI verb's host names its verb's module instead: a selection by manifest. */
+export type ModuleScope = "state" | "model";
+
+/** What the app decides for a module beyond its services: the chat's read-only policy and working directory. */
+export interface ModuleSettings {
+  memory?: MemoryModuleOptions;
+  chat?: { readOnly?: boolean; cwd?: string };
+}
+
+/** Each module's code, loaded on demand by id. */
+const loaders: Record<string, (settings: ModuleSettings) => Promise<AppModule>> = {
+  transcription: async () => (await import("@mlx-bun/module-transcription")).default,
+  datasets: async () => (await import("@mlx-bun/module-datasets")).default,
+  metrics: async () => (await import("@mlx-bun/module-metrics")).default,
+  quantize: async () => (await import("@mlx-bun/module-quantize")).default,
+  benchmarks: async () => (await import("@mlx-bun/module-benchmarks")).default,
+  train: async () => (await import("@mlx-bun/module-train")).default,
+  models: async () => (await import("@mlx-bun/module-models")).default,
+  memory: async settings => (await import("@mlx-bun/module-memory")).createMemoryModule(settings.memory),
+  chat: async settings => (await import("@mlx-bun/module-chat")).createChatModule(settings.chat),
+};
+
+/** The installed modules of a scope (or of a manifest selection), in activation order; a module's code loads only when selected. */
+export async function installedModules(scope: ModuleScope | ((manifest: Omit<AppModule, "activate">) => boolean) = "model", settings: ModuleSettings = {}): Promise<readonly AppModule[]> {
+  const select = typeof scope === "function" ? scope
+    : (manifest: Omit<AppModule, "activate">) => (manifest.placement ?? (manifest.requires.includes("jobs") || (manifest.sockets?.length ?? 0) > 0 ? "app" : "model")) === (scope === "state" ? "app" : "model");
+  return Promise.all(manifests.filter(select).map(manifest => loaders[manifest.id]!(settings)));
+}
+
+// Narrow public contracts and adapters used by this host's composition; module dependencies stay centralized here.
+export { runMemory as runModuleMemory } from "@mlx-bun/module-memory/cli";
+export type { MemoryDependencies, CommandArgs as MemoryCommandArgs } from "@mlx-bun/module-memory/cli";
+export { MEMORY_TASK_MODEL, adapterDirFor, locateTaskModel, memoryBatchSize, memoryPromptIds } from "@mlx-bun/module-memory/model";
+export type { MemoryCompletionClient, MemoryCompletionRequest } from "@mlx-bun/module-memory/model";
+export type { SynthesisClient } from "@mlx-bun/module-memory/synthesis";
+export { vaultRoot } from "@mlx-bun/module-memory/vault";
+export { memory as memoryManifest };

package/src/publishing/catalog-hub.ts

@@ -0,0 +1,18 @@
+import type { CatalogHub } from "@mlx-bun/app-services";
+import { createHfCredentials } from "./credentials";
+import { createPublisher } from "./upload";
+
+/** The Hub side of the model catalog: downloads through the hub library, transfers that outlive a request when the host owns them, and pushes with the app's saved token. */
+export function createCatalogHub(credentials: Pick<ReturnType<typeof createHfCredentials>, "get"> = createHfCredentials(), options: Pick<CatalogHub, "transfers"> = {}): CatalogHub {
+  const publish = createPublisher({ credentials, getJob: () => null });
+  return {
+    // The hub download honors the signal at every checkpoint and keeps the blob's .incomplete prefix resumable.
+    async download(id, { revision, signal, onProgress }) {
+      const { downloadModel } = await import("@mlx-bun/hub/download");
+      return downloadModel(id, { onProgress, signal, ...(revision !== undefined ? { revision } : {}) });
+    },
+    ...options,
+    canPublish: () => credentials.get() !== null,
+    publish: (directory, { repoId, signal, ...rest }) => publish({ kind: "quantize", repoId, sourcePath: directory, ...(signal ? { signal } : {}), ...rest }),
+  };
+}

package/src/publishing/credentials.ts

@@ -0,0 +1,31 @@
+import { chmodSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
+import { dirname } from "node:path";
+import { hfToken, type HfTokenOptions } from "@mlx-bun/hub/download";
+import { mlxBunHome, storagePath, userHome } from "../storage/paths";
+
+/** The app owns its saved token. Hub owns the shared environment/cache fallback.
+ * Explicit paths and environment isolate embedded consumers and tests. */
+export function createHfCredentials(options: HfTokenOptions & { tokenFile?: string } = {}) {
+  const environment = () => options.environment ?? process.env;
+  const file = () => options.tokenFile ?? storagePath("credentials", mlxBunHome(userHome(environment())));
+  return {
+    get(): string | null {
+      try {
+        const saved: unknown = JSON.parse(readFileSync(file(), "utf8"));
+        if (saved && typeof saved === "object" && "token" in saved && typeof saved.token === "string" && saved.token.trim())
+          return saved.token.trim();
+      } catch { /* Missing or malformed app settings fall back to shared credentials. */ }
+      const env = environment();
+      return hfToken({ environment: { HOME: userHome(env), HF_TOKEN: env.HF_TOKEN?.trim() || undefined },
+        cacheTokenPath: options.cacheTokenPath });
+    },
+    save(token: string): void {
+      const trimmed = token.trim();
+      if (!trimmed) throw new Error("saveHfToken: token is empty");
+      const target = file();
+      mkdirSync(dirname(target), { recursive: true });
+      writeFileSync(target, JSON.stringify({ token: trimmed, savedAt: new Date().toISOString() }, null, 2) + "\n", { mode: 0o600 });
+      chmodSync(target, 0o600);
+    },
+  };
+}

package/src/publishing/upload.ts

@@ -0,0 +1,37 @@
+import { uploadFolder } from "@mlx-bun/hub/upload";
+import type { createHfCredentials } from "./credentials";
+
+export interface PublishRequest {
+  kind: "quantize" | "finetune" | "dataset";
+  repoId: string;
+  private?: boolean;
+  sourcePath?: string;
+  jobId?: string;
+  commitMessage?: string;
+  /** Bytes of `file` sent so far out of `total`. */
+  onProgress?: (file: string, sent: number, total: number) => void;
+  /** Cancels the upload; see `@mlx-bun/hub/upload` for what an abort can and cannot undo. */
+  signal?: AbortSignal;
+}
+
+/** App policy chooses an artifact and credentials; hub owns the upload protocol.
+ * The job lookup is borrowed and read-only. No model or execution lease is needed. */
+export function createPublisher(options: {
+  credentials: Pick<ReturnType<typeof createHfCredentials>, "get">;
+  getJob(id: string): { output_path: string | null } | null | undefined | Promise<{ output_path: string | null } | null | undefined>;
+  upload?: typeof uploadFolder;
+}) {
+  return async (request: PublishRequest) => {
+    const token = options.credentials.get();
+    if (!token) throw new Error("no HF token saved — add one in Settings → Hugging Face");
+    let source = request.sourcePath;
+    if (!source && request.jobId) source = (await options.getJob(request.jobId))?.output_path ?? undefined;
+    if (!source) throw new Error("no source dir (pass job_id or source_path)");
+    return (options.upload ?? uploadFolder)(source, request.repoId, {
+      repoType: request.kind === "dataset" ? "dataset" : "model",
+      private: !!request.private, token, ...(request.signal ? { signal: request.signal } : {}),
+      ...(request.commitMessage !== undefined ? { commitMessage: request.commitMessage } : {}),
+      ...(request.onProgress ? { onProgress: request.onProgress } : {}),
+    });
+  };
+}

package/src/residency/leased-adapters.ts

@@ -0,0 +1,21 @@
+// The served model's adapter operation for a host that holds residency (serve-host.ts, serve-isolated.ts): every call leases
+// the current model, so it cannot be evicted while the call runs, and runs on the operation the model's own unit exposes.
+import type { AdapterOperation, ModelId } from "@mlx-bun/app-core";
+import type { ResidencyHost, ResidentUnit } from "./model-residency";
+
+export function leasedAdapters<U extends ResidentUnit>(host: Pick<ResidencyHost<U>, "acquire">, current: () => ModelId): AdapterOperation {
+  async function on<T>(signal: AbortSignal | undefined, use: (operation: AdapterOperation) => Promise<T>): Promise<T> {
+    const lease = await host.acquire(current(), signal ? { signal } : {});
+    try {
+      const operation = lease.operations.adapters;
+      if (!operation) throw new Error(`${lease.model.id} does not mount adapters`);
+      return await use(operation);
+    } finally { lease.release(); }
+  }
+  return {
+    list: signal => on(signal, operation => operation.list(signal)),
+    mount: (id, directory, signal) => on(signal, operation => operation.mount(id, directory, signal)),
+    unmount: (id, signal) => on(signal, operation => operation.unmount(id, signal)),
+    merge: (request, signal) => on(signal, operation => operation.merge(request, signal)),
+  };
+}

package/src/residency/local-records.ts

@@ -0,0 +1,16 @@
+// The local models a host may load, as the registry knows them: never a scan of the
+// hub cache beyond the first (an empty index) and never a download.
+import type { ModelOperation } from "@mlx-bun/app-core";
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+import { openRegistry } from "../storage/paths";
+
+/** Local checkpoints that declare `operation`: `generate` for chat models, `transcribe` for Whisper. */
+export async function listLocalRecords(operation: ModelOperation): Promise<readonly ModelRecord[]> {
+  const { declaredOperations } = await import("@mlx-bun/app-services/catalog");
+  const registry = openRegistry();
+  try {
+    // A fresh machine's index is empty until its first scan.
+    if (registry.list().length === 0) await registry.scan();
+    return registry.listCanonical().filter(record => declaredOperations(record.modelType).includes(operation));
+  } finally { registry.close(); }
+}

package/src/residency/model-residency.ts

@@ -0,0 +1,405 @@
+// Residency by memory fit: the `ModelHost` contract (@mlx-bun/app-core) for
+// models that generate. A model that fits the budget loads beside the others;
+// otherwise the least recently used unpinned, unleased model is drained, has
+// its state flushed durably, and is released before the newcomer loads, so
+// resident bytes never exceed the budget during a swap. A request for a model
+// that cannot fit while every other model is busy waits; it never evicts one
+// mid-request. Nothing here loads weights or names a model family: a source
+// resolves ids and loads units, and a unit says what it holds and how to close.
+import type {
+  AcquireOptions, CoreEvent, ModelHost, ModelHostError, ModelId, ModelLease, ModelOperation, ModelOperations, ModelRole, ModelStats,
+  ResidencyFailure, ResidencyPlan, ResidencyPolicy, ResidentModel,
+} from "@mlx-bun/app-core";
+import type { DisposableResource } from "@mlx-bun/inference/contracts/portable";
+
+export class ResidencyError extends Error implements ModelHostError {
+  constructor(readonly code: ResidencyFailure, message: string, options?: { cause?: unknown }) {
+    super(message, options);
+    this.name = "ResidencyError";
+  }
+}
+
+/** What closing a unit reported. */
+export interface UnitClosed {
+  /** Saved state reached durable storage before the weights were released. */
+  readonly flushed: boolean;
+  readonly drainMs?: number;
+  readonly flushMs?: number;
+}
+
+/** One loaded model the host serves. The host owns its release. */
+export interface ResidentUnit {
+  readonly id: ModelId;
+  readonly operations: readonly ModelOperation[];
+  /** Saved state from an earlier run was found when it loaded, so a request can resume from it. */
+  readonly resumed?: boolean;
+  /** It can no longer serve (a worker whose restarts are spent): the host releases it and loads the model again on the next request. */
+  broken?(): boolean;
+  /** Bytes it holds now: weights plus the state it carries. */
+  bytes(): number;
+  /** Re-read what it holds (a worker's measured memory) so the next `bytes()` is current; the host awaits it before it decides who fits. It never rejects the decision: a failure leaves the last reading. */
+  refresh?(): Promise<void>;
+  memory(): { readonly weightsBytes: number; readonly kvBytes: number; readonly prefixCacheBytes: number };
+  /** What a lease exposes; only what the model declared. */
+  operationsFor(): Partial<ModelOperations>;
+  /** Stop its execution until the resource is disposed (a job or a companion needs the GPU alone). */
+  pause(signal?: AbortSignal): Promise<DisposableResource>;
+  /** Stop admission, wait for work in flight, flush saved state durably, then release the weights. Once. */
+  close(options: { readonly flush: boolean }): Promise<UnitClosed>;
+}
+
+/** A model the host may load, with what residency needs to know before it does. */
+export interface ResidencyEntry {
+  readonly id: ModelId;
+  /** Bytes it needs resident (weights and working state), estimated before loading. */
+  readonly bytes: number;
+  readonly role?: ModelRole;
+  readonly operations: readonly ModelOperation[];
+}
+
+export interface ResidencySource<U extends ResidentUnit> {
+  /** Exact id only; never scans or downloads. */
+  resolve(id: ModelId): Promise<ResidencyEntry | undefined>;
+  load(entry: ResidencyEntry): Promise<U>;
+}
+
+export interface ResidencyOptions<U extends ResidentUnit> {
+  source: ResidencySource<U>;
+  /** Bytes all resident models may use together. */
+  budgetBytes: number;
+  /** Bytes the process actually holds (MLX active memory); residency never assumes less. */
+  measured?: () => number;
+  /** Bytes held by residents this host does not manage, such as a Whisper companion. */
+  external?: () => number;
+  pinned?: Iterable<ModelId>;
+  /** The model a caller who names none is served with, per operation. */
+  defaultFor?: (operation: ModelOperation) => Promise<ModelId | undefined>;
+  /** How this composition makes a model the one it serves (`ModelHost.serve`): it checks the id, loads the model through `acquire`, and
+   * records it as the default. Without it the host holds no switch and rejects with `not-switchable`. */
+  serve?: (id: ModelId, signal: AbortSignal | undefined) => Promise<void>;
+  /** The directories a resident model's unit reads (its snapshot and any it holds besides), reported as `ResidentModel.uses`. */
+  uses?: (unit: U) => readonly string[];
+  events?: { publish(event: CoreEvent): void };
+  log?: (line: string) => void;
+  now?: () => number;
+}
+
+/** A lease that also hands the composition the unit itself, for routes bound to it. */
+export interface ResidencyLease<U extends ResidentUnit> extends ModelLease { readonly unit: U }
+
+export interface ResidencyHost<U extends ResidentUnit> extends ModelHost {
+  acquire(id: ModelId, options?: AcquireOptions): Promise<ResidencyLease<U>>;
+  /** Register a unit loaded elsewhere (the startup model, a caller's own context) as resident. `pin` keeps it from ever being evicted. */
+  adopt(entry: ResidencyEntry, unit: U, options?: { readonly pin?: boolean; readonly loadMs?: number }): void;
+  /** The resident unit, without leasing it: for read-only views (stats, discovery) that never load a model. */
+  peek(id: ModelId): U | undefined;
+  /** Free `bytes` by draining least recently used models, waiting while every candidate is busy; resolves once they fit. */
+  makeRoom(bytes: number, signal?: AbortSignal): Promise<void>;
+  /** Stop every resident execution until disposed, and hold back new loads meanwhile (a managed job, a companion's decode). */
+  pauseAll(signal?: AbortSignal): Promise<DisposableResource>;
+  /** Stop admission, wait for loads in flight, close every unit. The results are each closed unit's. Idempotent. */
+  close(): Promise<readonly UnitClosed[]>;
+}
+
+/** What an eviction reports: the unit's own account, or the error its close failed with. */
+type Evicted = UnitClosed & { readonly error?: unknown };
+
+interface Slot<U extends ResidentUnit> {
+  readonly id: ModelId;
+  readonly entry: ResidencyEntry;
+  unit: U | undefined;
+  loading: Promise<U> | undefined;
+  closing: Promise<Evicted> | undefined;
+  leases: number;
+  lastUsedAt: number;
+  loads: number;
+  unloads: number;
+  lastLoadMs: number;
+}
+
+export function createResidencyHost<U extends ResidentUnit>(options: ResidencyOptions<U>): ResidencyHost<U> {
+  const { source } = options, budget = options.budgetBytes;
+  if (!(budget > 0)) throw new Error("the residency budget must be positive");
+  const log = options.log ?? (() => {}), now = options.now ?? (() => performance.now());
+  const measured = options.measured ?? (() => 0), external = options.external ?? (() => 0);
+  const slots = new Map<ModelId, Slot<U>>();
+  const pinned = new Set<ModelId>(options.pinned ?? []);
+  const waiters = new Set<() => void>();
+  let closed = false, closing: Promise<readonly UnitClosed[]> | undefined, paused = 0;
+  const closedError = () => new ResidencyError("closed", "the model host is closed");
+  const publish = (event: CoreEvent) => { try { options.events?.publish(event); } catch { /* a subscriber never slows a load */ } };
+  const notify = () => { for (const wake of [...waiters]) wake(); };
+  /** Wake on the next change (a release, a load, an unload, an unpause) or reject when the signal aborts. */
+  const changed = (signal?: AbortSignal) => new Promise<void>((resolve, reject) => {
+    const done = () => { waiters.delete(wake); signal?.removeEventListener("abort", abort); };
+    const wake = () => { done(); resolve(); };
+    const abort = () => { done(); reject(signal!.reason); };
+    waiters.add(wake);
+    if (signal?.aborted) abort(); else signal?.addEventListener("abort", abort, { once: true });
+  });
+
+  const held = (slot: Slot<U>) => slot.unit ? slot.unit.bytes() : slot.loading ? slot.entry.bytes : 0;
+  /** Bytes in use by everyone but `except`: the estimates, floored by what the process measurably holds. */
+  const used = (except?: Slot<U>) => {
+    let sum = external();
+    for (const slot of slots.values()) if (slot !== except) sum += held(slot);
+    return Math.max(sum, measured() - (except ? held(except) : 0));
+  };
+  /** Ask every resident unit for a current reading before a decision that depends on what they hold. */
+  const refreshed = async () => {
+    await Promise.all([...slots.values()].map(slot => slot.unit && !slot.closing ? slot.unit.refresh?.().catch(() => {}) : undefined));
+  };
+  const busy = (slot: Slot<U>) => slot.leases > 0 || !!slot.loading || !!slot.closing;
+  const evictable = (slot: Slot<U>) => !!slot.unit && !busy(slot) && !pinned.has(slot.id);
+
+  /** Least recently used first, the fewest that make `bytes` fit. `busy` is what waiting could still free. */
+  function decide(bytes: number, except?: Slot<U>): { fits: boolean; victims: Slot<U>[]; wait: boolean } {
+    const base = used(except);
+    if (base + bytes <= budget) return { fits: true, victims: [], wait: false };
+    const candidates = [...slots.values()].filter(slot => slot !== except && evictable(slot)).sort((a, b) => a.lastUsedAt - b.lastUsedAt);
+    const victims: Slot<U>[] = [];
+    let freed = 0;
+    for (const slot of candidates) {
+      victims.push(slot); freed += held(slot);
+      if (base - freed + bytes <= budget) return { fits: true, victims, wait: false };
+    }
+    // Even every candidate would not be enough. If busy models hold the rest, wait for one to finish rather
+    // than evict for nothing; if only pinned models and this one remain, a lone model is always served.
+    const others = [...slots.values()].some(slot => slot !== except && !pinned.has(slot.id) && !!(slot.unit || slot.loading) && busy(slot));
+    return { fits: false, victims: candidates, wait: others };
+  }
+
+  /** Drain, flush and release one resident model; it is `closing` for the duration so nothing else touches it. */
+  function evict(slot: Slot<U>, reason: "evicted" | "requested", quiet = false): Promise<Evicted> {
+    if (slot.closing) return slot.closing;
+    const unit = slot.unit!;
+    const started = now();
+    const weightsBytes = quiet ? 0 : unit.memory().weightsBytes;
+    if (!quiet) log(`[models] ${slot.id} draining (${reason}); flushing saved state before release`);
+    slot.closing = (async () => {
+      let result: Evicted = { flushed: false };
+      try { result = await unit.close({ flush: true }); }
+      catch (error) {
+        // The unit released what it could; whoever asked for the close decides what a failure means.
+        result = { flushed: false, error };
+        if (reason === "evicted") log(`[models] ${slot.id} closed with an error: ${error instanceof Error ? error.message : String(error)}`);
+      }
+      slot.unit = undefined;
+      slot.unloads++;
+      publish({ type: "model.unload", at: Date.now(), model: slot.id, reason, drainMs: result.drainMs ?? now() - started, flushed: result.flushed,
+        ...(result.flushMs !== undefined ? { flushMs: result.flushMs } : {}) });
+      publish({ type: "model.memory", at: Date.now(), model: slot.id, weightsBytes: 0, kvBytes: 0, prefixCacheBytes: 0 });
+      if (!quiet) log(`[models] ${slot.id} released (${(weightsBytes / 2 ** 30).toFixed(2)} GiB weights${result.flushed ? ", state saved" : ", state NOT durable"})`);
+      return result;
+    })().finally(() => { slot.closing = undefined; notify(); });
+    return slot.closing;
+  }
+
+  async function slotFor(id: ModelId, need: readonly ModelOperation[]): Promise<Slot<U>> {
+    const known = slots.get(id);
+    const entry = known?.entry ?? await source.resolve(id);
+    if (!entry) throw new ResidencyError("does-not-fit", `${id} is not a local model this host can serve`);
+    for (const operation of need)
+      if (!entry.operations.includes(operation)) throw new ResidencyError("does-not-fit", `${id} does not declare the ${operation} operation`);
+    let slot = slots.get(entry.id);
+    if (!slot) {
+      slot = { id: entry.id, entry, unit: undefined, loading: undefined, closing: undefined, leases: 0, lastUsedAt: 0, loads: 0, unloads: 0, lastLoadMs: 0 };
+      slots.set(entry.id, slot);
+    }
+    return slot;
+  }
+
+  function startLoad(slot: Slot<U>): Promise<U> {
+    const started = now();
+    publish({ type: "model.load", at: Date.now(), model: slot.id, phase: "started" });
+    return slot.loading = (async () => {
+      let unit: U;
+      try { unit = await source.load(slot.entry); }
+      catch (error) {
+        const message = error instanceof Error ? error.message : String(error);
+        publish({ type: "model.load", at: Date.now(), model: slot.id, phase: "failed", error: message });
+        throw new ResidencyError("load-failed", message, { cause: error });
+      }
+      slot.lastLoadMs = now() - started;
+      slot.loads++;
+      const memory = unit.memory();
+      publish({ type: "model.load", at: Date.now(), model: slot.id, phase: "finished", ms: slot.lastLoadMs, weightsBytes: memory.weightsBytes,
+        ...(unit.resumed ? { resumed: true } : {}) });
+      publish({ type: "model.memory", at: Date.now(), model: slot.id, ...memory });
+      log(`[models] ${slot.id} loaded in ${slot.lastLoadMs.toFixed(0)} ms (${(unit.bytes() / 2 ** 30).toFixed(2)} GiB resident)`);
+      if (closed) {
+        // The host closed while this loaded: nothing may serve it.
+        try { await unit.close({ flush: true }); } catch { /* released as far as it could be */ }
+        throw closedError();
+      }
+      slot.unit = unit;
+      return unit;
+    })().finally(() => { slot.loading = undefined; notify(); });
+  }
+
+  const lease = (slot: Slot<U>, unit: U, loadMs: number): ResidencyLease<U> => {
+    let released = false;
+    return {
+      unit, loadMs, operations: unit.operationsFor(),
+      get model() { return describe(slot); },
+      release() {
+        if (released) return;
+        released = true;
+        slot.leases--;
+        slot.lastUsedAt = now();
+        notify();
+      },
+    };
+  };
+  const describe = (slot: Slot<U>): ResidentModel => ({ id: slot.id, role: slot.entry.role ?? "primary",
+    state: slot.closing ? "draining" : slot.loading && !slot.unit ? "loading" : "ready", operations: slot.unit?.operations ?? slot.entry.operations,
+    bytes: held(slot), pinned: pinned.has(slot.id), leases: slot.leases, lastUsedAt: slot.lastUsedAt,
+    ...(slot.unit && options.uses ? { uses: options.uses(slot.unit) } : {}) });
+
+  const host: ResidencyHost<U> = {
+    get policy(): ResidencyPolicy { return { budgetBytes: budget, pinned: [...pinned], idleUnloadSec: Number.POSITIVE_INFINITY }; },
+
+    async acquire(id, acquireOptions = {}) {
+      const { signal } = acquireOptions, need = acquireOptions.need ?? [];
+      const waited = now();
+      for (;;) {
+        signal?.throwIfAborted();
+        if (closed) throw closedError();
+        const slot = await slotFor(id, need);
+        signal?.throwIfAborted();
+        if (closed) throw closedError();
+        if (slot.closing) { await slot.closing.catch(() => {}); continue; }
+        if (slot.unit?.broken?.() && !slot.leases) { await evict(slot, "evicted"); continue; }
+        // Reserve before anything suspends: a leased model is never evicted.
+        const reserve = () => { slot.leases++; slot.lastUsedAt = now(); };
+        if (slot.unit) { reserve(); return lease(slot, slot.unit, 0); }
+        if (slot.loading) {
+          reserve();
+          try { return lease(slot, await slot.loading, now() - waited); }
+          catch (error) { slot.leases--; notify(); throw error; }
+        }
+        // Not resident. A managed job or a companion's decode holds the GPU: no loads or evictions meanwhile.
+        if (paused > 0) { await changed(signal); continue; }
+        await refreshed();
+        if (slot.unit || slot.loading || slot.closing || closed || paused > 0) continue;
+        const plan = decide(slot.entry.bytes, slot);
+        if (!plan.fits && plan.wait) { await changed(signal); continue; }
+        if (plan.victims.length) {
+          // The largest set the plan named, released before the newcomer loads. A settled eviction re-decides.
+          await Promise.all(plan.victims.map(victim => evict(victim, "evicted")));
+          continue;
+        }
+        reserve();
+        try { return lease(slot, await startLoad(slot), now() - waited); }
+        catch (error) { slot.leases--; notify(); throw error; }
+      }
+    },
+
+    adopt(entry, unit, adoptOptions = {}) {
+      if (closed) throw closedError();
+      if (slots.get(entry.id)?.unit) throw new Error(`${entry.id} is already resident`);
+      const slot = slots.get(entry.id) ?? { id: entry.id, entry, unit: undefined, loading: undefined, closing: undefined, leases: 0,
+        lastUsedAt: 0, loads: 0, unloads: 0, lastLoadMs: 0 };
+      slot.unit = unit; slot.loads++; slot.lastUsedAt = now();
+      slots.set(entry.id, slot);
+      if (adoptOptions.pin) pinned.add(entry.id);
+      // The loader already loaded it: announce it as resident, with the time that took when the caller knows it.
+      const memory = unit.memory();
+      publish({ type: "model.load", at: Date.now(), model: slot.id, phase: "finished", ...(adoptOptions.loadMs !== undefined ? { ms: adoptOptions.loadMs } : {}),
+        weightsBytes: memory.weightsBytes, ...(unit.resumed ? { resumed: true } : {}) });
+    },
+    peek: id => slots.get(id)?.unit,
+    defaultFor: operation => options.defaultFor?.(operation) ?? Promise.resolve(undefined),
+    async serve(id, serveOptions = {}) {
+      if (!options.serve) throw new ResidencyError("not-switchable", "this host serves one model; switching needs a restart");
+      await options.serve(id, serveOptions.signal);
+    },
+
+    async plan(id) {
+      const slot = await slotFor(id, []);
+      await refreshed();
+      if (slot.unit || slot.loading) return { fits: true, requiredBytes: 0, freeBytes: Math.max(0, budget - used()), evict: [] } satisfies ResidencyPlan;
+      const plan = decide(slot.entry.bytes, slot);
+      return { fits: plan.fits, requiredBytes: slot.entry.bytes, freeBytes: Math.max(0, budget - used(slot)), evict: plan.victims.map(victim => victim.id) };
+    },
+
+    async unload(id, unloadOptions = {}) {
+      const slot = slots.get(id);
+      if (!slot) return;
+      if (slot.closing) { await slot.closing; return; }
+      if (slot.loading) await slot.loading.catch(() => {});
+      if (!slot.unit) return;
+      if (slot.leases > 0 && !unloadOptions.force) throw new ResidencyError("in-use", `${id} is in use`);
+      const result = await evict(slot, "requested");
+      if (result.error !== undefined) throw result.error;
+    },
+
+    async makeRoom(bytes, signal) {
+      for (;;) {
+        signal?.throwIfAborted();
+        if (closed) throw closedError();
+        if (paused > 0) { await changed(signal); continue; }
+        await refreshed();
+        if (closed || paused > 0) continue;
+        const plan = decide(bytes);
+        if (plan.fits && !plan.victims.length) return;
+        if (!plan.fits && plan.wait) { await changed(signal); continue; }
+        if (!plan.victims.length) return;
+        await Promise.all(plan.victims.map(victim => evict(victim, "evicted")));
+      }
+    },
+
+    async pauseAll(signal) {
+      paused++;
+      const held: DisposableResource[] = [];
+      try {
+        // A model that is draining is on its way out: wait for it, then pause the ones that stay.
+        await Promise.all([...slots.values()].map(slot => slot.closing?.catch(() => {})));
+        // A stable order, so two pausers never wait on each other.
+        for (const slot of [...slots.values()].filter(slot => slot.unit && !slot.closing).sort((a, b) => a.id < b.id ? -1 : 1))
+          held.push(await slot.unit!.pause(signal));
+      } catch (error) {
+        for (const resource of held.reverse()) resource.dispose();
+        paused--; notify();
+        throw error;
+      }
+      let released = false;
+      return { dispose() {
+        if (released) return;
+        released = true;
+        for (const resource of held.reverse()) resource.dispose();
+        paused--; notify();
+      } };
+    },
+
+    pin(id) { pinned.add(id); },
+    unpin(id) { pinned.delete(id); notify(); },
+    resident: () => [...slots.values()].filter(slot => slot.unit).map(describe),
+
+    stats(id): ModelStats {
+      const slot = slots.get(id);
+      return { resident: !!slot?.unit, loads: slot?.loads ?? 0, unloads: slot?.unloads ?? 0, lastLoadMs: slot?.lastLoadMs ?? 0,
+        idleUnloadSec: pinned.has(id) ? null : Number.POSITIVE_INFINITY };
+    },
+
+    close() {
+      if (closing) return closing;
+      closed = true;
+      notify();
+      return closing = (async () => {
+        const results: UnitClosed[] = [];
+        const errors: unknown[] = [];
+        // Loads in flight finish (and close themselves, seeing `closed`); evictions in flight finish.
+        for (const slot of slots.values()) await Promise.allSettled([slot.loading, slot.closing]);
+        for (const slot of [...slots.values()]) {
+          if (!slot.unit) continue;
+          const result = await evict(slot, "requested", true);
+          if (result.error !== undefined) errors.push(result.error); else results.push(result);
+        }
+        if (errors.length) throw new AggregateError(errors, "model host close failed");
+        return results;
+      })();
+    },
+  };
+  return host;
+}

package/src/residency/record-index.ts

@@ -0,0 +1,28 @@
+// The local models a host may load, by exact id, read from a source that lists
+// them (the registry). A miss refreshes the listing at most every few seconds,
+// so a model fetched from another shell becomes routable without a restart, and
+// naming a model that does not exist never rescans on every request.
+import type { ModelRecord } from "@mlx-bun/hub/registry";
+
+export interface RecordIndex {
+  /** Exact id only; the startup model is answered even when the listing does not know it (a path given to `--model`). */
+  find(id: string): Promise<ModelRecord | undefined>;
+  /** Forget the listing: a download or job finished. */
+  invalidate(): void;
+}
+
+export function createRecordIndex(list: () => readonly ModelRecord[] | Promise<readonly ModelRecord[]>, startup?: ModelRecord, refreshMs = 5_000): RecordIndex {
+  let index: { at: number; byId: Map<string, ModelRecord> } | undefined;
+  return {
+    async find(id) {
+      if (startup && id === startup.repoId) return startup;
+      for (const fresh of [false, true]) {
+        if (!index || fresh && Date.now() - index.at > refreshMs) index = { at: Date.now(), byId: new Map((await list()).map(record => [record.repoId, record])) };
+        const record = index.byId.get(id);
+        if (record) return record;
+      }
+      return undefined;
+    },
+    invalidate() { index = undefined; },
+  };
+}