@mlx-node/server 0.0.0 → 0.0.8

package/dist/chat-session-warm-reuse.d.ts

@@ -0,0 +1,51 @@
+/**
+ * Server-private `ChatSession` warm-reuse helper.
+ *
+ * Kept out of the `@mlx-node/lm` package exports entirely so downstream
+ * consumers cannot reach it: the module lives inside
+ * `@mlx-node/server`, its file path is not on the server's export map,
+ * and nothing re-exports it. Callers are the two server endpoints —
+ * `endpoints/responses.ts` (tier-1 / tier-2 HIT branch) and
+ * `endpoints/messages.ts` (`getOrCreateWarmAny` HIT branch) — each
+ * invoking the helper exclusively under a `SessionRegistry` HIT gate.
+ *
+ * Why this helper exists at all: `ChatSession.reset()` is the safe
+ * public wipe — it always calls `model.resetCaches()` because the
+ * underlying `SessionCapableModel` is shared across every session
+ * lifetime via the native `ModelRegistry`. A partial wipe that leaves
+ * the shared native KV cache intact without the HIT gate would leak a
+ * previous (unrelated) request's cached prefix into the next
+ * `chat_session_start_sync` call. The server's warm-lease replay path
+ * DOES have that HIT gate — the registry's own `getOrCreate` hit
+ * signal is the authoritative proof that the native cache genuinely
+ * belongs to this chain — so a JS-state-only reset that preserves the
+ * native cache is correct there and only there.
+ *
+ * Fields accessed: `inFlight`, `history`, `lastImagesKey`, `lastAudioKey`, `turnCount`,
+ * `unresolvedOkToolCallCount`, `needsFullReplay`. These are TypeScript `private` fields on
+ * `ChatSession` (compile-time only) — at runtime they are ordinary
+ * properties. The cast through {@link ChatSessionWarmReuseInternals}
+ * gives this helper a typed view of the instance without relaxing the
+ * class's `private` declarations. The field names MUST stay in sync
+ * with `packages/lm/src/chat-session.ts`; a mismatch would silently
+ * skip the intended state wipe and is covered by the chat-session
+ * unit tests that exercise this path through the endpoint handler.
+ */
+import type { ChatSession, SessionCapableModel } from '@mlx-node/lm';
+/**
+ * JS-state-only reset that DELIBERATELY preserves the underlying
+ * model's native KV cache and `cached_token_history`.
+ *
+ * @internal server-private — used only by `SessionRegistry` HIT
+ * branches in `endpoints/responses.ts` and `endpoints/messages.ts`.
+ * Never export from this package's `index.ts`.
+ *
+ * Wipes ONLY the JS-side session state (history array, image key, turn
+ * counter, tool-call fan-out guard). With this function, the JS session
+ * is fresh enough for `ChatSession.primeHistory()` (which requires
+ * `turnCount === 0`) while the native prefix verifier can still recover
+ * the reused prefix on the next `chatSessionStart` and skip the
+ * corresponding re-prefill.
+ */
+export declare function resetPreservingNativeCacheForWarmReuse<M extends SessionCapableModel>(session: ChatSession<M>): Promise<void>;
+//# sourceMappingURL=chat-session-warm-reuse.d.ts.map

package/dist/chat-session-warm-reuse.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"chat-session-warm-reuse.d.ts","sourceRoot":"","sources":["../src/chat-session-warm-reuse.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAgCG;AAEH,OAAO,KAAK,EAAE,WAAW,EAAE,mBAAmB,EAAE,MAAM,cAAc,CAAC;AAmBrE;;;;;;;;;;;;;;GAcG;AACH,wBAAsB,sCAAsC,CAAC,CAAC,SAAS,mBAAmB,EACxF,OAAO,EAAE,WAAW,CAAC,CAAC,CAAC,GACtB,OAAO,CAAC,IAAI,CAAC,CAqBf"}

package/dist/chat-session-warm-reuse.js

@@ -0,0 +1,68 @@
+/**
+ * Server-private `ChatSession` warm-reuse helper.
+ *
+ * Kept out of the `@mlx-node/lm` package exports entirely so downstream
+ * consumers cannot reach it: the module lives inside
+ * `@mlx-node/server`, its file path is not on the server's export map,
+ * and nothing re-exports it. Callers are the two server endpoints —
+ * `endpoints/responses.ts` (tier-1 / tier-2 HIT branch) and
+ * `endpoints/messages.ts` (`getOrCreateWarmAny` HIT branch) — each
+ * invoking the helper exclusively under a `SessionRegistry` HIT gate.
+ *
+ * Why this helper exists at all: `ChatSession.reset()` is the safe
+ * public wipe — it always calls `model.resetCaches()` because the
+ * underlying `SessionCapableModel` is shared across every session
+ * lifetime via the native `ModelRegistry`. A partial wipe that leaves
+ * the shared native KV cache intact without the HIT gate would leak a
+ * previous (unrelated) request's cached prefix into the next
+ * `chat_session_start_sync` call. The server's warm-lease replay path
+ * DOES have that HIT gate — the registry's own `getOrCreate` hit
+ * signal is the authoritative proof that the native cache genuinely
+ * belongs to this chain — so a JS-state-only reset that preserves the
+ * native cache is correct there and only there.
+ *
+ * Fields accessed: `inFlight`, `history`, `lastImagesKey`, `lastAudioKey`, `turnCount`,
+ * `unresolvedOkToolCallCount`, `needsFullReplay`. These are TypeScript `private` fields on
+ * `ChatSession` (compile-time only) — at runtime they are ordinary
+ * properties. The cast through {@link ChatSessionWarmReuseInternals}
+ * gives this helper a typed view of the instance without relaxing the
+ * class's `private` declarations. The field names MUST stay in sync
+ * with `packages/lm/src/chat-session.ts`; a mismatch would silently
+ * skip the intended state wipe and is covered by the chat-session
+ * unit tests that exercise this path through the endpoint handler.
+ */
+/**
+ * JS-state-only reset that DELIBERATELY preserves the underlying
+ * model's native KV cache and `cached_token_history`.
+ *
+ * @internal server-private — used only by `SessionRegistry` HIT
+ * branches in `endpoints/responses.ts` and `endpoints/messages.ts`.
+ * Never export from this package's `index.ts`.
+ *
+ * Wipes ONLY the JS-side session state (history array, image key, turn
+ * counter, tool-call fan-out guard). With this function, the JS session
+ * is fresh enough for `ChatSession.primeHistory()` (which requires
+ * `turnCount === 0`) while the native prefix verifier can still recover
+ * the reused prefix on the next `chatSessionStart` and skip the
+ * corresponding re-prefill.
+ */
+export async function resetPreservingNativeCacheForWarmReuse(session) {
+    // TypeScript `private` fields are only compile-time checks; at
+    // runtime they are ordinary properties. The cast through
+    // `ChatSessionWarmReuseInternals` preserves full static typing for
+    // this helper's mutations while bypassing the `private` gate — which
+    // is correct here because this helper is the designated server-side
+    // friend accessor. The cast is funneled through `unknown` because TS
+    // correctly rejects a direct `ChatSession → Internals` cast when the
+    // concrete class has other non-internals fields.
+    const internals = session;
+    if (internals.inFlight) {
+        throw new Error('ChatSession: cannot resetPreservingNativeCacheForWarmReuse() while a send() is in flight; await the previous call first');
+    }
+    internals.history = [];
+    internals.lastImagesKey = null;
+    internals.lastAudioKey = null;
+    internals.turnCount = 0;
+    internals.unresolvedOkToolCallCount = null;
+    internals.needsFullReplay = false;
+}

package/dist/endpoints/messages-count-tokens.d.ts

@@ -0,0 +1,8 @@
+/** POST /v1/messages/count_tokens — Anthropic Messages token-count endpoint. */
+import type { ServerResponse } from 'node:http';
+import type { IdleSweeper } from '../idle-sweeper.js';
+import type { ModelWorkCoordinator } from '../model-work-coordinator.js';
+import type { ModelRegistry } from '../registry.js';
+import type { AnthropicCountTokensRequest } from '../types-anthropic.js';
+export declare function handleCountMessageTokens(res: ServerResponse, body: AnthropicCountTokensRequest, registry: ModelRegistry, idleSweeper?: IdleSweeper | null, resolveModel?: (name: string) => Promise<void>, modelWorkCoordinator?: ModelWorkCoordinator): Promise<void>;
+//# sourceMappingURL=messages-count-tokens.d.ts.map

package/dist/endpoints/messages-count-tokens.d.ts.map

@@ -0,0 +1 @@
+{"version":3,"file":"messages-count-tokens.d.ts","sourceRoot":"","sources":["../../src/endpoints/messages-count-tokens.ts"],"names":[],"mappings":"AAAA,gFAAgF;AAEhF,OAAO,KAAK,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAUhD,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,oBAAoB,CAAC;AAEtD,OAAO,KAAK,EAAE,oBAAoB,EAAE,MAAM,8BAA8B,CAAC;AACzE,OAAO,KAAK,EAAE,aAAa,EAAiB,MAAM,gBAAgB,CAAC;AACnE,OAAO,KAAK,EAAE,2BAA2B,EAAgC,MAAM,uBAAuB,CAAC;AAqBvG,wBAAsB,wBAAwB,CAC5C,GAAG,EAAE,cAAc,EACnB,IAAI,EAAE,2BAA2B,EACjC,QAAQ,EAAE,aAAa,EACvB,WAAW,CAAC,EAAE,WAAW,GAAG,IAAI,EAChC,YAAY,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,OAAO,CAAC,IAAI,CAAC,EAC9C,oBAAoB,CAAC,EAAE,oBAAoB,GAC1C,OAAO,CAAC,IAAI,CAAC,CAuHf"}

package/dist/endpoints/messages-count-tokens.js

@@ -0,0 +1,121 @@
+/** POST /v1/messages/count_tokens — Anthropic Messages token-count endpoint. */
+import { sendAnthropicBadRequest, sendAnthropicInternalError, sendAnthropicNotFound, sendAnthropicNotImplemented, } from '../errors.js';
+import { mapAnthropicRequest } from '../mappers/anthropic-request.js';
+function getChatTemplateTokenCounter(model) {
+    const candidate = model;
+    return typeof candidate.applyChatTemplate === 'function' ? candidate : null;
+}
+function endJson(res, body) {
+    res.writeHead(200, { 'Content-Type': 'application/json' });
+    res.end(JSON.stringify(body));
+}
+export async function handleCountMessageTokens(res, body, registry, idleSweeper, resolveModel, modelWorkCoordinator) {
+    if (body == null || typeof body !== 'object') {
+        sendAnthropicBadRequest(res, 'Request body must be a JSON object');
+        return;
+    }
+    if (!body.model) {
+        sendAnthropicBadRequest(res, 'Missing required field: model');
+        return;
+    }
+    if (!body.messages || !Array.isArray(body.messages) || body.messages.length === 0) {
+        sendAnthropicBadRequest(res, 'Missing required field: messages');
+        return;
+    }
+    for (const msg of body.messages) {
+        if (msg == null || typeof msg !== 'object') {
+            sendAnthropicBadRequest(res, 'Each message must be a non-null object');
+            return;
+        }
+    }
+    let mapped;
+    try {
+        mapped = mapAnthropicRequest(body);
+    }
+    catch (err) {
+        sendAnthropicBadRequest(res, err instanceof Error ? err.message : 'Invalid request');
+        return;
+    }
+    if (resolveModel) {
+        try {
+            const runResolve = () => idleSweeper ? idleSweeper.withSuspendedDrains(() => resolveModel(body.model)) : resolveModel(body.model);
+            if (modelWorkCoordinator)
+                await modelWorkCoordinator.withModelLoad(runResolve);
+            else
+                await runResolve();
+        }
+        catch (err) {
+            sendAnthropicInternalError(res, err instanceof Error ? err.message : 'Failed to resolve model');
+            return;
+        }
+    }
+    const lease = registry.acquireDispatchLease(body.model);
+    if (!lease) {
+        if (registry.get(body.model) != null) {
+            sendAnthropicInternalError(res, 'session registry missing for registered model');
+            return;
+        }
+        sendAnthropicNotFound(res, `Model "${body.model}" not found`);
+        return;
+    }
+    const leaseModel = lease.model;
+    const sessionReg = lease.registry;
+    const preLockInstanceId = lease.instanceId;
+    try {
+        // Token counting is a pure-CPU tokenize-and-template operation; it must NOT
+        // queue behind the per-model generation FIFO (`sessionReg.withExclusive`)
+        // because that serializes against multi-minute decode passes and turns a
+        // millisecond call into a multi-hundred-second wait. The dispatch lease
+        // already pins the model object + binding for the duration of this call,
+        // and `withInference` (a shared reader lock against `withModelLoad`) is
+        // enough to keep the model from being swapped out mid-tokenize.
+        const bindingStillMatchesLease = () => {
+            const lockedSessionReg = registry.getSessionRegistry(body.model);
+            const lockedInstanceId = registry.getInstanceId(body.model);
+            return (lockedSessionReg !== undefined &&
+                lockedInstanceId !== undefined &&
+                lockedSessionReg === sessionReg &&
+                lockedInstanceId === preLockInstanceId);
+        };
+        const rejectChangedBinding = (phase) => {
+            sendAnthropicBadRequest(res, `Model "${body.model}" binding changed while the token-count request was ${phase}. ` +
+                `A concurrent register() re-pointed the name at a different model instance ` +
+                `(or released it entirely), so counting against the leased model would use ` +
+                `a stale model object. Retry the request — if the swap was intentional, the ` +
+                `new binding will service the retry cleanly.`);
+        };
+        const runCountWithModelRead = async () => {
+            if (!bindingStillMatchesLease()) {
+                rejectChangedBinding('waiting for the model-load reader gate');
+                return;
+            }
+            const counter = getChatTemplateTokenCounter(leaseModel);
+            if (!counter) {
+                sendAnthropicNotImplemented(res, `Model "${body.model}" does not expose applyChatTemplate(); token counting requires a ` +
+                    `non-generating chat-template tokenizer API on the registered model.`);
+                return;
+            }
+            try {
+                const tokens = await counter.applyChatTemplate(mapped.messages, true, mapped.config.tools ?? null);
+                if (!bindingStillMatchesLease()) {
+                    rejectChangedBinding('running');
+                    return;
+                }
+                endJson(res, { input_tokens: tokens.length });
+            }
+            catch (err) {
+                sendAnthropicInternalError(res, err instanceof Error ? err.message : 'Failed to count tokens');
+            }
+        };
+        if (modelWorkCoordinator)
+            await modelWorkCoordinator.withInference(runCountWithModelRead);
+        else
+            await runCountWithModelRead();
+    }
+    catch (err) {
+        sendAnthropicInternalError(res, err instanceof Error ? err.message : 'Failed to count tokens');
+    }
+    finally {
+        registry.releaseDispatchLease(leaseModel);
+    }
+}

package/dist/endpoints/messages.d.ts

@@ -1,13 +1,65 @@
 /**
  * POST /v1/messages — stateless Anthropic Messages API.
  *
- * Every request carries the full conversation in `req.messages`. We allocate
- * a fresh `ChatSession` per request via `SessionRegistry.getOrCreate(null)`,
- * prime with the mapped history, and run `startFromHistory[Stream]`. No
- * adopt/drop: the session's lifetime is this single call.
+ * Every request carries the full conversation in `req.messages`. The
+ * Anthropic Messages API is stateless on the wire: there is no
+ * `previous_response_id` to thread, and clients (e.g. Claude Code)
+ * also do NOT propagate `prompt_cache_key` back to the server. The
+ * cross-turn / cross-conversation prefix-reuse path is one of two
+ * mutually-exclusive mechanisms, picked at request time based on
+ * whether the underlying native model has the block-paged KV cache
+ * adapter active (`SessionCapableModel.hasBlockPagedCache?.()`):
+ *
+ *   * **Paged-active path** (Qwen3 + LFM2 + Gemma4 are paged-active
+ *     today; Qwen3.5 dense/MoE and Qianfan-OCR remain non-paged /
+ *     default-off pending a perf decision and adapter wiring
+ *     respectively). Each request allocates a fresh `ChatSession` via
+ *     `SessionRegistry.createFreshSession()` and runs a full
+ *     `session.reset()` + `primeHistory()` +
+ *     `startFromHistory[Stream]()`. The JS-side warm slot is
+ *     **not** consulted, **not** leased, and **not** adopted —
+ *     cross-request prefix reuse is handled entirely by the native
+ *     `BlockAllocator`'s content-addressed prefix-hash table, which
+ *     refcounts SYS blocks shared across requests transparently
+ *     (two parallel `/v1/messages` requests with the same system
+ *     prompt run on distinct `ChatSession` objects but reference
+ *     the same physical KV blocks). The non-streaming
+ *     `X-Session-Cache` header is promoted from `fresh` to
+ *     `prefix_hit` after dispatch when the engine reports
+ *     `cachedTokens > 0`.
+ *
+ *   * **Non-paged path** (Qwen3.5 dense + MoE — default-off pending a
+ *     perf decision; the Qianfan-OCR VLM — no adapter wired). Each
+ *     request looks up the warm slot via
+ *     `SessionRegistry.getOrCreateWarmAny(requestedSystem)`. On a
+ *     HIT we keep the underlying native KV cache alive
+ *     (`resetPreservingNativeCacheForWarmReuse` wipes only JS-side
+ *     session state) so the native `verify_cache_prefix_direct` can
+ *     recognize the cached prefix and re-prefill only the new
+ *     suffix. On a MISS we run a full `session.reset()` to wipe
+ *     both JS and native state — a fresh JS session does NOT imply
+ *     a fresh native cache (the underlying `SessionCapableModel` is
+ *     shared and its native `cached_token_history` persists across
+ *     requests). After the dispatch settles we adopt the session
+ *     back under the sentinel id `'__msg_warm__'` (or drop on
+ *     uncommitted streams / thrown errors) so the next turn can
+ *     lease it. The sentinel is never produced by either the OpenAI
+ *     or the Anthropic wire format, so cross-endpoint capture via
+ *     tier-1 is impossible by construction. The `/v1/responses` and
+ *     `/v1/messages` endpoints still SHARE the single warm slot
+ *     under the registry's single-warm invariant on this path — a
+ *     turn on one side can evict the other's slot.
+ *
+ * The `prompt_cache_key` request field is still NOT exposed on this
+ * endpoint. Cross-conversation block-level cache reuse on the
+ * paged path is now driven by native content-addressing instead of
+ * the JS warm slot, so adding the field is no longer a prerequisite
+ * for that use case.
  */
 import type { IncomingMessage, ServerResponse } from 'node:http';
+import type { IdleSweeper } from '../idle-sweeper.js';
+import type { ModelWorkCoordinator } from '../model-work-coordinator.js';
 import type { ModelRegistry } from '../registry.js';
 import type { AnthropicMessagesRequest } from '../types-anthropic.js';
-export declare function handleCreateMessage(res: ServerResponse, body: AnthropicMessagesRequest, registry: ModelRegistry, httpReq?: IncomingMessage): Promise<void>;
+export declare function handleCreateMessage(res: ServerResponse, body: AnthropicMessagesRequest, registry: ModelRegistry, httpReq?: IncomingMessage, idleSweeper?: IdleSweeper | null, resolveModel?: (name: string) => Promise<void>, modelWorkCoordinator?: ModelWorkCoordinator): Promise<void>;
 //# sourceMappingURL=messages.d.ts.map

package/dist/endpoints/messages.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"messages.d.ts","sourceRoot":"","sources":["../../src/endpoints/messages.ts"],"names":[],"mappings":"AAAA;;;;;;;GAOG;AAEH,OAAO,KAAK,EAAE,eAAe,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAuBjE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAYpD,OAAO,KAAK,EAAE,wBAAwB,EAAE,MAAM,uBAAuB,CAAC;AA0XtE,wBAAsB,mBAAmB,CACvC,GAAG,EAAE,cAAc,EACnB,IAAI,EAAE,wBAAwB,EAC9B,QAAQ,EAAE,aAAa,EACvB,OAAO,CAAC,EAAE,eAAe,GACxB,OAAO,CAAC,IAAI,CAAC,CA4Of"}
+{"version":3,"file":"messages.d.ts","sourceRoot":"","sources":["../../src/endpoints/messages.ts"],"names":[],"mappings":"AAAA;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;GAyDG;AAEH,OAAO,KAAK,EAAE,eAAe,EAAE,cAAc,EAAE,MAAM,WAAW,CAAC;AAajE,OAAO,KAAK,EAAE,WAAW,EAAE,MAAM,oBAAoB,CAAC;AAgBtD,OAAO,KAAK,EAAE,oBAAoB,EAAE,MAAM,8BAA8B,CAAC;AACzE,OAAO,KAAK,EAAE,aAAa,EAAE,MAAM,gBAAgB,CAAC;AAepD,OAAO,KAAK,EAAE,wBAAwB,EAAE,MAAM,uBAAuB,CAAC;AAg3BtE,wBAAsB,mBAAmB,CACvC,GAAG,EAAE,cAAc,EACnB,IAAI,EAAE,wBAAwB,EAC9B,QAAQ,EAAE,aAAa,EACvB,OAAO,CAAC,EAAE,eAAe,EACzB,WAAW,CAAC,EAAE,WAAW,GAAG,IAAI,EAChC,YAAY,CAAC,EAAE,CAAC,IAAI,EAAE,MAAM,KAAK,OAAO,CAAC,IAAI,CAAC,EAC9C,oBAAoB,CAAC,EAAE,oBAAoB,GAC1C,OAAO,CAAC,IAAI,CAAC,CAioBf"}