@mlx-node/server 0.0.7 → 0.0.8

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (74) hide show
  1. package/dist/chat-session-warm-reuse.d.ts +51 -0
  2. package/dist/chat-session-warm-reuse.d.ts.map +1 -0
  3. package/dist/chat-session-warm-reuse.js +68 -0
  4. package/dist/endpoints/messages-count-tokens.d.ts +8 -0
  5. package/dist/endpoints/messages-count-tokens.d.ts.map +1 -0
  6. package/dist/endpoints/messages-count-tokens.js +121 -0
  7. package/dist/endpoints/messages.d.ts +57 -5
  8. package/dist/endpoints/messages.d.ts.map +1 -1
  9. package/dist/endpoints/messages.js +1043 -147
  10. package/dist/endpoints/models.d.ts +2 -1
  11. package/dist/endpoints/models.d.ts.map +1 -1
  12. package/dist/endpoints/models.js +2 -2
  13. package/dist/endpoints/responses.d.ts +20 -7
  14. package/dist/endpoints/responses.d.ts.map +1 -1
  15. package/dist/endpoints/responses.js +572 -82
  16. package/dist/errors.d.ts +1 -0
  17. package/dist/errors.d.ts.map +1 -1
  18. package/dist/errors.js +3 -0
  19. package/dist/handler.d.ts +42 -0
  20. package/dist/handler.d.ts.map +1 -1
  21. package/dist/handler.js +6 -1
  22. package/dist/idle-sweeper.d.ts +245 -0
  23. package/dist/idle-sweeper.d.ts.map +1 -0
  24. package/dist/idle-sweeper.js +408 -0
  25. package/dist/index.d.ts +8 -2
  26. package/dist/index.d.ts.map +1 -1
  27. package/dist/index.js +10 -0
  28. package/dist/mappers/anthropic-request.d.ts +24 -2
  29. package/dist/mappers/anthropic-request.d.ts.map +1 -1
  30. package/dist/mappers/anthropic-request.js +222 -24
  31. package/dist/mappers/anthropic-response.d.ts +29 -4
  32. package/dist/mappers/anthropic-response.d.ts.map +1 -1
  33. package/dist/mappers/anthropic-response.js +143 -21
  34. package/dist/mappers/request.d.ts +48 -0
  35. package/dist/mappers/request.d.ts.map +1 -1
  36. package/dist/mappers/request.js +211 -35
  37. package/dist/mappers/response.d.ts.map +1 -1
  38. package/dist/mappers/response.js +13 -1
  39. package/dist/model-work-coordinator.d.ts +70 -0
  40. package/dist/model-work-coordinator.d.ts.map +1 -0
  41. package/dist/model-work-coordinator.js +120 -0
  42. package/dist/pending-writes.d.ts.map +1 -1
  43. package/dist/presets.d.ts +82 -0
  44. package/dist/presets.d.ts.map +1 -0
  45. package/dist/presets.js +98 -0
  46. package/dist/registry.d.ts +31 -1
  47. package/dist/registry.d.ts.map +1 -1
  48. package/dist/registry.js +33 -5
  49. package/dist/router.d.ts +4 -1
  50. package/dist/router.d.ts.map +1 -1
  51. package/dist/router.js +34 -4
  52. package/dist/server.d.ts +76 -0
  53. package/dist/server.d.ts.map +1 -1
  54. package/dist/server.js +48 -1
  55. package/dist/session-registry.d.ts +272 -18
  56. package/dist/session-registry.d.ts.map +1 -1
  57. package/dist/session-registry.js +509 -37
  58. package/dist/stop-sequence-buffer.d.ts +58 -0
  59. package/dist/stop-sequence-buffer.d.ts.map +1 -0
  60. package/dist/stop-sequence-buffer.js +148 -0
  61. package/dist/text-recovery.d.ts +35 -0
  62. package/dist/text-recovery.d.ts.map +1 -0
  63. package/dist/text-recovery.js +41 -0
  64. package/dist/timing.d.ts +80 -0
  65. package/dist/timing.d.ts.map +1 -0
  66. package/dist/timing.js +121 -0
  67. package/dist/tool-call-buffer.d.ts +5 -5
  68. package/dist/tool-call-buffer.d.ts.map +1 -1
  69. package/dist/tool-call-buffer.js +28 -8
  70. package/dist/types-anthropic.d.ts +161 -1
  71. package/dist/types-anthropic.d.ts.map +1 -1
  72. package/dist/types.d.ts +172 -2
  73. package/dist/types.d.ts.map +1 -1
  74. package/package.json +3 -3
@@ -8,12 +8,16 @@
8
8
  * `primeHistory` + `startFromHistory[Stream]`.
9
9
  */
10
10
  import { randomUUID } from 'node:crypto';
11
+ import { isContextCapacityError } from '@mlx-node/lm';
12
+ import { resetPreservingNativeCacheForWarmReuse } from '../chat-session-warm-reuse.js';
11
13
  import { sendBadRequest, sendInternalError, sendNotFound, sendRateLimit, sendStorageTimeout } from '../errors.js';
12
- import { mapRequest, reconstructMessagesFromChain } from '../mappers/request.js';
14
+ import { mapRequest, reconstructMessagesFromChain, stringifyStoredInputMessages } from '../mappers/request.js';
13
15
  import { buildPartialResponse, buildResponseObject, computeOutputText, genId, mapFinishReasonToStatus, } from '../mappers/response.js';
14
16
  import { getPendingWritesFor } from '../pending-writes.js';
15
- import { QueueFullError } from '../session-registry.js';
17
+ import { maybeWarnPromptCacheKeyIneligible, QueueFullError } from '../session-registry.js';
16
18
  import { beginSSE, endSSE, writeSSEEvent } from '../streaming.js';
19
+ import { longestSuffixPrefixOverlap } from '../text-recovery.js';
20
+ import { mergeTimingUsageExtensions, resolveServerTuningForUsage } from '../timing.js';
17
21
  import { ToolCallTagBuffer } from '../tool-call-buffer.js';
18
22
  import { createVisibility, endJson, flushTerminalSSE, markSSEMode, writeFallbackErrorSSE, } from '../transport-visibility.js';
19
23
  /**
@@ -23,6 +27,19 @@ import { createVisibility, endJson, flushTerminalSSE, markSSEMode, writeFallback
23
27
  * this 30-minute fallback is only used by legacy direct-invocation callers.
24
28
  */
25
29
  const RESPONSE_TTL_SECONDS = 1800;
30
+ /**
31
+ * Upper bound for a client-supplied output-token budget. The native
32
+ * `ChatConfig.max_new_tokens` is `Option<i32>`, and NAPI's
33
+ * `napi_get_value_int32` silently truncates a JS integer above `i32::MAX`
34
+ * to a NEGATIVE value — which the core clamp then turns into 0 (a silent
35
+ * empty completion). Reject anything above this bound at the edge so an
36
+ * over-large budget 400s instead of producing nothing. Shared with
37
+ * `/v1/messages` (`messages.ts`).
38
+ */
39
+ export const MAX_OUTPUT_TOKENS = 2147483647; // i32::MAX — native ChatConfig.max_new_tokens is i32
40
+ function withAdmissionControlledInference(sessionReg, modelWorkCoordinator, fn) {
41
+ return sessionReg.withExclusive(() => (modelWorkCoordinator ? modelWorkCoordinator.withInference(fn) : fn()));
42
+ }
26
43
  /**
27
44
  * Upper bound (ms) on how long the recovery path waits for an in-flight
28
45
  * `store.store(...)` to land. On timeout we re-probe `getChain` once to
@@ -136,8 +153,9 @@ export function getServerBootId() {
136
153
  export function __setServerBootIdForTesting(id) {
137
154
  serverBootId = id;
138
155
  }
139
- async function handleNonStreaming(res, result, req, responseId, previousResponseId, visibility) {
156
+ async function handleNonStreaming(res, result, req, responseId, previousResponseId, visibility, serverTiming) {
140
157
  const response = buildResponseObject(result, req, responseId, previousResponseId);
158
+ mergeTimingUsageExtensions(response.usage, result.performance, result.promptTokens, result.numTokens, result.cachedTokens, serverTiming);
141
159
  // `chatSession*` has no AbortSignal surface yet, so a mid-decode
142
160
  // client disconnect still burns the full decode budget — peer loss
143
161
  // is only observable when native decode resolves. Disconnect
@@ -161,7 +179,7 @@ async function handleNonStreaming(res, result, req, responseId, previousResponse
161
179
  * on a failed envelope cannot see success-shaped items inside it.
162
180
  * `ReasoningOutputItem` has no `status` field and is left alone.
163
181
  */
164
- function buildFailedTerminal(partial, outputItems, reason, usage) {
182
+ function buildFailedTerminal(partial, outputItems, reason, usage, errorMessage) {
165
183
  const normalized = outputItems.map((item) => {
166
184
  if (item.type === 'message') {
167
185
  const prev = item.status;
@@ -178,16 +196,25 @@ function buildFailedTerminal(partial, outputItems, reason, usage) {
178
196
  }
179
197
  return item;
180
198
  });
199
+ // Only the `reason: 'error'` path carries a diagnostic message —
200
+ // client_abort / stream_exhausted / finish_reason_error are caller
201
+ // or finite-state conditions, not server faults worth surfacing.
202
+ const error = errorMessage ? { type: 'server_error', message: errorMessage, code: null, param: null } : null;
181
203
  return {
182
204
  ...partial,
183
205
  status: 'failed',
184
206
  output: normalized,
185
207
  output_text: computeOutputText(normalized),
208
+ error,
186
209
  incomplete_details: { reason },
187
210
  usage,
188
211
  };
189
212
  }
190
- async function handleStreamingNative(res, chatStream, req, responseId, previousResponseId, wasCommitted, httpReq, visibility) {
213
+ async function handleStreamingNative(res, chatStream, req, responseId, previousResponseId, wasCommitted, httpReq, visibility, serverTiming) {
214
+ // `runSessionStreaming` completed the exact token/capacity preflight before
215
+ // handing us this iterator. Commit SSE immediately instead of entering the
216
+ // generator here: its first `next()` also starts image processing/prefill and
217
+ // may not resolve until the first generated token.
191
218
  beginSSE(res);
192
219
  // Commit to SSE wire format synchronously so the outer catch
193
220
  // branches on `responseMode` (not `headersSent`) and routes an
@@ -206,6 +233,13 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
206
233
  let messageText = '';
207
234
  let hasEmittedMessage = false;
208
235
  let hasEmittedReasoning = false;
236
+ // Tracks whether the reasoning output item's `response.output_item.done`
237
+ // has already been emitted. We close it eagerly on the reasoning→text
238
+ // transition (before opening the message item) so OpenAI Responses
239
+ // clients can populate their `thinkingSignature` via the `done`
240
+ // event's `currentBlock?.type === 'thinking'` guard. The terminal and
241
+ // failure paths check this flag to avoid double-emitting.
242
+ let hasClosedReasoning = false;
209
243
  let suppressedMessageIndex = -1;
210
244
  const tagBuffer = new ToolCallTagBuffer();
211
245
  // Terminal response is captured in the done branch but emitted AFTER
@@ -213,6 +247,14 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
213
247
  // `session.turns` once the producer's finally has run.
214
248
  let completedResponse = null;
215
249
  let sawDone = false;
250
+ // Lifted from the final stream event so the outer handler can set
251
+ // `X-Cached-Tokens` and promote the `X-Session-Cache` header to
252
+ // `prefix_hit` when tier-2 reuse actually happened. See
253
+ // `StreamingHandlerOutcome.cachedTokens` for the full rationale.
254
+ // Starts `undefined` — the field is only populated if the native
255
+ // terminal chunk carries `cachedTokens`. Today it never does; a
256
+ // future native plumbing change can lift it through.
257
+ let cachedTokens;
216
258
  // Fault state. `thrownError` sticks on a generator throw;
217
259
  // `clientAborted` sticks on any `close`/`error` from `httpReq`, `res`,
218
260
  // or `res.socket`. Either flips the post-loop block to the failure
@@ -286,18 +328,22 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
286
328
  delta: remainingText,
287
329
  });
288
330
  }
289
- // Close reasoning item if open
290
- if (hasEmittedReasoning && reasoningItemId) {
331
+ // Close reasoning item if still open (already closed eagerly on
332
+ // reasoning→text transition for most turns — this branch covers
333
+ // the reasoning-only shape where no text deltas ever arrived).
334
+ if (hasEmittedReasoning && !hasClosedReasoning && reasoningItemId) {
335
+ hasClosedReasoning = true;
336
+ const finalReasoningText = event.thinking ?? reasoningText;
291
337
  writeSSEEvent(res, 'response.reasoning_summary_text.done', {
292
338
  item_id: reasoningItemId,
293
339
  output_index: outputItems.length - (hasEmittedMessage ? 1 : 0) - 1,
294
340
  summary_index: 0,
295
- text: event.thinking ?? reasoningText,
341
+ text: finalReasoningText,
296
342
  });
297
343
  const reasoningItem = {
298
344
  id: reasoningItemId,
299
345
  type: 'reasoning',
300
- summary: [{ type: 'summary_text', text: event.thinking ?? reasoningText }],
346
+ summary: [{ type: 'summary_text', text: finalReasoningText }],
301
347
  };
302
348
  const riIndex = outputItems.findIndex((i) => i.id === reasoningItemId);
303
349
  if (riIndex >= 0) {
@@ -347,10 +393,38 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
347
393
  delta: finalText,
348
394
  });
349
395
  }
350
- else if (tagBuffer.suppressed && !hasToolCalls && finalText && hasEmittedMessage) {
351
- // Recovery: text was already being streamed but got cut off by a false-alarm
352
- // <tool_call> tag. Emit the unsent portion as a delta.
353
- const unsent = finalText.slice(messageText.length);
396
+ else if (tagBuffer.suppressed &&
397
+ !hasToolCalls &&
398
+ finalText &&
399
+ hasEmittedMessage &&
400
+ !messageText.includes(finalText)) {
401
+ // Recovery: streaming text was cut off by a false-alarm `<tool_call>` tag.
402
+ //
403
+ // The previous `finalText.slice(messageText.length)` is wrong: when the
404
+ // streamed text contains post-</think> whitespace (or any prefix the
405
+ // native side trimmed via `split_at_think_end` / `parse_tool_calls`),
406
+ // `messageText.length` indexes into the streamed buffer while
407
+ // `finalText` starts at the post-trim cleaned position — the two
408
+ // prefixes diverge (e.g. messageText=`"\n\n"`, finalText=`"<tool_call>..."`)
409
+ // and a length-based slice chops `<t` off `<tool_call>`, emitting
410
+ // `"ool_call>\n<function=..."` as visible text.
411
+ //
412
+ // Find the longest streamed-suffix == finalText-prefix overlap and emit
413
+ // whatever finalText has BEYOND that overlap.
414
+ //
415
+ // The `!messageText.includes(finalText)` guard distinguishes:
416
+ // (a) duplicate-trim case: streamed "Let me check. " + closed
417
+ // non-ok tool_call → finalText="Let me check." (trimmed). The
418
+ // trimmed text IS a substring of the streamed text → skip
419
+ // (otherwise we'd duplicate "Let me check.").
420
+ // (b) unclosed-tool case: streamed `\n\n` + unclosed
421
+ // `<tool_call>...` → finalText=`<tool_call>...`. The malformed
422
+ // tag is NOT a substring of the streamed whitespace → emit
423
+ // (this is the original `<t`-strip bug we're fixing).
424
+ // Length-based guards (`finalText.length > messageText.length`)
425
+ // misclassify case (b) when the streamed whitespace is long.
426
+ const overlap = longestSuffixPrefixOverlap(messageText, finalText);
427
+ const unsent = finalText.slice(overlap);
354
428
  if (unsent) {
355
429
  messageText += unsent;
356
430
  writeSSEEvent(res, 'response.output_text.delta', {
@@ -361,16 +435,29 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
361
435
  });
362
436
  }
363
437
  }
364
- // Emit any unsent suffix when final text is longer than what was streamed
365
- if (hasEmittedMessage && finalText && finalText.length > messageText.length && !tagBuffer.suppressed) {
366
- const unsent = finalText.slice(messageText.length);
367
- messageText += unsent;
368
- writeSSEEvent(res, 'response.output_text.delta', {
369
- item_id: messageItemId,
370
- output_index: outputItems.findIndex((i) => i.id === messageItemId),
371
- content_index: 0,
372
- delta: unsent,
373
- });
438
+ // Emit any unsent suffix when final text extends past what was
439
+ // streamed. Same divergence concern as above (post-</think> trim
440
+ // can leave `messageText` longer than the matching prefix of
441
+ // `finalText`), so we use the same overlap-based slice instead of
442
+ // a length-based one. When the overlap covers all of `finalText`
443
+ // (i.e. nothing more to emit) `unsent` is empty and we skip.
444
+ //
445
+ // The `!messageText.includes(finalText)` guard skips the
446
+ // duplicate-trim case where finalText is a substring of the
447
+ // streamed text (e.g. native `.trim()` shrinkage). See the
448
+ // companion comment above for the case-distinction rationale.
449
+ if (hasEmittedMessage && finalText && !tagBuffer.suppressed && !messageText.includes(finalText)) {
450
+ const overlap = longestSuffixPrefixOverlap(messageText, finalText);
451
+ const unsent = finalText.slice(overlap);
452
+ if (unsent) {
453
+ messageText += unsent;
454
+ writeSSEEvent(res, 'response.output_text.delta', {
455
+ item_id: messageItemId,
456
+ output_index: outputItems.findIndex((i) => i.id === messageItemId),
457
+ content_index: 0,
458
+ delta: unsent,
459
+ });
460
+ }
374
461
  }
375
462
  // Recovery: text was never emitted during streaming but final has text
376
463
  // (possible if all text arrived in the final event only)
@@ -493,12 +580,27 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
493
580
  // then the post-loop block handles emission + persistence.
494
581
  const promptTokens = event.promptTokens ?? 0;
495
582
  const reasoningTokens = event.reasoningTokens ?? 0;
583
+ cachedTokens = event.cachedTokens;
496
584
  const usage = {
497
585
  input_tokens: promptTokens,
498
586
  output_tokens: event.numTokens,
499
587
  output_tokens_details: { reasoning_tokens: reasoningTokens },
500
588
  total_tokens: promptTokens + event.numTokens,
501
589
  };
590
+ // Round 5 Fix #3: SSE headers flush before the native prefix
591
+ // verifier has reported cached-token counts, so streaming
592
+ // `X-Session-Cache` is documented as non-authoritative. The
593
+ // authoritative signal for streaming clients is this in-band
594
+ // `usage.input_tokens_details.cached_tokens` field on the
595
+ // terminal `response.completed` event — identical shape to
596
+ // the upstream OpenAI Responses API. Populated only when the
597
+ // native dispatch reports a non-zero reuse count so consumers
598
+ // cheaply distinguish "feature not active" from "active with
599
+ // zero reuse".
600
+ if (cachedTokens != null && cachedTokens > 0) {
601
+ usage.input_tokens_details = { cached_tokens: cachedTokens };
602
+ }
603
+ mergeTimingUsageExtensions(usage, event.performance, promptTokens, event.numTokens, cachedTokens, serverTiming);
502
604
  const finalOutput = outputItems.filter((_, idx) => idx !== suppressedMessageIndex);
503
605
  completedResponse = {
504
606
  ...partial,
@@ -538,6 +640,39 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
538
640
  });
539
641
  }
540
642
  else {
643
+ // Transition from reasoning to assistant text: close the reasoning
644
+ // output item BEFORE emitting any `response.output_item.added` for
645
+ // the message. OpenAI Responses clients (e.g. pi-mono) maintain a
646
+ // single `currentBlock` state — opening the message item while the
647
+ // reasoning item is still "in progress" overwrites that state, so
648
+ // the later `output_item.done` for reasoning fails its
649
+ // `currentBlock?.type === 'thinking'` guard and never sets
650
+ // `thinkingSignature`. Without the signature the next turn cannot
651
+ // echo the reasoning item back, and any thinking-model agent loses
652
+ // its chain of thought on each turn. Must run before the first
653
+ // message write on this branch.
654
+ if (hasEmittedReasoning && !hasClosedReasoning && reasoningItemId) {
655
+ hasClosedReasoning = true;
656
+ const riIndex = outputItems.findIndex((i) => i.id === reasoningItemId);
657
+ writeSSEEvent(res, 'response.reasoning_summary_text.done', {
658
+ item_id: reasoningItemId,
659
+ output_index: riIndex,
660
+ summary_index: 0,
661
+ text: reasoningText,
662
+ });
663
+ const reasoningItem = {
664
+ id: reasoningItemId,
665
+ type: 'reasoning',
666
+ summary: [{ type: 'summary_text', text: reasoningText }],
667
+ };
668
+ if (riIndex >= 0) {
669
+ outputItems[riIndex] = reasoningItem;
670
+ }
671
+ writeSSEEvent(res, 'response.output_item.done', {
672
+ output_index: riIndex >= 0 ? riIndex : 0,
673
+ item: reasoningItem,
674
+ });
675
+ }
541
676
  // Text delta with tool_call tag buffering
542
677
  const { safeText, tagFound, cleanPrefix } = tagBuffer.push(event.text);
543
678
  if (tagFound) {
@@ -618,6 +753,12 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
618
753
  // error would escape into the outer JSON error path with SSE
619
754
  // headers already on the wire.
620
755
  thrownError = err instanceof Error ? err : new Error(String(err));
756
+ // Surface the message to stderr even on the failure path — without
757
+ // this the native side (e.g. `Tokenizer encoded <turn|> to N
758
+ // tokens; expected 1`) is invisible to operators since the SSE
759
+ // `response.failed` payload only carries `incomplete_details.reason`
760
+ // and never the underlying exception text.
761
+ console.error(`[responses] native dispatch failed for ${req.model} (response ${responseId}):`, thrownError.message);
621
762
  }
622
763
  finally {
623
764
  if (httpReq) {
@@ -670,7 +811,7 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
670
811
  // adopt under an unseen responseId.
671
812
  await flushTerminalSSE(res, 'response.completed', { response: terminal }, visibility);
672
813
  endSSE(res);
673
- return { terminalToPersist: terminal, failureMode: null };
814
+ return { terminalToPersist: terminal, failureMode: null, cachedTokens };
674
815
  }
675
816
  // Failure epilogue. Close any dangling message items BEFORE the
676
817
  // terminal so clients tracking `output_index` see matching closes.
@@ -727,9 +868,10 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
727
868
  item: closedMessageItem,
728
869
  });
729
870
  }
730
- if (!sawDone && hasEmittedReasoning && reasoningItemId != null) {
871
+ if (!sawDone && hasEmittedReasoning && !hasClosedReasoning && reasoningItemId != null) {
731
872
  // No `status` field on reasoning items — just emit closes so
732
873
  // client-side output_index bookkeeping stays consistent.
874
+ hasClosedReasoning = true;
733
875
  writeSSEEvent(res, 'response.reasoning_summary_text.done', {
734
876
  item_id: reasoningItemId,
735
877
  output_index: outputItems.findIndex((i) => i.id === reasoningItemId),
@@ -748,13 +890,16 @@ async function handleStreamingNative(res, chatStream, req, responseId, previousR
748
890
  writeSSEEvent(res, 'response.output_item.done', { output_index: riIndex, item: reasoningItem });
749
891
  }
750
892
  }
751
- const failedTerminal = buildFailedTerminal(partial, finalOutput, reason, usage);
893
+ const failedTerminal = buildFailedTerminal(partial, finalOutput, reason, usage, reason === 'error' && thrownError ? thrownError.message : null);
752
894
  await flushTerminalSSE(res, 'response.failed', { response: failedTerminal }, visibility);
753
895
  endSSE(res);
754
896
  // No terminalToPersist on an uncommitted turn: a later continuation
755
897
  // that cold-replayed this record would silently resurrect failed
756
- // output as authoritative history.
757
- return { terminalToPersist: null, failureMode: reason };
898
+ // output as authoritative history. `cachedTokens` is meaningless on
899
+ // the failure path but returned verbatim to keep the type shape
900
+ // uniform — consumers already treat a non-null `failureMode` as the
901
+ // authoritative "do not use these numbers" signal.
902
+ return { terminalToPersist: null, failureMode: reason, cachedTokens };
758
903
  }
759
904
  // ---------------------------------------------------------------------------
760
905
  // Session routing
@@ -979,9 +1124,50 @@ export function validateAndCanonicalizeHistoryToolOrder(messages, apiSurface = '
979
1124
  * back to reset + cold re-prime on multi-message input. The caller
980
1125
  * is responsible for rejecting partial tool-result submissions
981
1126
  * against a fan-out (`handleCreateResponse` fan-out gate).
1127
+ *
1128
+ * `isFreshSession` is the MISS / HIT signal from `SessionRegistry`:
1129
+ * `true` when `lookup.hit === false` (a truly new session minted via
1130
+ * `newSession()`), `false` when a tier-1 or tier-2 warm lease was
1131
+ * handed out. On a MISS we must wipe the shared native model's
1132
+ * leftover `cached_token_history` + KV caches before re-priming, or
1133
+ * a previous UNRELATED request's cache could silently get reused as
1134
+ * a prefix (cross-request cache-affinity side channel). On a HIT we
1135
+ * must NOT wipe — the whole point of the warm lease is that
1136
+ * `verify_cache_prefix_direct` can recover the reused prefix on the
1137
+ * next `chat_session_start_sync`. The HIT branch calls the
1138
+ * server-private `resetPreservingNativeCacheForWarmReuse(session)`
1139
+ * helper from `../chat-session-warm-reuse.js` instead of the public
1140
+ * `reset()` to thread this distinction down to the JS-side state
1141
+ * clear. The helper lives inside `@mlx-node/server` and is never
1142
+ * re-exported from either `@mlx-node/lm` or `@mlx-node/server`'s
1143
+ * public surface, so downstream consumers cannot discover or invoke
1144
+ * it.
982
1145
  */
983
- async function runSessionNonStreaming(session, messages, newInputMessages, config) {
1146
+ async function runSessionNonStreaming(session, messages, newInputMessages, config, isFreshSession) {
984
1147
  if (session.turns === 0) {
1148
+ // Fresh JS session does NOT imply a fresh native cache — the
1149
+ // underlying `SessionCapableModel` is shared across every
1150
+ // `ChatSession` lifetime via `ModelRegistry`, and its native
1151
+ // `cached_token_history` + KV caches persist across requests.
1152
+ // After the native refactor moved the unconditional cache wipe
1153
+ // out of `chat_session_start_sync` into the miss branch of
1154
+ // `verify_cache_prefix_direct`, a MISS path that runs
1155
+ // `primeHistory() + startFromHistory()` on a fresh session would
1156
+ // inherit the PREVIOUS request's native cache and silently reuse
1157
+ // whatever prefix happened to overlap — a cross-request
1158
+ // cache-affinity side channel. Only registry HITS (tier-1 /
1159
+ // tier-2) are authorized for cache reuse, and a leased session
1160
+ // almost always has `turns > 0` so this branch is nearly always
1161
+ // a MISS. The `isFreshSession` flag is the authoritative signal
1162
+ // — on `false` we still need to clear JS-side state so
1163
+ // `primeHistory()` accepts the replay, but we keep the native
1164
+ // cache intact so the prefix verifier can recover it.
1165
+ if (isFreshSession) {
1166
+ await session.reset();
1167
+ }
1168
+ else {
1169
+ await resetPreservingNativeCacheForWarmReuse(session);
1170
+ }
985
1171
  session.primeHistory(messages);
986
1172
  const initialTurns = session.turns;
987
1173
  const result = await session.startFromHistory(config);
@@ -1010,7 +1196,12 @@ async function runSessionNonStreaming(session, messages, newInputMessages, confi
1010
1196
  throw new Error('tool message missing toolCallId');
1011
1197
  }
1012
1198
  const initialTurns = session.turns;
1013
- const result = await session.sendToolResult(last.toolCallId, last.content, { config });
1199
+ // Forward the structured `isError` field through to the native
1200
+ // renderer so the wire-format `[tool error]` marker stays in sync
1201
+ // with the Anthropic `tool_result.is_error === true` source field
1202
+ // (the structured channel is the authoritative signal — see
1203
+ // `ChatMessage.isError` rustdoc).
1204
+ const result = await session.sendToolResult(last.toolCallId, last.content, { config, isError: last.isError });
1014
1205
  return { result, committed: session.turns > initialTurns };
1015
1206
  }
1016
1207
  // Non-user / non-tool single-message continuation (assistant /
@@ -1021,20 +1212,57 @@ async function runSessionNonStreaming(session, messages, newInputMessages, confi
1021
1212
  // re-prime. `initialTurns` MUST be captured AFTER `session.reset()`
1022
1213
  // zeroes `turns`, otherwise the committed check reads stale.
1023
1214
  // Amortized: the caller re-keys this session under the new
1024
- // responseId on success.
1025
- await session.reset();
1215
+ // responseId on success. On a tier-1 / tier-2 HIT that landed here
1216
+ // (full-history replay is NOT a single-message delta, so even a
1217
+ // warm session falls through to this branch), we MUST keep the
1218
+ // native KV cache so the native `verify_cache_prefix_direct` can
1219
+ // recover the reused prefix — wiping it would neutralize the
1220
+ // entire warm-lease feature on multi-message hits. On a MISS we
1221
+ // still wipe to prevent cross-request cache-affinity leakage.
1222
+ if (isFreshSession) {
1223
+ await session.reset();
1224
+ }
1225
+ else {
1226
+ await resetPreservingNativeCacheForWarmReuse(session);
1227
+ }
1026
1228
  session.primeHistory(messages);
1027
1229
  const initialTurns = session.turns;
1028
1230
  const result = await session.startFromHistory(config);
1029
1231
  return { result, committed: session.turns > initialTurns };
1030
1232
  }
1031
1233
  /** Streaming counterpart to {@link runSessionNonStreaming}. */
1032
- async function runSessionStreaming(session, messages, newInputMessages, config, signal) {
1234
+ async function runSessionStreaming(session, messages, newInputMessages, config, signal, isFreshSession) {
1235
+ // Preserve the startFromHistoryStream precondition outside the lazy
1236
+ // generator. Without this guard an accepted `input: []` request would commit
1237
+ // SSE and only then throw when iteration begins; the former eager-first-item
1238
+ // path surfaced the same deterministic error before selecting wire format.
1239
+ if (messages.length === 0) {
1240
+ throw new Error('ChatSession: startFromHistoryStream() requires a primed history');
1241
+ }
1033
1242
  if (session.turns === 0) {
1243
+ // Fresh/replay requests carry their complete canonical history in
1244
+ // `messages`. Validate it before reset so context overflow cannot mutate
1245
+ // native state and still receives a JSON 400 before SSE begins.
1246
+ const constrainedConfig = await session.preflightContextCapacity(messages, config);
1247
+ // See `runSessionNonStreaming` for the full rationale. A fresh
1248
+ // JS session inherits the shared native model's KV cache from
1249
+ // prior requests; without an explicit `reset()` here the native
1250
+ // prefix verifier can silently reuse a previous request's cache
1251
+ // on any prompt-prefix overlap (a cross-request cache-affinity
1252
+ // side channel). On MISS (`isFreshSession === true`) we wipe
1253
+ // native cache + JS state. On tier-1 / tier-2 HIT we keep the
1254
+ // native cache so the prefix verifier can recover the reused
1255
+ // prefix, while still clearing JS-side state for `primeHistory`.
1256
+ if (isFreshSession) {
1257
+ await session.reset();
1258
+ }
1259
+ else {
1260
+ await resetPreservingNativeCacheForWarmReuse(session);
1261
+ }
1034
1262
  session.primeHistory(messages);
1035
1263
  const initialTurns = session.turns;
1036
1264
  return {
1037
- stream: session.startFromHistoryStream(config, signal),
1265
+ stream: session.startFromHistoryStream(constrainedConfig, signal),
1038
1266
  wasCommitted: () => session.turns > initialTurns,
1039
1267
  };
1040
1268
  }
@@ -1045,10 +1273,14 @@ async function runSessionStreaming(session, messages, newInputMessages, config,
1045
1273
  if (newInputMessages.length === 1) {
1046
1274
  const last = newInputMessages[0];
1047
1275
  if (last.role === 'user') {
1276
+ // Tier-2 prompt-cache hits may carry only this new message in the HTTP
1277
+ // request while the leased ChatSession owns the prior conversation.
1278
+ // Preflight against that authoritative private history, not `messages`.
1279
+ const constrainedConfig = await session.preflightPendingContextCapacity(last, config);
1048
1280
  const initialTurns = session.turns;
1049
1281
  const images = last.images ?? undefined;
1050
1282
  return {
1051
- stream: session.sendStream(last.content, images ? { images, config, signal } : { config, signal }),
1283
+ stream: session.sendStream(last.content, images ? { images, config: constrainedConfig, signal } : { config: constrainedConfig, signal }),
1052
1284
  wasCommitted: () => session.turns > initialTurns,
1053
1285
  };
1054
1286
  }
@@ -1056,9 +1288,18 @@ async function runSessionStreaming(session, messages, newInputMessages, config,
1056
1288
  if (!last.toolCallId) {
1057
1289
  throw new Error('tool message missing toolCallId');
1058
1290
  }
1291
+ const constrainedConfig = await session.preflightPendingContextCapacity(last, config);
1059
1292
  const initialTurns = session.turns;
1060
1293
  return {
1061
- stream: session.sendToolResultStream(last.toolCallId, last.content, { config, signal }),
1294
+ // Forward the structured `isError` field through to the native
1295
+ // renderer so the streaming wire-format `[tool error]` marker
1296
+ // stays in sync with the Anthropic `tool_result.is_error === true`
1297
+ // source field — same contract as the non-streaming path above.
1298
+ stream: session.sendToolResultStream(last.toolCallId, last.content, {
1299
+ config: constrainedConfig,
1300
+ signal,
1301
+ isError: last.isError,
1302
+ }),
1062
1303
  wasCommitted: () => session.turns > initialTurns,
1063
1304
  };
1064
1305
  }
@@ -1067,12 +1308,22 @@ async function runSessionStreaming(session, messages, newInputMessages, config,
1067
1308
  }
1068
1309
  // Multi-message (or single non-user/non-tool) hot path: same reset +
1069
1310
  // cold re-prime as the non-streaming variant. `initialTurns` must be
1070
- // captured AFTER reset.
1071
- await session.reset();
1311
+ // captured AFTER reset. On tier-1 / tier-2 HIT (warm lease that
1312
+ // cannot use the delta API because the input spans multiple
1313
+ // messages), keep the native KV cache so the prefix verifier can
1314
+ // reuse it on the replayed `chat_session_start_sync`; on MISS, wipe
1315
+ // to block cross-request cache-affinity leakage.
1316
+ const constrainedConfig = await session.preflightContextCapacity(messages, config);
1317
+ if (isFreshSession) {
1318
+ await session.reset();
1319
+ }
1320
+ else {
1321
+ await resetPreservingNativeCacheForWarmReuse(session);
1322
+ }
1072
1323
  session.primeHistory(messages);
1073
1324
  const initialTurns = session.turns;
1074
1325
  return {
1075
- stream: session.startFromHistoryStream(config, signal),
1326
+ stream: session.startFromHistoryStream(constrainedConfig, signal),
1076
1327
  wasCommitted: () => session.turns > initialTurns,
1077
1328
  };
1078
1329
  }
@@ -1106,7 +1357,7 @@ function buildResponseRecord(response, newInputMessages, previousResponseId, mod
1106
1357
  model: response.model,
1107
1358
  status: response.status,
1108
1359
  instructions: response.instructions ?? undefined,
1109
- inputJson: JSON.stringify(newInputMessages),
1360
+ inputJson: stringifyStoredInputMessages(newInputMessages),
1110
1361
  outputJson: JSON.stringify(response.output),
1111
1362
  outputText: response.output_text,
1112
1363
  usageJson: JSON.stringify(response.usage),
@@ -1161,7 +1412,8 @@ function readStoredModelIdentity(record) {
1161
1412
  // ---------------------------------------------------------------------------
1162
1413
  // Public handler
1163
1414
  // ---------------------------------------------------------------------------
1164
- export async function handleCreateResponse(res, body, registry, store, httpReq, responseRetentionSec) {
1415
+ export async function handleCreateResponse(res, body, registry, store, httpReq, responseRetentionSec, idleSweeper, modelWorkCoordinator) {
1416
+ const handlerStartedAt = Date.now();
1165
1417
  // Validate required fields
1166
1418
  if (body == null || typeof body !== 'object') {
1167
1419
  sendBadRequest(res, 'Request body must be a JSON object', 'body');
@@ -1179,6 +1431,22 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
1179
1431
  sendBadRequest(res, 'Field "input" must be a string or an array', 'input');
1180
1432
  return;
1181
1433
  }
1434
+ // A present `max_output_tokens` must be an integer in `[1, i32::MAX]`.
1435
+ // `null` / missing means "no explicit limit" and is fine. Mirrors the
1436
+ // Anthropic `/v1/messages` guard on `max_tokens`. Rejecting here (rather
1437
+ // than forwarding through the mapper) keeps a nonpositive budget from
1438
+ // reaching native chat, where a negative `i32` would size a cache /
1439
+ // allocation as a huge `usize`; core still clamps nonpositive to 0 as
1440
+ // a backstop. The upper bound matters because NAPI truncates a JS
1441
+ // integer above `i32::MAX` to a NEGATIVE `i32` (then clamped to 0 → a
1442
+ // silent empty completion) — reject it as a 400 instead.
1443
+ if (body.max_output_tokens != null &&
1444
+ (!Number.isInteger(body.max_output_tokens) ||
1445
+ body.max_output_tokens <= 0 ||
1446
+ body.max_output_tokens > MAX_OUTPUT_TOKENS)) {
1447
+ sendBadRequest(res, `Field "max_output_tokens" must be an integer between 1 and ${MAX_OUTPUT_TOKENS}`, 'max_output_tokens');
1448
+ return;
1449
+ }
1182
1450
  // Per-request retention override: `metadata.retention_seconds` lets a
1183
1451
  // client pin a single row to a longer (VIP / onboarding) or shorter
1184
1452
  // (one-shot PII) lifetime than the server-wide default. Bounds
@@ -1243,6 +1511,39 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
1243
1511
  // early-return path. These flags keep it a no-op when already run.
1244
1512
  let cleanupPerformed = false;
1245
1513
  let leaseReleased = false;
1514
+ // Idle-sweeper bracket flags. Hoisted to outer function scope so
1515
+ // the `finalizeIdleRequest` helper below can see them even though
1516
+ // the `beginRequest()` call lives inside the inner `try`.
1517
+ // Pre-dispatch validation failures (early returns that never called
1518
+ // `beginRequest`) observe `idleRequestStarted === false` and skip
1519
+ // the matching `endRequest()` entirely. `idleRequestEnded` is the
1520
+ // `done` flag that guarantees the decrement fires exactly once
1521
+ // regardless of which of the several finalize paths — outer
1522
+ // `finally`, `res.once('finish')`, `res.once('close')`,
1523
+ // `res.once('error')` — wins the race.
1524
+ //
1525
+ // Listeners are attached EAGERLY at `beginRequest()` time (not
1526
+ // lazily from the outer `finally`) to close the round-4 leak where
1527
+ // a terminal socket event fired *before* the outer `finally` ran:
1528
+ // the lazy attach saw `writableEnded === false && writableFinished
1529
+ // === false` at check time, attached listeners on a socket whose
1530
+ // final event had already been emitted, and `endRequest()` then
1531
+ // never fired, leaving `inFlight` pinned above zero and the
1532
+ // sweeper permanently armed.
1533
+ let idleRequestStarted = false;
1534
+ let idleRequestEnded = false;
1535
+ let idleListenersAttached = false;
1536
+ const finalizeIdleRequest = () => {
1537
+ if (!idleRequestStarted)
1538
+ return;
1539
+ if (idleRequestEnded)
1540
+ return;
1541
+ idleRequestEnded = true;
1542
+ idleSweeper?.endRequest();
1543
+ };
1544
+ const onFinalizeEvent = () => {
1545
+ finalizeIdleRequest();
1546
+ };
1246
1547
  try {
1247
1548
  // Initial snapshot of the live binding. On a continuation we
1248
1549
  // re-read after `await store.getChain()` and reject if the
@@ -1751,8 +2052,37 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
1751
2052
  // adopt when `failureMode === 'client_abort'` regardless of how
1752
2053
  // `committed` / `safeToSuppress` landed.
1753
2054
  let streamFailureMode = null;
2055
+ // Bracket native-model dispatch with the idle-sweeper counter.
2056
+ // Must happen AFTER request validation / store lookup but BEFORE
2057
+ // any native prefill or decode runs — so the pending-drain timer
2058
+ // is cancelled in time to avoid racing the allocator with a live
2059
+ // decode, and the post-dispatch `endRequest` arms a fresh drain
2060
+ // only after every native stream byte has been emitted.
2061
+ // Only inference traffic participates; `/v1/models`, health, and
2062
+ // CORS preflights intentionally do not touch the counter.
2063
+ //
2064
+ // Attach the terminal-event listeners BEFORE any `await` — this
2065
+ // is the round-4 fix for a sweeper leak where a fast terminal
2066
+ // event (e.g. `endJson()` rejecting after the socket already
2067
+ // emitted `close`) fired before the outer `finally` got a chance
2068
+ // to attach its listeners, leaving `inFlight` pinned above zero.
2069
+ // `finalizeIdleRequest` is the idempotency barrier — whichever
2070
+ // of the listener events, the outer `finally`, or a pre-dispatch
2071
+ // path fires first wins and the rest are no-ops.
2072
+ idleSweeper?.beginRequest();
2073
+ idleRequestStarted = true;
2074
+ res.once('finish', onFinalizeEvent);
2075
+ res.once('close', onFinalizeEvent);
2076
+ res.once('error', onFinalizeEvent);
2077
+ idleListenersAttached = true;
1754
2078
  try {
1755
- await sessionReg.withExclusive(async () => {
2079
+ const mutexQueuedAt = Date.now();
2080
+ const runInference = () => withAdmissionControlledInference(sessionReg, modelWorkCoordinator, async () => {
2081
+ const serverTiming = {
2082
+ server_queue_ms: Date.now() - mutexQueuedAt,
2083
+ server_pre_inference_ms: Date.now() - handlerStartedAt,
2084
+ ...resolveServerTuningForUsage(),
2085
+ };
1756
2086
  // Hot-swap race guard inside the mutex.
1757
2087
  //
1758
2088
  // `withExclusive` can park this waiter behind a long-running
@@ -1801,36 +2131,132 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
1801
2131
  // `send` / `sendToolResult` entry points cover exactly that
1802
2132
  // shape. A single `assistant` / `system` continuation cannot
1803
2133
  // be advanced incrementally against the warm KV cache and
1804
- // must be handled via reset + cold re-prime. Consuming a warm
1805
- // lease only to immediately `session.reset()` would destroy
1806
- // the cached prefix for no benefit (and mislabel the turn as
1807
- // `hit` when the client actually paid a full cold-replay
1808
- // prefill), so detect the case up front and force a fresh
1809
- // session lookup by passing `null` into `getOrCreate`. The
1810
- // subsequent `ResponseStore` reconstruction + `primeHistory` +
1811
- // `startFromHistory*` path below handles the rebuild. Multi-
1812
- // message continuations are left to the existing reset + cold
1813
- // re-prime fall-through inside `runSession*` (see comments
1814
- // there).
1815
- const hotPathIneligible = previousResponseId != null &&
1816
- newInputMessages.length === 1 &&
1817
- newInputMessages[0].role !== 'user' &&
1818
- newInputMessages[0].role !== 'tool';
1819
- const lookup = hotPathIneligible
1820
- ? sessionReg.getOrCreate(null, requestedInstructions)
1821
- : sessionReg.getOrCreate(previousResponseId ?? null, requestedInstructions);
2134
+ // must be handled via reset + cold re-prime. That branch is
2135
+ // still VALID it just routes through `runSession*`'s
2136
+ // `session.turns === 0` fall-through (`primeHistory` +
2137
+ // `startFromHistory*`) instead of the `send` / `sendToolResult`
2138
+ // delta path. Crucially, a tier-1 HIT on this branch is still
2139
+ // useful: `resetPreservingNativeCacheForWarmReuse(session)` keeps
2140
+ // the warm native KV cache, and the subsequent
2141
+ // `chat_session_start_sync` -> `verify_cache_prefix_direct`
2142
+ // recovers the reused prefix even across a full-history
2143
+ // replay. So we must NOT rewrite `previousResponseId` to null
2144
+ // to force a cold-replay lookup — the tier-1 lease is exactly
2145
+ // what makes warm reuse work. (Pre-Round 5 this branch passed
2146
+ // `null` to force a miss, but that was wrong: it threw away a
2147
+ // usable warm lease AND mislabeled the turn as `cold_replay`
2148
+ // when the native prefix verifier was about to reuse the
2149
+ // entire previous-turn prefix.) The hot-path-ineligibility
2150
+ // condition (single non-user/non-tool continuation) is no
2151
+ // longer consulted at lookup time — it's just the natural
2152
+ // fall-through to the `session.turns === 0 || multi-message`
2153
+ // cold-re-prime branch in `runSession*`, which correctly
2154
+ // preserves the warm native cache on HIT via
2155
+ // `resetPreservingNativeCacheForWarmReuse(session)`.
2156
+ // Normalize the caller-supplied `prompt_cache_key` into the
2157
+ // `string | null` shape the registry expects. `undefined` and
2158
+ // missing both map to `null`; an explicit empty string is
2159
+ // preserved distinct from `null` so the registry's tier-2
2160
+ // scan treats "no key" and "empty key" as different tenants
2161
+ // (prevents an unkeyed client from accidentally colliding
2162
+ // with one that explicitly empty-keyed).
2163
+ const promptCacheKey = typeof body.prompt_cache_key === 'string' ? body.prompt_cache_key : null;
2164
+ // Precedence gate: when `previous_response_id` is present, tier-2
2165
+ // (prompt-cache-key) lookup is DISABLED — even if the request is
2166
+ // hot-path ineligible (single `assistant` / `system` continuation).
2167
+ // The documented precedence is "prev-id wins; tier-1 miss falls
2168
+ // through to FRESH, not tier-2", and that rule has to hold whether
2169
+ // we take the hot path or the ineligible cold-replay branch. If we
2170
+ // let the ineligible branch fall through to tier-2, a mis-routed
2171
+ // prev-id request could lease an UNRELATED warm session that
2172
+ // happens to share `prompt_cache_key`, then cold-replay on top of
2173
+ // it — which `session.reset()` + `primeHistory()` would destroy,
2174
+ // corrupting an unrelated chain. Force `null` for the cache key on
2175
+ // both branches whenever a prev-id is set; tier-2 only runs for
2176
+ // requests with no prev-id at all.
2177
+ const effectivePromptCacheKey = previousResponseId != null ? null : promptCacheKey;
2178
+ // Integrator nudge: if the caller supplied a non-empty
2179
+ // `prompt_cache_key` but tier-2 prerequisites are missing
2180
+ // (env gate off or key below the min-length floor) the turn
2181
+ // silently cold-starts. Emit a once-per-distinct-raw-key
2182
+ // stderr warning so `X-Session-Cache: fresh` on every request
2183
+ // can be diagnosed without reading source. Gated on
2184
+ // `effectivePromptCacheKey` (not raw `promptCacheKey`) so a
2185
+ // request that suppresses the key via `previous_response_id`
2186
+ // precedence does not also log a misleading "key ignored"
2187
+ // message — that case is documented precedence, not a
2188
+ // misconfiguration.
2189
+ if (effectivePromptCacheKey !== null) {
2190
+ maybeWarnPromptCacheKeyIneligible(effectivePromptCacheKey);
2191
+ }
2192
+ const lookup = sessionReg.getOrCreate(previousResponseId ?? null, requestedInstructions, effectivePromptCacheKey);
1822
2193
  const session = lookup.session;
1823
2194
  // `X-Session-Cache` observability header: classify this turn as
1824
- // `fresh` (no `previous_response_id` on the request), `hit`
1825
- // (warm-cache lease consumed), or `cold_replay` (request carried
1826
- // `previous_response_id` but the warm entry was missing / expired
1827
- // / instructions-mismatched / already leased, OR the request
1828
- // shape is ineligible for the hot path the endpoint will
1829
- // rebuild the session from the `ResponseStore` below). Set
1830
- // before any `writeHead` / SSE `beginSSE` so both JSON and SSE
1831
- // responses carry it. See `endpoints/messages.ts` for the
1832
- // matching always-`fresh` emission on `/v1/messages`.
1833
- const sessionCacheStatus = previousResponseId == null ? 'fresh' : lookup.hit && !hotPathIneligible ? 'hit' : 'cold_replay';
2195
+ // `fresh` (no `previous_response_id` on the request and tier-2
2196
+ // prompt-cache-key did not hit), `hit` (prev-id warm-cache
2197
+ // lease consumed on tier 1), `prefix_hit` (tier-2 warm-cache
2198
+ // lease consumed only promoted from `fresh` later once the
2199
+ // native `cachedTokens > 0` confirms real prefix reuse), or
2200
+ // `cold_replay` (request carried `previous_response_id` but
2201
+ // the warm entry was missing / expired / instructions-
2202
+ // mismatched / already leased, OR the request shape is
2203
+ // ineligible for the hot path — the endpoint will rebuild the
2204
+ // session from the `ResponseStore` below). Set before any
2205
+ // `writeHead` / SSE `beginSSE` so both JSON and SSE responses
2206
+ // carry it. See `endpoints/messages.ts` for the matching
2207
+ // emission on `/v1/messages`.
2208
+ //
2209
+ // The `prefix_hit` promotion and the companion
2210
+ // `X-Cached-Tokens: N` header both depend on the native
2211
+ // ChatResult's `cachedTokens` field, which is only authoritative
2212
+ // AFTER the native dispatch completes. We therefore emit the
2213
+ // initial `fresh` / `hit` / `cold_replay` value here and let
2214
+ // the post-dispatch branch below promote a `fresh`+tier2-hit
2215
+ // classification to `prefix_hit` once the cached-tokens count
2216
+ // is known.
2217
+ const tier2Hit = previousResponseId == null && lookup.hit;
2218
+ // Optimistic pre-dispatch classification. SSE flushes headers on
2219
+ // `beginSSE` inside `handleStreamingNative`, so the streaming
2220
+ // path has exactly one shot to commit the header value — before
2221
+ // the dispatch runs, i.e. BEFORE the native prefix verifier has
2222
+ // reported whether any tokens were actually reused.
2223
+ //
2224
+ // For non-streaming we still commit optimistically to
2225
+ // `prefix_hit` on tier-2 hit and demote to `fresh` post-dispatch
2226
+ // if `cachedTokens === 0` (template drift, tokenizer change,
2227
+ // image-set change, etc.) — `res.end` has not fired yet so the
2228
+ // header is still settable.
2229
+ //
2230
+ // For streaming we deliberately do NOT promote to `prefix_hit`
2231
+ // on tier-2 hit. Once SSE headers flush they cannot be
2232
+ // corrected, and a false-positive `prefix_hit` would contradict
2233
+ // consumers that read `cachedTokens` from the terminal event.
2234
+ // Approach B (this path): emit `fresh` on streaming even when
2235
+ // tier-2 found a warm session — the cache reuse still happens,
2236
+ // only the observability header is conservative. A future
2237
+ // refactor can thread `cached_tokens` through the native
2238
+ // streaming `start` chunk so the server knows authoritatively
2239
+ // before `beginSSE()` flushes, at which point streaming can
2240
+ // commit `prefix_hit` too (Approach A). Until then,
2241
+ // `prefix_hit` is a non-streaming-only signal.
2242
+ const isStreaming = mappedBody.stream === true;
2243
+ // Prev-id branch classification. A tier-1 HIT is labeled `hit`
2244
+ // regardless of whether the request is hot-path-eligible: the
2245
+ // warm native KV cache is reused in both paths (the cheap
2246
+ // `send` / `sendToolResult` delta on the eligible branch; the
2247
+ // full-history `primeHistory` + `startFromHistory*` replay on
2248
+ // the ineligible branch, where `resetPreservingNativeCacheForWarmReuse`
2249
+ // keeps the cache alive for `verify_cache_prefix_direct` to
2250
+ // recover). Only a registry MISS on the prev-id branch
2251
+ // downgrades to `cold_replay` (no warm cache to reuse, the
2252
+ // request must rebuild from `ResponseStore`).
2253
+ let sessionCacheStatus = previousResponseId == null
2254
+ ? tier2Hit && !isStreaming
2255
+ ? 'prefix_hit'
2256
+ : 'fresh'
2257
+ : lookup.hit
2258
+ ? 'hit'
2259
+ : 'cold_replay';
1834
2260
  res.setHeader('X-Session-Cache', sessionCacheStatus);
1835
2261
  // Multi-tool-call fan-out gate.
1836
2262
  //
@@ -2067,10 +2493,10 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2067
2493
  // (JSON error, SSE `error` frame, or socket destroy).
2068
2494
  let handlerError = null;
2069
2495
  if (mappedBody.stream) {
2070
- const outcome = await runSessionStreaming(session, messages, newInputMessages, config, streamSignal);
2496
+ const outcome = await runSessionStreaming(session, messages, newInputMessages, config, streamSignal, !lookup.hit);
2071
2497
  const streamingWasCommitted = () => outcome.wasCommitted();
2072
2498
  try {
2073
- const handlerOutcome = await handleStreamingNative(res, outcome.stream, mappedBody, responseId, previousResponseId, streamingWasCommitted, httpReq, visibility);
2499
+ const handlerOutcome = await handleStreamingNative(res, outcome.stream, mappedBody, responseId, previousResponseId, streamingWasCommitted, httpReq, visibility, serverTiming);
2074
2500
  streamFailureMode = handlerOutcome.failureMode;
2075
2501
  if (handlerOutcome.terminalToPersist != null && store && body.store !== false) {
2076
2502
  // Initiate the write SYNCHRONOUSLY inside the mutex so
@@ -2294,9 +2720,30 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2294
2720
  // `handleNonStreaming` (short-circuits `endJson` and
2295
2721
  // signals the outer persist gate) plus this documented
2296
2722
  // limitation.
2297
- const outcome = await runSessionNonStreaming(session, messages, newInputMessages, config);
2723
+ const outcome = await runSessionNonStreaming(session, messages, newInputMessages, config, !lookup.hit);
2724
+ // Prefix-cache observability headers for the non-streaming
2725
+ // path. `res.end` has not fired yet (the handler's
2726
+ // `endJson` call below is what flushes), so `setHeader`
2727
+ // still lands on the wire. We re-classify the
2728
+ // `X-Session-Cache` header here so a tier-2 lookup that
2729
+ // did NOT actually produce native prefix reuse
2730
+ // (`cachedTokens === 0`) gets demoted from the optimistic
2731
+ // `prefix_hit` back to `fresh` — matching the plan's
2732
+ // contract that `prefix_hit` only fires when the registry
2733
+ // served a match via `promptCacheKey` AND the ChatResult
2734
+ // reports `cachedTokens > 0`. The companion
2735
+ // `X-Cached-Tokens: N` header reports the exact count for
2736
+ // operators and downstream telemetry whenever reuse
2737
+ // happened.
2738
+ if (tier2Hit && outcome.result.cachedTokens === 0) {
2739
+ sessionCacheStatus = 'fresh';
2740
+ res.setHeader('X-Session-Cache', sessionCacheStatus);
2741
+ }
2742
+ if (outcome.result.cachedTokens > 0) {
2743
+ res.setHeader('X-Cached-Tokens', String(outcome.result.cachedTokens));
2744
+ }
2298
2745
  try {
2299
- const handlerOutcome = await handleNonStreaming(res, outcome.result, mappedBody, responseId, previousResponseId, visibility);
2746
+ const handlerOutcome = await handleNonStreaming(res, outcome.result, mappedBody, responseId, previousResponseId, visibility, serverTiming);
2300
2747
  if (store && body.store !== false) {
2301
2748
  // Same in-lock-initiate / off-lock-await split as the
2302
2749
  // streaming branch. The non-streaming handler only
@@ -2531,7 +2978,21 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2531
2978
  // stream path completed cleanly and the adopted session
2532
2979
  // is genuinely reachable via the responseId.
2533
2980
  if (committed && (handlerError == null || safeToSuppress) && streamFailureMode === null) {
2534
- sessionReg.adopt(responseId, session, requestedInstructions);
2981
+ // Adopt under the SAME `effectivePromptCacheKey` that was
2982
+ // used for `getOrCreate` above — not the raw
2983
+ // `promptCacheKey` from the request body. When a request
2984
+ // carries `previous_response_id` (tier-1 path), tier-2 is
2985
+ // deliberately disabled on the lookup side by forcing
2986
+ // `effectivePromptCacheKey = null`; the adopt side must
2987
+ // follow the same rule or a mixed-mode request
2988
+ // (`previous_response_id=rA + prompt_cache_key=K`) would
2989
+ // store the adopted session under `K` even though the
2990
+ // lookup was resolved via `rA`. A subsequent keyless/
2991
+ // prev-idless request with `prompt_cache_key=K` would then
2992
+ // tier-2 hit and lease rA's chain session — a cross-chain
2993
+ // corruption the precedence rule was explicitly designed
2994
+ // to prevent. Keep adopt's key aligned with lookup's key.
2995
+ sessionReg.adopt(responseId, session, requestedInstructions, effectivePromptCacheKey);
2535
2996
  }
2536
2997
  // Rethrow handler errors when the client hasn't seen a
2537
2998
  // terminal yet, regardless of commit state. The outer
@@ -2560,9 +3021,15 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2560
3021
  // already received — or no output at all if the terminal
2561
3022
  // already landed.
2562
3023
  if (visibility.responseMode === null) {
2563
- // Headers never went out. Safe to emit a clean 500 JSON
2564
- // error.
2565
- sendInternalError(res, message);
3024
+ // Capacity failures are deterministic request errors, raised
3025
+ // before native cache mutation. Keep them out of the generic
3026
+ // 500 path so clients can compact/truncate and retry.
3027
+ if (isContextCapacityError(err)) {
3028
+ sendBadRequest(res, message);
3029
+ }
3030
+ else {
3031
+ sendInternalError(res, message);
3032
+ }
2566
3033
  }
2567
3034
  else if (visibility.responseMode === 'json') {
2568
3035
  // We already wrote `Content-Type: application/json` and
@@ -2602,6 +3069,7 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2602
3069
  }
2603
3070
  }
2604
3071
  });
3072
+ await runInference();
2605
3073
  }
2606
3074
  catch (err) {
2607
3075
  // Admission-control rejection from the per-model queue cap
@@ -2765,11 +3233,16 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2765
3233
  // never ran (early-return validation failure, or an exception
2766
3234
  // raised inside the outer `try` block between lease
2767
3235
  // acquisition and the `withExclusive` call), make sure the
2768
- // abort listeners are detached and the dispatch lease is
3236
+ // abort listeners + idle-sweeper listeners are detached, the
3237
+ // in-flight counter is decremented, and the dispatch lease is
2769
3238
  // released here. `runPostDispatchCleanup` is safe to re-invoke
2770
- // — the `abortListenersAttached` check and
2771
- // `releaseDispatchLease`'s `inFlight < 0` floor make it a
2772
- // no-op when the happy-path already fired it.
3239
+ // — the `abortListenersAttached` / `idleListenersAttached` /
3240
+ // `leaseReleased` flags make every sub-step idempotent on a
3241
+ // second pass. `finalizeIdleRequest` inside the helper is also
3242
+ // guarded by `idleRequestEnded`, so the decrement fires exactly
3243
+ // once regardless of which path wins the race between the
3244
+ // eagerly-attached terminal listener firing, the happy-path
3245
+ // cleanup on `withExclusive` return, and this outer fallback.
2773
3246
  if (!cleanupPerformed) {
2774
3247
  runPostDispatchCleanup();
2775
3248
  }
@@ -2793,6 +3266,23 @@ export async function handleCreateResponse(res, body, registry, store, httpReq,
2793
3266
  }
2794
3267
  abortListenersAttached = false;
2795
3268
  }
3269
+ // Drop the idle-sweeper's finalize listeners AND decrement the
3270
+ // in-flight counter. Post-dispatch cleanup runs AFTER
3271
+ // `handleStreamingNative` / `handleNonStreaming` have awaited
3272
+ // their terminal `res.end()` — the native dispatch is done at
3273
+ // this point, so the sweeper can safely arm a new pending
3274
+ // drain. Firing here (rather than only in the outer `finally`)
3275
+ // means the post-commit persist wait does not keep `inFlight`
3276
+ // pinned above zero on a wedged store. `finalizeIdleRequest`
3277
+ // is idempotent via the `done` flag so a subsequent fire from
3278
+ // the outer finally / a stray listener is a no-op.
3279
+ if (idleListenersAttached) {
3280
+ res.removeListener('finish', onFinalizeEvent);
3281
+ res.removeListener('close', onFinalizeEvent);
3282
+ res.removeListener('error', onFinalizeEvent);
3283
+ idleListenersAttached = false;
3284
+ }
3285
+ finalizeIdleRequest();
2796
3286
  // Release the dispatch lease on the ORIGINAL model object the
2797
3287
  // lease was acquired against (not a re-read of `body.model`,
2798
3288
  // which may have been hot-swapped while we held the mutex). A