@mlx-node/server 0.0.7 → 0.0.9

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.
Files changed (110) hide show
  1. package/dist/auth.d.ts +56 -0
  2. package/dist/auth.d.ts.map +1 -0
  3. package/dist/auth.js +106 -0
  4. package/dist/chat-session-warm-reuse.d.ts +51 -0
  5. package/dist/chat-session-warm-reuse.d.ts.map +1 -0
  6. package/dist/chat-session-warm-reuse.js +72 -0
  7. package/dist/endpoints/messages-count-tokens.d.ts +8 -0
  8. package/dist/endpoints/messages-count-tokens.d.ts.map +1 -0
  9. package/dist/endpoints/messages-count-tokens.js +121 -0
  10. package/dist/endpoints/messages.d.ts +57 -5
  11. package/dist/endpoints/messages.d.ts.map +1 -1
  12. package/dist/endpoints/messages.js +1043 -147
  13. package/dist/endpoints/models.d.ts +2 -1
  14. package/dist/endpoints/models.d.ts.map +1 -1
  15. package/dist/endpoints/models.js +2 -2
  16. package/dist/endpoints/responses.d.ts +22 -7
  17. package/dist/endpoints/responses.d.ts.map +1 -1
  18. package/dist/endpoints/responses.js +608 -85
  19. package/dist/errors.d.ts +1 -0
  20. package/dist/errors.d.ts.map +1 -1
  21. package/dist/errors.js +3 -0
  22. package/dist/handler.d.ts +69 -1
  23. package/dist/handler.d.ts.map +1 -1
  24. package/dist/handler.js +70 -16
  25. package/dist/health.d.ts +146 -0
  26. package/dist/health.d.ts.map +1 -0
  27. package/dist/health.js +107 -0
  28. package/dist/host/discover.d.ts +19 -0
  29. package/dist/host/discover.d.ts.map +1 -0
  30. package/dist/host/discover.js +50 -0
  31. package/dist/host/env-policy.d.ts +62 -0
  32. package/dist/host/env-policy.d.ts.map +1 -0
  33. package/dist/host/env-policy.js +69 -0
  34. package/dist/host/index.d.ts +202 -0
  35. package/dist/host/index.d.ts.map +1 -0
  36. package/dist/host/index.js +325 -0
  37. package/dist/host/logger.d.ts +36 -0
  38. package/dist/host/logger.d.ts.map +1 -0
  39. package/dist/host/logger.js +376 -0
  40. package/dist/host/net.d.ts +65 -0
  41. package/dist/host/net.d.ts.map +1 -0
  42. package/dist/host/net.js +97 -0
  43. package/dist/host/paths.d.ts +28 -0
  44. package/dist/host/paths.d.ts.map +1 -0
  45. package/dist/host/paths.js +71 -0
  46. package/dist/host/swap.d.ts +27 -0
  47. package/dist/host/swap.d.ts.map +1 -0
  48. package/dist/host/swap.js +178 -0
  49. package/dist/host/temp-root.d.ts +57 -0
  50. package/dist/host/temp-root.d.ts.map +1 -0
  51. package/dist/host/temp-root.js +99 -0
  52. package/dist/idle-sweeper.d.ts +245 -0
  53. package/dist/idle-sweeper.d.ts.map +1 -0
  54. package/dist/idle-sweeper.js +408 -0
  55. package/dist/index.d.ts +19 -4
  56. package/dist/index.d.ts.map +1 -1
  57. package/dist/index.js +17 -1
  58. package/dist/load-model.d.ts +69 -0
  59. package/dist/load-model.d.ts.map +1 -0
  60. package/dist/load-model.js +63 -0
  61. package/dist/mappers/anthropic-request.d.ts +24 -2
  62. package/dist/mappers/anthropic-request.d.ts.map +1 -1
  63. package/dist/mappers/anthropic-request.js +222 -24
  64. package/dist/mappers/anthropic-response.d.ts +29 -4
  65. package/dist/mappers/anthropic-response.d.ts.map +1 -1
  66. package/dist/mappers/anthropic-response.js +143 -21
  67. package/dist/mappers/request.d.ts +48 -0
  68. package/dist/mappers/request.d.ts.map +1 -1
  69. package/dist/mappers/request.js +211 -35
  70. package/dist/mappers/response.d.ts.map +1 -1
  71. package/dist/mappers/response.js +13 -1
  72. package/dist/model-work-coordinator.d.ts +95 -0
  73. package/dist/model-work-coordinator.d.ts.map +1 -0
  74. package/dist/model-work-coordinator.js +201 -0
  75. package/dist/pending-writes.d.ts.map +1 -1
  76. package/dist/presets.d.ts +82 -0
  77. package/dist/presets.d.ts.map +1 -0
  78. package/dist/presets.js +98 -0
  79. package/dist/registry.d.ts +31 -1
  80. package/dist/registry.d.ts.map +1 -1
  81. package/dist/registry.js +33 -5
  82. package/dist/router.d.ts +37 -1
  83. package/dist/router.d.ts.map +1 -1
  84. package/dist/router.js +79 -7
  85. package/dist/server.d.ts +193 -3
  86. package/dist/server.d.ts.map +1 -1
  87. package/dist/server.js +173 -10
  88. package/dist/session-registry.d.ts +279 -18
  89. package/dist/session-registry.d.ts.map +1 -1
  90. package/dist/session-registry.js +518 -37
  91. package/dist/stop-sequence-buffer.d.ts +58 -0
  92. package/dist/stop-sequence-buffer.d.ts.map +1 -0
  93. package/dist/stop-sequence-buffer.js +148 -0
  94. package/dist/streaming.d.ts +14 -0
  95. package/dist/streaming.d.ts.map +1 -1
  96. package/dist/streaming.js +45 -0
  97. package/dist/text-recovery.d.ts +35 -0
  98. package/dist/text-recovery.d.ts.map +1 -0
  99. package/dist/text-recovery.js +41 -0
  100. package/dist/timing.d.ts +80 -0
  101. package/dist/timing.d.ts.map +1 -0
  102. package/dist/timing.js +121 -0
  103. package/dist/tool-call-buffer.d.ts +5 -5
  104. package/dist/tool-call-buffer.d.ts.map +1 -1
  105. package/dist/tool-call-buffer.js +28 -8
  106. package/dist/types-anthropic.d.ts +161 -1
  107. package/dist/types-anthropic.d.ts.map +1 -1
  108. package/dist/types.d.ts +172 -2
  109. package/dist/types.d.ts.map +1 -1
  110. package/package.json +15 -3
@@ -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 `RESPONSE_TTL_SECONDS`
54
- * in `packages/server/src/endpoints/responses.ts` so the cached
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,40 @@ export class SessionRegistry {
171
406
  get queueDepth() {
172
407
  return this.queuedCount;
173
408
  }
409
+ /**
410
+ * Configured waiter cap for this model's execution mutex, or `undefined`
411
+ * when unbounded. Paired with {@link queueDepth} so a readiness probe can
412
+ * tell "3 waiters, unbounded" (fine) from "3 waiters, cap of 3" (the next
413
+ * request gets a 429) without reaching into private state.
414
+ */
415
+ get queueDepthLimit() {
416
+ return this.maxQueueDepth;
417
+ }
418
+ /**
419
+ * Current sampling defaults applied to every new `ChatSession` this
420
+ * registry allocates. Exposed primarily for tests and diagnostics.
421
+ */
422
+ get defaultSamplingConfig() {
423
+ return this.samplingDefaults;
424
+ }
425
+ get outputTokenLimit() {
426
+ return this.maxOutputTokens;
427
+ }
428
+ /**
429
+ * Replace the sampling defaults forwarded into every future
430
+ * `ChatSession` this registry allocates. Called by `ModelRegistry`
431
+ * on a `register(name, model, { samplingDefaults })` refresh so a
432
+ * fresh registration's defaults immediately apply to the next
433
+ * cache-miss cold-start. Sessions already cached at call time keep
434
+ * the defaults they were constructed with — they settle naturally
435
+ * through the single-warm cache rotation.
436
+ */
437
+ setSamplingDefaults(defaults) {
438
+ this.samplingDefaults = defaults;
439
+ }
440
+ setMaxOutputTokens(limit) {
441
+ this.maxOutputTokens = limit;
442
+ }
174
443
  /** Number of sessions currently cached. Primarily for tests and diagnostics. Always 0 or 1. */
175
444
  get size() {
176
445
  return this.entries.size;
@@ -180,60 +449,256 @@ export class SessionRegistry {
180
449
  * Always returns a `SessionLookupResult` and always leaves the cache
181
450
  * empty after return (single-warm invariant).
182
451
  *
183
- * On a null id, missing key, expired entry, or prefix-state
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.
452
+ * Lookup proceeds in two tiers:
453
+ *
454
+ * 1. **Tier 1 `previousResponseId`.** The existing hot path:
455
+ * exact id match on a live, non-expired entry whose stored
456
+ * `instructions` are byte-equal to `requestedInstructions`. On
457
+ * a match the entry is leased out (single-use: removed from the
458
+ * map so a concurrent second request cannot share the live
459
+ * `ChatSession`). On a miss — unknown id, expired, or
460
+ * instructions drift — the method falls through to a FRESH
461
+ * session regardless of whether tier 2 would have hit.
187
462
  *
188
- * On a hit: the entry is removed and its live session is returned
189
- * alongside `hit: true`. Overlapping requests against the same
190
- * `previous_response_id` cannot share the same live `ChatSession`
191
- * the first wins, the second misses and cold-replays.
463
+ * `previousResponseId` wins unconditionally when supplied. The
464
+ * two keys could legitimately identify different conversation
465
+ * branches (e.g. a client fork where one arm chose the prev-id
466
+ * path and the other arm chose to set `prompt_cache_key`
467
+ * without one), so routing the prev-id branch through tier 2
468
+ * on miss risks splicing the wrong warm state into the wrong
469
+ * chain. Cold-replay is the safe default.
192
470
  *
193
- * `requestedInstructions` is the caller's prefix/system state
194
- * (OpenAI `instructions`, Anthropic `system`, or `null`); byte-for-
195
- * byte mismatch against the cached entry forces cold replay so
196
- * the new prefix is re-primed.
471
+ * 2. **Tier 2 `promptCacheKey`.** Only runs when
472
+ * `previousResponseId` is `null`. Stateless agent clients
473
+ * (pi-mono, Aider, Codex CLI, Continue, etc.) never use
474
+ * `previous_response_id` they own the conversation history
475
+ * client-side and resend the full transcript on every turn —
476
+ * so the only way to reuse a warm session across those turns
477
+ * is to key on the client-supplied `prompt_cache_key`. Scans
478
+ * for any live, non-expired entry whose stored
479
+ * `promptCacheKey` is non-null AND byte-equal to the caller's
480
+ * `promptCacheKey` AND whose stored `instructions` are byte-
481
+ * equal. Empty string is treated as a distinct key from
482
+ * `null` — an opt-out sentinel from a client that forgot to
483
+ * thread the key must NOT collide with another client that
484
+ * did set it to empty. On a match the entry is leased out
485
+ * (same single-use semantics as tier 1). On a miss, fall
486
+ * through to a fresh session.
197
487
  *
198
488
  * The `hit` flag drives the `X-Session-Cache` observability header
199
489
  * emitted by both `/v1/responses` and `/v1/messages`: when the caller
200
490
  * supplied a `previous_response_id`, `hit === true` yields `hit` and
201
491
  * `hit === false` yields `cold_replay` (the endpoint then rebuilds
202
492
  * from the `ResponseStore` on a fresh session). Requests with no
203
- * `previous_response_id` (or the stateless `/v1/messages` endpoint,
204
- * which always passes `null`) yield `fresh` regardless of this flag.
493
+ * `previous_response_id` yield either `fresh` (tier-2 miss) or
494
+ * `prefix_hit` (tier-2 hit only classified as such once the
495
+ * native `cachedTokens > 0` confirms the prefix-cache machinery
496
+ * actually reused the cached tokens).
205
497
  */
206
- getOrCreate(previousResponseId, requestedInstructions) {
498
+ getOrCreate(previousResponseId, requestedInstructions, promptCacheKey = null) {
499
+ // Tier 1: previousResponseId exact match.
500
+ //
207
501
  // Every call is about to overwrite native KV state, so drop any
208
502
  // other cached entry now — a later `getOrCreate` must not hand
209
503
  // out a wrapper whose assumed state has been stomped. Under the
210
504
  // single-warm invariant the map holds at most one entry, so the
211
505
  // common case is either "the entry we want" or "nothing".
212
- if (previousResponseId === null) {
213
- this.entries.clear();
214
- return { session: new ChatSession(this.model), hit: false };
215
- }
216
- const entry = this.entries.get(previousResponseId);
217
- if (entry === undefined) {
506
+ if (previousResponseId !== null) {
507
+ const entry = this.entries.get(previousResponseId);
508
+ if (entry === undefined) {
509
+ this.entries.clear();
510
+ return { session: this.newSession(), hit: false };
511
+ }
512
+ if (entry.expiresAt < nowSec()) {
513
+ this.entries.clear();
514
+ return { session: this.newSession(), hit: false };
515
+ }
516
+ // Prefix-state mismatch forces cold replay so the new
517
+ // instructions are re-primed; without this guard, output would
518
+ // silently depend on cache state instead of request contents.
519
+ if (entry.instructions !== requestedInstructions) {
520
+ this.entries.clear();
521
+ return { session: this.newSession(), hit: false };
522
+ }
523
+ // Tier-1 hit: clear and hand the session out as a single-use
524
+ // lease so a concurrent second request against the same id
525
+ // cold-replays instead of sharing this live ChatSession. Note
526
+ // that even on a prev-id tier-1 MISS we do NOT fall through to
527
+ // tier 2 — see the docstring above for the precedence
528
+ // rationale.
218
529
  this.entries.clear();
219
- return { session: new ChatSession(this.model), hit: false };
530
+ return { session: entry.session, hit: true };
220
531
  }
221
- if (entry.expiresAt < nowSec()) {
222
- this.entries.clear();
223
- return { session: new ChatSession(this.model), hit: false };
532
+ // Tier 2: promptCacheKey scan (only reached when previousResponseId is null).
533
+ //
534
+ // The registry holds at most one entry under the single-warm
535
+ // invariant, so the "scan" is actually a single lookup — walk the
536
+ // map, check the one entry if present, hit or miss. A non-null
537
+ // scoped key on both the request and the entry plus byte-equal
538
+ // instructions is the match condition.
539
+ //
540
+ // SECURITY: raw caller-supplied keys never touch the map. They
541
+ // are run through {@link scopePromptCacheKey}, which (a) returns
542
+ // `null` when the tier-2 opt-in env var is unset, (b) enforces a
543
+ // minimum length, and (c) HMACs the key with a boot-time nonce
544
+ // held only in this process's memory. Without the opt-in every
545
+ // tier-2 lookup below immediately misses; with the opt-in,
546
+ // attackers who cannot read the process-local nonce cannot craft
547
+ // a lookup that matches a stored entry by guessing the raw key.
548
+ const scopedKey = scopePromptCacheKey(promptCacheKey);
549
+ if (scopedKey !== null) {
550
+ for (const entry of this.entries.values()) {
551
+ if (entry.expiresAt < nowSec())
552
+ continue;
553
+ if (entry.promptCacheKey === null)
554
+ continue;
555
+ if (entry.promptCacheKey !== scopedKey)
556
+ continue;
557
+ if (entry.instructions !== requestedInstructions)
558
+ continue;
559
+ // Tier-2 hit: clear and lease (same single-warm / single-use
560
+ // semantics as tier 1).
561
+ this.entries.clear();
562
+ return { session: entry.session, hit: true };
563
+ }
224
564
  }
225
- // Prefix-state mismatch forces cold replay so the new
226
- // instructions are re-primed; without this guard, output would
227
- // silently depend on cache state instead of request contents.
228
- if (entry.instructions !== requestedInstructions) {
565
+ // Fall through: fresh session. Clear any leftover entry so a
566
+ // later lookup cannot hand out a wrapper whose assumed state has
567
+ // been overwritten by this dispatch.
568
+ this.entries.clear();
569
+ return { session: this.newSession(), hit: false };
570
+ }
571
+ /**
572
+ * Allocate a fresh `ChatSession` bound to this registry's model
573
+ * without touching the warm slot. Intended for the `/v1/messages`
574
+ * endpoint when the underlying model has a block-paged KV cache
575
+ * active: the native cache already reuses SYS blocks across requests
576
+ * via content-addressing in `BlockAllocator`'s prefix-hash table, so
577
+ * the JS-side warm slot in
578
+ * {@link SessionRegistry.getOrCreateWarmAny} is redundant.
579
+ *
580
+ * Crucially, this call is purely additive — it does **NOT** clear,
581
+ * read, or evict the warm slot. Two parallel `/v1/messages` requests
582
+ * sharing a system prompt both call `createFreshSession` and both
583
+ * get distinct sessions; the native cache transparently refcounts
584
+ * the shared SYS blocks across them. This is the routing decision
585
+ * the long block comment in `packages/server/src/endpoints/messages.ts`
586
+ * documents: paged → fresh session, non-paged → warm-any lookup.
587
+ *
588
+ * The returned session is pre-seeded with the operator-configured
589
+ * `samplingDefaults` (matching every other cache-miss branch) so a
590
+ * client that picks the paged path does not silently stray from the
591
+ * server's pinned sampling knobs.
592
+ *
593
+ * Returned with `hit: false` to keep the result shape uniform with
594
+ * {@link SessionRegistry.getOrCreate} and
595
+ * {@link SessionRegistry.getOrCreateWarmAny}; callers that care
596
+ * about the cache header semantics should observe
597
+ * `result.cachedTokens` from the dispatch instead — that's the
598
+ * authoritative signal for whether the native engine recovered any
599
+ * prefix on this turn (paged or otherwise).
600
+ */
601
+ createFreshSession() {
602
+ return { session: this.newSession(), hit: false };
603
+ }
604
+ /**
605
+ * @deprecated **Redundant on `/v1/messages` for paged-active models.**
606
+ * There is no call site for paged-active full-attention models (Qwen3 +
607
+ * LFM2 + Gemma4 today): the native block-paged KV adapter (`PagedKVCacheAdapter` +
608
+ * `BlockAllocator` + `LayerKVPool`) recovers a turn's prefix from
609
+ * refcounted KV blocks keyed by token-prefix hash, so the JS-side
610
+ * single-warm slot this method walks is redundant — the native
611
+ * cache picks up the same cross-turn reuse without the
612
+ * byte-equal-`instructions` gate, and additionally supports
613
+ * cross-conversation prefix sharing the warm slot cannot.
614
+ *
615
+ * The `/v1/messages` endpoint now branches at request time on
616
+ * {@link SessionCapableModel.hasBlockPagedCache}: paged-active models
617
+ * call {@link SessionRegistry.createFreshSession} per request and
618
+ * never touch the warm slot; non-paged models (Qwen3.5 dense + MoE —
619
+ * default-OFF pending a perf decision; the `QianfanOCRModel` VLM —
620
+ * no adapter wired) still call this method because the JS-side warm
621
+ * slot is the ONLY cross-conversation reuse mechanism available to
622
+ * them. Removing this method would silently disable cross-turn reuse
623
+ * on every non-paged model, so it stays load-bearing until ALL
624
+ * session-capable models have paged enabled by default. Treat
625
+ * `@deprecated` as an intent signal that paged-active callers should
626
+ * use `createFreshSession` instead.
627
+ *
628
+ * Third lookup mode — for STATELESS full-history endpoints that have
629
+ * no `previous_response_id` to thread and do not propagate
630
+ * `prompt_cache_key` back to the server. The Anthropic
631
+ * `/v1/messages` endpoint is the canonical caller: clients (e.g.
632
+ * Claude Code) POST the entire conversation each turn, so the only
633
+ * remaining signal that a turn N continues turn N-1's prefix is the
634
+ * registry's own warm slot.
635
+ *
636
+ * Behaviour: walk the registry's at-most-one warm entry. If it is
637
+ * non-expired AND its stored `instructions` are byte-equal to
638
+ * `requestedInstructions`, lease it out (single-use — `entries.clear()`
639
+ * before return, mirroring the tier-1 / tier-2 lease-on-hit
640
+ * semantics). Otherwise clear the map and return a fresh session.
641
+ *
642
+ * Crucially, this lookup IGNORES `entry.promptCacheKey` and ignores
643
+ * the entry's prior `previousResponseId` keying — any warm slot is
644
+ * fair game for `/v1/messages` reuse. The byte-equal `instructions`
645
+ * compare is the SOLE correctness gate: a system prompt change
646
+ * forces cold replay so the new prefix state is re-primed instead
647
+ * of silently reusing a stale warmed prompt.
648
+ *
649
+ * **Adoption sentinel.** `/v1/messages` adopts back under the literal
650
+ * sentinel id `'__msg_warm__'`. That sentinel will never appear as a
651
+ * `previous_response_id` on a `/v1/responses` request — the
652
+ * Anthropic Messages API does not produce a `previous_response_id`
653
+ * value clients could echo back, and the OpenAI side mints fresh
654
+ * `resp_*` ids — so cross-endpoint capture via tier-1 is impossible
655
+ * by construction. The two endpoints still SHARE the single warm
656
+ * slot under the registry's single-warm invariant: a
657
+ * `/v1/messages` turn that follows a `/v1/responses` turn can evict
658
+ * (and vice versa). That is the explicit trade-off of holding at
659
+ * most one warm entry per model.
660
+ *
661
+ * **Trust model.** Multi-tenant isolation on this endpoint requires
662
+ * fronting the server with an auth proxy that scopes warm-slot
663
+ * visibility per tenant — same trust boundary documented at the top
664
+ * of this file for the tier-2 `prompt_cache_key` path. The single-
665
+ * warm invariant plus `withExclusive`'s per-model serialization make
666
+ * the lookup safe under SINGLE-tenant assumptions: no two requests
667
+ * race the slot, and there is at most one slot to lease.
668
+ *
669
+ * **Caller contract on miss.** If `instructions` drifts between
670
+ * turns (system prompt changed) this returns `hit: false` and a
671
+ * fresh session — and the caller MUST then run a full
672
+ * `session.reset()` before priming history, NOT the JS-only
673
+ * `resetPreservingNativeCacheForWarmReuse` path. A fresh JS session
674
+ * does NOT imply a fresh native cache (the underlying
675
+ * `SessionCapableModel` is shared and its native
676
+ * `cached_token_history` persists across requests), so skipping the
677
+ * native wipe on a miss would let the next `chatSessionStart` reuse
678
+ * an unrelated previous request's prefix — the cross-request
679
+ * cache-affinity side channel that the long block comment in
680
+ * `responses.ts` (around the `runSessionNonStreaming` /
681
+ * `runSessionStreaming` branches) describes.
682
+ */
683
+ getOrCreateWarmAny(requestedInstructions) {
684
+ // Single-warm invariant: at most one entry. Walk it once, lease
685
+ // on a fresh + instructions-matched hit, otherwise clear and
686
+ // cold-start. The ignored fields (promptCacheKey,
687
+ // previousResponseId-keying) are deliberate — see the docstring.
688
+ for (const entry of this.entries.values()) {
689
+ if (entry.expiresAt < nowSec())
690
+ continue;
691
+ if (entry.instructions !== requestedInstructions)
692
+ continue;
693
+ // Hit: clear and lease (single-use semantics, same as tiers 1/2).
229
694
  this.entries.clear();
230
- return { session: new ChatSession(this.model), hit: false };
695
+ return { session: entry.session, hit: true };
231
696
  }
232
- // Hit: clear and hand the session out as a single-use lease so
233
- // a concurrent second request against the same id cold-replays
234
- // instead of sharing this live ChatSession.
697
+ // Miss (no entry, expired, or instructions drift). Clear the map
698
+ // so a stale wrapper cannot leak into a later lookup, and return
699
+ // a fresh session.
235
700
  this.entries.clear();
236
- return { session: entry.session, hit: true };
701
+ return { session: this.newSession(), hit: false };
237
702
  }
238
703
  /**
239
704
  * Insert a session under a newly allocated response id. Clears the
@@ -243,12 +708,28 @@ export class SessionRegistry {
243
708
  * `instructions` is the prefix/system state used for this turn;
244
709
  * stored on the entry and compared on the next `getOrCreate` to
245
710
  * detect prefix changes that must force a cold replay.
711
+ *
712
+ * `promptCacheKey` is the client-supplied conversation-chain key
713
+ * that enables the registry's tier-2 lookup for stateless agent
714
+ * turns that do not carry a `previous_response_id`. `null` /
715
+ * `undefined` means "no key supplied" — stored verbatim so a
716
+ * subsequent stateless lookup that also omits the key does NOT
717
+ * accidentally pick up this entry (only explicit non-null
718
+ * key-equality on both sides can hit tier 2). See
719
+ * {@link SessionRegistry.getOrCreate} for the precedence rules.
246
720
  */
247
- adopt(responseId, session, instructions) {
721
+ adopt(responseId, session, instructions, promptCacheKey = null) {
722
+ // Scope the caller-supplied key BEFORE storing so a later
723
+ // `getOrCreate` can only resolve entries via the same opt-in +
724
+ // HMAC path. When tier-2 reuse is disabled (or the key is too
725
+ // short / absent) `scopePromptCacheKey` returns `null`, which
726
+ // disables this entry from ever matching a tier-2 lookup — the
727
+ // raw caller-supplied key is NEVER stored.
248
728
  this.entries.clear();
249
729
  this.entries.set(responseId, {
250
730
  session,
251
731
  instructions,
732
+ promptCacheKey: scopePromptCacheKey(promptCacheKey ?? null),
252
733
  expiresAt: nowSec() + this.ttlSec,
253
734
  });
254
735
  }