mlx-bun 0.0.13 → 0.3.0

package/src/serve/isolate.ts

@@ -1,4 +1,11 @@
-// Runtime isolation — the parent half (docs/design/runtime-isolation.md).
+import { ExecutionCoordinator } from "../engine/execution-coordinator";
+import { AdmissionPool, AdmissionRejected } from "../engine/admission";
+import { acquireReservation } from "./preparation";
+import { createParentApplication } from "./parent-application";
+import type { DisposableResource } from "../contracts/resources";
+import { retainResponseLease } from "./response-lease";
+import type { EngineHost } from "../contracts/host";
+// Runtime isolation — the parent half (docs/reference/server-config.md).
 //
 // ARCHITECTURE (decision 2026-07-05, deviating from the doc's original
 // structured-IPC sketch — recorded there): the inference ENGINE is the
@@ -33,11 +40,17 @@ import { join } from "node:path";
 
 /** How to (re)start the engine child. `argv[0]` is the executable. */
 export interface EngineSpec {
+  /** Parent-owned activity lease; covers startup and complete response bodies. */
+  acquire?: (signal: AbortSignal) => Promise<DisposableResource>;
   argv: string[];
   socketPath: string;
   /** Max ms to wait for the child's /health after spawn (weights load —
    *  large models take a while). Default 15 min. */
   readyTimeoutMs?: number;
+  /** Bound automatic restarts in a rolling window, including startup failure. */
+  maxRestarts?: number;
+  restartWindowMs?: number;
+  restartDelayMs?: number;
 }
 
 const HOP_BY_HOP = new Set([
@@ -50,17 +63,23 @@ export function defaultSocketPath(): string {
 }
 
 /** The persistent engine child: spawn, health-gate, respawn-on-crash. */
-export class EngineChild {
+export class EngineChild implements EngineHost<Request, Response> {
   readonly spec: EngineSpec;
   #proc: ReturnType<typeof Bun.spawn> | null = null;
   #ready: Promise<void>;
   #stopping = false;
+  #lifetime = new AbortController();
   restarts = 0;
   #lastSpawnAt = 0;
+  #restartTimes: number[] = [];
+  #backoff: ReturnType<typeof setTimeout> | undefined;
+  #wakeBackoff: (() => void) | undefined;
+  #startups = new Set<Promise<void>>();
 
   constructor(spec: EngineSpec) {
-    this.spec = spec;
+    this.spec = Object.freeze({ ...spec, argv: [...spec.argv] });
     this.#ready = this.#spawn();
+    void this.#ready.catch(() => {}); // failures remain observable through ready/forward
   }
 
   get ready(): Promise<void> {
@@ -71,38 +90,78 @@ export class EngineChild {
     return this.#proc?.pid ?? null;
   }
 
-  async #spawn(): Promise<void> {
-    this.#lastSpawnAt = Date.now();
-    try { unlinkSync(this.spec.socketPath); } catch {}
-    const proc = Bun.spawn(this.spec.argv, {
-      stdio: ["ignore", "inherit", "inherit"], // load progress → user's terminal
-      env: process.env,
-    });
-    this.#proc = proc;
-    void proc.exited.then((code) => {
-      if (this.#stopping || this.#proc !== proc) return;
-      console.error(`[isolate] engine exited (code ${code}) — respawning`);
-      this.restarts++;
-      // Crash-loop backoff: an engine that dies within 10 s of spawning
-      // (bad flags, OOM on load) waits 5 s before the retry.
-      const delay = Date.now() - this.#lastSpawnAt < 10_000 ? 5_000 : 0;
-      this.#ready = new Promise((r) => setTimeout(r, delay)).then(() => this.#spawn());
-    });
-    // Health-gate: poll the child's /health over the socket until it
-    // answers — model load happens behind this.
-    const deadline = Date.now() + (this.spec.readyTimeoutMs ?? 15 * 60_000);
-    while (Date.now() < deadline) {
-      if (proc.killed || this.#proc !== proc) throw new Error("engine died during startup");
-      try {
-        const r = await fetch("http://engine/health", {
-          unix: this.spec.socketPath,
-          signal: AbortSignal.timeout(2_000),
-        } as RequestInit & { unix: string });
-        if (r.ok) return;
-      } catch { /* not up yet */ }
-      await new Promise((r) => setTimeout(r, 250));
+  #spawn(): Promise<void> {
+    const work = this.#start();
+    this.#startups.add(work);
+    void work.finally(() => this.#startups.delete(work)).catch(() => {});
+    return work;
+  }
+
+  async #start(): Promise<void> {
+    if (this.#stopping) throw new Error("engine host is closed");
+    const lease = await this.spec.acquire?.(this.#lifetime.signal);
+    let starting: ReturnType<typeof Bun.spawn> | undefined;
+    let healthy = false;
+    try {
+      if (this.#stopping) throw new Error("engine host is closed");
+      this.#lastSpawnAt = Date.now();
+      try { unlinkSync(this.spec.socketPath); } catch {}
+      const proc = Bun.spawn(this.spec.argv, {
+        stdio: ["ignore", "inherit", "inherit"], // load progress → user's terminal
+        env: process.env,
+      });
+      this.#proc = proc;
+      starting = proc;
+      void proc.exited.then((code) => {
+        if (this.#stopping || this.#proc !== proc) return;
+        const now = Date.now();
+        const window = this.spec.restartWindowMs ?? 60_000;
+        this.#restartTimes = this.#restartTimes.filter((time) => now - time < window);
+        if (this.#restartTimes.length >= (this.spec.maxRestarts ?? 3)) {
+          this.#ready = Promise.reject(new Error(`engine restart limit reached after exit ${code}`));
+          void this.#ready.catch(() => {});
+          return;
+        }
+        this.#restartTimes.push(now);
+        console.error(`[isolate] engine exited (code ${code}) — respawning`);
+        this.restarts++;
+        // Crash-loop backoff: an engine that dies within 10 s of spawning
+        // (bad flags, OOM on load) waits 5 s before the retry.
+        const delay = this.spec.restartDelayMs ?? (Date.now() - this.#lastSpawnAt < 10_000 ? 5_000 : 0);
+        this.#ready = new Promise<void>((resolve) => {
+          this.#wakeBackoff = resolve;
+          this.#backoff = setTimeout(resolve, delay);
+        }).then(() => {
+          this.#backoff = undefined; this.#wakeBackoff = undefined;
+          return this.#spawn();
+        });
+        void this.#ready.catch(() => {});
+      });
+      // Health-gate: poll the child's /health over the socket until it
+      // answers — model load happens behind this.
+      const deadline = Date.now() + (this.spec.readyTimeoutMs ?? 15 * 60_000);
+      while (Date.now() < deadline) {
+        if (this.#stopping || proc.killed || proc.exitCode !== null || this.#proc !== proc)
+          throw new Error("engine died during startup");
+        try {
+          const r = await fetch("http://engine/health", {
+            unix: this.spec.socketPath,
+            signal: AbortSignal.timeout(2_000),
+          } as RequestInit & { unix: string });
+          if (r.ok) { healthy = true; return; }
+        } catch { /* not up yet */ }
+        await new Promise((r) => setTimeout(r, 250));
+      }
+      proc.kill();
+      throw new Error("engine did not become healthy in time");
+    } finally {
+      if (!healthy && starting && starting.exitCode === null) {
+        starting.kill();
+        const force = setTimeout(() => { if (starting!.exitCode === null) starting!.kill("SIGKILL"); }, 3000);
+        try { await starting.exited; } finally { clearTimeout(force); }
+      }
+      lease?.dispose();
     }
-    throw new Error("engine did not become healthy in time");
   }
 
   /** Forward one request to the child; the response streams through.
@@ -124,33 +183,81 @@ export class EngineChild {
   }
 
   async #forwardOnce(request: Request): Promise<Response> {
-    await this.#ready;
-    const headers = new Headers();
-    request.headers.forEach((v, k) => {
-      if (!HOP_BY_HOP.has(k.toLowerCase())) headers.set(k, v);
+    if (this.#stopping) throw new Error("engine host is closed");
+    request.signal.throwIfAborted();
+    await new Promise<void>((resolve, reject) => {
+      const abort = () => { request.signal.removeEventListener("abort", abort); reject(request.signal.reason); };
+      request.signal.addEventListener("abort", abort, { once: true });
+      this.#ready.then(() => { request.signal.removeEventListener("abort", abort); resolve(); },
+        (error) => { request.signal.removeEventListener("abort", abort); reject(error); });
     });
-    const proxied = await fetch(request.url, {
-      method: request.method,
-      headers,
-      body: request.body,
-      signal: request.signal, // client abort → child sees the disconnect
-      redirect: "manual",
-      unix: this.spec.socketPath,
-      // Streaming request bodies need half-duplex; buffered bodies ignore it.
-      duplex: "half",
-    } as RequestInit & { unix: string; duplex: string });
-    const outHeaders = new Headers();
-    proxied.headers.forEach((v, k) => {
-      if (!HOP_BY_HOP.has(k.toLowerCase())) outHeaders.set(k, v);
-    });
-    return new Response(proxied.body, { status: proxied.status, headers: outHeaders });
+    if (this.#stopping) throw new Error("engine host is closed");
+    request.signal.throwIfAborted();
+    const signal = AbortSignal.any([request.signal, this.#lifetime.signal]);
+    const lease = await this.spec.acquire?.(signal);
+    try {
+      signal.throwIfAborted();
+      const headers = new Headers();
+      request.headers.forEach((v, k) => {
+        if (!HOP_BY_HOP.has(k.toLowerCase())) headers.set(k, v);
+      });
+      const proxied = await fetch(request.url, {
+        method: request.method,
+        headers,
+        body: request.body,
+        signal, // client abort → child sees the disconnect
+        redirect: "manual",
+        unix: this.spec.socketPath,
+        // Streaming request bodies need half-duplex; buffered bodies ignore it.
+        duplex: "half",
+      } as RequestInit & { unix: string; duplex: string });
+      const outHeaders = new Headers();
+      proxied.headers.forEach((v, k) => {
+        if (!HOP_BY_HOP.has(k.toLowerCase())) outHeaders.set(k, v);
+      });
+      const response = new Response(proxied.body, { status: proxied.status, headers: outHeaders });
+      return lease ? retainResponseLease(response, lease) : response;
+    } catch (error) { lease?.dispose(); throw error; }
+  }
+
+  /** A live UDS response owns the worker's native lease. Disconnect releases it
+   * even if the parent crashes. External cancellation applies to acquisition;
+   * the returned owner alone releases an admitted lease after job death. */
+  async reserveExecution(signal: AbortSignal): Promise<DisposableResource> {
+    if (!this.#proc || this.#proc.exitCode !== null) return { dispose() {} };
+    const abort = new AbortController();
+    const cancel = () => abort.abort(signal.reason);
+    signal.throwIfAborted();
+    signal.addEventListener("abort", cancel, { once: true });
+    try {
+      const response = await fetch("http://engine/admin/lease", {
+        method: "POST", unix: this.spec.socketPath, signal: abort.signal,
+      } as RequestInit & { unix: string });
+      if (!response.ok) { await response.body?.cancel(); throw new Error(`worker lease failed (${response.status})`); }
+      return { dispose() { abort.abort(); void response.body?.cancel().catch(() => {}); } };
+    } catch (error) { abort.abort(); throw error; }
+    finally { signal.removeEventListener("abort", cancel); }
   }
 
   stop(): void {
     this.#stopping = true;
+    this.#lifetime.abort(new Error("engine host is closed"));
+    clearTimeout(this.#backoff);
+    this.#wakeBackoff?.();
     this.#proc?.kill();
     try { unlinkSync(this.spec.socketPath); } catch {}
   }
+
+  async close(): Promise<void> {
+    this.stop();
+    const process = this.#proc;
+    const force = process ? setTimeout(() => { if (process.exitCode === null) process.kill("SIGKILL"); }, 3000) : undefined;
+    try {
+      // Startup owns an admission lease even before a child exists. Earlier
+      // crashed attempts can still be unwinding after #ready switches to a retry.
+      await Promise.allSettled([this.#ready, ...this.#startups, ...(process ? [process.exited] : [])]);
+    } finally { if (force) clearTimeout(force); }
+  }
 }
 
 /** Model resolution the parent injects (Registry-backed in production;
@@ -158,7 +265,7 @@ export class EngineChild {
  *  drop-in ignore semantics (served by the default model, like mlx-lm). */
 export type ModelResolver = (query: string) => { repoId: string; path: string } | null;
 
-/** Child-per-model pool (P2, runtime-isolation.md): LRU residency over
+/** Child-per-model pool (P2, docs/reference/server-config.md): LRU residency over
  *  engine children. Eviction is graceful and lossless — POST /admin/drain
  *  (gateway quiesce + demote the prompt cache to the SSD tier, keyed by
  *  that model's own fingerprint dir) and only then stop the child; the
@@ -170,6 +277,9 @@ export class ModelPool {
   #children = new Map<string, EngineChild>(); // repoId → child
   #lru: string[] = []; // most-recent last
   #spawning = new Map<string, Promise<EngineChild>>();
+  #coldStarts = new AdmissionPool(1);
+  #starting = new Set<EngineChild>();
+  #stopping = false;
   readonly defaultKey: string;
 
   constructor(
@@ -210,40 +320,55 @@ export class ModelPool {
   /** Route a request's `model` field to a child, spawning on first use.
    *  undefined/empty/unknown → the default child (drop-in semantics). */
   async childFor(modelField: string | undefined | null): Promise<EngineChild> {
-    if (!modelField) return this.#use(this.defaultKey);
+    if (this.#stopping) throw new Error("model pool is closed");
+    modelField ||= this.defaultKey;
     // Already resident under this exact id?
     if (this.#children.has(modelField)) return this.#use(modelField);
     const rec = (() => {
       try { return this.opts.resolve(modelField); } catch { return null; }
     })();
-    if (!rec) return this.#use(this.defaultKey); // unknown → ignore, like mlx-lm
-    if (this.#children.has(rec.repoId)) return this.#use(rec.repoId);
-    const inflight = this.#spawning.get(rec.repoId);
-    if (inflight) return inflight.then(() => this.#use(rec.repoId));
+    if (!rec && modelField !== this.defaultKey) return this.childFor(this.defaultKey);
+    const resolved = rec ?? { repoId: this.defaultKey, path: "" };
+    if (this.#children.has(resolved.repoId)) return this.#use(resolved.repoId);
+    const inflight = this.#spawning.get(resolved.repoId);
+    if (inflight) return inflight.then(() => this.#use(resolved.repoId));
     const spawnP = (async () => {
-      const sock = this.#socketFor(rec.repoId);
-      const child = new EngineChild({
-        argv: [...this.opts.selfArgv, ...engineArgvForModel(this.opts.rawArgs, sock, rec.path)],
+      const admission = await this.#coldStarts.acquire();
+      try {
+      if (this.#stopping) throw new Error("model pool is closed");
+      const sock = this.#socketFor(resolved.repoId);
+      const child = new EngineChild(resolved.repoId === this.defaultKey ? this.opts.defaultChild.spec : {
+        argv: [...this.opts.selfArgv, ...engineArgvForModel(this.opts.rawArgs, sock, resolved.path)],
         socketPath: sock,
+        acquire: this.opts.defaultChild.spec.acquire,
       });
-      // SPAWN-OVERLAP: the old model keeps serving while this loads.
-      await child.ready;
-      this.#children.set(rec.repoId, child);
-      this.#lru.push(rec.repoId);
-      await this.#evictOverCap();
-      return child;
+      this.#starting.add(child);
+      try {
+        await child.ready;
+        if (this.#stopping) throw new Error("model pool is closed");
+        this.#children.set(resolved.repoId, child);
+        this.#lru.push(resolved.repoId);
+        await this.#evictOverCap();
+        return child;
+      } catch (error) {
+        await child.close();
+        throw error;
+      } finally { this.#starting.delete(child); }
+      } finally { admission.dispose(); }
     })();
-    this.#spawning.set(rec.repoId, spawnP);
+    this.#spawning.set(resolved.repoId, spawnP);
     try {
       return await spawnP;
     } finally {
-      this.#spawning.delete(rec.repoId);
+      this.#spawning.delete(resolved.repoId);
     }
   }
 
   #use(key: string): EngineChild {
     this.#bump(key);
-    return this.#children.get(key)!;
+    const child = this.#children.get(key);
+    if (!child) throw new Error(`model ${key} is no longer resident`);
+    return child;
   }
 
   /** Evict least-recently-used children over the cap: drain + demote (the
@@ -253,23 +378,38 @@ export class ModelPool {
       const victim = this.#lru.find((k) => this.#children.has(k));
       if (!victim || this.#children.size <= 1) return;
       const child = this.#children.get(victim)!;
-      this.#children.delete(victim);
-      this.#lru = this.#lru.filter((k) => k !== victim);
+      const lease = await child.spec.acquire?.(AbortSignal.timeout(120_000));
       try {
-        await fetch("http://engine/admin/drain", {
-          method: "POST",
-          unix: child.spec.socketPath,
-          signal: AbortSignal.timeout(120_000),
-        } as RequestInit & { unix: string });
-      } catch { /* best-effort — state also persists via evict-spill */ }
-      child.stop();
-      console.log(`[isolate] evicted engine ${victim} (pool cap ${this.opts.poolMax})`);
+        // Another overlapping eviction may have completed while we waited.
+        if (this.#children.get(victim) !== child) continue;
+        this.#children.delete(victim);
+        this.#lru = this.#lru.filter((k) => k !== victim);
+        try {
+          const response = await fetch("http://engine/admin/drain", {
+            method: "POST", unix: child.spec.socketPath,
+            signal: AbortSignal.timeout(120_000),
+          } as RequestInit & { unix: string });
+          await response.arrayBuffer();
+        } catch { /* best-effort — state also persists via evict-spill */ }
+        await child.close();
+        console.log(`[isolate] evicted engine ${victim} (pool cap ${this.opts.poolMax})`);
+      } finally { lease?.dispose(); }
     }
   }
 
   stopAll(): void {
-    for (const c of this.#children.values()) c.stop();
+    this.#stopping = true;
+    this.#coldStarts.close();
+    for (const child of new Set([...this.#children.values(), ...this.#starting])) child.stop();
     this.#children.clear();
+    this.#lru = [];
+  }
+
+  async close(): Promise<void> {
+    const children = new Set([...this.#children.values(), ...this.#starting]);
+    this.stopAll();
+    await Promise.all([...children].map((child) => child.close()));
+    await Promise.allSettled(this.#spawning.values());
   }
 }
 
@@ -280,6 +420,8 @@ const MODEL_ROUTED = new Set([
 ]);
 
 export interface ProxyServerOptions {
+  /** Application-store injection for embedded hosts and model-free tests. */
+  createJobStore?: () => Promise<import("../jobs/db").JobStore>;
   port: number;
   hostname?: string;
   engine: EngineSpec;
@@ -301,11 +443,61 @@ export function startProxyServer(opts: ProxyServerOptions): {
   server: ReturnType<typeof Bun.serve>;
   engine: EngineChild;
   pool: ModelPool | null;
+  close(): Promise<void>;
 } {
-  const engine = new EngineChild(opts.engine);
+  const coordinator = new ExecutionCoordinator();
+  const acquire = (mode: "shared" | "exclusive", signal: AbortSignal) =>
+    acquireReservation({ acquire: (cancellation) => coordinator.acquire(mode, cancellation) }, signal);
+  const engine = new EngineChild({ ...opts.engine, async acquire(signal) {
+    const parent = await acquire("shared", signal);
+    try {
+      const supplied = await opts.engine.acquire?.(signal);
+      return { dispose() { try { supplied?.dispose(); } finally { parent.dispose(); } } };
+    } catch (error) { parent.dispose(); throw error; }
+  } });
   const pool = opts.pool
     ? new ModelPool({ ...opts.pool, defaultChild: engine })
     : null;
+  const children = () => pool ? pool.residentKeys.map((key) => pool.child(key)!).filter(Boolean) : [engine];
+  const forwardRequest = async (request: Request): Promise<Response> => {
+    const url = new URL(request.url);
+    // P2 routing: generation endpoints carry `model` in the JSON body —
+    // buffer it (small), pick/spawn that model's child, forward the
+    // buffered body. Everything else rides the default engine.
+    if (pool && request.method === "POST" && MODEL_ROUTED.has(url.pathname)) {
+      const raw = await request.text();
+      let modelField: string | undefined;
+      try { modelField = (JSON.parse(raw) as { model?: string }).model; } catch { /* malformed → default child answers with its own 400 */ }
+      const target = await pool.childFor(modelField);
+      return await target.forward(new Request(request.url, {
+        method: request.method,
+        headers: request.headers,
+        body: raw,
+        signal: request.signal,
+      }));
+    }
+    return await (pool ? await pool.childFor(undefined) : engine).forward(request);
+  };
+  const application = createParentApplication({
+    createJobStore: opts.createJobStore,
+    serverPort: () => server.port,
+    async acquireGpu(signal) {
+      const lease = await acquire("exclusive", signal);
+      try {
+        // Network cancellation is not a native completion fence. Explicitly
+        // drain every resident worker before the subprocess may touch the GPU.
+        const results = await Promise.allSettled(children().map((child) => child.reserveExecution(signal)));
+        const held = results.filter((result) => result.status === "fulfilled").map((result) => result.value);
+        const failed = results.find((result) => result.status === "rejected");
+        if (failed) { for (const worker of held) worker.dispose(); throw failed.reason; }
+        return { dispose() { for (const worker of held) worker.dispose(); lease.dispose(); } };
+      } catch (error) { lease.dispose(); throw error; }
+    },
+    invalidateLibrary() {
+      for (const child of children()) void child.forward(new Request("http://engine/library?refresh=1"))
+        .then((response) => response.arrayBuffer()).catch(() => {});
+    },
+  });
   const server = Bun.serve({
     port: opts.port,
     ...(opts.hostname ? { hostname: opts.hostname } : {}),
@@ -317,32 +509,26 @@ export function startProxyServer(opts: ProxyServerOptions): {
           { error: { message: "WebSocket chat is not proxied under --isolate yet; run without --isolate for the web chat UI" } },
           { status: 501 },
         );
-      if (url.pathname === "/engine" && request.method === "GET")
+      if (url.pathname === "/engine" && request.method === "GET") {
+        // A pool may have evicted and recreated the default worker. Inspect
+        // current residency without starting a worker merely for diagnostics.
+        const current = pool ? pool.child(pool.defaultKey) : engine;
         return Response.json({
           isolated: true,
-          pid: engine.pid,
-          restarts: engine.restarts,
-          socket: engine.spec.socketPath,
+          response_store: application.responseStats,
+          pid: current?.pid ?? null,
+          restarts: current?.restarts ?? null,
+          socket: current?.spec.socketPath ?? null,
           ...(pool ? { pool: { resident: pool.residentKeys, default: pool.defaultKey } } : {}),
         });
+      }
       try {
-        // P2 routing: generation endpoints carry `model` in the JSON body —
-        // buffer it (small), pick/spawn that model's child, forward the
-        // buffered body. Everything else rides the default engine.
-        if (pool && request.method === "POST" && MODEL_ROUTED.has(url.pathname)) {
-          const raw = await request.text();
-          let modelField: string | undefined;
-          try { modelField = (JSON.parse(raw) as { model?: string }).model; } catch { /* malformed → default child answers with its own 400 */ }
-          const target = await pool.childFor(modelField);
-          return await target.forward(new Request(request.url, {
-            method: request.method,
-            headers: request.headers,
-            body: raw,
-            signal: request.signal,
-          }));
-        }
-        return await engine.forward(request);
+        const local = await application.handle(request, forwardRequest);
+        if (local) return local;
+        return await forwardRequest(request);
       } catch (e) {
+        if (e instanceof AdmissionRejected) return Response.json(
+          { error: { message: e.message, type: "resource_admission", code: "queue_full" } }, { status: 429 });
         if (request.signal.aborted) return new Response(null, { status: 499 });
         return Response.json(
           {
@@ -356,7 +542,13 @@ export function startProxyServer(opts: ProxyServerOptions): {
       }
     },
   });
-  return { server, engine, pool };
+  let closing: Promise<void> | undefined;
+  return { server, engine, pool, close: () => closing ??= (async () => {
+    await server.stop(true);
+    coordinator.close();
+    await application.close();
+    if (pool) await pool.close(); else await engine.close();
+  })() };
 }
 
 /** serve's value-taking flags — needed to tell a flag VALUE from the bare
@@ -368,7 +560,7 @@ const SERVE_VALUE_FLAGS = new Set([
   "--batch", "--decode-concurrency", "--kv-quant",
   "--kv-group-size", "--quantized-kv-start", "--kv-budget",
   "--memory-budget", "--prompt-cache", "--ssd-cache", "--ssd-cache-max",
-  "--ssd-demote-idle", "--adapter", "--max-tokens", "--unix",
+  "--ssd-demote-idle", "--generation-checkpoint", "--adapter", "--max-tokens", "--unix",
   "--model-pool", "--compiled-decode", "--compiled-activations",
   "--fused-sdpa", "--thinking",
 ]);

package/src/serve/kv-budget.ts

@@ -0,0 +1,14 @@
+import type { ModelConfig } from "../config";
+import { kvBytesAt, KvScheme } from "../kv-scheme";
+
+/** Project one batch row's worst-case KV bytes using the same per-layer
+ * scheme that the server-wide admission ceiling uses. */
+export function batchRowKvBytes(
+  config: ModelConfig,
+  promptTokens: number,
+  maxTokens: number,
+  scheme?: KvScheme,
+): number {
+  const tokens = promptTokens + maxTokens;
+  return scheme ? scheme.bytesAt(config, tokens) : kvBytesAt(config, tokens);
+}