@mlx-node/server 0.0.13 → 0.0.15

package/src/host/swap.ts

@@ -0,0 +1,200 @@
+/**
+ * Single-resident lazy-load policy for an inference host.
+ *
+ * The host discovers every local model up-front but loads at most one into
+ * the `ModelRegistry` at a time. Switching models (e.g. via Claude Code's
+ * `/model` picker, or the desktop app's model menu) unregisters the previous
+ * instance, letting GC + native destructors reclaim memory, before loading
+ * the new one.
+ */
+
+import type { LoadableModel, SessionCapableModel } from '@mlx-node/lm';
+
+import type { PublicModelEntry } from '../handler.js';
+import type { ModelRegistry } from '../registry.js';
+import type { DiscoveredModel } from './discover.js';
+
+export interface SwapController {
+  resolveModel: (name: string) => Promise<void>;
+  listModels: () => PublicModelEntry[];
+}
+
+/**
+ * Build the `resolveModel` + `listModels` callbacks for the handler.
+ *
+ * `loadModelFn` is injected so tests can stub it without touching native code.
+ * The controller serializes every `resolveModel` invocation on a single
+ * promise chain so two concurrent requests for different-but-currently-
+ * unloaded models cannot race on the native compiled-path globals.
+ */
+export function makeSwapController(
+  discovered: DiscoveredModel[],
+  registry: ModelRegistry,
+  loadModelFn: (path: string) => Promise<LoadableModel>,
+  defaultName?: string,
+): SwapController {
+  const byName = new Map<string, DiscoveredModel>();
+  for (const entry of discovered) byName.set(entry.name, entry);
+
+  const ordered = [...discovered].sort((a, b) => (a.name < b.name ? -1 : a.name > b.name ? 1 : 0));
+
+  // Which entry unknown names (haiku subagent dispatches, etc.) fall back
+  // to before anything is resident. Defaults to discovered[0], but the
+  // caller can pin it to the user's `--model` pick so the first haiku
+  // title-gen doesn't trigger a load of the alphabetically-first model
+  // followed by an immediate swap to the user's real choice.
+  const fallbackEntry = (defaultName != null ? byName.get(defaultName) : undefined) ?? discovered[0];
+
+  let resident: { name: string } | null = null;
+  // Names we registered as aliases to the current resident (e.g.
+  // Claude Code's hardcoded `claude-haiku-*` for subagent dispatches /
+  // title generation). Tracked so we can unregister them on `/model`
+  // swap — otherwise an alias's refcount would keep the old binding
+  // alive past the user's swap.
+  const aliases = new Set<string>();
+  let currentOp: Promise<unknown> = Promise.resolve();
+
+  async function resolveModel(name: string): Promise<void> {
+    // Fast path: already registered under this name (either as a real
+    // resident or as an alias we previously installed). Avoid chaining.
+    if (registry.get(name)) return;
+
+    const next = currentOp.then(async () => {
+      // Re-check under the serialized section — a prior waiter may have loaded it.
+      if (registry.get(name)) return;
+
+      // Pick the discovered entry to resolve against. If the requested
+      // name matches a discovered model, use it. Otherwise (unknown
+      // name — Claude Code's hardcoded small-fast-model, etc.) fall
+      // through to the current resident so subagent dispatches don't
+      // 404, loading discovered[0] on first boot if nothing is resident
+      // yet.
+      //
+      // CRITICAL: this must read `resident` at RUN time, not QUEUE time.
+      // If we capture it before chaining onto `currentOp`, a swap that
+      // ran ahead of us will leave us with a stale target — e.g. a haiku
+      // alias request that arrived during a `/model a → b` switch would
+      // capture `targetEntry = a`, then re-bind itself to `a` and undo
+      // the user's switch when its turn finally comes around.
+      const knownEntry = byName.get(name);
+      const targetEntry = knownEntry ?? (resident ? (byName.get(resident.name) ?? fallbackEntry) : fallbackEntry);
+      const isAlias = targetEntry.name !== name;
+
+      // Swap out any stale resident that isn't the target.
+      //
+      // We do NOT unregister the aliases here: in-flight messages.ts
+      // requests may be microtask-racing between "resolveModel returned"
+      // and "registry.get(body.model)" and dropping the alias in that
+      // window yields a spurious 404. Instead we carry the alias set
+      // across the swap and re-point them to the new resident below,
+      // so the name always resolves to *some* live instance.
+      const oldResident = resident;
+      const carriedAliases = new Set(aliases);
+      if (oldResident && oldResident.name !== targetEntry.name) {
+        // Drop our local alias bookkeeping AND the old resident's primary
+        // name binding, but leave the alias *names* in the registry pointed
+        // at the old model — they hold the only refcount preventing GC,
+        // and an in-flight `registry.get(alias)` must keep resolving to
+        // *some* live instance until the new model is in hand.
+        aliases.clear();
+        registry.unregister(oldResident.name);
+        resident = null;
+      }
+
+      // Ensure the target is resident. If the load throws, restore the
+      // pre-swap controller state so future swap attempts know about the
+      // aliases we just cleared — otherwise the alias *names* stay bound
+      // in the registry to the old model object forever (alias bindings
+      // hold their own refcount), but the controller forgets they exist
+      // and never repoints them, leaving alias-routed traffic permanently
+      // pinned to a stale model.
+      let instance = registry.get(targetEntry.name);
+      if (!instance) {
+        let loaded: LoadableModel;
+        try {
+          loaded = await loadModelFn(targetEntry.path);
+        } catch (err) {
+          // Recovery: re-populate the controller's alias set so the next
+          // resolveModel call still owns them. The alias→old-model bindings
+          // are still live in the registry (we never unregistered them), so
+          // we can recover the old model object via any surviving alias and
+          // re-bind the old resident's primary name.
+          //
+          // If `carriedAliases` is empty (no aliases ever existed) AND we
+          // unregistered `oldResident.name`, the binding's refcount may have
+          // hit zero and the model is gone. There's nothing the controller
+          // can do to recover in that case — the user will need to /model-
+          // pick again. This is acceptable: alias-less load failures are
+          // rare, and the user-facing symptom is "prior model gone, please
+          // re-pick", not silently-wrong responses.
+          for (const aliasName of carriedAliases) aliases.add(aliasName);
+          if (oldResident) {
+            let oldInstance: SessionCapableModel | undefined;
+            for (const aliasName of carriedAliases) {
+              const probe = registry.get(aliasName);
+              if (probe) {
+                oldInstance = probe;
+                break;
+              }
+            }
+            if (oldInstance) {
+              const oldEntry = byName.get(oldResident.name);
+              registry.register(oldResident.name, oldInstance, {
+                samplingDefaults: oldEntry?.preset.sampling,
+                maxOutputTokens: oldEntry?.preset.maxOutputTokens,
+              });
+              resident = oldResident;
+            }
+          }
+          throw err;
+        }
+        instance = loaded as unknown as SessionCapableModel;
+        registry.register(targetEntry.name, instance, {
+          samplingDefaults: targetEntry.preset.sampling,
+          maxOutputTokens: targetEntry.preset.maxOutputTokens,
+        });
+        resident = { name: targetEntry.name };
+      } else if (!resident) {
+        resident = { name: targetEntry.name };
+      }
+
+      // Re-point any aliases carried across the swap onto the new
+      // resident. `registry.register(sameName, differentModel)` drops
+      // the old binding's refcount and installs the new one atomically,
+      // so any concurrent `registry.get(alias)` either sees the old or
+      // new instance — never null.
+      for (const aliasName of carriedAliases) {
+        if (aliasName === targetEntry.name) continue;
+        registry.register(aliasName, instance, {
+          samplingDefaults: targetEntry.preset.sampling,
+          maxOutputTokens: targetEntry.preset.maxOutputTokens,
+        });
+        aliases.add(aliasName);
+      }
+
+      // For unknown names, register an alias on the resident instance so
+      // the endpoint's `registry.get(name)` lookup succeeds.
+      if (isAlias) {
+        registry.register(name, instance, {
+          samplingDefaults: targetEntry.preset.sampling,
+          maxOutputTokens: targetEntry.preset.maxOutputTokens,
+        });
+        aliases.add(name);
+      }
+    });
+    currentOp = next.catch(() => undefined);
+    await next;
+  }
+
+  function listModels(): PublicModelEntry[] {
+    const created = Math.floor(Date.now() / 1000);
+    return ordered.map((entry) => ({
+      id: entry.name,
+      object: 'model',
+      created,
+      owned_by: 'mlx-node',
+    }));
+  }
+
+  return { resolveModel, listModels };
+}

package/src/host/temp-root.ts

@@ -0,0 +1,110 @@
+/**
+ * Pid-scoped temp roots for the host's paged-config overrides, plus the
+ * startup sweep that reclaims the ones a killed host left behind.
+ *
+ * `PagedConfigOverrideManager` clones a checkpoint directory (config.json
+ * rewritten, everything else symlinked) into a temp root and removes that root
+ * in `cleanup()`. `cleanup()` runs on a normal shutdown — but the host's
+ * headline deployment is an Electron `utilityProcess`, and a `utilityProcess`
+ * can be SIGKILLed (app force-quit, OOM killer, `kill -9`). SIGKILL runs no
+ * handler, so without a sweep EVERY hard kill leaks a root. The clones are
+ * symlink farms rather than copies, so the leak is inodes and directory
+ * entries rather than model-sized bytes — but the roots accumulate forever and
+ * a partially-written clone can hold a real config.json.
+ *
+ * The reclaim strategy is "name the owner in the directory name": the manager
+ * gets a `tempDirPrefix` of `mlx-inference-host-<pid>-`, `mkdtemp` appends its
+ * own random suffix, and {@link sweepOrphanHostTempRoots} parses the pid back
+ * out and removes any root whose owner is gone.
+ *
+ * Known limitation: pid reuse. A long-dead host's root whose pid has since
+ * been recycled by an unrelated process is SPARED (never wrongly deleted), so
+ * the failure mode is a leaked directory, not data loss.
+ */
+
+import { readdir, rm } from 'node:fs/promises';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+
+/** Shared stem. A directory is host-owned iff its name starts with this. */
+export const HOST_TEMP_DIR_STEM = 'mlx-inference-host-';
+
+/** Matches `mlx-inference-host-<pid>-<mkdtemp suffix>`. */
+const HOST_TEMP_DIR_RE = /^mlx-inference-host-(\d+)-/;
+
+/**
+ * `tempDirPrefix` to hand `PagedConfigOverrideManager` so the root it creates
+ * carries its owner's pid.
+ */
+export function hostTempDirPrefix(pid: number = process.pid): string {
+  return `${HOST_TEMP_DIR_STEM}${pid}-`;
+}
+
+/**
+ * Does a process with this pid exist?
+ *
+ * `process.kill(pid, 0)` sends no signal; it only performs the permission +
+ * existence check. `ESRCH` is the sole "gone" answer — `EPERM` means the
+ * process exists but belongs to another user, which must count as ALIVE so a
+ * multi-user box never has one user's sweep delete another user's live root.
+ */
+export function isProcessAlive(pid: number): boolean {
+  // pid 0 addresses the caller's whole process group on POSIX and pid < 0 a
+  // group by id; neither is ever a real owner, and signalling them would be
+  // actively dangerous. Treat as alive so they are never swept.
+  if (!Number.isInteger(pid) || pid <= 0) return true;
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch (err) {
+    return (err as NodeJS.ErrnoException).code !== 'ESRCH';
+  }
+}
+
+export interface SweepOrphanTempRootsOptions {
+  /** Directory to scan. Defaults to the OS temp dir. */
+  root?: string;
+  /** Our own pid — never swept, however the liveness probe answers. */
+  selfPid?: number;
+  /** Injectable liveness probe. Defaults to {@link isProcessAlive}. */
+  isAlive?: (pid: number) => boolean;
+}
+
+/**
+ * Remove every host temp root whose owning pid is no longer alive.
+ *
+ * Best effort by contract: a scan or unlink failure (permissions, a root a
+ * concurrently-exiting host is removing under us) is swallowed, because a
+ * housekeeping step must never be the reason a host refuses to start.
+ * Returns the absolute paths actually removed.
+ */
+export async function sweepOrphanHostTempRoots(opts: SweepOrphanTempRootsOptions = {}): Promise<string[]> {
+  const root = opts.root ?? tmpdir();
+  const selfPid = opts.selfPid ?? process.pid;
+  const isAlive = opts.isAlive ?? isProcessAlive;
+
+  let entries: string[];
+  try {
+    entries = await readdir(root);
+  } catch {
+    return [];
+  }
+
+  const removed: string[] = [];
+  for (const name of entries) {
+    const match = HOST_TEMP_DIR_RE.exec(name);
+    if (match === null) continue;
+    const pid = Number.parseInt(match[1], 10);
+    if (pid === selfPid) continue;
+    if (isAlive(pid)) continue;
+
+    const full = join(root, name);
+    try {
+      await rm(full, { recursive: true, force: true });
+      removed.push(full);
+    } catch {
+      /* another sweeper won the race, or we lack permission; leave it */
+    }
+  }
+  return removed;
+}