@mlx-node/server 0.0.0 → 0.0.8
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/chat-session-warm-reuse.d.ts +51 -0
- package/dist/chat-session-warm-reuse.d.ts.map +1 -0
- package/dist/chat-session-warm-reuse.js +68 -0
- package/dist/endpoints/messages-count-tokens.d.ts +8 -0
- package/dist/endpoints/messages-count-tokens.d.ts.map +1 -0
- package/dist/endpoints/messages-count-tokens.js +121 -0
- package/dist/endpoints/messages.d.ts +57 -5
- package/dist/endpoints/messages.d.ts.map +1 -1
- package/dist/endpoints/messages.js +1043 -147
- package/dist/endpoints/models.d.ts +2 -1
- package/dist/endpoints/models.d.ts.map +1 -1
- package/dist/endpoints/models.js +2 -2
- package/dist/endpoints/responses.d.ts +20 -7
- package/dist/endpoints/responses.d.ts.map +1 -1
- package/dist/endpoints/responses.js +572 -82
- package/dist/errors.d.ts +1 -0
- package/dist/errors.d.ts.map +1 -1
- package/dist/errors.js +3 -0
- package/dist/handler.d.ts +42 -0
- package/dist/handler.d.ts.map +1 -1
- package/dist/handler.js +6 -1
- package/dist/idle-sweeper.d.ts +245 -0
- package/dist/idle-sweeper.d.ts.map +1 -0
- package/dist/idle-sweeper.js +408 -0
- package/dist/index.d.ts +8 -2
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +10 -0
- package/dist/mappers/anthropic-request.d.ts +24 -2
- package/dist/mappers/anthropic-request.d.ts.map +1 -1
- package/dist/mappers/anthropic-request.js +222 -24
- package/dist/mappers/anthropic-response.d.ts +29 -4
- package/dist/mappers/anthropic-response.d.ts.map +1 -1
- package/dist/mappers/anthropic-response.js +143 -21
- package/dist/mappers/request.d.ts +48 -0
- package/dist/mappers/request.d.ts.map +1 -1
- package/dist/mappers/request.js +211 -35
- package/dist/mappers/response.d.ts.map +1 -1
- package/dist/mappers/response.js +13 -1
- package/dist/model-work-coordinator.d.ts +70 -0
- package/dist/model-work-coordinator.d.ts.map +1 -0
- package/dist/model-work-coordinator.js +120 -0
- package/dist/pending-writes.d.ts.map +1 -1
- package/dist/presets.d.ts +82 -0
- package/dist/presets.d.ts.map +1 -0
- package/dist/presets.js +98 -0
- package/dist/registry.d.ts +31 -1
- package/dist/registry.d.ts.map +1 -1
- package/dist/registry.js +33 -5
- package/dist/router.d.ts +4 -1
- package/dist/router.d.ts.map +1 -1
- package/dist/router.js +34 -4
- package/dist/server.d.ts +76 -0
- package/dist/server.d.ts.map +1 -1
- package/dist/server.js +48 -1
- package/dist/session-registry.d.ts +272 -18
- package/dist/session-registry.d.ts.map +1 -1
- package/dist/session-registry.js +509 -37
- package/dist/stop-sequence-buffer.d.ts +58 -0
- package/dist/stop-sequence-buffer.d.ts.map +1 -0
- package/dist/stop-sequence-buffer.js +148 -0
- package/dist/text-recovery.d.ts +35 -0
- package/dist/text-recovery.d.ts.map +1 -0
- package/dist/text-recovery.js +41 -0
- package/dist/timing.d.ts +80 -0
- package/dist/timing.d.ts.map +1 -0
- package/dist/timing.js +121 -0
- package/dist/tool-call-buffer.d.ts +5 -5
- package/dist/tool-call-buffer.d.ts.map +1 -1
- package/dist/tool-call-buffer.js +28 -8
- package/dist/types-anthropic.d.ts +161 -1
- package/dist/types-anthropic.d.ts.map +1 -1
- package/dist/types.d.ts +172 -2
- package/dist/types.d.ts.map +1 -1
- package/package.json +5 -5
package/dist/session-registry.js
CHANGED
|
@@ -2,6 +2,18 @@
|
|
|
2
2
|
* SessionRegistry -- per-model cache holding AT MOST one live
|
|
3
3
|
* `ChatSession` whose native KV state is currently valid.
|
|
4
4
|
*
|
|
5
|
+
* **Tier-2 `prompt_cache_key` reuse is ON by default** so the server
|
|
6
|
+
* is compatible with any stateless LLM agent that sends the full
|
|
7
|
+
* conversation history each turn. The key is caller-controlled and
|
|
8
|
+
* HMAC-scoped with a boot-time nonce (raw value never stored on the
|
|
9
|
+
* entry), but two clients that pick the same raw key will still
|
|
10
|
+
* lease the same warm session — a session-hijack surface in
|
|
11
|
+
* multi-tenant settings. For multi-tenant deployments, opt out via
|
|
12
|
+
* `MLX_DISABLE_PROMPT_CACHE_KEY=1` or front the server with an auth
|
|
13
|
+
* proxy that rewrites or namespaces `prompt_cache_key` per tenant
|
|
14
|
+
* before it reaches this process. See also the comments on
|
|
15
|
+
* {@link scopePromptCacheKey}.
|
|
16
|
+
*
|
|
5
17
|
* Design notes:
|
|
6
18
|
*
|
|
7
19
|
* - **One registry per model.** Composed alongside each registered
|
|
@@ -50,8 +62,8 @@
|
|
|
50
62
|
* atomically appends the new user turn, so cold replay is
|
|
51
63
|
* indistinguishable from a hot hit.
|
|
52
64
|
*
|
|
53
|
-
* - **TTL.** Default 1800 seconds mirrors `
|
|
54
|
-
* in `packages/server/src/
|
|
65
|
+
* - **TTL.** Default 1800 seconds mirrors `DEFAULT_RESPONSE_RETENTION_SECONDS`
|
|
66
|
+
* in `packages/server/src/server.ts` so the cached
|
|
55
67
|
* entry ages out alongside its stored response metadata. With
|
|
56
68
|
* at most one entry there is no LRU bookkeeping — just a single
|
|
57
69
|
* expiry check on lookup.
|
|
@@ -79,7 +91,204 @@
|
|
|
79
91
|
* time. A weaker epoch-token scheme would let the losing
|
|
80
92
|
* `adopt()` no-op but the native KV would already be wrong.
|
|
81
93
|
*/
|
|
94
|
+
import { createHash, createHmac, randomBytes } from 'node:crypto';
|
|
82
95
|
import { ChatSession } from '@mlx-node/lm';
|
|
96
|
+
/**
|
|
97
|
+
* Tier-2 `prompt_cache_key` reuse is **ON by default** so the server
|
|
98
|
+
* is immediately compatible with any stateless LLM agent (pi-mono,
|
|
99
|
+
* Aider, Codex CLI, Claude Code, Cline, Continue) that sends the full
|
|
100
|
+
* transcript every turn. Agents that set OpenAI's standard
|
|
101
|
+
* `prompt_cache_key` field — which most do — get automatic KV cache
|
|
102
|
+
* reuse across turns of the same logical session.
|
|
103
|
+
*
|
|
104
|
+
* Opt out via `MLX_DISABLE_PROMPT_CACHE_KEY=1` for **multi-tenant
|
|
105
|
+
* deployments**, where the tier-2 lookup becomes unsafe: two clients
|
|
106
|
+
* that pick the same raw `prompt_cache_key` (by accident or on
|
|
107
|
+
* purpose) would share a warm `ChatSession`, leaking conversation
|
|
108
|
+
* history and sampling state across principals.
|
|
109
|
+
*
|
|
110
|
+
* Multi-tenant isolation is out of scope for this registry. The
|
|
111
|
+
* HMAC-scoping applied below only hides the raw key from memory /
|
|
112
|
+
* dumps — it does NOT protect against two clients sharing the same
|
|
113
|
+
* raw input. Operators who need multi-tenant isolation must either
|
|
114
|
+
* disable the feature or front the server with an auth proxy that
|
|
115
|
+
* rewrites `prompt_cache_key` per-tenant before it reaches the
|
|
116
|
+
* process.
|
|
117
|
+
*
|
|
118
|
+
* Read at call time (not cached at module load) so tests can flip the
|
|
119
|
+
* env via `vi.stubEnv()` between cases without re-importing the module.
|
|
120
|
+
* The check is a single env-var read plus a string compare — negligible
|
|
121
|
+
* against the rest of the lookup work.
|
|
122
|
+
*/
|
|
123
|
+
function isPromptCacheKeyEnabled() {
|
|
124
|
+
return process.env.MLX_DISABLE_PROMPT_CACHE_KEY !== '1';
|
|
125
|
+
}
|
|
126
|
+
/**
|
|
127
|
+
* Minimum length accepted for a caller-supplied `prompt_cache_key`
|
|
128
|
+
* before tier-2 scoping. Short keys make trivial guessing collisions
|
|
129
|
+
* plausible; reject anything shorter than this as if the caller had
|
|
130
|
+
* not supplied a key at all. Chosen to reject one- / two- / few-byte
|
|
131
|
+
* values a client might accidentally pass through while still allowing
|
|
132
|
+
* any reasonable opaque id (UUID prefix, short client token, etc.).
|
|
133
|
+
*/
|
|
134
|
+
const PROMPT_CACHE_KEY_MIN_LENGTH = 8;
|
|
135
|
+
/**
|
|
136
|
+
* Lazily-initialized boot-time nonce used to HMAC every caller-supplied
|
|
137
|
+
* `prompt_cache_key` before it is stored or looked up. Held in memory
|
|
138
|
+
* only — never persisted to disk. A process restart invalidates every
|
|
139
|
+
* tier-2 entry because the next module instance produces a fresh
|
|
140
|
+
* nonce.
|
|
141
|
+
*
|
|
142
|
+
* The nonce makes pre-existing entries unmatchable from outside the
|
|
143
|
+
* process: an attacker who knows a victim's raw `prompt_cache_key` but
|
|
144
|
+
* cannot read the nonce from the server's memory also cannot craft a
|
|
145
|
+
* lookup that collides with the stored HMAC'd key. Combined with the
|
|
146
|
+
* opt-in gate above, the tier-2 surface is off-by-default and bound to
|
|
147
|
+
* a server-instance secret when enabled.
|
|
148
|
+
*
|
|
149
|
+
* Populated lazily on first use so the cost is not paid when tier-2 is
|
|
150
|
+
* disabled. Module-scope so it is shared across every `SessionRegistry`
|
|
151
|
+
* in the process (each per-model registry's HMAC'd keys are still
|
|
152
|
+
* distinct via their per-registry `entries` map — there is no
|
|
153
|
+
* cross-model leakage).
|
|
154
|
+
*/
|
|
155
|
+
let cachedNonce = null;
|
|
156
|
+
/** Lazily obtain the module-scoped HMAC nonce. */
|
|
157
|
+
function getNonce() {
|
|
158
|
+
if (cachedNonce === null) {
|
|
159
|
+
cachedNonce = randomBytes(32);
|
|
160
|
+
}
|
|
161
|
+
return cachedNonce;
|
|
162
|
+
}
|
|
163
|
+
/**
|
|
164
|
+
* Test-only hook used by the scoping unit tests to simulate a
|
|
165
|
+
* server restart: resets the module-scoped HMAC nonce (so every
|
|
166
|
+
* previously stored tier-2 key misses) and clears the silent-miss
|
|
167
|
+
* dedupe cache so tests can re-exercise the once-per-key diagnostic
|
|
168
|
+
* path.
|
|
169
|
+
*
|
|
170
|
+
* **Not exported from the package's public `index.ts` surface** —
|
|
171
|
+
* exporting it there would let downstream consumers nuke tier-2
|
|
172
|
+
* state in production (every stored entry would go unreachable).
|
|
173
|
+
* Tests reach the function via the deep path
|
|
174
|
+
* `packages/server/src/session-registry.ts` instead; that import is
|
|
175
|
+
* deliberately noisy to signal "test-only, do not use from app
|
|
176
|
+
* code". The `__` prefix is a loud-enough convention for the
|
|
177
|
+
* ergonomic test path but NOT sufficient for a public package
|
|
178
|
+
* export.
|
|
179
|
+
*/
|
|
180
|
+
export function __resetPromptCacheKeyNonceForTests() {
|
|
181
|
+
cachedNonce = null;
|
|
182
|
+
loggedSilentMissKeys.clear();
|
|
183
|
+
}
|
|
184
|
+
/**
|
|
185
|
+
* Test-only probe for the silent-miss dedupe Map's live size. Same
|
|
186
|
+
* "do not use from app code" rationale as
|
|
187
|
+
* {@link __resetPromptCacheKeyNonceForTests}; the `__` prefix and the
|
|
188
|
+
* test-only deep-import path are the load-bearing signals. Exists so
|
|
189
|
+
* the flooding regression test can assert the FIFO cap actually bounds
|
|
190
|
+
* the stored set — an invariant that is otherwise unobservable from
|
|
191
|
+
* outside the module and is NOT implied by the warning count (each
|
|
192
|
+
* call emits at most one warning regardless of whether the cap works).
|
|
193
|
+
*/
|
|
194
|
+
export function __loggedSilentMissKeysSizeForTests() {
|
|
195
|
+
return loggedSilentMissKeys.size;
|
|
196
|
+
}
|
|
197
|
+
/**
|
|
198
|
+
* Normalize and HMAC-scope a caller-supplied `prompt_cache_key` before
|
|
199
|
+
* it is stored or used for lookup.
|
|
200
|
+
*
|
|
201
|
+
* **Single-tenant trust boundary.** HMAC-scoping hides the raw key
|
|
202
|
+
* from memory dumps and keeps one process instance's stored keys
|
|
203
|
+
* unreachable from another instance (a restart rerolls the nonce),
|
|
204
|
+
* but it does NOT protect against two clients supplying the same raw
|
|
205
|
+
* key — by construction both lookups HMAC to the same scoped key and
|
|
206
|
+
* share the entry. Multi-tenant isolation is out of scope; see the
|
|
207
|
+
* module docstring.
|
|
208
|
+
*
|
|
209
|
+
* Returns `null` when:
|
|
210
|
+
* - The tier-2 feature is disabled
|
|
211
|
+
* (`MLX_DISABLE_PROMPT_CACHE_KEY` is set to `"1"`).
|
|
212
|
+
* - `rawKey` is `null`, `undefined`, or the empty string (callers
|
|
213
|
+
* that forget to thread the key must not accidentally opt into
|
|
214
|
+
* tier-2 reuse).
|
|
215
|
+
* - `rawKey` is shorter than {@link PROMPT_CACHE_KEY_MIN_LENGTH}
|
|
216
|
+
* characters, which keeps trivial guessing collisions off the
|
|
217
|
+
* table.
|
|
218
|
+
*
|
|
219
|
+
* Otherwise returns the first 32 hex chars of
|
|
220
|
+
* `HMAC-SHA256(cachedNonce, rawKey)` — opaque, server-instance-scoped,
|
|
221
|
+
* and long enough to preserve the 64-bit entropy floor that the
|
|
222
|
+
* pre-scoped path relied on for key uniqueness.
|
|
223
|
+
*/
|
|
224
|
+
function scopePromptCacheKey(rawKey) {
|
|
225
|
+
if (!isPromptCacheKeyEnabled())
|
|
226
|
+
return null;
|
|
227
|
+
if (rawKey == null || rawKey.length < PROMPT_CACHE_KEY_MIN_LENGTH)
|
|
228
|
+
return null;
|
|
229
|
+
return createHmac('sha256', getNonce()).update(rawKey).digest('hex').slice(0, 32);
|
|
230
|
+
}
|
|
231
|
+
/**
|
|
232
|
+
* Bounded set of SHA-256-digest prefixes tracking which
|
|
233
|
+
* `prompt_cache_key` values have already had a silent-miss warning
|
|
234
|
+
* emitted in this process. Stored as 64-bit hex digests (not raw
|
|
235
|
+
* strings) so an attacker flooding the endpoint with distinct
|
|
236
|
+
* attacker-controlled keys cannot drive unbounded memory growth.
|
|
237
|
+
* FIFO-evicted at {@link LOGGED_SILENT_MISS_KEYS_MAX} entries via the
|
|
238
|
+
* Map insertion-order guarantee. Reset alongside the nonce / warning
|
|
239
|
+
* flag in {@link __resetPromptCacheKeyNonceForTests} so unit tests
|
|
240
|
+
* can re-exercise the once-per-key path.
|
|
241
|
+
*/
|
|
242
|
+
const LOGGED_SILENT_MISS_KEYS_MAX = 256;
|
|
243
|
+
const loggedSilentMissKeys = new Map();
|
|
244
|
+
/** Hash a raw key to a bounded digest for dedupe storage. */
|
|
245
|
+
function digestSilentMissKey(rawKey) {
|
|
246
|
+
// 64-bit prefix is enough for dedupe across a 256-entry window.
|
|
247
|
+
return createHash('sha256').update(rawKey).digest('hex').slice(0, 16);
|
|
248
|
+
}
|
|
249
|
+
/**
|
|
250
|
+
* Emit a once-per-raw-key stderr debug warning when the caller
|
|
251
|
+
* supplies a non-empty `prompt_cache_key` but at least one tier-2
|
|
252
|
+
* prerequisite is missing — the env gate is off, or the key is
|
|
253
|
+
* shorter than {@link PROMPT_CACHE_KEY_MIN_LENGTH}. The silent-miss
|
|
254
|
+
* fallback (cold-start as if no key were supplied) is the documented
|
|
255
|
+
* behaviour but is easy to miss during integration; this nudge
|
|
256
|
+
* surfaces the cause once per distinct key so operators don't have
|
|
257
|
+
* to grep source to diagnose a flat `X-Session-Cache: fresh`.
|
|
258
|
+
*
|
|
259
|
+
* No-op when `rawKey` is null / undefined / empty (the caller did
|
|
260
|
+
* not ask for tier-2 at all) or when scoping would succeed (the
|
|
261
|
+
* gate already accepted the key). Called by the endpoint layer
|
|
262
|
+
* right after it has decided `effectivePromptCacheKey`.
|
|
263
|
+
*/
|
|
264
|
+
export function maybeWarnPromptCacheKeyIneligible(rawKey) {
|
|
265
|
+
if (rawKey == null || rawKey.length === 0)
|
|
266
|
+
return;
|
|
267
|
+
// Happy-path branch: scoping would succeed, no nudge needed.
|
|
268
|
+
if (isPromptCacheKeyEnabled() && rawKey.length >= PROMPT_CACHE_KEY_MIN_LENGTH)
|
|
269
|
+
return;
|
|
270
|
+
const digest = digestSilentMissKey(rawKey);
|
|
271
|
+
if (loggedSilentMissKeys.has(digest))
|
|
272
|
+
return;
|
|
273
|
+
// FIFO eviction via Map insertion order — bounds memory under
|
|
274
|
+
// adversarial key flooding while preserving once-per-key semantics
|
|
275
|
+
// within the recent-key window.
|
|
276
|
+
if (loggedSilentMissKeys.size >= LOGGED_SILENT_MISS_KEYS_MAX) {
|
|
277
|
+
const oldest = loggedSilentMissKeys.keys().next().value;
|
|
278
|
+
if (oldest !== undefined)
|
|
279
|
+
loggedSilentMissKeys.delete(oldest);
|
|
280
|
+
}
|
|
281
|
+
loggedSilentMissKeys.set(digest, true);
|
|
282
|
+
if (!isPromptCacheKeyEnabled()) {
|
|
283
|
+
console.warn(`[mlx-node] prompt_cache_key supplied but tier-2 reuse is disabled ` +
|
|
284
|
+
`(MLX_DISABLE_PROMPT_CACHE_KEY=1). The key will be ignored and this ` +
|
|
285
|
+
`turn will cold-start. This message is logged once per distinct key.`);
|
|
286
|
+
return;
|
|
287
|
+
}
|
|
288
|
+
console.warn(`[mlx-node] prompt_cache_key is shorter than ${PROMPT_CACHE_KEY_MIN_LENGTH} chars; tier-2 reuse requires ` +
|
|
289
|
+
`at least ${PROMPT_CACHE_KEY_MIN_LENGTH} characters. The key will be ignored and this turn will ` +
|
|
290
|
+
`cold-start. This message is logged once per distinct key.`);
|
|
291
|
+
}
|
|
83
292
|
/**
|
|
84
293
|
* Thrown synchronously by {@link SessionRegistry.withExclusive} when
|
|
85
294
|
* the per-model queue cap (`maxQueueDepth`) is exceeded. The error is
|
|
@@ -104,6 +313,14 @@ export class SessionRegistry {
|
|
|
104
313
|
model;
|
|
105
314
|
ttlSec;
|
|
106
315
|
maxQueueDepth;
|
|
316
|
+
/**
|
|
317
|
+
* Per-model sampling defaults forwarded into every new `ChatSession`
|
|
318
|
+
* via its `defaultConfig` constructor option. `undefined` preserves
|
|
319
|
+
* the pre-defaults behaviour (empty `defaultConfig`). See
|
|
320
|
+
* {@link SessionRegistryOptions.samplingDefaults}.
|
|
321
|
+
*/
|
|
322
|
+
samplingDefaults;
|
|
323
|
+
maxOutputTokens;
|
|
107
324
|
/**
|
|
108
325
|
* Number of callers that are currently WAITING for the per-model
|
|
109
326
|
* execution mutex — i.e. have entered `withExclusive` but have not
|
|
@@ -162,6 +379,24 @@ export class SessionRegistry {
|
|
|
162
379
|
this.model = opts.model;
|
|
163
380
|
this.ttlSec = opts.ttlSec ?? 1800;
|
|
164
381
|
this.maxQueueDepth = opts.maxQueueDepth;
|
|
382
|
+
this.samplingDefaults = opts.samplingDefaults;
|
|
383
|
+
this.maxOutputTokens = opts.maxOutputTokens;
|
|
384
|
+
}
|
|
385
|
+
/**
|
|
386
|
+
* Construct a fresh `ChatSession` bound to this registry's model and
|
|
387
|
+
* pre-seeded with the operator-configured `samplingDefaults` (if any).
|
|
388
|
+
* Centralized so every cache-miss branch of `getOrCreate` produces a
|
|
389
|
+
* session whose per-call overlay will merge on top of the same
|
|
390
|
+
* defaults — clients cannot accidentally stray from the server's
|
|
391
|
+
* pinned sampling knobs by picking a cold-replay path.
|
|
392
|
+
*/
|
|
393
|
+
newSession() {
|
|
394
|
+
if (this.samplingDefaults === undefined) {
|
|
395
|
+
return new ChatSession(this.model);
|
|
396
|
+
}
|
|
397
|
+
return new ChatSession(this.model, {
|
|
398
|
+
defaultConfig: this.samplingDefaults,
|
|
399
|
+
});
|
|
165
400
|
}
|
|
166
401
|
/**
|
|
167
402
|
* Number of requests currently WAITING to acquire the per-model
|
|
@@ -171,6 +406,31 @@ export class SessionRegistry {
|
|
|
171
406
|
get queueDepth() {
|
|
172
407
|
return this.queuedCount;
|
|
173
408
|
}
|
|
409
|
+
/**
|
|
410
|
+
* Current sampling defaults applied to every new `ChatSession` this
|
|
411
|
+
* registry allocates. Exposed primarily for tests and diagnostics.
|
|
412
|
+
*/
|
|
413
|
+
get defaultSamplingConfig() {
|
|
414
|
+
return this.samplingDefaults;
|
|
415
|
+
}
|
|
416
|
+
get outputTokenLimit() {
|
|
417
|
+
return this.maxOutputTokens;
|
|
418
|
+
}
|
|
419
|
+
/**
|
|
420
|
+
* Replace the sampling defaults forwarded into every future
|
|
421
|
+
* `ChatSession` this registry allocates. Called by `ModelRegistry`
|
|
422
|
+
* on a `register(name, model, { samplingDefaults })` refresh so a
|
|
423
|
+
* fresh registration's defaults immediately apply to the next
|
|
424
|
+
* cache-miss cold-start. Sessions already cached at call time keep
|
|
425
|
+
* the defaults they were constructed with — they settle naturally
|
|
426
|
+
* through the single-warm cache rotation.
|
|
427
|
+
*/
|
|
428
|
+
setSamplingDefaults(defaults) {
|
|
429
|
+
this.samplingDefaults = defaults;
|
|
430
|
+
}
|
|
431
|
+
setMaxOutputTokens(limit) {
|
|
432
|
+
this.maxOutputTokens = limit;
|
|
433
|
+
}
|
|
174
434
|
/** Number of sessions currently cached. Primarily for tests and diagnostics. Always 0 or 1. */
|
|
175
435
|
get size() {
|
|
176
436
|
return this.entries.size;
|
|
@@ -180,60 +440,256 @@ export class SessionRegistry {
|
|
|
180
440
|
* Always returns a `SessionLookupResult` and always leaves the cache
|
|
181
441
|
* empty after return (single-warm invariant).
|
|
182
442
|
*
|
|
183
|
-
*
|
|
184
|
-
* mismatch: clear and return `{ session: new ChatSession(model), hit: false }`.
|
|
185
|
-
* The caller primes / cold-replays from the `ResponseStore` and
|
|
186
|
-
* re-adopts after the turn commits.
|
|
443
|
+
* Lookup proceeds in two tiers:
|
|
187
444
|
*
|
|
188
|
-
*
|
|
189
|
-
*
|
|
190
|
-
*
|
|
191
|
-
* the
|
|
445
|
+
* 1. **Tier 1 — `previousResponseId`.** The existing hot path:
|
|
446
|
+
* exact id match on a live, non-expired entry whose stored
|
|
447
|
+
* `instructions` are byte-equal to `requestedInstructions`. On
|
|
448
|
+
* a match the entry is leased out (single-use: removed from the
|
|
449
|
+
* map so a concurrent second request cannot share the live
|
|
450
|
+
* `ChatSession`). On a miss — unknown id, expired, or
|
|
451
|
+
* instructions drift — the method falls through to a FRESH
|
|
452
|
+
* session regardless of whether tier 2 would have hit.
|
|
192
453
|
*
|
|
193
|
-
*
|
|
194
|
-
*
|
|
195
|
-
*
|
|
196
|
-
* the
|
|
454
|
+
* `previousResponseId` wins unconditionally when supplied. The
|
|
455
|
+
* two keys could legitimately identify different conversation
|
|
456
|
+
* branches (e.g. a client fork where one arm chose the prev-id
|
|
457
|
+
* path and the other arm chose to set `prompt_cache_key`
|
|
458
|
+
* without one), so routing the prev-id branch through tier 2
|
|
459
|
+
* on miss risks splicing the wrong warm state into the wrong
|
|
460
|
+
* chain. Cold-replay is the safe default.
|
|
461
|
+
*
|
|
462
|
+
* 2. **Tier 2 — `promptCacheKey`.** Only runs when
|
|
463
|
+
* `previousResponseId` is `null`. Stateless agent clients
|
|
464
|
+
* (pi-mono, Aider, Codex CLI, Continue, etc.) never use
|
|
465
|
+
* `previous_response_id` — they own the conversation history
|
|
466
|
+
* client-side and resend the full transcript on every turn —
|
|
467
|
+
* so the only way to reuse a warm session across those turns
|
|
468
|
+
* is to key on the client-supplied `prompt_cache_key`. Scans
|
|
469
|
+
* for any live, non-expired entry whose stored
|
|
470
|
+
* `promptCacheKey` is non-null AND byte-equal to the caller's
|
|
471
|
+
* `promptCacheKey` AND whose stored `instructions` are byte-
|
|
472
|
+
* equal. Empty string is treated as a distinct key from
|
|
473
|
+
* `null` — an opt-out sentinel from a client that forgot to
|
|
474
|
+
* thread the key must NOT collide with another client that
|
|
475
|
+
* did set it to empty. On a match the entry is leased out
|
|
476
|
+
* (same single-use semantics as tier 1). On a miss, fall
|
|
477
|
+
* through to a fresh session.
|
|
197
478
|
*
|
|
198
479
|
* The `hit` flag drives the `X-Session-Cache` observability header
|
|
199
480
|
* emitted by both `/v1/responses` and `/v1/messages`: when the caller
|
|
200
481
|
* supplied a `previous_response_id`, `hit === true` yields `hit` and
|
|
201
482
|
* `hit === false` yields `cold_replay` (the endpoint then rebuilds
|
|
202
483
|
* from the `ResponseStore` on a fresh session). Requests with no
|
|
203
|
-
* `previous_response_id`
|
|
204
|
-
*
|
|
484
|
+
* `previous_response_id` yield either `fresh` (tier-2 miss) or
|
|
485
|
+
* `prefix_hit` (tier-2 hit — only classified as such once the
|
|
486
|
+
* native `cachedTokens > 0` confirms the prefix-cache machinery
|
|
487
|
+
* actually reused the cached tokens).
|
|
205
488
|
*/
|
|
206
|
-
getOrCreate(previousResponseId, requestedInstructions) {
|
|
489
|
+
getOrCreate(previousResponseId, requestedInstructions, promptCacheKey = null) {
|
|
490
|
+
// Tier 1: previousResponseId exact match.
|
|
491
|
+
//
|
|
207
492
|
// Every call is about to overwrite native KV state, so drop any
|
|
208
493
|
// other cached entry now — a later `getOrCreate` must not hand
|
|
209
494
|
// out a wrapper whose assumed state has been stomped. Under the
|
|
210
495
|
// single-warm invariant the map holds at most one entry, so the
|
|
211
496
|
// common case is either "the entry we want" or "nothing".
|
|
212
|
-
if (previousResponseId
|
|
213
|
-
this.entries.
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
|
|
497
|
+
if (previousResponseId !== null) {
|
|
498
|
+
const entry = this.entries.get(previousResponseId);
|
|
499
|
+
if (entry === undefined) {
|
|
500
|
+
this.entries.clear();
|
|
501
|
+
return { session: this.newSession(), hit: false };
|
|
502
|
+
}
|
|
503
|
+
if (entry.expiresAt < nowSec()) {
|
|
504
|
+
this.entries.clear();
|
|
505
|
+
return { session: this.newSession(), hit: false };
|
|
506
|
+
}
|
|
507
|
+
// Prefix-state mismatch forces cold replay so the new
|
|
508
|
+
// instructions are re-primed; without this guard, output would
|
|
509
|
+
// silently depend on cache state instead of request contents.
|
|
510
|
+
if (entry.instructions !== requestedInstructions) {
|
|
511
|
+
this.entries.clear();
|
|
512
|
+
return { session: this.newSession(), hit: false };
|
|
513
|
+
}
|
|
514
|
+
// Tier-1 hit: clear and hand the session out as a single-use
|
|
515
|
+
// lease so a concurrent second request against the same id
|
|
516
|
+
// cold-replays instead of sharing this live ChatSession. Note
|
|
517
|
+
// that even on a prev-id tier-1 MISS we do NOT fall through to
|
|
518
|
+
// tier 2 — see the docstring above for the precedence
|
|
519
|
+
// rationale.
|
|
218
520
|
this.entries.clear();
|
|
219
|
-
return { session:
|
|
521
|
+
return { session: entry.session, hit: true };
|
|
220
522
|
}
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
523
|
+
// Tier 2: promptCacheKey scan (only reached when previousResponseId is null).
|
|
524
|
+
//
|
|
525
|
+
// The registry holds at most one entry under the single-warm
|
|
526
|
+
// invariant, so the "scan" is actually a single lookup — walk the
|
|
527
|
+
// map, check the one entry if present, hit or miss. A non-null
|
|
528
|
+
// scoped key on both the request and the entry plus byte-equal
|
|
529
|
+
// instructions is the match condition.
|
|
530
|
+
//
|
|
531
|
+
// SECURITY: raw caller-supplied keys never touch the map. They
|
|
532
|
+
// are run through {@link scopePromptCacheKey}, which (a) returns
|
|
533
|
+
// `null` when the tier-2 opt-in env var is unset, (b) enforces a
|
|
534
|
+
// minimum length, and (c) HMACs the key with a boot-time nonce
|
|
535
|
+
// held only in this process's memory. Without the opt-in every
|
|
536
|
+
// tier-2 lookup below immediately misses; with the opt-in,
|
|
537
|
+
// attackers who cannot read the process-local nonce cannot craft
|
|
538
|
+
// a lookup that matches a stored entry by guessing the raw key.
|
|
539
|
+
const scopedKey = scopePromptCacheKey(promptCacheKey);
|
|
540
|
+
if (scopedKey !== null) {
|
|
541
|
+
for (const entry of this.entries.values()) {
|
|
542
|
+
if (entry.expiresAt < nowSec())
|
|
543
|
+
continue;
|
|
544
|
+
if (entry.promptCacheKey === null)
|
|
545
|
+
continue;
|
|
546
|
+
if (entry.promptCacheKey !== scopedKey)
|
|
547
|
+
continue;
|
|
548
|
+
if (entry.instructions !== requestedInstructions)
|
|
549
|
+
continue;
|
|
550
|
+
// Tier-2 hit: clear and lease (same single-warm / single-use
|
|
551
|
+
// semantics as tier 1).
|
|
552
|
+
this.entries.clear();
|
|
553
|
+
return { session: entry.session, hit: true };
|
|
554
|
+
}
|
|
224
555
|
}
|
|
225
|
-
//
|
|
226
|
-
//
|
|
227
|
-
//
|
|
228
|
-
|
|
556
|
+
// Fall through: fresh session. Clear any leftover entry so a
|
|
557
|
+
// later lookup cannot hand out a wrapper whose assumed state has
|
|
558
|
+
// been overwritten by this dispatch.
|
|
559
|
+
this.entries.clear();
|
|
560
|
+
return { session: this.newSession(), hit: false };
|
|
561
|
+
}
|
|
562
|
+
/**
|
|
563
|
+
* Allocate a fresh `ChatSession` bound to this registry's model
|
|
564
|
+
* without touching the warm slot. Intended for the `/v1/messages`
|
|
565
|
+
* endpoint when the underlying model has a block-paged KV cache
|
|
566
|
+
* active: the native cache already reuses SYS blocks across requests
|
|
567
|
+
* via content-addressing in `BlockAllocator`'s prefix-hash table, so
|
|
568
|
+
* the JS-side warm slot in
|
|
569
|
+
* {@link SessionRegistry.getOrCreateWarmAny} is redundant.
|
|
570
|
+
*
|
|
571
|
+
* Crucially, this call is purely additive — it does **NOT** clear,
|
|
572
|
+
* read, or evict the warm slot. Two parallel `/v1/messages` requests
|
|
573
|
+
* sharing a system prompt both call `createFreshSession` and both
|
|
574
|
+
* get distinct sessions; the native cache transparently refcounts
|
|
575
|
+
* the shared SYS blocks across them. This is the routing decision
|
|
576
|
+
* the long block comment in `packages/server/src/endpoints/messages.ts`
|
|
577
|
+
* documents: paged → fresh session, non-paged → warm-any lookup.
|
|
578
|
+
*
|
|
579
|
+
* The returned session is pre-seeded with the operator-configured
|
|
580
|
+
* `samplingDefaults` (matching every other cache-miss branch) so a
|
|
581
|
+
* client that picks the paged path does not silently stray from the
|
|
582
|
+
* server's pinned sampling knobs.
|
|
583
|
+
*
|
|
584
|
+
* Returned with `hit: false` to keep the result shape uniform with
|
|
585
|
+
* {@link SessionRegistry.getOrCreate} and
|
|
586
|
+
* {@link SessionRegistry.getOrCreateWarmAny}; callers that care
|
|
587
|
+
* about the cache header semantics should observe
|
|
588
|
+
* `result.cachedTokens` from the dispatch instead — that's the
|
|
589
|
+
* authoritative signal for whether the native engine recovered any
|
|
590
|
+
* prefix on this turn (paged or otherwise).
|
|
591
|
+
*/
|
|
592
|
+
createFreshSession() {
|
|
593
|
+
return { session: this.newSession(), hit: false };
|
|
594
|
+
}
|
|
595
|
+
/**
|
|
596
|
+
* @deprecated **Redundant on `/v1/messages` for paged-active models.**
|
|
597
|
+
* There is no call site for paged-active full-attention models (Qwen3 +
|
|
598
|
+
* LFM2 + Gemma4 today): the native block-paged KV adapter (`PagedKVCacheAdapter` +
|
|
599
|
+
* `BlockAllocator` + `LayerKVPool`) recovers a turn's prefix from
|
|
600
|
+
* refcounted KV blocks keyed by token-prefix hash, so the JS-side
|
|
601
|
+
* single-warm slot this method walks is redundant — the native
|
|
602
|
+
* cache picks up the same cross-turn reuse without the
|
|
603
|
+
* byte-equal-`instructions` gate, and additionally supports
|
|
604
|
+
* cross-conversation prefix sharing the warm slot cannot.
|
|
605
|
+
*
|
|
606
|
+
* The `/v1/messages` endpoint now branches at request time on
|
|
607
|
+
* {@link SessionCapableModel.hasBlockPagedCache}: paged-active models
|
|
608
|
+
* call {@link SessionRegistry.createFreshSession} per request and
|
|
609
|
+
* never touch the warm slot; non-paged models (Qwen3.5 dense + MoE —
|
|
610
|
+
* default-OFF pending a perf decision; the `QianfanOCRModel` VLM —
|
|
611
|
+
* no adapter wired) still call this method because the JS-side warm
|
|
612
|
+
* slot is the ONLY cross-conversation reuse mechanism available to
|
|
613
|
+
* them. Removing this method would silently disable cross-turn reuse
|
|
614
|
+
* on every non-paged model, so it stays load-bearing until ALL
|
|
615
|
+
* session-capable models have paged enabled by default. Treat
|
|
616
|
+
* `@deprecated` as an intent signal that paged-active callers should
|
|
617
|
+
* use `createFreshSession` instead.
|
|
618
|
+
*
|
|
619
|
+
* Third lookup mode — for STATELESS full-history endpoints that have
|
|
620
|
+
* no `previous_response_id` to thread and do not propagate
|
|
621
|
+
* `prompt_cache_key` back to the server. The Anthropic
|
|
622
|
+
* `/v1/messages` endpoint is the canonical caller: clients (e.g.
|
|
623
|
+
* Claude Code) POST the entire conversation each turn, so the only
|
|
624
|
+
* remaining signal that a turn N continues turn N-1's prefix is the
|
|
625
|
+
* registry's own warm slot.
|
|
626
|
+
*
|
|
627
|
+
* Behaviour: walk the registry's at-most-one warm entry. If it is
|
|
628
|
+
* non-expired AND its stored `instructions` are byte-equal to
|
|
629
|
+
* `requestedInstructions`, lease it out (single-use — `entries.clear()`
|
|
630
|
+
* before return, mirroring the tier-1 / tier-2 lease-on-hit
|
|
631
|
+
* semantics). Otherwise clear the map and return a fresh session.
|
|
632
|
+
*
|
|
633
|
+
* Crucially, this lookup IGNORES `entry.promptCacheKey` and ignores
|
|
634
|
+
* the entry's prior `previousResponseId` keying — any warm slot is
|
|
635
|
+
* fair game for `/v1/messages` reuse. The byte-equal `instructions`
|
|
636
|
+
* compare is the SOLE correctness gate: a system prompt change
|
|
637
|
+
* forces cold replay so the new prefix state is re-primed instead
|
|
638
|
+
* of silently reusing a stale warmed prompt.
|
|
639
|
+
*
|
|
640
|
+
* **Adoption sentinel.** `/v1/messages` adopts back under the literal
|
|
641
|
+
* sentinel id `'__msg_warm__'`. That sentinel will never appear as a
|
|
642
|
+
* `previous_response_id` on a `/v1/responses` request — the
|
|
643
|
+
* Anthropic Messages API does not produce a `previous_response_id`
|
|
644
|
+
* value clients could echo back, and the OpenAI side mints fresh
|
|
645
|
+
* `resp_*` ids — so cross-endpoint capture via tier-1 is impossible
|
|
646
|
+
* by construction. The two endpoints still SHARE the single warm
|
|
647
|
+
* slot under the registry's single-warm invariant: a
|
|
648
|
+
* `/v1/messages` turn that follows a `/v1/responses` turn can evict
|
|
649
|
+
* (and vice versa). That is the explicit trade-off of holding at
|
|
650
|
+
* most one warm entry per model.
|
|
651
|
+
*
|
|
652
|
+
* **Trust model.** Multi-tenant isolation on this endpoint requires
|
|
653
|
+
* fronting the server with an auth proxy that scopes warm-slot
|
|
654
|
+
* visibility per tenant — same trust boundary documented at the top
|
|
655
|
+
* of this file for the tier-2 `prompt_cache_key` path. The single-
|
|
656
|
+
* warm invariant plus `withExclusive`'s per-model serialization make
|
|
657
|
+
* the lookup safe under SINGLE-tenant assumptions: no two requests
|
|
658
|
+
* race the slot, and there is at most one slot to lease.
|
|
659
|
+
*
|
|
660
|
+
* **Caller contract on miss.** If `instructions` drifts between
|
|
661
|
+
* turns (system prompt changed) this returns `hit: false` and a
|
|
662
|
+
* fresh session — and the caller MUST then run a full
|
|
663
|
+
* `session.reset()` before priming history, NOT the JS-only
|
|
664
|
+
* `resetPreservingNativeCacheForWarmReuse` path. A fresh JS session
|
|
665
|
+
* does NOT imply a fresh native cache (the underlying
|
|
666
|
+
* `SessionCapableModel` is shared and its native
|
|
667
|
+
* `cached_token_history` persists across requests), so skipping the
|
|
668
|
+
* native wipe on a miss would let the next `chatSessionStart` reuse
|
|
669
|
+
* an unrelated previous request's prefix — the cross-request
|
|
670
|
+
* cache-affinity side channel that the long block comment in
|
|
671
|
+
* `responses.ts` (around the `runSessionNonStreaming` /
|
|
672
|
+
* `runSessionStreaming` branches) describes.
|
|
673
|
+
*/
|
|
674
|
+
getOrCreateWarmAny(requestedInstructions) {
|
|
675
|
+
// Single-warm invariant: at most one entry. Walk it once, lease
|
|
676
|
+
// on a fresh + instructions-matched hit, otherwise clear and
|
|
677
|
+
// cold-start. The ignored fields (promptCacheKey,
|
|
678
|
+
// previousResponseId-keying) are deliberate — see the docstring.
|
|
679
|
+
for (const entry of this.entries.values()) {
|
|
680
|
+
if (entry.expiresAt < nowSec())
|
|
681
|
+
continue;
|
|
682
|
+
if (entry.instructions !== requestedInstructions)
|
|
683
|
+
continue;
|
|
684
|
+
// Hit: clear and lease (single-use semantics, same as tiers 1/2).
|
|
229
685
|
this.entries.clear();
|
|
230
|
-
return { session:
|
|
686
|
+
return { session: entry.session, hit: true };
|
|
231
687
|
}
|
|
232
|
-
//
|
|
233
|
-
// a
|
|
234
|
-
//
|
|
688
|
+
// Miss (no entry, expired, or instructions drift). Clear the map
|
|
689
|
+
// so a stale wrapper cannot leak into a later lookup, and return
|
|
690
|
+
// a fresh session.
|
|
235
691
|
this.entries.clear();
|
|
236
|
-
return { session:
|
|
692
|
+
return { session: this.newSession(), hit: false };
|
|
237
693
|
}
|
|
238
694
|
/**
|
|
239
695
|
* Insert a session under a newly allocated response id. Clears the
|
|
@@ -243,12 +699,28 @@ export class SessionRegistry {
|
|
|
243
699
|
* `instructions` is the prefix/system state used for this turn;
|
|
244
700
|
* stored on the entry and compared on the next `getOrCreate` to
|
|
245
701
|
* detect prefix changes that must force a cold replay.
|
|
702
|
+
*
|
|
703
|
+
* `promptCacheKey` is the client-supplied conversation-chain key
|
|
704
|
+
* that enables the registry's tier-2 lookup for stateless agent
|
|
705
|
+
* turns that do not carry a `previous_response_id`. `null` /
|
|
706
|
+
* `undefined` means "no key supplied" — stored verbatim so a
|
|
707
|
+
* subsequent stateless lookup that also omits the key does NOT
|
|
708
|
+
* accidentally pick up this entry (only explicit non-null
|
|
709
|
+
* key-equality on both sides can hit tier 2). See
|
|
710
|
+
* {@link SessionRegistry.getOrCreate} for the precedence rules.
|
|
246
711
|
*/
|
|
247
|
-
adopt(responseId, session, instructions) {
|
|
712
|
+
adopt(responseId, session, instructions, promptCacheKey = null) {
|
|
713
|
+
// Scope the caller-supplied key BEFORE storing so a later
|
|
714
|
+
// `getOrCreate` can only resolve entries via the same opt-in +
|
|
715
|
+
// HMAC path. When tier-2 reuse is disabled (or the key is too
|
|
716
|
+
// short / absent) `scopePromptCacheKey` returns `null`, which
|
|
717
|
+
// disables this entry from ever matching a tier-2 lookup — the
|
|
718
|
+
// raw caller-supplied key is NEVER stored.
|
|
248
719
|
this.entries.clear();
|
|
249
720
|
this.entries.set(responseId, {
|
|
250
721
|
session,
|
|
251
722
|
instructions,
|
|
723
|
+
promptCacheKey: scopePromptCacheKey(promptCacheKey ?? null),
|
|
252
724
|
expiresAt: nowSec() + this.ttlSec,
|
|
253
725
|
});
|
|
254
726
|
}
|