@mlx-node/server 0.0.0

package/dist/registry.d.ts

@@ -0,0 +1,363 @@
+/**
+ * ModelRegistry -- maps friendly model names to loaded model instances.
+ *
+ * All models exposing the chat-session surface (see `SessionCapableModel`
+ * from `@mlx-node/lm`) are eligible for serving. Every registered model
+ * is paired with a `SessionRegistry` — an LRU+TTL cache of live
+ * `ChatSession` instances keyed by server-allocated response id.
+ *
+ * **Model-instance identity, not name.** Session registries are keyed
+ * by MODEL OBJECT identity. The single-warm-session invariant
+ * enforced by `SessionRegistry` is a property of the underlying
+ * `SessionCapableModel` (one shared native KV cache per instance), so
+ * registering the SAME model object under two names MUST yield the
+ * SAME `SessionRegistry` — otherwise each alias's local single-warm
+ * cache would hand out warm wrappers while the other alias silently
+ * stomps them via the shared native state. `register()` looks the
+ * model up in an identity-keyed map and reuses the existing registry
+ * on alias, or allocates a fresh one on first sight.
+ *
+ * **Monotonic per-instance ids.** Every distinct model object gets a
+ * monotonic `instanceId` on first registration, reused across aliases,
+ * and dropped when the binding is fully torn down. The responses
+ * endpoint persists this id alongside each stored record and, on a
+ * `previous_response_id` continuation, compares the stored id against
+ * the live id for `body.model`. This closes two holes a friendly-name
+ * check leaves open:
+ *
+ *  1. A name hot-swap — `register("foo", modelA)` then
+ *     `register("foo", modelB)` — would pass a string check, so a
+ *     chain produced by `modelA` could be silently replayed through
+ *     `modelB`'s tokenizer / chat template / KV layout. With instance
+ *     ids the stored id (modelA's) no longer matches the live id
+ *     (modelB's) and the continuation is rejected with 400.
+ *  2. Two NAMES aliasing the SAME model object would be spuriously
+ *     rejected by a string check comparing the stored name against
+ *     `body.model`. Instance ids recognise them as the same binding
+ *     and the continuation is accepted.
+ */
+import type { SessionCapableModel } from '@mlx-node/lm';
+import { SessionRegistry } from './session-registry.js';
+/** Minimal contract for a model that can be served via chat sessions. */
+export type ServableModel = SessionCapableModel;
+/** Model entry stored in the registry. */
+export interface ModelEntry {
+    id: string;
+    model: ServableModel;
+    createdAt: number;
+    /**
+     * Per-model-instance session cache, shared across every name that
+     * points at this exact model object. See the module-level rustdoc.
+     */
+    sessionRegistry: SessionRegistry;
+}
+/**
+ * Constructor options for {@link ModelRegistry}.
+ */
+export interface ModelRegistryOptions {
+    /**
+     * Maximum queue depth (waiters-only) per-model for the session
+     * registry's execution mutex. Forwarded into every
+     * `SessionRegistry` this registry allocates. See
+     * {@link SessionRegistryOptions.maxQueueDepth}. Default: `undefined`
+     * (unbounded — current behaviour).
+     */
+    maxQueueDepth?: number;
+}
+export declare class ModelRegistry {
+    private readonly maxQueueDepth;
+    private readonly models;
+    /**
+     * Identity-keyed (WeakMap semantics, but strong refs because the
+     * registry already holds the model through its ModelEntry) map
+     * from a model instance to its shared `SessionRegistry` binding.
+     * Every name that references the same model object resolves to
+     * the same binding — an alias of a registered model shares its
+     * session cache and therefore its single-warm invariant.
+     */
+    private readonly sessionRegistriesByModel;
+    /**
+     * Identity-keyed map from a model instance to its monotonic
+     * instance id. Entries are allocated on first registration,
+     * reused across aliasing, and dropped when the last binding
+     * releases (mirrors `sessionRegistriesByModel` lifetime exactly).
+     */
+    private readonly instanceIds;
+    /** Monotonic counter for `instanceIds`. Never reused. */
+    private nextInstanceId;
+    /**
+     * Tombstone map for instance ids retired by the hard-timeout
+     * breaker. When the responses endpoint force-releases a wedged
+     * persist's `retainBinding`, the breaker calls
+     * `retireInstanceIdForForceRelease(model)` BEFORE dropping the
+     * retain so the live id (already stamped into the pending record)
+     * is preserved here. A subsequent `register()` of the SAME model
+     * object that arrives AFTER the binding has fully torn down
+     * inherits the retired id instead of minting a fresh one — so a
+     * late-landing persist's row stays chainable. A true hot-swap
+     * (different model object) has no tombstone for the new model,
+     * so a fresh id is minted and the stale stored row is correctly
+     * rejected with 400 instance-mismatch.
+     *
+     * `WeakMap`-keyed on the model object so entries do not keep the
+     * model alive; if the model is GC'd the tombstone is cleaned up
+     * automatically.
+     *
+     * Lifetime is refcounted: store ONE `{ instanceId, outstandingCount }`
+     * entry per model. `retireInstanceIdForForceRelease` increments
+     * (creating the entry on first retire); `releaseTombstone`
+     * decrements and drops the entry when count hits zero. Because
+     * `register()` inherits the retired id whenever the tombstone
+     * exists, concurrent breakers on the same model all target the
+     * SAME numeric `instanceId` — one shared refcount keeps the
+     * tombstone alive as long as ANY pending persist still needs it,
+     * and memory is bounded at O(1) per model regardless of how many
+     * hard-timeouts have fired.
+     */
+    private readonly retiredInstanceIds;
+    constructor(opts?: ModelRegistryOptions);
+    /**
+     * Register a model under a given name.
+     *
+     * If the name is already registered and the new model is a
+     * DIFFERENT instance, the old binding's refcount is decremented
+     * (and dropped if no other alias references it) before the new
+     * binding is taken. Re-registering with the SAME model instance
+     * leaves the binding unchanged.
+     *
+     * On first sight of a model object a fresh `SessionRegistry` is
+     * allocated. On alias the existing registry is reused so the
+     * single-warm invariant spans both names.
+     *
+     * Tombstone-inherit path: if the binding was previously torn down
+     * AND the hard-timeout breaker called
+     * `retireInstanceIdForForceRelease(model)` before teardown fired,
+     * the fresh binding inherits the retired instance id from
+     * `retiredInstanceIds` — so a late-landing persist's record stays
+     * chainable. A hot-swap (different model object) has no tombstone,
+     * so a fresh id is minted and the stale record fails
+     * `previous_response_id` with 400. The aliasing path naturally
+     * preserves the id because `instanceIds.has(model)` is already
+     * true.
+     */
+    register(name: string, model: ServableModel): void;
+    /**
+     * Unregister a model by name.
+     *
+     * Drops the name -> ModelEntry mapping and decrements the shared
+     * session-registry binding's refcount. When the refcount hits zero
+     * (no other alias references this model object) the binding — and
+     * the `SessionRegistry` it owns — is dropped entirely so cached
+     * sessions for the now-unreferenced model are released.
+     *
+     * @returns true if the model was removed.
+     */
+    unregister(name: string): boolean;
+    /**
+     * Decrement the refcount on a model binding; drop it at zero iff
+     * no dispatch holds a lease AND no post-commit persist is still
+     * retaining it. When either counter is non-zero the teardown is
+     * deferred via `pendingTeardown` so the `SessionRegistry` (and
+     * its `execLock` FIFO) stays alive until the last holder releases.
+     *
+     * A concurrent `register(sameModel)` between `dropNameReference()`
+     * and the final release clears `pendingTeardown` and reuses the
+     * still-live binding, preserving `modelInstanceId` so any row the
+     * pending persist is about to land still resolves to a live id
+     * when the next continuation arrives.
+     */
+    private dropNameReference;
+    /**
+     * Drop `model`'s binding and instance id from the registry.
+     * Shared teardown step; a subsequent re-registration usually mints
+     * a FRESH instance id — once the last alias, lease, and persist
+     * retention all release, any previously stored record referencing
+     * this id belongs to a logically dead binding and a continuation
+     * against it must fall through to the `currentInstanceId === undefined`
+     * rejection path so the stale chain cannot be replayed.
+     *
+     * Tombstone exception: if the hard-timeout breaker retired the
+     * previous id via `retireInstanceIdForForceRelease` before the
+     * forced release, a subsequent same-object `register()` inherits
+     * the retired id from `retiredInstanceIds` instead of minting
+     * fresh — this preserves chain continuity for late-landing
+     * persists that crossed the safety breaker.
+     */
+    private finalizeBindingTeardown;
+    /**
+     * Acquire a dispatch lease on the session registry bound to `name`.
+     * Returns the live `SessionRegistry` and the binding's instance id,
+     * or `undefined` if the name is not registered. Every successful
+     * acquisition MUST be balanced with exactly one
+     * `releaseDispatchLease(model)` call (typically via try/finally).
+     *
+     * The lease keeps the binding alive past a concurrent
+     * `unregister()` / `register(differentModel)` sequence: the
+     * `SessionRegistry` and its `execLock` FIFO remain valid while any
+     * lease is outstanding, so a newly registered same-model alias
+     * will rebind to the SAME registry and its `withExclusive` will
+     * serialize behind the in-flight dispatch.
+     *
+     * The returned `model` handle is what the caller passes to
+     * `releaseDispatchLease()` — the lease binds to the model OBJECT
+     * (not the friendly name) because the name can be hot-swapped
+     * while the lease is held.
+     */
+    acquireDispatchLease(name: string): {
+        model: ServableModel;
+        registry: SessionRegistry;
+        instanceId: number;
+    } | undefined;
+    /**
+     * Release a dispatch lease previously obtained via
+     * `acquireDispatchLease()`. Decrements the binding's in-flight
+     * counter and, if the binding has been flagged for teardown (its
+     * refcount hit zero while the lease was held), drops it once the
+     * last lease releases. Safe to call exactly once per acquired
+     * lease; calling it on a model whose binding has already been
+     * fully torn down is a no-op.
+     */
+    releaseDispatchLease(model: ServableModel): void;
+    /**
+     * Retain the binding for the duration of a post-commit persist.
+     *
+     * The responses endpoint starts `store.store(record)` synchronously
+     * inside `withExclusive` so the pending-writes tracker observes
+     * the in-flight write before the mutex releases, but does NOT
+     * await it on the critical path. The write still carries the
+     * binding's `modelInstanceId` (stamped into `configJson` by
+     * `buildResponseRecord`); without a retention, a same-model
+     * unregister + re-register completing while the write is in flight
+     * would delete the instance id and the re-registration would mint
+     * a fresh one, so the row — when it finally lands — would
+     * reference a dead id and the next continuation would be rejected
+     * with 400 instance-mismatch.
+     *
+     * The retention counter is CHECKED in every teardown gate
+     * (`dropNameReference`, `releaseDispatchLease`, `releaseBinding`).
+     * It is orthogonal to `inFlight` so the dispatch lease can release
+     * eagerly after `withExclusive` returns while the binding stays
+     * pinned long enough for the backgrounded `store.store(...)` to
+     * settle.
+     *
+     * Safe to call on a model whose binding has already been torn down
+     * (no-op). The matching `releaseBinding(model)` MUST still run in
+     * the persist's `.finally(...)` so the counter stays balanced.
+     */
+    retainBinding(model: ServableModel): void;
+    /**
+     * Balance a prior `retainBinding()` call. Decrements the persist
+     * retention counter and, if the binding has been flagged for
+     * teardown (refcount hit zero while the retention was held AND
+     * every dispatch lease has already released), drops it once the
+     * last retention releases. Safe to call exactly once per retain;
+     * calling it on a model whose binding has already been fully torn
+     * down is a no-op.
+     */
+    releaseBinding(model: ServableModel): void;
+    /**
+     * Tombstone installer, invoked exclusively by the responses
+     * endpoint's hard-timeout breaker (see
+     * `getPostCommitPersistHardTimeoutMs` in `endpoints/responses.ts`)
+     * when it force-releases the `retainBinding` on a wedged persist.
+     * Must be called BEFORE the idempotent
+     * `persistRetainBox.release?.()` so `instanceIds.get(model)`
+     * still returns the live id that the already-stamped record
+     * carries.
+     *
+     * Returns the retired id so the caller can capture it and scope
+     * the tombstone's lifetime to the specific wedged persist via
+     * `releaseTombstone(model)` inside that persist's `.finally(...)`.
+     * Returns `undefined` when the model has no current instance id
+     * assignment (caller raced the natural teardown path).
+     *
+     * Refcounted: each call increments a shared
+     * `{ instanceId, outstandingCount }` entry per model. Overlapping
+     * breakers share one slot (they all target the same numeric id
+     * because `register()` inherits the retired id whenever the
+     * tombstone is present), so memory stays O(1) per model.
+     */
+    retireInstanceIdForForceRelease(model: ServableModel): {
+        instanceId: number;
+    } | undefined;
+    /**
+     * Tombstone cleanup. Called from the post-commit persist's
+     * `.finally(...)` to balance exactly one prior
+     * `retireInstanceIdForForceRelease(model)` call. Decrements the
+     * shared refcount and drops the entry at zero so the next natural
+     * teardown mints a fresh id.
+     *
+     * Safe to call on a model whose tombstone has already been drained
+     * (no-op). The counter is clamped non-negative so spurious
+     * releases cannot underflow and re-enable inheritance.
+     */
+    releaseTombstone(model: ServableModel): void;
+    /**
+     * Retrieve a model instance by name.
+     */
+    get(name: string): ServableModel | undefined;
+    /**
+     * Retrieve the monotonic instance id for the model currently bound
+     * to `name`, or `undefined` if the name isn't registered.
+     *
+     * Two names that alias the same model object return the SAME id
+     * (they share a binding), and a name that has been hot-swapped to
+     * a different model object returns a DIFFERENT id than before the
+     * swap (the prior binding's id was dropped by `dropNameReference`
+     * and a fresh id was minted for the new model on re-registration).
+     *
+     * The responses endpoint uses this to key the
+     * `previous_response_id` cross-chain guard on instance identity
+     * instead of friendly name, so hot swaps are caught and safe
+     * aliases are not spuriously rejected.
+     */
+    getInstanceId(name: string): number | undefined;
+    /**
+     * Retrieve the session registry for a given model name, or
+     * `undefined` if the name is not registered.
+     *
+     * Every name that points at the same model instance returns the
+     * SAME `SessionRegistry` object. Two aliases `a` and `b` of one
+     * model therefore satisfy
+     * `registry.getSessionRegistry('a') === registry.getSessionRegistry('b')`,
+     * which is what the single-warm invariant requires: any turn
+     * through either alias advances the same cache's state, so a later
+     * lookup via either alias sees the current warm wrapper (if
+     * freshly adopted) or misses and cold-replays (if it was leased
+     * out by the other alias) — never a stale wrapper pointing at
+     * stomped native state.
+     */
+    getSessionRegistry(name: string): SessionRegistry | undefined;
+    /**
+     * Iterate every DISTINCT session registry currently in use.
+     *
+     * Two aliases of the same model share one `SessionRegistry`, so
+     * naively walking every `ModelEntry` would yield duplicates. We
+     * walk the identity-keyed bindings instead so each registry
+     * appears exactly once, which is what the periodic `sweep()`
+     * scheduler in `server.ts` needs to avoid redundantly sweeping
+     * the same cache multiple times per tick.
+     */
+    listSessionRegistries(): SessionRegistry[];
+    /**
+     * List all registered models in the OpenAI /v1/models format.
+     */
+    list(): {
+        id: string;
+        object: string;
+        created: number;
+        owned_by: string;
+    }[];
+    /**
+     * Check whether a model supports streaming.
+     *
+     * Every `SessionCapableModel` structurally exposes
+     * `chatStreamSessionStart`, so this is universally `true` for any
+     * properly-typed model registered through the session-capable
+     * interface. Kept as a belt-and-suspenders duck-type so a partially
+     * stubbed test double (pre-migration or intentionally non-streaming)
+     * can still opt out by omitting the method.
+     */
+    hasStreamSupport(model: ServableModel): boolean;
+}
+//# sourceMappingURL=registry.d.ts.map

package/dist/registry.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"registry.d.ts","sourceRoot":"","sources":["../src/registry.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAqCG;AAEH,OAAO,KAAK,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAExD,OAAO,EAAE,eAAe,EAAE,MAAM,uBAAuB,CAAC;AAExD,yEAAyE;AACzE,MAAM,MAAM,aAAa,GAAG,mBAAmB,CAAC;AAEhD,0CAA0C;AAC1C,MAAM,WAAW,UAAU;IACzB,EAAE,EAAE,MAAM,CAAC;IACX,KAAK,EAAE,aAAa,CAAC;IACrB,SAAS,EAAE,MAAM,CAAC;IAClB;;;OAGG;IACH,eAAe,EAAE,eAAe,CAAC;CAClC;AAsCD;;GAEG;AACH,MAAM,WAAW,oBAAoB;IACnC;;;;;;OAMG;IACH,aAAa,CAAC,EAAE,MAAM,CAAC;CACxB;AAED,qBAAa,aAAa;IACxB,OAAO,CAAC,QAAQ,CAAC,aAAa,CAAqB;IACnD,OAAO,CAAC,QAAQ,CAAC,MAAM,CAAiC;IACxD;;;;;;;OAOG;IACH,OAAO,CAAC,QAAQ,CAAC,wBAAwB,CAAoD;IAC7F;;;;;OAKG;IACH,OAAO,CAAC,QAAQ,CAAC,WAAW,CAAoC;IAChE,yDAAyD;IACzD,OAAO,CAAC,cAAc,CAAK;IAC3B;;;;;;;;;;;;;;;;;;;;;;;;;;;;OA4BG;IACH,OAAO,CAAC,QAAQ,CAAC,kBAAkB,CAAkF;gBAEzG,IAAI,CAAC,EAAE,oBAAoB;IAIvC;;;;;;;;;;;;;;;;;;;;;;;OAuBG;IACH,QAAQ,CAAC,IAAI,EAAE,MAAM,EAAE,KAAK,EAAE,aAAa,GAAG,IAAI;IA8DlD;;;;;;;;;;OAUG;IACH,UAAU,CAAC,IAAI,EAAE,MAAM,GAAG,OAAO;IAQjC;;;;;;;;;;;;OAYG;IACH,OAAO,CAAC,iBAAiB;IAiBzB;;;;;;;;;;;;;;;OAeG;IACH,OAAO,CAAC,uBAAuB;IAK/B;;;;;;;;;;;;;;;;;;OAkBG;IACH,oBAAoB,CAClB,IAAI,EAAE,MAAM,GACX;QAAE,KAAK,EAAE,aAAa,CAAC;QAAC,QAAQ,EAAE,eAAe,CAAC;QAAC,UAAU,EAAE,MAAM,CAAA;KAAE,GAAG,SAAS;IAWtF;;;;;;;;OAQG;IACH,oBAAoB,CAAC,KAAK,EAAE,aAAa,GAAG,IAAI;IAUhD;;;;;;;;;;;;;;;;;;;;;;;;;OAyBG;IACH,aAAa,CAAC,KAAK,EAAE,aAAa,GAAG,IAAI;IAMzC;;;;;;;;OAQG;IACH,cAAc,CAAC,KAAK,EAAE,aAAa,GAAG,IAAI;IAU1C;;;;;;;;;;;;;;;;;;;;;OAqBG;IACH,+BAA+B,CAAC,KAAK,EAAE,aAAa,GAAG;QAAE,UAAU,EAAE,MAAM,CAAA;KAAE,GAAG,SAAS;IAYzF;;;;;;;;;;OAUG;IACH,gBAAgB,CAAC,KAAK,EAAE,aAAa,GAAG,IAAI;IAS5C;;OAEG;IACH,GAAG,CAAC,IAAI,EAAE,MAAM,GAAG,aAAa,GAAG,SAAS;IAI5C;;;;;;;;;;;;;;OAcG;IACH,aAAa,CAAC,IAAI,EAAE,MAAM,GAAG,MAAM,GAAG,SAAS;IAM/C;;;;;;;;;;;;;;OAcG;IACH,kBAAkB,CAAC,IAAI,EAAE,MAAM,GAAG,eAAe,GAAG,SAAS;IAI7D;;;;;;;;;OASG;IACH,qBAAqB,IAAI,eAAe,EAAE;IAQ1C;;OAEG;IACH,IAAI,IAAI;QAAE,EAAE,EAAE,MAAM,CAAC;QAAC,MAAM,EAAE,MAAM,CAAC;QAAC,OAAO,EAAE,MAAM,CAAC;QAAC,QAAQ,EAAE,MAAM,CAAA;KAAE,EAAE;IAa3E;;;;;;;;;OASG;IACH,gBAAgB,CAAC,KAAK,EAAE,aAAa,GAAG,OAAO;CAIhD"}