@mlx-node/server 0.0.13 → 0.0.15
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/host/discover.d.ts +3 -6
- package/dist/host/discover.d.ts.map +1 -1
- package/dist/host/discover.js +9 -42
- package/dist/host/index.d.ts +2 -2
- package/dist/host/index.d.ts.map +1 -1
- package/dist/host/index.js +8 -1
- package/package.json +9 -4
- package/src/auth.ts +111 -0
- package/src/chat-session-warm-reuse.ts +96 -0
- package/src/endpoints/messages-count-tokens.ts +164 -0
- package/src/endpoints/messages.ts +1802 -0
- package/src/endpoints/models.ts +20 -0
- package/src/endpoints/responses.ts +3928 -0
- package/src/errors.ts +120 -0
- package/src/handler.ts +195 -0
- package/src/health.ts +213 -0
- package/src/host/discover.ts +25 -0
- package/src/host/env-policy.ts +81 -0
- package/src/host/index.ts +496 -0
- package/src/host/logger.ts +419 -0
- package/src/host/net.ts +100 -0
- package/src/host/paths.ts +77 -0
- package/src/host/swap.ts +200 -0
- package/src/host/temp-root.ts +110 -0
- package/src/idle-sweeper.ts +555 -0
- package/src/index.ts +114 -0
- package/src/load-model.ts +92 -0
- package/src/mappers/anthropic-request.ts +485 -0
- package/src/mappers/anthropic-response.ts +306 -0
- package/src/mappers/request.ts +456 -0
- package/src/mappers/response.ts +163 -0
- package/src/model-work-coordinator.ts +416 -0
- package/src/pending-writes.ts +481 -0
- package/src/registry.ts +691 -0
- package/src/router.ts +220 -0
- package/src/server.ts +579 -0
- package/src/session-registry.ts +1371 -0
- package/src/stop-sequence-buffer.ts +161 -0
- package/src/streaming.ts +205 -0
- package/src/text-recovery.ts +41 -0
- package/src/timing.ts +236 -0
- package/src/tool-call-buffer.ts +78 -0
- package/src/transport-visibility.ts +185 -0
- package/src/types-anthropic.ts +409 -0
- package/src/types.ts +470 -0
|
@@ -0,0 +1,1371 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* SessionRegistry -- per-model cache holding AT MOST one live
|
|
3
|
+
* `ChatSession` whose native KV state is currently valid.
|
|
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
|
+
*
|
|
17
|
+
* Design notes:
|
|
18
|
+
*
|
|
19
|
+
* - **One registry per model.** Composed alongside each registered
|
|
20
|
+
* `ServableModel` in `ModelRegistry`. Sessions are keyed purely
|
|
21
|
+
* by response id — no secondary keying on model name because the
|
|
22
|
+
* registry is already scoped per model.
|
|
23
|
+
*
|
|
24
|
+
* - **Single-warm-session invariant.** The JS warm registry retains at
|
|
25
|
+
* most ONE `ChatSession` entry — both `getOrCreate` and `adopt` clear
|
|
26
|
+
* the map before returning or inserting. This remains the sole safe
|
|
27
|
+
* reuse mechanism for flat-cache models whose native cache is one
|
|
28
|
+
* mutable vector. Block-paged schedulers instead isolate live turns by
|
|
29
|
+
* cache owner and reuse verified physical blocks through the native
|
|
30
|
+
* prefix table; those models may run fresh JS sessions concurrently.
|
|
31
|
+
*
|
|
32
|
+
* - **Lease semantics on hit.** Clear-on-hit also gives single-
|
|
33
|
+
* flight lease semantics: two overlapping requests referencing
|
|
34
|
+
* the same `previous_response_id` cannot share the same live
|
|
35
|
+
* `ChatSession`. The first wins the cleared entry; the second
|
|
36
|
+
* finds the map empty and cold-replays from `ResponseStore` on
|
|
37
|
+
* a fresh session. Without this, the second would hit
|
|
38
|
+
* `ChatSession`'s single-flight "concurrent send() not allowed"
|
|
39
|
+
* guard.
|
|
40
|
+
*
|
|
41
|
+
* - **Prefix compatibility changes miss.** Each entry records its
|
|
42
|
+
* `instructions` plus an opaque fingerprint of the cache salt used to
|
|
43
|
+
* adopt it. `getOrCreate` compares both against the new request;
|
|
44
|
+
* mismatch forces owner release and cold replay instead of reusing a
|
|
45
|
+
* stale prompt or changing the security domain of one live native
|
|
46
|
+
* request. OpenAI `instructions` and Anthropic `system` share the same
|
|
47
|
+
* parameter; both endpoints also thread their mapped `cache_salt`.
|
|
48
|
+
*
|
|
49
|
+
* - **Cache miss fallback.** On a miss (eviction, interleaved turn
|
|
50
|
+
* on a different chain, restart, lease-on-hit) the endpoint
|
|
51
|
+
* layer reconstructs the conversation from the `ResponseStore`
|
|
52
|
+
* history, primes a fresh `ChatSession` via `primeHistory()`,
|
|
53
|
+
* and resumes through `startFromHistory()` /
|
|
54
|
+
* `startFromHistoryStream()`. That pair dispatches one
|
|
55
|
+
* `chatSessionStart*` call that rebuilds the full KV cache and
|
|
56
|
+
* atomically appends the new user turn, so cold replay is
|
|
57
|
+
* indistinguishable from a hot hit.
|
|
58
|
+
*
|
|
59
|
+
* - **TTL.** Default 1800 seconds mirrors `DEFAULT_RESPONSE_RETENTION_SECONDS`
|
|
60
|
+
* in `packages/server/src/server.ts` so the cached
|
|
61
|
+
* entry ages out alongside its stored response metadata. With
|
|
62
|
+
* at most one entry there is no LRU bookkeeping — just a single
|
|
63
|
+
* expiry check on lookup.
|
|
64
|
+
*
|
|
65
|
+
* - **Thread safety.** Node.js is single-threaded within one
|
|
66
|
+
* event-loop tick, so the internal `Map` is safe against
|
|
67
|
+
* concurrent mutation by design. `sweep()` can be scheduled
|
|
68
|
+
* via `setInterval` without colliding with in-flight calls.
|
|
69
|
+
*
|
|
70
|
+
* - **Per-model admission lane.** Flat and not-yet-batched families use
|
|
71
|
+
* `withExclusive(fn)`, the original FIFO mutex across the full dispatch.
|
|
72
|
+
* A model that explicitly reports `maxConcurrentSequences() > 1` uses
|
|
73
|
+
* `withAdmission(fn)`, a counting semaphore sized to that native
|
|
74
|
+
* scheduler. Per-session serialization still lives in `ChatSession`;
|
|
75
|
+
* only independent sessions share the model lane.
|
|
76
|
+
*/
|
|
77
|
+
|
|
78
|
+
import { createHash, createHmac, randomBytes } from 'node:crypto';
|
|
79
|
+
|
|
80
|
+
import type { ChatConfig } from '@mlx-node/core';
|
|
81
|
+
import { ChatSession, type SessionCapableModel } from '@mlx-node/lm';
|
|
82
|
+
|
|
83
|
+
/**
|
|
84
|
+
* Tier-2 `prompt_cache_key` reuse is **ON by default** so the server
|
|
85
|
+
* is immediately compatible with any stateless LLM agent (pi-mono,
|
|
86
|
+
* Aider, Codex CLI, Claude Code, Cline, Continue) that sends the full
|
|
87
|
+
* transcript every turn. Agents that set OpenAI's standard
|
|
88
|
+
* `prompt_cache_key` field — which most do — get automatic KV cache
|
|
89
|
+
* reuse across turns of the same logical session.
|
|
90
|
+
*
|
|
91
|
+
* Opt out via `MLX_DISABLE_PROMPT_CACHE_KEY=1` for **multi-tenant
|
|
92
|
+
* deployments**, where the tier-2 lookup becomes unsafe: two clients
|
|
93
|
+
* that pick the same raw `prompt_cache_key` (by accident or on
|
|
94
|
+
* purpose) would share a warm `ChatSession`, leaking conversation
|
|
95
|
+
* history and sampling state across principals.
|
|
96
|
+
*
|
|
97
|
+
* Multi-tenant isolation is out of scope for this registry. The
|
|
98
|
+
* HMAC-scoping applied below only hides the raw key from memory /
|
|
99
|
+
* dumps — it does NOT protect against two clients sharing the same
|
|
100
|
+
* raw input. Operators who need multi-tenant isolation must either
|
|
101
|
+
* disable the feature or front the server with an auth proxy that
|
|
102
|
+
* rewrites `prompt_cache_key` per-tenant before it reaches the
|
|
103
|
+
* process.
|
|
104
|
+
*
|
|
105
|
+
* Read at call time (not cached at module load) so tests can flip the
|
|
106
|
+
* env via `vi.stubEnv()` between cases without re-importing the module.
|
|
107
|
+
* The check is a single env-var read plus a string compare — negligible
|
|
108
|
+
* against the rest of the lookup work.
|
|
109
|
+
*/
|
|
110
|
+
function isPromptCacheKeyEnabled(): boolean {
|
|
111
|
+
return process.env.MLX_DISABLE_PROMPT_CACHE_KEY !== '1';
|
|
112
|
+
}
|
|
113
|
+
|
|
114
|
+
/**
|
|
115
|
+
* Minimum length accepted for a caller-supplied `prompt_cache_key`
|
|
116
|
+
* before tier-2 scoping. Short keys make trivial guessing collisions
|
|
117
|
+
* plausible; reject anything shorter than this as if the caller had
|
|
118
|
+
* not supplied a key at all. Chosen to reject one- / two- / few-byte
|
|
119
|
+
* values a client might accidentally pass through while still allowing
|
|
120
|
+
* any reasonable opaque id (UUID prefix, short client token, etc.).
|
|
121
|
+
*/
|
|
122
|
+
const PROMPT_CACHE_KEY_MIN_LENGTH = 8;
|
|
123
|
+
|
|
124
|
+
/**
|
|
125
|
+
* Lazily-initialized boot-time nonce used to HMAC every caller-supplied
|
|
126
|
+
* `prompt_cache_key` before it is stored or looked up. Held in memory
|
|
127
|
+
* only — never persisted to disk. A process restart invalidates every
|
|
128
|
+
* tier-2 entry because the next module instance produces a fresh
|
|
129
|
+
* nonce.
|
|
130
|
+
*
|
|
131
|
+
* The nonce makes pre-existing entries unmatchable from outside the
|
|
132
|
+
* process: an attacker who knows a victim's raw `prompt_cache_key` but
|
|
133
|
+
* cannot read the nonce from the server's memory also cannot craft a
|
|
134
|
+
* lookup that collides with the stored HMAC'd key. Combined with the
|
|
135
|
+
* opt-in gate above, the tier-2 surface is off-by-default and bound to
|
|
136
|
+
* a server-instance secret when enabled.
|
|
137
|
+
*
|
|
138
|
+
* Populated lazily on first use so the cost is not paid when tier-2 is
|
|
139
|
+
* disabled. Module-scope so it is shared across every `SessionRegistry`
|
|
140
|
+
* in the process (each per-model registry's HMAC'd keys are still
|
|
141
|
+
* distinct via their per-registry `entries` map — there is no
|
|
142
|
+
* cross-model leakage).
|
|
143
|
+
*/
|
|
144
|
+
let cachedNonce: Buffer | null = null;
|
|
145
|
+
|
|
146
|
+
/** Lazily obtain the module-scoped HMAC nonce. */
|
|
147
|
+
function getNonce(): Buffer {
|
|
148
|
+
if (cachedNonce === null) {
|
|
149
|
+
cachedNonce = randomBytes(32);
|
|
150
|
+
}
|
|
151
|
+
return cachedNonce;
|
|
152
|
+
}
|
|
153
|
+
|
|
154
|
+
/** Opaque equality token for a cache salt; the caller's raw value is never retained. */
|
|
155
|
+
function fingerprintCacheSalt(cacheSalt: string | null | undefined): string | null {
|
|
156
|
+
if (cacheSalt == null) return null;
|
|
157
|
+
return createHmac('sha256', getNonce()).update(cacheSalt).digest('hex').slice(0, 32);
|
|
158
|
+
}
|
|
159
|
+
|
|
160
|
+
/**
|
|
161
|
+
* Test-only hook used by the scoping unit tests to simulate a server
|
|
162
|
+
* restart: resets the module-scoped HMAC nonce (so every previously stored
|
|
163
|
+
* tier-2 key and cache-salt fingerprint misses) and clears the silent-miss
|
|
164
|
+
* dedupe cache so tests can re-exercise the once-per-key diagnostic path.
|
|
165
|
+
*
|
|
166
|
+
* **Not exported from the package's public `index.ts` surface** —
|
|
167
|
+
* exporting it there would let downstream consumers nuke tier-2
|
|
168
|
+
* state in production (every stored entry would go unreachable).
|
|
169
|
+
* Tests reach the function via the deep path
|
|
170
|
+
* `packages/server/src/session-registry.ts` instead; that import is
|
|
171
|
+
* deliberately noisy to signal "test-only, do not use from app
|
|
172
|
+
* code". The `__` prefix is a loud-enough convention for the
|
|
173
|
+
* ergonomic test path but NOT sufficient for a public package
|
|
174
|
+
* export.
|
|
175
|
+
*/
|
|
176
|
+
export function __resetPromptCacheKeyNonceForTests(): void {
|
|
177
|
+
cachedNonce = null;
|
|
178
|
+
loggedSilentMissKeys.clear();
|
|
179
|
+
}
|
|
180
|
+
|
|
181
|
+
/**
|
|
182
|
+
* Test-only probe for the silent-miss dedupe Map's live size. Same
|
|
183
|
+
* "do not use from app code" rationale as
|
|
184
|
+
* {@link __resetPromptCacheKeyNonceForTests}; the `__` prefix and the
|
|
185
|
+
* test-only deep-import path are the load-bearing signals. Exists so
|
|
186
|
+
* the flooding regression test can assert the FIFO cap actually bounds
|
|
187
|
+
* the stored set — an invariant that is otherwise unobservable from
|
|
188
|
+
* outside the module and is NOT implied by the warning count (each
|
|
189
|
+
* call emits at most one warning regardless of whether the cap works).
|
|
190
|
+
*/
|
|
191
|
+
export function __loggedSilentMissKeysSizeForTests(): number {
|
|
192
|
+
return loggedSilentMissKeys.size;
|
|
193
|
+
}
|
|
194
|
+
|
|
195
|
+
/**
|
|
196
|
+
* Normalize and HMAC-scope a caller-supplied `prompt_cache_key` before
|
|
197
|
+
* it is stored or used for lookup.
|
|
198
|
+
*
|
|
199
|
+
* **Single-tenant trust boundary.** HMAC-scoping hides the raw key
|
|
200
|
+
* from memory dumps and keeps one process instance's stored keys
|
|
201
|
+
* unreachable from another instance (a restart rerolls the nonce),
|
|
202
|
+
* but it does NOT protect against two clients supplying the same raw
|
|
203
|
+
* key — by construction both lookups HMAC to the same scoped key and
|
|
204
|
+
* share the entry. Multi-tenant isolation is out of scope; see the
|
|
205
|
+
* module docstring.
|
|
206
|
+
*
|
|
207
|
+
* Returns `null` when:
|
|
208
|
+
* - The tier-2 feature is disabled
|
|
209
|
+
* (`MLX_DISABLE_PROMPT_CACHE_KEY` is set to `"1"`).
|
|
210
|
+
* - `rawKey` is `null`, `undefined`, or the empty string (callers
|
|
211
|
+
* that forget to thread the key must not accidentally opt into
|
|
212
|
+
* tier-2 reuse).
|
|
213
|
+
* - `rawKey` is shorter than {@link PROMPT_CACHE_KEY_MIN_LENGTH}
|
|
214
|
+
* characters, which keeps trivial guessing collisions off the
|
|
215
|
+
* table.
|
|
216
|
+
*
|
|
217
|
+
* Otherwise returns the first 32 hex chars of
|
|
218
|
+
* `HMAC-SHA256(cachedNonce, rawKey)` — opaque, server-instance-scoped,
|
|
219
|
+
* and long enough to preserve the 64-bit entropy floor that the
|
|
220
|
+
* pre-scoped path relied on for key uniqueness.
|
|
221
|
+
*/
|
|
222
|
+
function scopePromptCacheKey(rawKey: string | null | undefined): string | null {
|
|
223
|
+
if (!isPromptCacheKeyEnabled()) return null;
|
|
224
|
+
if (rawKey == null || rawKey.length < PROMPT_CACHE_KEY_MIN_LENGTH) return null;
|
|
225
|
+
return createHmac('sha256', getNonce()).update(rawKey).digest('hex').slice(0, 32);
|
|
226
|
+
}
|
|
227
|
+
|
|
228
|
+
/**
|
|
229
|
+
* Bounded set of SHA-256-digest prefixes tracking which
|
|
230
|
+
* `prompt_cache_key` values have already had a silent-miss warning
|
|
231
|
+
* emitted in this process. Stored as 64-bit hex digests (not raw
|
|
232
|
+
* strings) so an attacker flooding the endpoint with distinct
|
|
233
|
+
* attacker-controlled keys cannot drive unbounded memory growth.
|
|
234
|
+
* FIFO-evicted at {@link LOGGED_SILENT_MISS_KEYS_MAX} entries via the
|
|
235
|
+
* Map insertion-order guarantee. Reset alongside the nonce / warning
|
|
236
|
+
* flag in {@link __resetPromptCacheKeyNonceForTests} so unit tests
|
|
237
|
+
* can re-exercise the once-per-key path.
|
|
238
|
+
*/
|
|
239
|
+
const LOGGED_SILENT_MISS_KEYS_MAX = 256;
|
|
240
|
+
const loggedSilentMissKeys = new Map<string, true>();
|
|
241
|
+
|
|
242
|
+
/** Hash a raw key to a bounded digest for dedupe storage. */
|
|
243
|
+
function digestSilentMissKey(rawKey: string): string {
|
|
244
|
+
// 64-bit prefix is enough for dedupe across a 256-entry window.
|
|
245
|
+
return createHash('sha256').update(rawKey).digest('hex').slice(0, 16);
|
|
246
|
+
}
|
|
247
|
+
|
|
248
|
+
/**
|
|
249
|
+
* Emit a once-per-raw-key stderr debug warning when the caller
|
|
250
|
+
* supplies a non-empty `prompt_cache_key` but at least one tier-2
|
|
251
|
+
* prerequisite is missing — the env gate is off, or the key is
|
|
252
|
+
* shorter than {@link PROMPT_CACHE_KEY_MIN_LENGTH}. The silent-miss
|
|
253
|
+
* fallback (cold-start as if no key were supplied) is the documented
|
|
254
|
+
* behaviour but is easy to miss during integration; this nudge
|
|
255
|
+
* surfaces the cause once per distinct key so operators don't have
|
|
256
|
+
* to grep source to diagnose a flat `X-Session-Cache: fresh`.
|
|
257
|
+
*
|
|
258
|
+
* No-op when `rawKey` is null / undefined / empty (the caller did
|
|
259
|
+
* not ask for tier-2 at all) or when scoping would succeed (the
|
|
260
|
+
* gate already accepted the key). Called by the endpoint layer
|
|
261
|
+
* right after it has decided `effectivePromptCacheKey`.
|
|
262
|
+
*/
|
|
263
|
+
export function maybeWarnPromptCacheKeyIneligible(rawKey: string | null | undefined): void {
|
|
264
|
+
if (rawKey == null || rawKey.length === 0) return;
|
|
265
|
+
// Happy-path branch: scoping would succeed, no nudge needed.
|
|
266
|
+
if (isPromptCacheKeyEnabled() && rawKey.length >= PROMPT_CACHE_KEY_MIN_LENGTH) return;
|
|
267
|
+
const digest = digestSilentMissKey(rawKey);
|
|
268
|
+
if (loggedSilentMissKeys.has(digest)) return;
|
|
269
|
+
// FIFO eviction via Map insertion order — bounds memory under
|
|
270
|
+
// adversarial key flooding while preserving once-per-key semantics
|
|
271
|
+
// within the recent-key window.
|
|
272
|
+
if (loggedSilentMissKeys.size >= LOGGED_SILENT_MISS_KEYS_MAX) {
|
|
273
|
+
const oldest = loggedSilentMissKeys.keys().next().value;
|
|
274
|
+
if (oldest !== undefined) loggedSilentMissKeys.delete(oldest);
|
|
275
|
+
}
|
|
276
|
+
loggedSilentMissKeys.set(digest, true);
|
|
277
|
+
if (!isPromptCacheKeyEnabled()) {
|
|
278
|
+
console.warn(
|
|
279
|
+
`[mlx-node] prompt_cache_key supplied but tier-2 reuse is disabled ` +
|
|
280
|
+
`(MLX_DISABLE_PROMPT_CACHE_KEY=1). The key will be ignored and this ` +
|
|
281
|
+
`turn will cold-start. This message is logged once per distinct key.`,
|
|
282
|
+
);
|
|
283
|
+
return;
|
|
284
|
+
}
|
|
285
|
+
console.warn(
|
|
286
|
+
`[mlx-node] prompt_cache_key is shorter than ${PROMPT_CACHE_KEY_MIN_LENGTH} chars; tier-2 reuse requires ` +
|
|
287
|
+
`at least ${PROMPT_CACHE_KEY_MIN_LENGTH} characters. The key will be ignored and this turn will ` +
|
|
288
|
+
`cold-start. This message is logged once per distinct key.`,
|
|
289
|
+
);
|
|
290
|
+
}
|
|
291
|
+
|
|
292
|
+
/** Constructor options for {@link SessionRegistry}. */
|
|
293
|
+
export interface SessionRegistryOptions {
|
|
294
|
+
/** The model that every session in this registry wraps. Single-model per registry. */
|
|
295
|
+
model: SessionCapableModel;
|
|
296
|
+
/** TTL in seconds before an unused session is evicted. Default: 1800 (30 min). */
|
|
297
|
+
ttlSec?: number;
|
|
298
|
+
/**
|
|
299
|
+
* Maximum number of requests that may be WAITING for the
|
|
300
|
+
* per-model execution mutex at the same time (the in-flight holder
|
|
301
|
+
* does NOT count toward this). When set and the cap is exceeded at
|
|
302
|
+
* `withExclusive` entry, the call throws {@link QueueFullError}
|
|
303
|
+
* synchronously so the endpoint layer can emit HTTP 429 and the
|
|
304
|
+
* client can retry later.
|
|
305
|
+
*
|
|
306
|
+
* Default at THIS layer: `undefined` (unbounded) — a directly
|
|
307
|
+
* constructed `SessionRegistry` has no cap. `createServer` supplies
|
|
308
|
+
* its own default of 16 per {@link ServerConfig.maxQueueDepthPerModel}
|
|
309
|
+
* (or the `MLX_MAX_QUEUE_DEPTH_PER_MODEL` env var; config
|
|
310
|
+
* `'unbounded'` opts out), so server-allocated registries are capped
|
|
311
|
+
* unless the operator explicitly opted out.
|
|
312
|
+
*/
|
|
313
|
+
maxQueueDepth?: number;
|
|
314
|
+
/**
|
|
315
|
+
* Number of independent dispatches the native model scheduler can advance
|
|
316
|
+
* concurrently. Values below two retain the legacy exclusive FIFO. The
|
|
317
|
+
* model registry derives this from `maxConcurrentSequences()`; direct
|
|
318
|
+
* construction defaults to one.
|
|
319
|
+
*/
|
|
320
|
+
maxConcurrentDispatches?: number;
|
|
321
|
+
/**
|
|
322
|
+
* Optional sampling defaults applied to every `ChatSession` this
|
|
323
|
+
* registry allocates. Forwarded verbatim into `new ChatSession(model,
|
|
324
|
+
* { defaultConfig })` so the session's `mergeConfig(overlay)` shallow-
|
|
325
|
+
* merges per-call config on top. Intended for server operators who
|
|
326
|
+
* want to pin per-model sampling knobs (temperature, topK, penalties,
|
|
327
|
+
* etc.) without client cooperation — per-request values from the
|
|
328
|
+
* OpenAI `/v1/responses` or Anthropic `/v1/messages` body still win
|
|
329
|
+
* where present because `ChatSession` treats them as an overlay.
|
|
330
|
+
*
|
|
331
|
+
* When `undefined`, behaviour is unchanged from the pre-defaults era
|
|
332
|
+
* (each `new ChatSession(model)` uses an empty `defaultConfig`).
|
|
333
|
+
*/
|
|
334
|
+
samplingDefaults?: ChatConfig;
|
|
335
|
+
/**
|
|
336
|
+
* Optional per-model cap for generated output tokens. Endpoint handlers
|
|
337
|
+
* apply it after request mapping, before dispatching into native decode.
|
|
338
|
+
*/
|
|
339
|
+
maxOutputTokens?: number;
|
|
340
|
+
}
|
|
341
|
+
|
|
342
|
+
/**
|
|
343
|
+
* Thrown synchronously by {@link SessionRegistry.withExclusive} when
|
|
344
|
+
* the per-model queue cap (`maxQueueDepth`) is exceeded. The error is
|
|
345
|
+
* raised BEFORE awaiting the previous lock holder so endpoint handlers
|
|
346
|
+
* can reliably catch it without racing the chain.
|
|
347
|
+
*/
|
|
348
|
+
export class QueueFullError extends Error {
|
|
349
|
+
readonly queueDepth: number;
|
|
350
|
+
readonly preDispatchAdmissions: number;
|
|
351
|
+
readonly admissionFootprint: number;
|
|
352
|
+
readonly limit: number;
|
|
353
|
+
|
|
354
|
+
constructor(queueDepth: number, preDispatchAdmissions: number, limit: number) {
|
|
355
|
+
const admissionFootprint = queueDepth + preDispatchAdmissions;
|
|
356
|
+
super(
|
|
357
|
+
`Model queue full: ${queueDepth} queued, ${preDispatchAdmissions} pre-dispatch ` +
|
|
358
|
+
`(${admissionFootprint} admitted outside the active runner; waiter limit ${limit})`,
|
|
359
|
+
);
|
|
360
|
+
this.name = 'QueueFullError';
|
|
361
|
+
this.queueDepth = queueDepth;
|
|
362
|
+
this.preDispatchAdmissions = preDispatchAdmissions;
|
|
363
|
+
this.admissionFootprint = admissionFootprint;
|
|
364
|
+
this.limit = limit;
|
|
365
|
+
}
|
|
366
|
+
}
|
|
367
|
+
|
|
368
|
+
/**
|
|
369
|
+
* Admission permit returned by
|
|
370
|
+
* {@link SessionRegistry.beginPreDispatchAdmission}. Represents exactly
|
|
371
|
+
* ONE unit of the per-model admission budget, and ends in exactly one of
|
|
372
|
+
* two terminal states:
|
|
373
|
+
*
|
|
374
|
+
* - **handed off**: passed as the second argument of
|
|
375
|
+
* `withExclusive(fn, permit)` or `withAdmission(fn, permit)`, which
|
|
376
|
+
* consumes it atomically as that
|
|
377
|
+
* call's admission — the unit converts into the waiter charge (or is
|
|
378
|
+
* retired when the caller wins the runner slot). `release()` becomes
|
|
379
|
+
* a no-op afterwards.
|
|
380
|
+
* - **released**: `release()` frees the unit without a dispatch. Every
|
|
381
|
+
* exit between admission and lock placement must do this — the
|
|
382
|
+
* recommended pattern is one unconditional `release()` in a
|
|
383
|
+
* `finally`, which is safe on every path because `release()` is
|
|
384
|
+
* idempotent and a no-op after handoff.
|
|
385
|
+
*/
|
|
386
|
+
export interface PreDispatchAdmission {
|
|
387
|
+
/** Idempotent: free the budget unit unless already handed off. */
|
|
388
|
+
release(): void;
|
|
389
|
+
}
|
|
390
|
+
|
|
391
|
+
/**
|
|
392
|
+
* Result of {@link SessionRegistry.getOrCreate}. `hit` reflects whether
|
|
393
|
+
* the call consumed a live warm entry (single-use lease) or returned a
|
|
394
|
+
* fresh `ChatSession` on a miss. The endpoint layer uses `hit` to
|
|
395
|
+
* classify the per-request session-cache status emitted to clients via
|
|
396
|
+
* the `X-Session-Cache` observability header.
|
|
397
|
+
*/
|
|
398
|
+
export interface SessionLookupResult {
|
|
399
|
+
session: ChatSession<SessionCapableModel>;
|
|
400
|
+
hit: boolean;
|
|
401
|
+
}
|
|
402
|
+
|
|
403
|
+
interface SessionEntry {
|
|
404
|
+
session: ChatSession<SessionCapableModel>;
|
|
405
|
+
/**
|
|
406
|
+
* The `instructions` / `system` string the caller adopted this
|
|
407
|
+
* session with. `null` if the caller did not supply any. Compared
|
|
408
|
+
* byte-for-byte against the caller's `requestedInstructions` in
|
|
409
|
+
* `getOrCreate` to detect prefix/system-state changes that would
|
|
410
|
+
* otherwise let a hit silently reuse a stale warmed prompt.
|
|
411
|
+
*/
|
|
412
|
+
instructions: string | null;
|
|
413
|
+
/**
|
|
414
|
+
* Prefix-cache security domain used when this warm session was adopted.
|
|
415
|
+
* Omission is represented as `null` (native salt 0) and must remain
|
|
416
|
+
* distinct from every explicit salt because a live native request cannot
|
|
417
|
+
* change domains.
|
|
418
|
+
*/
|
|
419
|
+
cacheSaltFingerprint: string | null;
|
|
420
|
+
/**
|
|
421
|
+
* Stable caller-supplied key identifying the logical conversation
|
|
422
|
+
* chain for warm-session reuse across stateless turns that do NOT
|
|
423
|
+
* carry a `previous_response_id`. `null` when the adopting caller
|
|
424
|
+
* supplied no key (or when the request came through an endpoint
|
|
425
|
+
* that does not honor the key, e.g. the chain terminated on an
|
|
426
|
+
* in-`previous_response_id` hop). See
|
|
427
|
+
* {@link SessionRegistry.getOrCreate} for the tier-2 lookup
|
|
428
|
+
* semantics.
|
|
429
|
+
*
|
|
430
|
+
* The distinction between `null` (no key supplied) and the empty
|
|
431
|
+
* string `""` (key explicitly set to empty) is load-bearing — tier-2
|
|
432
|
+
* lookup treats them as different keys so a client that forgets to
|
|
433
|
+
* thread the key does not accidentally collide with another client
|
|
434
|
+
* that did set it to empty.
|
|
435
|
+
*/
|
|
436
|
+
promptCacheKey: string | null;
|
|
437
|
+
/** Unix seconds at which this entry becomes eligible for eviction. */
|
|
438
|
+
expiresAt: number;
|
|
439
|
+
}
|
|
440
|
+
|
|
441
|
+
/** Current time in unix seconds. Kept as a helper so tests can patch `Date.now` via fake timers. */
|
|
442
|
+
function nowSec(): number {
|
|
443
|
+
return Math.floor(Date.now() / 1000);
|
|
444
|
+
}
|
|
445
|
+
|
|
446
|
+
export class SessionRegistry {
|
|
447
|
+
private readonly model: SessionCapableModel;
|
|
448
|
+
private readonly ttlSec: number;
|
|
449
|
+
private readonly maxQueueDepth: number | undefined;
|
|
450
|
+
private readonly maxConcurrentDispatches: number;
|
|
451
|
+
/**
|
|
452
|
+
* Per-model sampling defaults forwarded into every new `ChatSession`
|
|
453
|
+
* via its `defaultConfig` constructor option. `undefined` preserves
|
|
454
|
+
* the pre-defaults behaviour (empty `defaultConfig`). See
|
|
455
|
+
* {@link SessionRegistryOptions.samplingDefaults}.
|
|
456
|
+
*/
|
|
457
|
+
private samplingDefaults: ChatConfig | undefined;
|
|
458
|
+
private maxOutputTokens: number | undefined;
|
|
459
|
+
/**
|
|
460
|
+
* Number of callers that are currently WAITING for the per-model
|
|
461
|
+
* execution mutex — i.e. have entered `withExclusive` but have not
|
|
462
|
+
* yet started running their closure. The caller that is actively
|
|
463
|
+
* running inside `fn()` is NOT counted here, so a cap of
|
|
464
|
+
* `maxQueueDepth = N` means "1 running + up to N waiting".
|
|
465
|
+
*
|
|
466
|
+
* Mutated strictly inside `withExclusive`: the admitting caller is
|
|
467
|
+
* counted as a waiter ONLY when the execution chain is already
|
|
468
|
+
* non-idle (i.e. some earlier caller still holds the mutex). The
|
|
469
|
+
* first caller into an idle chain is admitted directly as the
|
|
470
|
+
* runner slot and never contributes to `queuedCount`. Waiters
|
|
471
|
+
* decrement exactly once as they transition from waiting to
|
|
472
|
+
* running (after `await prev`). The counter is intentionally
|
|
473
|
+
* NEVER touched on cap-reject paths (the caller never queued) so
|
|
474
|
+
* the cap check is stable across concurrent entries, and runner-
|
|
475
|
+
* slot admissions leave it alone so a synchronous burst
|
|
476
|
+
* (e.g. `Promise.all([fn, fn])`) does not spuriously bill the
|
|
477
|
+
* runner-slot caller against the waiter cap.
|
|
478
|
+
*/
|
|
479
|
+
private queuedCount = 0;
|
|
480
|
+
/**
|
|
481
|
+
* Requests admitted by {@link beginPreDispatchAdmission} whose permit
|
|
482
|
+
* is still outstanding — parked in the `ModelWorkCoordinator` writer
|
|
483
|
+
* queue (host mode), blocked in pre-lock store lookups
|
|
484
|
+
* (`previous_response_id` continuations), or anywhere else between the
|
|
485
|
+
* endpoint gate and resident-lane placement. None of that parking is
|
|
486
|
+
* visible to `queuedCount`; this counter is what lets the gate bound
|
|
487
|
+
* it. Decremented ONLY by the permit itself: `release()` on a bail-out
|
|
488
|
+
* or the atomic consume inside the selected lane on handoff.
|
|
489
|
+
*/
|
|
490
|
+
private preDispatchAdmits = 0;
|
|
491
|
+
/**
|
|
492
|
+
* Consume hooks for outstanding permits, keyed by permit identity.
|
|
493
|
+
* Registry-scoped on purpose: both execution lanes consult THIS map, so a
|
|
494
|
+
* permit minted by a different registry is simply not found and the
|
|
495
|
+
* call falls back to normal waiter charging — a cross-registry handoff
|
|
496
|
+
* cannot corrupt either registry's counters.
|
|
497
|
+
*/
|
|
498
|
+
private readonly permitConsumers = new WeakMap<PreDispatchAdmission, () => boolean>();
|
|
499
|
+
/**
|
|
500
|
+
* Holds AT MOST ONE entry under the single-warm invariant (see the
|
|
501
|
+
* module-level rustdoc). `getOrCreate` and `adopt` both clear the
|
|
502
|
+
* map as part of their contract so a later lookup cannot hand out
|
|
503
|
+
* a wrapper whose assumed native state has been overwritten by a
|
|
504
|
+
* turn on another cached entry.
|
|
505
|
+
*/
|
|
506
|
+
private readonly entries: Map<string, SessionEntry> = new Map();
|
|
507
|
+
/**
|
|
508
|
+
* Eviction is synchronous at the map boundary, but releasing a native
|
|
509
|
+
* scheduler owner is asynchronous. Start every disposal immediately and
|
|
510
|
+
* retain its promise so endpoint admission lanes can wait for command-order
|
|
511
|
+
* visibility before dispatching a replacement turn.
|
|
512
|
+
*/
|
|
513
|
+
private readonly pendingDisposals = new Set<Promise<void>>();
|
|
514
|
+
private readonly disposalBySession = new WeakMap<ChatSession<SessionCapableModel>, Promise<void>>();
|
|
515
|
+
private readonly failedDisposals = new Set<ChatSession<SessionCapableModel>>();
|
|
516
|
+
/**
|
|
517
|
+
* Shared sentinel representing "the execution chain is idle" — a
|
|
518
|
+
* pre-resolved promise. `execLock` starts at this value and is
|
|
519
|
+
* reset to it whenever the last holder releases without a
|
|
520
|
+
* successor chained behind it. `withExclusive` uses reference
|
|
521
|
+
* equality against this sentinel (`execLock === initialLock`) to
|
|
522
|
+
* tell "I am the runner slot on an idle chain" apart from "I am a
|
|
523
|
+
* waiter behind someone else", which is how the burst
|
|
524
|
+
* (`Promise.all([fn, fn])`) admission bug is avoided.
|
|
525
|
+
*/
|
|
526
|
+
private readonly initialLock: Promise<void> = Promise.resolve();
|
|
527
|
+
/**
|
|
528
|
+
* Tail of the per-model execution FIFO. Every `withExclusive` call
|
|
529
|
+
* captures this value as its predecessor, then overwrites it with
|
|
530
|
+
* its own pending promise so the next waiter chains after it. The
|
|
531
|
+
* chain is resolved only when the current holder's `fn` has
|
|
532
|
+
* settled (success or failure), guaranteeing that at most one
|
|
533
|
+
* dispatch runs through this registry's native model at a time.
|
|
534
|
+
* Initialized to `initialLock` so the first caller proceeds
|
|
535
|
+
* without waiting AND is recognised as the runner slot (no waiter
|
|
536
|
+
* increment). When a holder releases as the current chain tail it
|
|
537
|
+
* restores `execLock` to `initialLock` so the next burst starts
|
|
538
|
+
* cleanly from the idle state.
|
|
539
|
+
*/
|
|
540
|
+
private execLock: Promise<void> = this.initialLock;
|
|
541
|
+
/** Active holders in the continuous-batching admission lane. */
|
|
542
|
+
private activeAdmissions = 0;
|
|
543
|
+
/** FIFO waiters parked behind the continuous-batching admission limit. */
|
|
544
|
+
private readonly admissionWaiters: Array<() => void> = [];
|
|
545
|
+
|
|
546
|
+
constructor(opts: SessionRegistryOptions) {
|
|
547
|
+
this.model = opts.model;
|
|
548
|
+
this.ttlSec = opts.ttlSec ?? 1800;
|
|
549
|
+
this.maxQueueDepth = opts.maxQueueDepth;
|
|
550
|
+
const requestedConcurrency = opts.maxConcurrentDispatches ?? 1;
|
|
551
|
+
this.maxConcurrentDispatches =
|
|
552
|
+
Number.isSafeInteger(requestedConcurrency) && requestedConcurrency > 1 ? requestedConcurrency : 1;
|
|
553
|
+
this.samplingDefaults = opts.samplingDefaults;
|
|
554
|
+
this.maxOutputTokens = opts.maxOutputTokens;
|
|
555
|
+
}
|
|
556
|
+
|
|
557
|
+
/**
|
|
558
|
+
* Construct a fresh `ChatSession` bound to this registry's model and
|
|
559
|
+
* pre-seeded with the operator-configured `samplingDefaults` (if any).
|
|
560
|
+
* Centralized so every cache-miss branch of `getOrCreate` produces a
|
|
561
|
+
* session whose per-call overlay will merge on top of the same
|
|
562
|
+
* defaults — clients cannot accidentally stray from the server's
|
|
563
|
+
* pinned sampling knobs by picking a cold-replay path.
|
|
564
|
+
*/
|
|
565
|
+
private newSession(): ChatSession<SessionCapableModel> {
|
|
566
|
+
if (this.samplingDefaults === undefined) {
|
|
567
|
+
return new ChatSession(this.model);
|
|
568
|
+
}
|
|
569
|
+
return new ChatSession(this.model, {
|
|
570
|
+
defaultConfig: this.samplingDefaults,
|
|
571
|
+
});
|
|
572
|
+
}
|
|
573
|
+
|
|
574
|
+
private scheduleDispose(session: ChatSession<SessionCapableModel>): void {
|
|
575
|
+
if (this.disposalBySession.has(session)) return;
|
|
576
|
+
|
|
577
|
+
const disposal = this.disposeSession(session)
|
|
578
|
+
.catch((error: unknown) => {
|
|
579
|
+
console.error('[server] failed to release an evicted chat-session cache owner:', error);
|
|
580
|
+
})
|
|
581
|
+
.finally(() => {
|
|
582
|
+
this.pendingDisposals.delete(disposal);
|
|
583
|
+
this.disposalBySession.delete(session);
|
|
584
|
+
});
|
|
585
|
+
this.disposalBySession.set(session, disposal);
|
|
586
|
+
this.pendingDisposals.add(disposal);
|
|
587
|
+
}
|
|
588
|
+
|
|
589
|
+
/**
|
|
590
|
+
* Dispose one leased session while retaining failed cleanup for a later
|
|
591
|
+
* registry flush. `ChatSession.dispose()` removes successful owners as it
|
|
592
|
+
* goes, so a retry only revisits owners whose native release failed.
|
|
593
|
+
*/
|
|
594
|
+
async disposeSession(session: ChatSession<SessionCapableModel>): Promise<void> {
|
|
595
|
+
try {
|
|
596
|
+
await session.dispose();
|
|
597
|
+
this.failedDisposals.delete(session);
|
|
598
|
+
} catch (error) {
|
|
599
|
+
this.failedDisposals.add(session);
|
|
600
|
+
throw error;
|
|
601
|
+
}
|
|
602
|
+
}
|
|
603
|
+
|
|
604
|
+
/** Remove every cached entry except an optional session being leased. */
|
|
605
|
+
private evictEntriesExcept(keep?: ChatSession<SessionCapableModel>): void {
|
|
606
|
+
for (const entry of this.entries.values()) {
|
|
607
|
+
if (entry.session !== keep) this.scheduleDispose(entry.session);
|
|
608
|
+
}
|
|
609
|
+
this.entries.clear();
|
|
610
|
+
}
|
|
611
|
+
|
|
612
|
+
/**
|
|
613
|
+
* Wait until every disposal scheduled so far has settled, retrying each
|
|
614
|
+
* failure discovered during this flush once. A persistent failure remains
|
|
615
|
+
* recorded for a later flush instead of spinning forever. Disposals log and
|
|
616
|
+
* absorb their own failures so cleanup cannot rewrite a response that has
|
|
617
|
+
* already reached the client.
|
|
618
|
+
*/
|
|
619
|
+
async flushPendingDisposals(): Promise<void> {
|
|
620
|
+
const retried = new Set<ChatSession<SessionCapableModel>>();
|
|
621
|
+
while (true) {
|
|
622
|
+
const retries = Array.from(this.failedDisposals).filter((session) => !retried.has(session));
|
|
623
|
+
for (const session of retries) {
|
|
624
|
+
this.failedDisposals.delete(session);
|
|
625
|
+
retried.add(session);
|
|
626
|
+
this.scheduleDispose(session);
|
|
627
|
+
}
|
|
628
|
+
if (this.pendingDisposals.size === 0) return;
|
|
629
|
+
await Promise.all(Array.from(this.pendingDisposals));
|
|
630
|
+
}
|
|
631
|
+
}
|
|
632
|
+
|
|
633
|
+
/**
|
|
634
|
+
* Disposals started but not yet settled — each one is an in-flight
|
|
635
|
+
* native `releaseCacheOwner` round-trip — plus disposals whose initial
|
|
636
|
+
* attempt and bounded retry both failed and remain owed to the native
|
|
637
|
+
* scheduler ({@link failedDisposals}), retried by the next
|
|
638
|
+
* {@link flushPendingDisposals}. Endpoints await that flush before
|
|
639
|
+
* leaving the admission lane, but it runs after the response has
|
|
640
|
+
* finished, so an observer keyed on request completion can still beat
|
|
641
|
+
* the release. `adopt`/`drop`/`sweep` schedule synchronously, so once
|
|
642
|
+
* the request counters read zero any disposal those requests will ever
|
|
643
|
+
* cause is already counted here. Primarily for diagnostics/tests.
|
|
644
|
+
*/
|
|
645
|
+
get pendingDisposalCount(): number {
|
|
646
|
+
return this.pendingDisposals.size + this.failedDisposals.size;
|
|
647
|
+
}
|
|
648
|
+
|
|
649
|
+
/**
|
|
650
|
+
* Number of requests waiting for this model's selected admission lane.
|
|
651
|
+
* Active dispatches are not included. Primarily for diagnostics/tests.
|
|
652
|
+
*/
|
|
653
|
+
get queueDepth(): number {
|
|
654
|
+
return this.queuedCount;
|
|
655
|
+
}
|
|
656
|
+
|
|
657
|
+
/**
|
|
658
|
+
* Configured waiter cap for this model's admission lane, or `undefined`
|
|
659
|
+
* when unbounded. Paired with {@link queueDepth} so a readiness probe can
|
|
660
|
+
* tell "3 waiters, unbounded" (fine) from "3 waiters, cap of 3" (the next
|
|
661
|
+
* request gets a 429) without reaching into private state.
|
|
662
|
+
*/
|
|
663
|
+
get queueDepthLimit(): number | undefined {
|
|
664
|
+
return this.maxQueueDepth;
|
|
665
|
+
}
|
|
666
|
+
|
|
667
|
+
/** Native continuous-batching capacity used by the endpoint route switch. */
|
|
668
|
+
get concurrentAdmissionLimit(): number {
|
|
669
|
+
return this.maxConcurrentDispatches;
|
|
670
|
+
}
|
|
671
|
+
|
|
672
|
+
/**
|
|
673
|
+
* Outstanding pre-dispatch permits — requests admitted by
|
|
674
|
+
* {@link beginPreDispatchAdmission} that have neither handed their
|
|
675
|
+
* permit to the selected execution lane nor released it yet. These permits,
|
|
676
|
+
* queued callers, and active dispatches share one bounded budget; see
|
|
677
|
+
* {@link assertAdmissionCapacity}. For probes and diagnostics.
|
|
678
|
+
*/
|
|
679
|
+
get preDispatchAdmitCount(): number {
|
|
680
|
+
return this.preDispatchAdmits;
|
|
681
|
+
}
|
|
682
|
+
|
|
683
|
+
/**
|
|
684
|
+
* Single source of truth for the per-model admission budget — the
|
|
685
|
+
* ONLY place the cap arithmetic lives. Every admission path calls
|
|
686
|
+
* this: {@link beginPreDispatchAdmission} before minting a permit,
|
|
687
|
+
* and both execution lanes for every non-handed-off caller. A handed-off
|
|
688
|
+
* permit skips the call for its OWN
|
|
689
|
+
* token only: that token was charged here at acquisition and its
|
|
690
|
+
* conversion keeps the total constant, so re-checking would
|
|
691
|
+
* double-charge an already-admitted request. It never exempts anyone
|
|
692
|
+
* else — all other outstanding state stays counted for every caller
|
|
693
|
+
* that did not pay.
|
|
694
|
+
*
|
|
695
|
+
* Budget invariant: active dispatches + queued callers + outstanding
|
|
696
|
+
* permits never exceed `maxConcurrentDispatches + maxQueueDepth`.
|
|
697
|
+
* On the exclusive lane `maxConcurrentDispatches` is one, exactly the
|
|
698
|
+
* original runner entitlement. On the batched lane it is the native
|
|
699
|
+
* scheduler's sequence capacity.
|
|
700
|
+
*
|
|
701
|
+
* Charging stays at the call sites (`preDispatchAdmits += 1` at the
|
|
702
|
+
* gate, `queuedCount += 1` for waiters, and the selected lane's active
|
|
703
|
+
* count for runners); every admitted unit is counted by exactly one
|
|
704
|
+
* at any time, which is what makes the footprint sum complete across
|
|
705
|
+
* any interleaving of permitted and permitless callers.
|
|
706
|
+
*
|
|
707
|
+
* Throws {@link QueueFullError} — reporting the footprint and the
|
|
708
|
+
* cap — when the caller does not fit; returns normally otherwise.
|
|
709
|
+
* No-op when the registry is unbounded.
|
|
710
|
+
*/
|
|
711
|
+
private assertAdmissionCapacity(): void {
|
|
712
|
+
if (this.maxQueueDepth === undefined) return;
|
|
713
|
+
const activeDispatches =
|
|
714
|
+
this.maxConcurrentDispatches > 1 ? this.activeAdmissions : this.execLock === this.initialLock ? 0 : 1;
|
|
715
|
+
const footprint = activeDispatches + this.queuedCount + this.preDispatchAdmits;
|
|
716
|
+
if (footprint >= this.maxQueueDepth + this.maxConcurrentDispatches) {
|
|
717
|
+
throw new QueueFullError(this.queuedCount, this.preDispatchAdmits, this.maxQueueDepth);
|
|
718
|
+
}
|
|
719
|
+
}
|
|
720
|
+
|
|
721
|
+
/**
|
|
722
|
+
* Endpoint-side early admission against this registry's cap, taken BEFORE
|
|
723
|
+
* the request enters any pre-dispatch parking spot the resident lane cannot
|
|
724
|
+
* see. Resident host traffic deliberately bypasses the model-load writer so
|
|
725
|
+
* continuous batching remains reachable, but a continuation can still block
|
|
726
|
+
* in `await store.getChain(...)` (or other pre-lock work) with `queuedCount`
|
|
727
|
+
* unchanged. This permit keeps all such work inside the same bounded budget.
|
|
728
|
+
*
|
|
729
|
+
* Accounting: pre-dispatch permits, active dispatches, and queued callers
|
|
730
|
+
* draw from ONE budget: `maxQueueDepth` waiter slots plus the selected
|
|
731
|
+
* lane's active capacity. Both lanes and this early gate route through
|
|
732
|
+
* {@link assertAdmissionCapacity}, so permitted and permitless arrivals
|
|
733
|
+
* cannot double-spend a slot.
|
|
734
|
+
*
|
|
735
|
+
* Throws {@link QueueFullError} synchronously when over cap (the
|
|
736
|
+
* caller maps it to the same 429 envelope as resident-lane rejection
|
|
737
|
+
* reject). On admission returns a {@link PreDispatchAdmission} permit
|
|
738
|
+
* the caller must RETAIN through ALL pre-lock asynchronous work and
|
|
739
|
+
* then hand to `withExclusive(fn, permit)` or `withAdmission(fn, permit)`,
|
|
740
|
+
* which consumes it
|
|
741
|
+
* atomically as that call's admission — one budget, one token per
|
|
742
|
+
* request, never double-counted. Releasing the permit early instead
|
|
743
|
+
* of handing it off re-opens the hole this gate closes: the request
|
|
744
|
+
* would be counted by NEITHER counter while parked, arrivals would
|
|
745
|
+
* refill the budget, and the resident lane would then admit a second
|
|
746
|
+
* full waiter budget on top. `release()` belongs on bail-out paths
|
|
747
|
+
* only (idempotent, no-op after handoff — an unconditional `finally`
|
|
748
|
+
* release is the recommended shape).
|
|
749
|
+
*/
|
|
750
|
+
beginPreDispatchAdmission(): PreDispatchAdmission {
|
|
751
|
+
this.assertAdmissionCapacity();
|
|
752
|
+
this.preDispatchAdmits += 1;
|
|
753
|
+
let settled = false;
|
|
754
|
+
const settle = (): boolean => {
|
|
755
|
+
if (settled) return false;
|
|
756
|
+
settled = true;
|
|
757
|
+
this.preDispatchAdmits -= 1;
|
|
758
|
+
if (this.preDispatchAdmits < 0) this.preDispatchAdmits = 0;
|
|
759
|
+
this.permitConsumers.delete(permit);
|
|
760
|
+
return true;
|
|
761
|
+
};
|
|
762
|
+
const permit: PreDispatchAdmission = {
|
|
763
|
+
release: (): void => {
|
|
764
|
+
void settle();
|
|
765
|
+
},
|
|
766
|
+
};
|
|
767
|
+
this.permitConsumers.set(permit, settle);
|
|
768
|
+
return permit;
|
|
769
|
+
}
|
|
770
|
+
|
|
771
|
+
/**
|
|
772
|
+
* Current sampling defaults applied to every new `ChatSession` this
|
|
773
|
+
* registry allocates. Exposed primarily for tests and diagnostics.
|
|
774
|
+
*/
|
|
775
|
+
get defaultSamplingConfig(): ChatConfig | undefined {
|
|
776
|
+
return this.samplingDefaults;
|
|
777
|
+
}
|
|
778
|
+
|
|
779
|
+
get outputTokenLimit(): number | undefined {
|
|
780
|
+
return this.maxOutputTokens;
|
|
781
|
+
}
|
|
782
|
+
|
|
783
|
+
/**
|
|
784
|
+
* Replace the sampling defaults forwarded into every future
|
|
785
|
+
* `ChatSession` this registry allocates. Called by `ModelRegistry`
|
|
786
|
+
* on a `register(name, model, { samplingDefaults })` refresh so a
|
|
787
|
+
* fresh registration's defaults immediately apply to the next
|
|
788
|
+
* cache-miss cold-start. Sessions already cached at call time keep
|
|
789
|
+
* the defaults they were constructed with — they settle naturally
|
|
790
|
+
* through the single-warm cache rotation.
|
|
791
|
+
*/
|
|
792
|
+
setSamplingDefaults(defaults: ChatConfig | undefined): void {
|
|
793
|
+
this.samplingDefaults = defaults;
|
|
794
|
+
}
|
|
795
|
+
|
|
796
|
+
setMaxOutputTokens(limit: number | undefined): void {
|
|
797
|
+
this.maxOutputTokens = limit;
|
|
798
|
+
}
|
|
799
|
+
|
|
800
|
+
/** Number of sessions currently cached. Primarily for tests and diagnostics. Always 0 or 1. */
|
|
801
|
+
get size(): number {
|
|
802
|
+
return this.entries.size;
|
|
803
|
+
}
|
|
804
|
+
|
|
805
|
+
/**
|
|
806
|
+
* Look up or allocate a session for the given previous response id.
|
|
807
|
+
* Always returns a `SessionLookupResult` and always leaves the cache
|
|
808
|
+
* empty after return (single-warm invariant).
|
|
809
|
+
*
|
|
810
|
+
* Lookup proceeds in two tiers:
|
|
811
|
+
*
|
|
812
|
+
* 1. **Tier 1 — `previousResponseId`.** The existing hot path:
|
|
813
|
+
* exact id match on a live, non-expired entry whose stored
|
|
814
|
+
* `instructions` and cache-salt fingerprint match the request. On
|
|
815
|
+
* a match the entry is leased out (single-use: removed from the
|
|
816
|
+
* map so a concurrent second request cannot share the live
|
|
817
|
+
* `ChatSession`). On a miss — unknown id, expired, instructions
|
|
818
|
+
* drift, or cache-salt drift — the method falls through to a FRESH
|
|
819
|
+
* session regardless of whether tier 2 would have hit.
|
|
820
|
+
*
|
|
821
|
+
* `previousResponseId` wins unconditionally when supplied. The
|
|
822
|
+
* two keys could legitimately identify different conversation
|
|
823
|
+
* branches (e.g. a client fork where one arm chose the prev-id
|
|
824
|
+
* path and the other arm chose to set `prompt_cache_key`
|
|
825
|
+
* without one), so routing the prev-id branch through tier 2
|
|
826
|
+
* on miss risks splicing the wrong warm state into the wrong
|
|
827
|
+
* chain. Cold-replay is the safe default.
|
|
828
|
+
*
|
|
829
|
+
* 2. **Tier 2 — `promptCacheKey`.** Only runs when
|
|
830
|
+
* `previousResponseId` is `null`. Stateless agent clients
|
|
831
|
+
* (pi-mono, Aider, Codex CLI, Continue, etc.) never use
|
|
832
|
+
* `previous_response_id` — they own the conversation history
|
|
833
|
+
* client-side and resend the full transcript on every turn —
|
|
834
|
+
* so the only way to reuse a warm session across those turns
|
|
835
|
+
* is to key on the client-supplied `prompt_cache_key`. Scans
|
|
836
|
+
* for any live, non-expired entry whose stored
|
|
837
|
+
* `promptCacheKey` is non-null AND byte-equal to the caller's
|
|
838
|
+
* `promptCacheKey`, whose stored `instructions` are byte-equal,
|
|
839
|
+
* AND whose cache-salt fingerprint matches. Empty string is treated
|
|
840
|
+
* as a distinct key from `null` — an opt-out sentinel from a client that forgot to
|
|
841
|
+
* thread the key must NOT collide with another client that
|
|
842
|
+
* did set it to empty. On a match the entry is leased out
|
|
843
|
+
* (same single-use semantics as tier 1). On a miss, fall
|
|
844
|
+
* through to a fresh session.
|
|
845
|
+
*
|
|
846
|
+
* The `hit` flag drives the `X-Session-Cache` observability header
|
|
847
|
+
* emitted by both `/v1/responses` and `/v1/messages`: when the caller
|
|
848
|
+
* supplied a `previous_response_id`, `hit === true` yields `hit` and
|
|
849
|
+
* `hit === false` yields `cold_replay` (the endpoint then rebuilds
|
|
850
|
+
* from the `ResponseStore` on a fresh session). Requests with no
|
|
851
|
+
* `previous_response_id` yield either `fresh` (tier-2 miss) or
|
|
852
|
+
* `prefix_hit` (tier-2 hit — only classified as such once the
|
|
853
|
+
* native `cachedTokens > 0` confirms the prefix-cache machinery
|
|
854
|
+
* actually reused the cached tokens).
|
|
855
|
+
*/
|
|
856
|
+
getOrCreate(
|
|
857
|
+
previousResponseId: string | null,
|
|
858
|
+
requestedInstructions: string | null,
|
|
859
|
+
promptCacheKey: string | null = null,
|
|
860
|
+
requestedCacheSalt: string | null = null,
|
|
861
|
+
): SessionLookupResult {
|
|
862
|
+
const requestedCacheSaltFingerprint = fingerprintCacheSalt(requestedCacheSalt);
|
|
863
|
+
// Tier 1: previousResponseId exact match.
|
|
864
|
+
//
|
|
865
|
+
// Every call is about to overwrite native KV state, so drop any
|
|
866
|
+
// other cached entry now — a later `getOrCreate` must not hand
|
|
867
|
+
// out a wrapper whose assumed state has been stomped. Under the
|
|
868
|
+
// single-warm invariant the map holds at most one entry, so the
|
|
869
|
+
// common case is either "the entry we want" or "nothing".
|
|
870
|
+
if (previousResponseId !== null) {
|
|
871
|
+
const entry = this.entries.get(previousResponseId);
|
|
872
|
+
if (entry === undefined) {
|
|
873
|
+
this.evictEntriesExcept();
|
|
874
|
+
return { session: this.newSession(), hit: false };
|
|
875
|
+
}
|
|
876
|
+
if (entry.expiresAt < nowSec()) {
|
|
877
|
+
this.evictEntriesExcept();
|
|
878
|
+
return { session: this.newSession(), hit: false };
|
|
879
|
+
}
|
|
880
|
+
// Prefix-state mismatch forces cold replay so new instructions are
|
|
881
|
+
// re-primed; cache-salt mismatch below likewise prevents a live native
|
|
882
|
+
// request from crossing prefix-cache security domains.
|
|
883
|
+
if (entry.instructions !== requestedInstructions) {
|
|
884
|
+
this.evictEntriesExcept();
|
|
885
|
+
return { session: this.newSession(), hit: false };
|
|
886
|
+
}
|
|
887
|
+
if (entry.cacheSaltFingerprint !== requestedCacheSaltFingerprint) {
|
|
888
|
+
this.evictEntriesExcept();
|
|
889
|
+
return { session: this.newSession(), hit: false };
|
|
890
|
+
}
|
|
891
|
+
// Tier-1 hit: clear and hand the session out as a single-use
|
|
892
|
+
// lease so a concurrent second request against the same id
|
|
893
|
+
// cold-replays instead of sharing this live ChatSession. Note
|
|
894
|
+
// that even on a prev-id tier-1 MISS we do NOT fall through to
|
|
895
|
+
// tier 2 — see the docstring above for the precedence
|
|
896
|
+
// rationale.
|
|
897
|
+
this.evictEntriesExcept(entry.session);
|
|
898
|
+
return { session: entry.session, hit: true };
|
|
899
|
+
}
|
|
900
|
+
|
|
901
|
+
// Tier 2: promptCacheKey scan (only reached when previousResponseId is null).
|
|
902
|
+
//
|
|
903
|
+
// The registry holds at most one entry under the single-warm
|
|
904
|
+
// invariant, so the "scan" is actually a single lookup — walk the
|
|
905
|
+
// map, check the one entry if present, hit or miss. A non-null
|
|
906
|
+
// scoped key on both the request and the entry plus byte-equal
|
|
907
|
+
// instructions and cache-salt fingerprints is the match condition.
|
|
908
|
+
//
|
|
909
|
+
// SECURITY: raw caller-supplied keys never touch the map. They
|
|
910
|
+
// are run through {@link scopePromptCacheKey}, which (a) returns
|
|
911
|
+
// `null` when the tier-2 opt-in env var is unset, (b) enforces a
|
|
912
|
+
// minimum length, and (c) HMACs the key with a boot-time nonce
|
|
913
|
+
// held only in this process's memory. Without the opt-in every
|
|
914
|
+
// tier-2 lookup below immediately misses; with the opt-in,
|
|
915
|
+
// attackers who cannot read the process-local nonce cannot craft
|
|
916
|
+
// a lookup that matches a stored entry by guessing the raw key.
|
|
917
|
+
const scopedKey = scopePromptCacheKey(promptCacheKey);
|
|
918
|
+
if (scopedKey !== null) {
|
|
919
|
+
for (const entry of this.entries.values()) {
|
|
920
|
+
if (entry.expiresAt < nowSec()) continue;
|
|
921
|
+
if (entry.promptCacheKey === null) continue;
|
|
922
|
+
if (entry.promptCacheKey !== scopedKey) continue;
|
|
923
|
+
if (entry.instructions !== requestedInstructions) continue;
|
|
924
|
+
if (entry.cacheSaltFingerprint !== requestedCacheSaltFingerprint) continue;
|
|
925
|
+
// Tier-2 hit: clear and lease (same single-warm / single-use
|
|
926
|
+
// semantics as tier 1).
|
|
927
|
+
this.evictEntriesExcept(entry.session);
|
|
928
|
+
return { session: entry.session, hit: true };
|
|
929
|
+
}
|
|
930
|
+
}
|
|
931
|
+
|
|
932
|
+
// Fall through: fresh session. Clear any leftover entry so a
|
|
933
|
+
// later lookup cannot hand out a wrapper whose assumed state has
|
|
934
|
+
// been overwritten by this dispatch.
|
|
935
|
+
this.evictEntriesExcept();
|
|
936
|
+
return { session: this.newSession(), hit: false };
|
|
937
|
+
}
|
|
938
|
+
|
|
939
|
+
/**
|
|
940
|
+
* Allocate a fresh `ChatSession` bound to this registry's model
|
|
941
|
+
* without touching the warm slot. Intended for the `/v1/messages`
|
|
942
|
+
* endpoint when the underlying model has a block-paged KV cache
|
|
943
|
+
* active: the native cache already reuses SYS blocks across requests
|
|
944
|
+
* via content-addressing in `BlockAllocator`'s prefix-hash table, so
|
|
945
|
+
* the JS-side warm slot in
|
|
946
|
+
* {@link SessionRegistry.getOrCreateWarmAny} is redundant.
|
|
947
|
+
*
|
|
948
|
+
* Crucially, this call is purely additive — it does **NOT** clear,
|
|
949
|
+
* read, or evict the warm slot. Two parallel `/v1/messages` requests
|
|
950
|
+
* sharing a system prompt both call `createFreshSession` and both
|
|
951
|
+
* get distinct sessions; the native cache transparently refcounts
|
|
952
|
+
* the shared SYS blocks across them. This is the routing decision
|
|
953
|
+
* the long block comment in `packages/server/src/endpoints/messages.ts`
|
|
954
|
+
* documents: paged → fresh session, non-paged → warm-any lookup.
|
|
955
|
+
*
|
|
956
|
+
* The returned session is pre-seeded with the operator-configured
|
|
957
|
+
* `samplingDefaults` (matching every other cache-miss branch) so a
|
|
958
|
+
* client that picks the paged path does not silently stray from the
|
|
959
|
+
* server's pinned sampling knobs.
|
|
960
|
+
*
|
|
961
|
+
* Returned with `hit: false` to keep the result shape uniform with
|
|
962
|
+
* {@link SessionRegistry.getOrCreate} and
|
|
963
|
+
* {@link SessionRegistry.getOrCreateWarmAny}; callers that care
|
|
964
|
+
* about the cache header semantics should observe
|
|
965
|
+
* `result.cachedTokens` from the dispatch instead — that's the
|
|
966
|
+
* authoritative signal for whether the native engine recovered any
|
|
967
|
+
* prefix on this turn (paged or otherwise).
|
|
968
|
+
*/
|
|
969
|
+
createFreshSession(): SessionLookupResult {
|
|
970
|
+
return { session: this.newSession(), hit: false };
|
|
971
|
+
}
|
|
972
|
+
|
|
973
|
+
/**
|
|
974
|
+
* @deprecated **Redundant on `/v1/messages` for paged-active models.**
|
|
975
|
+
* There is no call site for paged-active full-attention models (Qwen3 +
|
|
976
|
+
* LFM2 + Gemma4 today): the native block-paged KV adapter (`PagedKVCacheAdapter` +
|
|
977
|
+
* `BlockAllocator` + `LayerKVPool`) recovers a turn's prefix from
|
|
978
|
+
* refcounted KV blocks keyed by token-prefix hash, so the JS-side
|
|
979
|
+
* single-warm slot this method walks is redundant — the native
|
|
980
|
+
* cache picks up the same cross-turn reuse without the
|
|
981
|
+
* byte-equal-`instructions` gate, and additionally supports
|
|
982
|
+
* cross-conversation prefix sharing the warm slot cannot.
|
|
983
|
+
*
|
|
984
|
+
* The `/v1/messages` endpoint now branches at request time on
|
|
985
|
+
* {@link SessionCapableModel.hasBlockPagedCache}: paged-active models
|
|
986
|
+
* call {@link SessionRegistry.createFreshSession} per request and
|
|
987
|
+
* never touch the warm slot; non-paged models (Qwen3.5 dense + MoE —
|
|
988
|
+
* default-OFF pending a perf decision; the `QianfanOCRModel` VLM —
|
|
989
|
+
* no adapter wired) still call this method because the JS-side warm
|
|
990
|
+
* slot is the ONLY cross-conversation reuse mechanism available to
|
|
991
|
+
* them. Removing this method would silently disable cross-turn reuse
|
|
992
|
+
* on every non-paged model, so it stays load-bearing until ALL
|
|
993
|
+
* session-capable models have paged enabled by default. Treat
|
|
994
|
+
* `@deprecated` as an intent signal that paged-active callers should
|
|
995
|
+
* use `createFreshSession` instead.
|
|
996
|
+
*
|
|
997
|
+
* Third lookup mode — for STATELESS full-history endpoints that have
|
|
998
|
+
* no `previous_response_id` to thread and do not propagate
|
|
999
|
+
* `prompt_cache_key` back to the server. The Anthropic
|
|
1000
|
+
* `/v1/messages` endpoint is the canonical caller: clients (e.g.
|
|
1001
|
+
* Claude Code) POST the entire conversation each turn, so the only
|
|
1002
|
+
* remaining signal that a turn N continues turn N-1's prefix is the
|
|
1003
|
+
* registry's own warm slot.
|
|
1004
|
+
*
|
|
1005
|
+
* Behaviour: walk the registry's at-most-one warm entry. If it is
|
|
1006
|
+
* non-expired, its stored `instructions` are byte-equal to
|
|
1007
|
+
* `requestedInstructions`, AND its stored cache salt equals
|
|
1008
|
+
* `requestedCacheSalt`, lease it out (single-use — `entries.clear()`
|
|
1009
|
+
* before return, mirroring the tier-1 / tier-2 lease-on-hit semantics).
|
|
1010
|
+
* Otherwise clear the map and return a fresh session.
|
|
1011
|
+
*
|
|
1012
|
+
* Crucially, this lookup IGNORES `entry.promptCacheKey` and ignores
|
|
1013
|
+
* the entry's prior `previousResponseId` keying — any warm slot is
|
|
1014
|
+
* fair game for `/v1/messages` reuse. Byte-equal instructions and cache
|
|
1015
|
+
* salt are the correctness gates: a system prompt or prefix-cache
|
|
1016
|
+
* security-domain change forces cold replay instead of reusing stale state
|
|
1017
|
+
* or asking the native adapter to mutate a live request's salt.
|
|
1018
|
+
*
|
|
1019
|
+
* **Adoption sentinel.** `/v1/messages` adopts back under the literal
|
|
1020
|
+
* sentinel id `'__msg_warm__'`. That sentinel will never appear as a
|
|
1021
|
+
* `previous_response_id` on a `/v1/responses` request — the
|
|
1022
|
+
* Anthropic Messages API does not produce a `previous_response_id`
|
|
1023
|
+
* value clients could echo back, and the OpenAI side mints fresh
|
|
1024
|
+
* `resp_*` ids — so cross-endpoint capture via tier-1 is impossible
|
|
1025
|
+
* by construction. The two endpoints still SHARE the single warm
|
|
1026
|
+
* slot under the registry's single-warm invariant: a
|
|
1027
|
+
* `/v1/messages` turn that follows a `/v1/responses` turn can evict
|
|
1028
|
+
* (and vice versa). That is the explicit trade-off of holding at
|
|
1029
|
+
* most one warm entry per model.
|
|
1030
|
+
*
|
|
1031
|
+
* **Trust model.** Multi-tenant isolation on this endpoint requires
|
|
1032
|
+
* fronting the server with an auth proxy that scopes warm-slot
|
|
1033
|
+
* visibility per tenant — same trust boundary documented at the top
|
|
1034
|
+
* of this file for the tier-2 `prompt_cache_key` path. The single-
|
|
1035
|
+
* warm invariant plus `withExclusive`'s per-model serialization make
|
|
1036
|
+
* the lookup safe under SINGLE-tenant assumptions: no two requests
|
|
1037
|
+
* race the slot, and there is at most one slot to lease.
|
|
1038
|
+
*
|
|
1039
|
+
* **Caller contract on miss.** If `instructions` drifts between
|
|
1040
|
+
* turns (system prompt changed) this returns `hit: false` and a
|
|
1041
|
+
* fresh session — and the caller MUST then run a full
|
|
1042
|
+
* `session.reset()` before priming history, NOT the JS-only
|
|
1043
|
+
* `resetPreservingNativeCacheForWarmReuse` path. A fresh JS session
|
|
1044
|
+
* does NOT imply a fresh native cache (the underlying
|
|
1045
|
+
* `SessionCapableModel` is shared and its native
|
|
1046
|
+
* `cached_token_history` persists across requests), so skipping the
|
|
1047
|
+
* native wipe on a miss would let the next `chatSessionStart` reuse
|
|
1048
|
+
* an unrelated previous request's prefix — the cross-request
|
|
1049
|
+
* cache-affinity side channel that the long block comment in
|
|
1050
|
+
* `responses.ts` (around the `runSessionNonStreaming` /
|
|
1051
|
+
* `runSessionStreaming` branches) describes.
|
|
1052
|
+
*/
|
|
1053
|
+
getOrCreateWarmAny(
|
|
1054
|
+
requestedInstructions: string | null,
|
|
1055
|
+
requestedCacheSalt: string | null = null,
|
|
1056
|
+
): SessionLookupResult {
|
|
1057
|
+
const requestedCacheSaltFingerprint = fingerprintCacheSalt(requestedCacheSalt);
|
|
1058
|
+
// Single-warm invariant: at most one entry. Walk it once, lease
|
|
1059
|
+
// on a fresh + instructions-and-salt-matched hit, otherwise clear and
|
|
1060
|
+
// cold-start. The ignored fields (promptCacheKey,
|
|
1061
|
+
// previousResponseId-keying) are deliberate — see the docstring.
|
|
1062
|
+
for (const entry of this.entries.values()) {
|
|
1063
|
+
if (entry.expiresAt < nowSec()) continue;
|
|
1064
|
+
if (entry.instructions !== requestedInstructions) continue;
|
|
1065
|
+
if (entry.cacheSaltFingerprint !== requestedCacheSaltFingerprint) continue;
|
|
1066
|
+
// Hit: clear and lease (single-use semantics, same as tiers 1/2).
|
|
1067
|
+
this.evictEntriesExcept(entry.session);
|
|
1068
|
+
return { session: entry.session, hit: true };
|
|
1069
|
+
}
|
|
1070
|
+
// Miss (no entry, expired, instructions drift, or salt drift). Clear the map
|
|
1071
|
+
// so a stale wrapper cannot leak into a later lookup, and return
|
|
1072
|
+
// a fresh session.
|
|
1073
|
+
this.evictEntriesExcept();
|
|
1074
|
+
return { session: this.newSession(), hit: false };
|
|
1075
|
+
}
|
|
1076
|
+
|
|
1077
|
+
/**
|
|
1078
|
+
* Insert a session under a newly allocated response id. Clears the
|
|
1079
|
+
* map before inserting to keep the single-warm invariant explicit
|
|
1080
|
+
* regardless of caller ordering.
|
|
1081
|
+
*
|
|
1082
|
+
* `instructions` is the prefix/system state used for this turn;
|
|
1083
|
+
* stored on the entry and compared on the next `getOrCreate` to
|
|
1084
|
+
* detect prefix changes that must force a cold replay.
|
|
1085
|
+
* `cacheSalt` is HMAC-fingerprinted before storage and compared alongside
|
|
1086
|
+
* the prefix so the raw security-domain key is not retained.
|
|
1087
|
+
*
|
|
1088
|
+
* `promptCacheKey` is the client-supplied conversation-chain key
|
|
1089
|
+
* that enables the registry's tier-2 lookup for stateless agent
|
|
1090
|
+
* turns that do not carry a `previous_response_id`. `null` /
|
|
1091
|
+
* `undefined` means "no key supplied" — stored verbatim so a
|
|
1092
|
+
* subsequent stateless lookup that also omits the key does NOT
|
|
1093
|
+
* accidentally pick up this entry (only explicit non-null
|
|
1094
|
+
* key-equality on both sides can hit tier 2). See
|
|
1095
|
+
* {@link SessionRegistry.getOrCreate} for the precedence rules.
|
|
1096
|
+
*/
|
|
1097
|
+
adopt(
|
|
1098
|
+
responseId: string,
|
|
1099
|
+
session: ChatSession<SessionCapableModel>,
|
|
1100
|
+
instructions: string | null,
|
|
1101
|
+
promptCacheKey: string | null | undefined = null,
|
|
1102
|
+
cacheSalt: string | null | undefined = null,
|
|
1103
|
+
): void {
|
|
1104
|
+
// Scope the caller-supplied key BEFORE storing so a later
|
|
1105
|
+
// `getOrCreate` can only resolve entries via the same opt-in +
|
|
1106
|
+
// HMAC path. When tier-2 reuse is disabled (or the key is too
|
|
1107
|
+
// short / absent) `scopePromptCacheKey` returns `null`, which
|
|
1108
|
+
// disables this entry from ever matching a tier-2 lookup — the
|
|
1109
|
+
// raw caller-supplied key is NEVER stored.
|
|
1110
|
+
this.evictEntriesExcept(session);
|
|
1111
|
+
this.entries.set(responseId, {
|
|
1112
|
+
session,
|
|
1113
|
+
instructions,
|
|
1114
|
+
cacheSaltFingerprint: fingerprintCacheSalt(cacheSalt),
|
|
1115
|
+
promptCacheKey: scopePromptCacheKey(promptCacheKey ?? null),
|
|
1116
|
+
expiresAt: nowSec() + this.ttlSec,
|
|
1117
|
+
});
|
|
1118
|
+
}
|
|
1119
|
+
|
|
1120
|
+
/**
|
|
1121
|
+
* Remove a session by response id. No-op if the key is not present.
|
|
1122
|
+
*/
|
|
1123
|
+
drop(responseId: string): void {
|
|
1124
|
+
const entry = this.entries.get(responseId);
|
|
1125
|
+
if (entry === undefined) return;
|
|
1126
|
+
this.entries.delete(responseId);
|
|
1127
|
+
this.scheduleDispose(entry.session);
|
|
1128
|
+
}
|
|
1129
|
+
|
|
1130
|
+
/**
|
|
1131
|
+
* Walk the map and drop the entry if its TTL has expired.
|
|
1132
|
+
* Intended for periodic cleanup via `setInterval`. Under the
|
|
1133
|
+
* single-warm invariant the map holds at most one entry.
|
|
1134
|
+
*/
|
|
1135
|
+
sweep(): void {
|
|
1136
|
+
const cutoff = nowSec();
|
|
1137
|
+
for (const [key, entry] of this.entries) {
|
|
1138
|
+
if (entry.expiresAt < cutoff) {
|
|
1139
|
+
this.entries.delete(key);
|
|
1140
|
+
this.scheduleDispose(entry.session);
|
|
1141
|
+
}
|
|
1142
|
+
}
|
|
1143
|
+
}
|
|
1144
|
+
|
|
1145
|
+
/** Empty the registry. Useful at shutdown and in tests. */
|
|
1146
|
+
clear(): void {
|
|
1147
|
+
this.evictEntriesExcept();
|
|
1148
|
+
}
|
|
1149
|
+
|
|
1150
|
+
/**
|
|
1151
|
+
* Serialize `fn` against every other dispatch through this
|
|
1152
|
+
* registry's model. The caller must hold the lock across the
|
|
1153
|
+
* entire per-model dispatch span — `getOrCreate` ->
|
|
1154
|
+
* `primeHistory`/`send*` -> `adopt`/`drop`. Without it, two
|
|
1155
|
+
* concurrent `primeHistory()` / `send*()` calls would race on
|
|
1156
|
+
* the single mutable native KV cache and whichever finished last
|
|
1157
|
+
* would corrupt the other's chain.
|
|
1158
|
+
*
|
|
1159
|
+
* FIFO chaining via a rolling `execLock` promise: each caller
|
|
1160
|
+
* captures the current tail, publishes a fresh pending promise as
|
|
1161
|
+
* the new tail, awaits the old tail, then runs `fn`. The
|
|
1162
|
+
* `finally` releases regardless of whether `fn` threw.
|
|
1163
|
+
*
|
|
1164
|
+
* **Admission control.** Every non-handed-off call — waiter AND
|
|
1165
|
+
* runner — routes through {@link assertAdmissionCapacity}, the single
|
|
1166
|
+
* source of truth for the budget: the combined admission footprint
|
|
1167
|
+
* (`queuedCount + preDispatchAdmits`) plus this caller must fit
|
|
1168
|
+
* within `maxQueueDepth` plus the idle-chain runner entitlement.
|
|
1169
|
+
* Over-budget calls throw {@link QueueFullError} — SYNCHRONOUSLY
|
|
1170
|
+
* from the caller's perspective, not merely before `await prev`. The
|
|
1171
|
+
* wrapper is deliberately NOT declared `async` so the admission gate
|
|
1172
|
+
* throws on the caller's stack frame, letting endpoint handlers wrap
|
|
1173
|
+
* the call site in a plain try/catch without racing promise
|
|
1174
|
+
* microtasks. On acceptance the async body takes over via the
|
|
1175
|
+
* returned `Promise<T>`.
|
|
1176
|
+
*
|
|
1177
|
+
* A cap of N permits one running dispatch plus N admitted-but-not-
|
|
1178
|
+
* running requests (queued waiters and outstanding permits
|
|
1179
|
+
* combined), rejecting the next. The default (undefined) preserves
|
|
1180
|
+
* the original unbounded behaviour.
|
|
1181
|
+
*
|
|
1182
|
+
* **Runner-slot admission.** Whether a given caller counts as the
|
|
1183
|
+
* runner slot or as a waiter is decided up front by comparing
|
|
1184
|
+
* `execLock` against the idle sentinel `initialLock`. If they are
|
|
1185
|
+
* identical, nobody is currently in-flight and this caller is
|
|
1186
|
+
* admitted against the extra runner entitlement rather than a waiter
|
|
1187
|
+
* slot — it never touches `queuedCount`. Otherwise it is a waiter
|
|
1188
|
+
* and the normal charge / increment / decrement cycle applies. This
|
|
1189
|
+
* is what keeps a synchronous burst such as `Promise.all([fn, fn])`
|
|
1190
|
+
* admissible under `maxQueueDepth = 1` — Call 1 is the runner,
|
|
1191
|
+
* Call 2 is the one allowed waiter, Call 3 would throw. The runner
|
|
1192
|
+
* path is NOT exempt from the budget: while the chain is idle,
|
|
1193
|
+
* {@link beginPreDispatchAdmission} lends out `maxQueueDepth + 1`
|
|
1194
|
+
* permits precisely because one of them is entitled to become the
|
|
1195
|
+
* runner, so a permitless call that would take that seat while the
|
|
1196
|
+
* whole runner-plus-waiter capacity is already spoken for must
|
|
1197
|
+
* reject — otherwise the outstanding permits would later convert on
|
|
1198
|
+
* top of it and breach the cap.
|
|
1199
|
+
*
|
|
1200
|
+
* **Permit handoff.** A caller that was already admitted by
|
|
1201
|
+
* {@link beginPreDispatchAdmission} passes its permit as the second
|
|
1202
|
+
* argument; the permit is consumed atomically as this call's
|
|
1203
|
+
* admission instead of charging `queuedCount` a second time — one
|
|
1204
|
+
* budget, one token per request. See the handoff comment in the
|
|
1205
|
+
* body and {@link PreDispatchAdmission}.
|
|
1206
|
+
*/
|
|
1207
|
+
withExclusive<T>(fn: () => Promise<T>, permit?: PreDispatchAdmission): Promise<T> {
|
|
1208
|
+
// Distinguish runner-slot from waiter admission. If the chain is
|
|
1209
|
+
// idle (`execLock === initialLock`) the current caller is about
|
|
1210
|
+
// to become the active holder on its very first `await prev`
|
|
1211
|
+
// microtask — it must NOT be billed against the waiter cap and
|
|
1212
|
+
// must NOT touch `queuedCount`. Only chained callers (someone
|
|
1213
|
+
// else still holds or is ahead in the FIFO) count as waiters.
|
|
1214
|
+
const asWaiter = this.execLock !== this.initialLock;
|
|
1215
|
+
|
|
1216
|
+
// Atomic permit handoff (see `beginPreDispatchAdmission`). A caller
|
|
1217
|
+
// handing in a still-outstanding permit from THIS registry already
|
|
1218
|
+
// owns one unit of the shared budget: consuming it here — in the
|
|
1219
|
+
// same synchronous block as the waiter accounting below — converts
|
|
1220
|
+
// that unit into this call's admission (waiter path increments
|
|
1221
|
+
// `queuedCount`, keeping the total constant) or retires it (runner
|
|
1222
|
+
// path). The handoff skips only the caller's OWN token: its budget
|
|
1223
|
+
// share was charged at acquisition and conversion keeps the total
|
|
1224
|
+
// constant, so the handoff itself can never create a breach —
|
|
1225
|
+
// whereas re-checking would double-charge and could reject a
|
|
1226
|
+
// request that was already admitted at the endpoint gate. A permit
|
|
1227
|
+
// minted by a DIFFERENT registry is not in `permitConsumers` and
|
|
1228
|
+
// falls through to normal charging (its own budget stays balanced
|
|
1229
|
+
// by the caller's `finally` release).
|
|
1230
|
+
const consume = permit === undefined ? undefined : this.permitConsumers.get(permit);
|
|
1231
|
+
const handedOff = consume !== undefined && consume();
|
|
1232
|
+
|
|
1233
|
+
// Admission check — raised synchronously so endpoint handlers
|
|
1234
|
+
// can reliably catch `QueueFullError` without racing any `await`.
|
|
1235
|
+
// Every non-handed-off caller — runner AND waiter — runs the same
|
|
1236
|
+
// shared predicate as `beginPreDispatchAdmission`: the combined
|
|
1237
|
+
// footprint (queued waiters PLUS outstanding pre-dispatch permits)
|
|
1238
|
+
// against the cap plus the idle-chain runner entitlement. Checking
|
|
1239
|
+
// `queuedCount` alone would let a permitless caller spend a slot
|
|
1240
|
+
// an outstanding permit already owns (e.g. a continuation parked
|
|
1241
|
+
// in `store.getChain`); exempting the runner path would let a
|
|
1242
|
+
// permitless call take the runner seat the gate already lent to
|
|
1243
|
+
// one of `maxQueueDepth + 1` idle-chain permits — either way the
|
|
1244
|
+
// cap breaches once those permits convert. Nothing is mutated on
|
|
1245
|
+
// the reject path; the request never queued.
|
|
1246
|
+
if (!handedOff) {
|
|
1247
|
+
this.assertAdmissionCapacity();
|
|
1248
|
+
}
|
|
1249
|
+
|
|
1250
|
+
const prev = this.execLock;
|
|
1251
|
+
let release!: () => void;
|
|
1252
|
+
const myLock = new Promise<void>((resolve) => {
|
|
1253
|
+
release = resolve;
|
|
1254
|
+
});
|
|
1255
|
+
this.execLock = myLock;
|
|
1256
|
+
if (asWaiter) {
|
|
1257
|
+
this.queuedCount += 1;
|
|
1258
|
+
}
|
|
1259
|
+
return this._runExclusive(prev, myLock, release, fn, asWaiter);
|
|
1260
|
+
}
|
|
1261
|
+
|
|
1262
|
+
/**
|
|
1263
|
+
* Admit one dispatch to the model's continuous-batching lane.
|
|
1264
|
+
*
|
|
1265
|
+
* Admission rejection is synchronous, matching {@link withExclusive} and
|
|
1266
|
+
* preserving the endpoint's existing QueueFullError-to-429 mapping. Once
|
|
1267
|
+
* accepted, at most {@link concurrentAdmissionLimit} closures run at once;
|
|
1268
|
+
* excess accepted callers wait FIFO and contribute to {@link queueDepth}.
|
|
1269
|
+
* A pre-dispatch permit is consumed atomically into either an active slot or
|
|
1270
|
+
* a queued slot, so the early endpoint gate and this semaphore share one
|
|
1271
|
+
* bounded budget.
|
|
1272
|
+
*/
|
|
1273
|
+
withAdmission<T>(fn: () => Promise<T>, permit?: PreDispatchAdmission): Promise<T> {
|
|
1274
|
+
const consume = permit === undefined ? undefined : this.permitConsumers.get(permit);
|
|
1275
|
+
const handedOff = consume !== undefined && consume();
|
|
1276
|
+
if (!handedOff) {
|
|
1277
|
+
this.assertAdmissionCapacity();
|
|
1278
|
+
}
|
|
1279
|
+
|
|
1280
|
+
let acquire!: () => void;
|
|
1281
|
+
const acquired = new Promise<void>((resolve) => {
|
|
1282
|
+
acquire = resolve;
|
|
1283
|
+
});
|
|
1284
|
+
if (this.activeAdmissions < this.maxConcurrentDispatches) {
|
|
1285
|
+
this.activeAdmissions += 1;
|
|
1286
|
+
acquire();
|
|
1287
|
+
} else {
|
|
1288
|
+
this.queuedCount += 1;
|
|
1289
|
+
this.admissionWaiters.push(() => {
|
|
1290
|
+
this.queuedCount -= 1;
|
|
1291
|
+
if (this.queuedCount < 0) this.queuedCount = 0;
|
|
1292
|
+
this.activeAdmissions += 1;
|
|
1293
|
+
acquire();
|
|
1294
|
+
});
|
|
1295
|
+
}
|
|
1296
|
+
return this._runAdmission(acquired, fn);
|
|
1297
|
+
}
|
|
1298
|
+
|
|
1299
|
+
private async _runAdmission<T>(acquired: Promise<void>, fn: () => Promise<T>): Promise<T> {
|
|
1300
|
+
await acquired;
|
|
1301
|
+
try {
|
|
1302
|
+
return await fn();
|
|
1303
|
+
} finally {
|
|
1304
|
+
this.activeAdmissions -= 1;
|
|
1305
|
+
if (this.activeAdmissions < 0) this.activeAdmissions = 0;
|
|
1306
|
+
const next = this.admissionWaiters.shift();
|
|
1307
|
+
next?.();
|
|
1308
|
+
}
|
|
1309
|
+
}
|
|
1310
|
+
|
|
1311
|
+
/**
|
|
1312
|
+
* Async tail of {@link withExclusive}. Kept separate so the public
|
|
1313
|
+
* wrapper stays a plain (non-async) function whose admission-gate
|
|
1314
|
+
* throw lands on the caller's stack synchronously. This helper
|
|
1315
|
+
* owns the post-acceptance bookkeeping: awaiting the predecessor
|
|
1316
|
+
* lock, transitioning from waiter to holder (`queuedCount`
|
|
1317
|
+
* decrement for waiters only), running `fn`, releasing the FIFO
|
|
1318
|
+
* tail, and resetting the chain to `initialLock` when this caller
|
|
1319
|
+
* is still the tail (so a future burst admits its first entry as
|
|
1320
|
+
* a runner-slot rather than as a waiter).
|
|
1321
|
+
*/
|
|
1322
|
+
private async _runExclusive<T>(
|
|
1323
|
+
prev: Promise<void>,
|
|
1324
|
+
myLock: Promise<void>,
|
|
1325
|
+
release: () => void,
|
|
1326
|
+
fn: () => Promise<T>,
|
|
1327
|
+
asWaiter: boolean,
|
|
1328
|
+
): Promise<T> {
|
|
1329
|
+
// Track whether the waiting-counter has already been balanced so
|
|
1330
|
+
// an error raised by `await prev` (should never happen today but
|
|
1331
|
+
// is cheap to defend against) cannot double-decrement via the
|
|
1332
|
+
// outer `finally` below. Runner-slot admissions never touch the
|
|
1333
|
+
// counter, so the flag starts already-balanced for them.
|
|
1334
|
+
let waitingDecremented = !asWaiter;
|
|
1335
|
+
try {
|
|
1336
|
+
try {
|
|
1337
|
+
await prev;
|
|
1338
|
+
} finally {
|
|
1339
|
+
// Transition from "waiting" to "running" — the counter must
|
|
1340
|
+
// drop exactly here regardless of whether `prev` fulfilled
|
|
1341
|
+
// or rejected, because from this point forward the caller is
|
|
1342
|
+
// the active holder and no longer part of the queue depth.
|
|
1343
|
+
if (!waitingDecremented) {
|
|
1344
|
+
this.queuedCount -= 1;
|
|
1345
|
+
if (this.queuedCount < 0) this.queuedCount = 0;
|
|
1346
|
+
waitingDecremented = true;
|
|
1347
|
+
}
|
|
1348
|
+
}
|
|
1349
|
+
return await fn();
|
|
1350
|
+
} finally {
|
|
1351
|
+
// Belt-and-suspenders: if `await prev` managed to throw before
|
|
1352
|
+
// reaching the inner `finally` (extremely unlikely given the
|
|
1353
|
+
// chain is always resolved with `undefined`), still balance the
|
|
1354
|
+
// queued counter so a future cap check doesn't drift upward.
|
|
1355
|
+
if (!waitingDecremented) {
|
|
1356
|
+
this.queuedCount -= 1;
|
|
1357
|
+
if (this.queuedCount < 0) this.queuedCount = 0;
|
|
1358
|
+
waitingDecremented = true;
|
|
1359
|
+
}
|
|
1360
|
+
release();
|
|
1361
|
+
// Reset the chain to the idle sentinel ONLY when this caller
|
|
1362
|
+
// is still the tail — if someone else has already extended the
|
|
1363
|
+
// FIFO behind us, leave their tail in place. Reference-equality
|
|
1364
|
+
// gate here is what lets the next burst see `execLock ===
|
|
1365
|
+
// initialLock` and admit its first caller as a runner slot.
|
|
1366
|
+
if (this.execLock === myLock) {
|
|
1367
|
+
this.execLock = this.initialLock;
|
|
1368
|
+
}
|
|
1369
|
+
}
|
|
1370
|
+
}
|
|
1371
|
+
}
|