@mlx-node/lm 0.0.10 → 0.0.13
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/chat-session.d.ts +175 -68
- package/dist/chat-session.d.ts.map +1 -1
- package/dist/chat-session.js +271 -122
- package/dist/family-data.d.ts +407 -0
- package/dist/family-data.d.ts.map +1 -0
- package/dist/family-data.js +381 -0
- package/dist/index.d.ts +6 -2
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +15 -1
- package/dist/models/model-loader.d.ts +91 -135
- package/dist/models/model-loader.d.ts.map +1 -1
- package/dist/models/model-loader.js +74 -179
- package/dist/models/paged-config-override.d.ts +0 -2
- package/dist/models/paged-config-override.d.ts.map +1 -1
- package/dist/models/paged-config-override.js +29 -6
- package/dist/models/qwen3_5-configs.d.ts +1 -3
- package/dist/models/qwen3_5-configs.d.ts.map +1 -1
- package/dist/stream.d.ts +29 -5
- package/dist/stream.d.ts.map +1 -1
- package/dist/stream.js +127 -24
- package/package.json +7 -3
- package/dist/interfaces.d.ts +0 -3
- package/dist/interfaces.d.ts.map +0 -1
- package/dist/interfaces.js +0 -1
package/dist/chat-session.js
CHANGED
|
@@ -21,8 +21,10 @@
|
|
|
21
21
|
* - An image hash (`lastImagesKey`) tracks the images bound to the
|
|
22
22
|
* current cache. A `send()` call whose image set has changed
|
|
23
23
|
* (different bytes or different ordering) triggers a full
|
|
24
|
-
* restart:
|
|
25
|
-
* images) to history → `chatSessionStart(history)`.
|
|
24
|
+
* restart: release this session's native cache owner → push the new user
|
|
25
|
+
* message (with images) to history → `chatSessionStart(history)`. Native
|
|
26
|
+
* models without owner-scoped release retain the exclusive-model
|
|
27
|
+
* `resetCaches()` fallback.
|
|
26
28
|
*
|
|
27
29
|
* - Text-only `send()` on turn >= 1 still gets incremental prefill
|
|
28
30
|
* on a token-prefix hit. Prompt structure is never reconstructed
|
|
@@ -83,7 +85,7 @@
|
|
|
83
85
|
* await session.reset();
|
|
84
86
|
* ```
|
|
85
87
|
*/
|
|
86
|
-
import { createHash } from 'node:crypto';
|
|
88
|
+
import { createHash, randomUUID } from 'node:crypto';
|
|
87
89
|
/**
|
|
88
90
|
* Typed prefix native media guards use when a history cannot be continued
|
|
89
91
|
* from the held image/audio state. The session layer recognizes this exact
|
|
@@ -106,6 +108,23 @@ const IMAGE_CHANGE_RESTART_PREFIX = 'IMAGE_CHANGE_REQUIRES_SESSION_RESTART:';
|
|
|
106
108
|
* `crates/mlx-core/src/engine/params.rs`.
|
|
107
109
|
*/
|
|
108
110
|
const NATIVE_DEFAULT_MAX_NEW_TOKENS = 2048;
|
|
111
|
+
/**
|
|
112
|
+
* Model wrapper class names whose in-checkpoint MTP head must NOT be
|
|
113
|
+
* auto-enabled. Fallback only: {@link ChatSession#mtpAutoDefaultAllowed}
|
|
114
|
+
* prefers the native {@link SessionCapableModel.mtpAutoEnabled} getter in BOTH
|
|
115
|
+
* directions whenever the binding exposes it.
|
|
116
|
+
*
|
|
117
|
+
* NemotronH ships a complete MTP head on every checkpoint, so
|
|
118
|
+
* `hasMtpWeights()` is unconditionally `true`. Setting `enable_mtp` forces the
|
|
119
|
+
* turn into the exclusive/barrier scheduler lane, which takes the session OUT
|
|
120
|
+
* of continuous batching so concurrent sessions serialize; a streaming MTP turn
|
|
121
|
+
* additionally has no flat-core streaming arm and falls back to paged AR — no
|
|
122
|
+
* speculation and no batching. The BARRIER is the reason, not the head's speed:
|
|
123
|
+
* MTP is roughly perf-neutral on this family, so an explicit per-session
|
|
124
|
+
* `enableMtp: true` costs nothing. Drop the family from this set once an MTP
|
|
125
|
+
* turn can share the continuous-batching lane.
|
|
126
|
+
*/
|
|
127
|
+
const MTP_AUTO_DEFAULT_SUPPRESSED_MODELS = new Set(['NemotronHModel']);
|
|
109
128
|
/**
|
|
110
129
|
* Stable, provider-neutral error raised before native inference when a
|
|
111
130
|
* rendered prompt cannot fit in the model's physically available hot KV
|
|
@@ -300,39 +319,6 @@ function publicStreamEvent(event, config) {
|
|
|
300
319
|
const safeRaw = event.publicRawText;
|
|
301
320
|
return { ...event, thinking: null, rawText: safeRaw ?? event.text };
|
|
302
321
|
}
|
|
303
|
-
/**
|
|
304
|
-
* Compute a stable hex-encoded identity key for a list of image
|
|
305
|
-
* byte buffers.
|
|
306
|
-
*
|
|
307
|
-
* Returns `null` when no images are provided so `send()` can
|
|
308
|
-
* distinguish "no-images" from "image set changed". The key is
|
|
309
|
-
* order-sensitive: `[A, B]` and `[B, A]` produce different keys,
|
|
310
|
-
* matching the positional semantics of the underlying VLM chat
|
|
311
|
-
* template.
|
|
312
|
-
*
|
|
313
|
-
* This is a byte-identity check — callers use the key solely to
|
|
314
|
-
* decide whether to restart the server-side session, so any
|
|
315
|
-
* collision-resistant digest is sufficient. We use SHA-256 (native
|
|
316
|
-
* `node:crypto`) with a length-prefixed framing so different image
|
|
317
|
-
* counts and different byte lengths cannot collide by accident.
|
|
318
|
-
*
|
|
319
|
-
* Implementation note: kept fully sync + self-contained so
|
|
320
|
-
* `send()` can stay synchronous in its routing decision. `node:crypto`
|
|
321
|
-
* is a Node built-in, so this adds no external runtime dependency
|
|
322
|
-
* beyond `@mlx-node/core` and the existing stream bridge.
|
|
323
|
-
*/
|
|
324
|
-
function computeImagesKey(images) {
|
|
325
|
-
return computeByteListKey(images);
|
|
326
|
-
}
|
|
327
|
-
/**
|
|
328
|
-
* Audio counterpart of {@link computeImagesKey}: a stable, order-sensitive
|
|
329
|
-
* byte-identity key for a list of encoded audio buffers. Used by `send()` /
|
|
330
|
-
* `sendStream()` to decide whether a new audio set must cold-restart the
|
|
331
|
-
* server-side session. Shares the exact SHA-256 framing as the image key.
|
|
332
|
-
*/
|
|
333
|
-
function computeAudioKey(audio) {
|
|
334
|
-
return computeByteListKey(audio);
|
|
335
|
-
}
|
|
336
322
|
/**
|
|
337
323
|
* SHA-256 byte-identity key for a length-framed list of byte buffers.
|
|
338
324
|
* Returns `null` for an empty/absent list so callers can distinguish
|
|
@@ -345,6 +331,10 @@ function computeAudioKey(audio) {
|
|
|
345
331
|
* before any `await` in `send()`/`sendStream()`, so the JS loop's cost
|
|
346
332
|
* was a real head-of-line-blocking stall for every other request
|
|
347
333
|
* handled by the same process.
|
|
334
|
+
*
|
|
335
|
+
* The key is order-sensitive: `[A, B]` and `[B, A]` produce different keys,
|
|
336
|
+
* matching the positional semantics of the underlying VLM chat template.
|
|
337
|
+
* Kept fully sync so `send()` can stay synchronous in its routing decision.
|
|
348
338
|
*/
|
|
349
339
|
function computeByteListKey(buffers) {
|
|
350
340
|
if (!buffers || buffers.length === 0)
|
|
@@ -376,6 +366,19 @@ export class ChatSession {
|
|
|
376
366
|
model;
|
|
377
367
|
system;
|
|
378
368
|
defaultConfig;
|
|
369
|
+
/**
|
|
370
|
+
* Stable native cache/scheduler identity for this JS session.
|
|
371
|
+
*
|
|
372
|
+
* Direct HTTP callers do not carry the agent provider's cacheOwnerId. If
|
|
373
|
+
* they reach a paged model without an owner, native must conservatively use
|
|
374
|
+
* the legacy exclusive lane because sequence zero cannot represent two live
|
|
375
|
+
* requests. Give every ChatSession its own identity while still allowing an
|
|
376
|
+
* explicit provider identity to override it.
|
|
377
|
+
*/
|
|
378
|
+
cacheOwnerId = null;
|
|
379
|
+
/** The native owner used by this session, retained as a set for retryable release. */
|
|
380
|
+
nativeCacheOwnerIds = new Set();
|
|
381
|
+
disposed = false;
|
|
379
382
|
/** Tool definitions are conversation state for deterministic template replay. */
|
|
380
383
|
activeTools;
|
|
381
384
|
/**
|
|
@@ -386,14 +389,14 @@ export class ChatSession {
|
|
|
386
389
|
history = [];
|
|
387
390
|
/**
|
|
388
391
|
* Hex-encoded byte-identity key of the image set currently bound
|
|
389
|
-
* to the server's KV cache (SHA-256; see `
|
|
392
|
+
* to the server's KV cache (SHA-256; see `computeByteListKey`).
|
|
390
393
|
* `null` when no images are cached. A `send()` whose new key
|
|
391
394
|
* differs triggers a full `chatSessionStart` restart.
|
|
392
395
|
*/
|
|
393
396
|
lastImagesKey = null;
|
|
394
397
|
/**
|
|
395
398
|
* Hex-encoded byte-identity key of the audio set currently bound to the
|
|
396
|
-
* server's KV cache (see {@link
|
|
399
|
+
* server's KV cache (see {@link computeByteListKey}). `null` when no audio is
|
|
397
400
|
* cached. A `send()` whose new key differs triggers a full
|
|
398
401
|
* `chatSessionStart` restart — the audio counterpart of `lastImagesKey`.
|
|
399
402
|
*/
|
|
@@ -464,7 +467,7 @@ export class ChatSession {
|
|
|
464
467
|
if (this.inFlight) {
|
|
465
468
|
throw new Error('ChatSession: cannot preflight context capacity while a send() is in flight');
|
|
466
469
|
}
|
|
467
|
-
return await this.constrainToContextCapacity(messages.slice(), this.mergeConfig(config));
|
|
470
|
+
return await this.constrainToContextCapacity(messages.slice(), this.mergeConfig(config, false));
|
|
468
471
|
}
|
|
469
472
|
/**
|
|
470
473
|
* Capacity-preflight one pending user/tool message against this session's
|
|
@@ -487,7 +490,7 @@ export class ChatSession {
|
|
|
487
490
|
else {
|
|
488
491
|
throw new Error('ChatSession: pending context capacity preflight requires a user or tool message');
|
|
489
492
|
}
|
|
490
|
-
return await this.constrainToContextCapacity(this.historyWithPending(pending), this.mergeConfig(config));
|
|
493
|
+
return await this.constrainToContextCapacity(this.historyWithPending(pending), this.mergeConfig(config, false));
|
|
491
494
|
}
|
|
492
495
|
/**
|
|
493
496
|
* Count of `ok` tool calls from the most recent assistant turn, or
|
|
@@ -525,8 +528,8 @@ export class ChatSession {
|
|
|
525
528
|
this.inFlight = true;
|
|
526
529
|
try {
|
|
527
530
|
const mergedConfig = this.mergeConfig(opts.config);
|
|
528
|
-
const newImagesKey =
|
|
529
|
-
const newAudioKey =
|
|
531
|
+
const newImagesKey = computeByteListKey(opts.images);
|
|
532
|
+
const newAudioKey = computeByteListKey(opts.audio);
|
|
530
533
|
// Only an explicit NEW image/audio set can trigger a forced restart. Omitting
|
|
531
534
|
// `images`/`audio` (key === null) is interpreted as "keep the current
|
|
532
535
|
// media cache state" — the server-side cache already holds any prior
|
|
@@ -537,7 +540,7 @@ export class ChatSession {
|
|
|
537
540
|
const isFirstTurn = this.turnCount === 0;
|
|
538
541
|
const replayRequired = this.needsFullReplay;
|
|
539
542
|
if (isFirstTurn || imageChanged || audioChanged || replayRequired) {
|
|
540
|
-
return await this.runStartPath(userMessage, opts.images, opts.audio, imageChanged || audioChanged || replayRequired, isFirstTurn, mergedConfig);
|
|
543
|
+
return await this.runStartPath(userMessage, opts.images, opts.audio, imageChanged || audioChanged || replayRequired, isFirstTurn, mergedConfig, opts.signal);
|
|
541
544
|
}
|
|
542
545
|
// Role-aware continuation: pass the complete structured transcript.
|
|
543
546
|
// Native code renders it with the checkpoint template and only reuses
|
|
@@ -547,7 +550,7 @@ export class ChatSession {
|
|
|
547
550
|
const constrainedConfig = await this.constrainToContextCapacity(pendingHistory, mergedConfig);
|
|
548
551
|
let result;
|
|
549
552
|
try {
|
|
550
|
-
result = await this.
|
|
553
|
+
result = await this.runNonStreamingNative('continue', pendingHistory, withReplayReasoning(constrainedConfig, this.model), opts.signal);
|
|
551
554
|
}
|
|
552
555
|
catch (err) {
|
|
553
556
|
if (!isMediaHeldRestartError(err)) {
|
|
@@ -561,7 +564,7 @@ export class ChatSession {
|
|
|
561
564
|
// the trailing-media keys keep `lastImagesKey`/`lastAudioKey`
|
|
562
565
|
// consistent across the replay. The continuation path has NOT pushed
|
|
563
566
|
// `userMessage` yet, so `runStartPath` pushing it adds no duplicate.
|
|
564
|
-
return await this.runStartPath(userMessage, undefined, undefined, true, false, constrainedConfig);
|
|
567
|
+
return await this.runStartPath(userMessage, undefined, undefined, true, false, constrainedConfig, opts.signal);
|
|
565
568
|
}
|
|
566
569
|
this.history.push(pendingUser);
|
|
567
570
|
this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
|
|
@@ -593,8 +596,8 @@ export class ChatSession {
|
|
|
593
596
|
this.inFlight = true;
|
|
594
597
|
try {
|
|
595
598
|
const mergedConfig = this.mergeConfig(opts.config);
|
|
596
|
-
const newImagesKey =
|
|
597
|
-
const newAudioKey =
|
|
599
|
+
const newImagesKey = computeByteListKey(opts.images);
|
|
600
|
+
const newAudioKey = computeByteListKey(opts.audio);
|
|
598
601
|
// Only an explicit NEW image/audio set can trigger a restart. Omitting
|
|
599
602
|
// `images`/`audio` (key === null) is interpreted as "keep the current
|
|
600
603
|
// media cache state" — the server-side cache already holds any prior
|
|
@@ -732,7 +735,7 @@ export class ChatSession {
|
|
|
732
735
|
this.assertCanSendToolResult('sendToolResult');
|
|
733
736
|
this.inFlight = true;
|
|
734
737
|
try {
|
|
735
|
-
const { isError, config } = opts;
|
|
738
|
+
const { isError, config, signal } = opts;
|
|
736
739
|
const mergedConfig = this.mergeConfig(config);
|
|
737
740
|
const toolMsg = {
|
|
738
741
|
role: 'tool',
|
|
@@ -753,10 +756,10 @@ export class ChatSession {
|
|
|
753
756
|
// tool-call turn (turnCount>=1), so this never fires on the happy
|
|
754
757
|
// path.
|
|
755
758
|
if (this.turnCount === 0 || this.needsFullReplay) {
|
|
756
|
-
return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig);
|
|
759
|
+
return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig, signal);
|
|
757
760
|
}
|
|
758
761
|
try {
|
|
759
|
-
const result = await this.
|
|
762
|
+
const result = await this.runNonStreamingNative('continueTool', pendingHistory, withReplayReasoning(constrainedConfig, this.model), signal);
|
|
760
763
|
this.history.push({ role: 'tool', content, toolCallId, isError });
|
|
761
764
|
this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
|
|
762
765
|
this.turnCount++;
|
|
@@ -779,7 +782,7 @@ export class ChatSession {
|
|
|
779
782
|
// restart core pushes it — `isError` rides on that message so the
|
|
780
783
|
// wire-format error marker is re-rendered, and a tool result
|
|
781
784
|
// always follows >=1 prior turn so `isFirstTurn` is false.
|
|
782
|
-
return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig);
|
|
785
|
+
return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig, signal);
|
|
783
786
|
}
|
|
784
787
|
}
|
|
785
788
|
finally {
|
|
@@ -793,12 +796,12 @@ export class ChatSession {
|
|
|
793
796
|
* when the native session is cold (turnCount===0 — e.g. after an
|
|
794
797
|
* interrupted media-held replay rolled the cache back) and by the
|
|
795
798
|
* media-held rejection catch. `mediaChanged=true` forces a
|
|
796
|
-
*
|
|
797
|
-
*
|
|
798
|
-
*
|
|
799
|
+
* native cache invalidation so the prefill starts from a guaranteed-clean
|
|
800
|
+
* owner; `isFirstTurn=false` because a tool result always follows a prior
|
|
801
|
+
* tool-call turn.
|
|
799
802
|
*/
|
|
800
|
-
async replayToolResultThroughStartPath(toolMsg, config) {
|
|
801
|
-
return await this.runStartPathWithMessage(toolMsg, true, false, config);
|
|
803
|
+
async replayToolResultThroughStartPath(toolMsg, config, signal) {
|
|
804
|
+
return await this.runStartPathWithMessage(toolMsg, true, false, config, signal);
|
|
802
805
|
}
|
|
803
806
|
/**
|
|
804
807
|
* Streaming variant of {@link ChatSession#sendToolResult}.
|
|
@@ -938,42 +941,106 @@ export class ChatSession {
|
|
|
938
941
|
/**
|
|
939
942
|
* Reset the session state.
|
|
940
943
|
*
|
|
941
|
-
*
|
|
942
|
-
*
|
|
943
|
-
* `chatSessionStart` again.
|
|
944
|
-
*
|
|
945
|
-
*
|
|
946
|
-
*
|
|
947
|
-
*
|
|
948
|
-
* is shared across every `ChatSession` lifetime via the native
|
|
949
|
-
* `ModelRegistry`: a partial wipe that leaves the shared native
|
|
950
|
-
* KV cache intact would leak a previous (unrelated) request's
|
|
951
|
-
* cached prefix into the next `chat_session_start_sync` call. The
|
|
952
|
-
* server-side warm-lease replay path (where preserving the native
|
|
953
|
-
* cache is correct) uses its own server-private helper gated by
|
|
954
|
-
* the `SessionRegistry` HIT signal — the only authoritative proof
|
|
955
|
-
* that the native cache genuinely belongs to this chain. That
|
|
956
|
-
* helper lives inside `@mlx-node/server`, never touches the
|
|
957
|
-
* `@mlx-node/lm` export map, and is not reachable from downstream
|
|
958
|
-
* consumers. Public consumers of `@mlx-node/lm` have no such HIT
|
|
959
|
-
* signal, so the public API intentionally offers only the full-wipe
|
|
960
|
-
* option.
|
|
944
|
+
* Invalidates this session's native KV state and wipes local history,
|
|
945
|
+
* media keys, and turn counter so the next `send()` goes through
|
|
946
|
+
* `chatSessionStart` again. Block-paged models release only this stable
|
|
947
|
+
* session owner; clearing the model-wide scheduler would invalidate every
|
|
948
|
+
* other live `ChatSession`. Exclusive/flat models retain the model-wide
|
|
949
|
+
* `resetCaches()` barrier because they have no independently releasable
|
|
950
|
+
* owner state.
|
|
961
951
|
*
|
|
962
|
-
*
|
|
963
|
-
*
|
|
952
|
+
* Native invalidation is async: the release/reset command is awaited so it
|
|
953
|
+
* has fully drained through the model thread before this promise resolves —
|
|
954
|
+
* callers that `await reset()` and
|
|
955
|
+
* then start a turn keep strict command-queue ordering. Because the
|
|
956
|
+
* await genuinely suspends, `reset()` RESERVES the same in-flight
|
|
957
|
+
* guard as the send entry points for its whole duration: any
|
|
958
|
+
* concurrent `send*()`, `reset()`, `primeHistory()`, or
|
|
959
|
+
* `startFromHistory*()` on this session rejects until the reset
|
|
960
|
+
* settles, so a racing turn can never commit against the pre-wipe
|
|
961
|
+
* history (or have its state erased mid-commit). If native invalidation
|
|
962
|
+
* rejects, no JS state is wiped and the guard is released.
|
|
964
963
|
*/
|
|
965
964
|
async reset() {
|
|
965
|
+
if (this.disposed) {
|
|
966
|
+
throw new Error('ChatSession: session has been disposed');
|
|
967
|
+
}
|
|
966
968
|
if (this.inFlight) {
|
|
967
969
|
throw new Error('ChatSession: cannot reset() while a send() is in flight; await the previous call first');
|
|
968
970
|
}
|
|
969
|
-
this.
|
|
970
|
-
|
|
971
|
-
|
|
972
|
-
|
|
973
|
-
|
|
974
|
-
|
|
975
|
-
|
|
976
|
-
|
|
971
|
+
this.inFlight = true;
|
|
972
|
+
try {
|
|
973
|
+
if (this.model.hasBlockPagedCache?.() === true && this.model.releaseCacheOwner) {
|
|
974
|
+
for (const ownerId of Array.from(this.nativeCacheOwnerIds)) {
|
|
975
|
+
await this.model.releaseCacheOwner(ownerId);
|
|
976
|
+
this.nativeCacheOwnerIds.delete(ownerId);
|
|
977
|
+
}
|
|
978
|
+
}
|
|
979
|
+
else {
|
|
980
|
+
await this.model.resetCaches();
|
|
981
|
+
this.nativeCacheOwnerIds.clear();
|
|
982
|
+
}
|
|
983
|
+
this.history = [];
|
|
984
|
+
this.lastImagesKey = null;
|
|
985
|
+
this.lastAudioKey = null;
|
|
986
|
+
this.turnCount = 0;
|
|
987
|
+
this.unresolvedOkToolCallCount = null;
|
|
988
|
+
this.needsFullReplay = false;
|
|
989
|
+
this.activeTools = this.defaultConfig.tools;
|
|
990
|
+
}
|
|
991
|
+
finally {
|
|
992
|
+
this.inFlight = false;
|
|
993
|
+
}
|
|
994
|
+
}
|
|
995
|
+
/**
|
|
996
|
+
* Permanently dispose this JS session and release its native scheduler
|
|
997
|
+
* owner. The operation is awaited and idempotent; it rejects while
|
|
998
|
+
* a turn is in flight so native state cannot be torn down mid-decode.
|
|
999
|
+
*
|
|
1000
|
+
* A caller-supplied first `cacheOwnerId` becomes this session's stable owner
|
|
1001
|
+
* just like the generated default. Do not share an explicit owner id between
|
|
1002
|
+
* independently disposed sessions: disposing either session releases that
|
|
1003
|
+
* owner's native scheduler state for both.
|
|
1004
|
+
*/
|
|
1005
|
+
async dispose() {
|
|
1006
|
+
if (this.disposed)
|
|
1007
|
+
return;
|
|
1008
|
+
if (this.inFlight) {
|
|
1009
|
+
throw new Error('ChatSession: cannot dispose() while a send() is in flight; await the previous call first');
|
|
1010
|
+
}
|
|
1011
|
+
this.inFlight = true;
|
|
1012
|
+
try {
|
|
1013
|
+
let firstReleaseError;
|
|
1014
|
+
let hadReleaseError = false;
|
|
1015
|
+
if (this.model.releaseCacheOwner) {
|
|
1016
|
+
for (const ownerId of Array.from(this.nativeCacheOwnerIds)) {
|
|
1017
|
+
try {
|
|
1018
|
+
await this.model.releaseCacheOwner(ownerId);
|
|
1019
|
+
this.nativeCacheOwnerIds.delete(ownerId);
|
|
1020
|
+
}
|
|
1021
|
+
catch (error) {
|
|
1022
|
+
if (!hadReleaseError)
|
|
1023
|
+
firstReleaseError = error;
|
|
1024
|
+
hadReleaseError = true;
|
|
1025
|
+
}
|
|
1026
|
+
}
|
|
1027
|
+
}
|
|
1028
|
+
else {
|
|
1029
|
+
this.nativeCacheOwnerIds.clear();
|
|
1030
|
+
}
|
|
1031
|
+
if (hadReleaseError)
|
|
1032
|
+
throw firstReleaseError;
|
|
1033
|
+
this.disposed = true;
|
|
1034
|
+
this.history = [];
|
|
1035
|
+
this.lastImagesKey = null;
|
|
1036
|
+
this.lastAudioKey = null;
|
|
1037
|
+
this.turnCount = 0;
|
|
1038
|
+
this.unresolvedOkToolCallCount = null;
|
|
1039
|
+
this.needsFullReplay = false;
|
|
1040
|
+
}
|
|
1041
|
+
finally {
|
|
1042
|
+
this.inFlight = false;
|
|
1043
|
+
}
|
|
977
1044
|
}
|
|
978
1045
|
/**
|
|
979
1046
|
* Prime the session history without running inference.
|
|
@@ -1026,7 +1093,7 @@ export class ChatSession {
|
|
|
1026
1093
|
* stay on the delta path, and subsequent image turns correctly
|
|
1027
1094
|
* trigger restart).
|
|
1028
1095
|
*/
|
|
1029
|
-
async startFromHistory(config) {
|
|
1096
|
+
async startFromHistory(config, opts = {}) {
|
|
1030
1097
|
if (this.inFlight) {
|
|
1031
1098
|
throw new Error('ChatSession: cannot startFromHistory() while a send() is in flight');
|
|
1032
1099
|
}
|
|
@@ -1041,7 +1108,7 @@ export class ChatSession {
|
|
|
1041
1108
|
const mergedConfig = this.mergeConfig(config);
|
|
1042
1109
|
const historySnapshot = this.history.slice();
|
|
1043
1110
|
const constrainedConfig = await this.constrainToContextCapacity(historySnapshot, mergedConfig);
|
|
1044
|
-
const result = await this.
|
|
1111
|
+
const result = await this.runNonStreamingNative('start', historySnapshot, withReplayReasoning(constrainedConfig, this.model), opts.signal);
|
|
1045
1112
|
this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
|
|
1046
1113
|
this.turnCount++;
|
|
1047
1114
|
this.needsFullReplay = false;
|
|
@@ -1137,6 +1204,34 @@ export class ChatSession {
|
|
|
1137
1204
|
// -------------------------------------------------------------------
|
|
1138
1205
|
// Internal helpers
|
|
1139
1206
|
// -------------------------------------------------------------------
|
|
1207
|
+
/**
|
|
1208
|
+
* Dispatch one non-streaming native turn through the public AbortSignal
|
|
1209
|
+
* surface (H2). The model wrapper owns the attach-race-safe translation
|
|
1210
|
+
* from AbortSignal to the lower-level native operation handle; ChatSession
|
|
1211
|
+
* never exposes that two-phase primitive to callers.
|
|
1212
|
+
*
|
|
1213
|
+
* Callers sit inside the entry points' existing try/finally blocks,
|
|
1214
|
+
* so a rejection here releases `inFlight` and (on the delta paths)
|
|
1215
|
+
* sets `needsFullReplay` exactly like any other native failure.
|
|
1216
|
+
*/
|
|
1217
|
+
async runNonStreamingNative(kind, messages, config, signal) {
|
|
1218
|
+
const model = this.model;
|
|
1219
|
+
if (signal?.aborted === true)
|
|
1220
|
+
throw new Error('chat session cancelled');
|
|
1221
|
+
if (kind === 'start') {
|
|
1222
|
+
return signal == null
|
|
1223
|
+
? await model.chatSessionStart(messages, config)
|
|
1224
|
+
: await model.chatSessionStart(messages, config, signal);
|
|
1225
|
+
}
|
|
1226
|
+
if (kind === 'continue') {
|
|
1227
|
+
return signal == null
|
|
1228
|
+
? await model.chatSessionContinue(messages, config)
|
|
1229
|
+
: await model.chatSessionContinue(messages, config, signal);
|
|
1230
|
+
}
|
|
1231
|
+
return signal == null
|
|
1232
|
+
? await model.chatSessionContinueTool(messages, config)
|
|
1233
|
+
: await model.chatSessionContinueTool(messages, config, signal);
|
|
1234
|
+
}
|
|
1140
1235
|
/**
|
|
1141
1236
|
* Gate plain-text continuation entry points (`send`, `sendStream`)
|
|
1142
1237
|
* on the tool-call resolution invariant. Any outstanding `ok` tool
|
|
@@ -1242,25 +1337,46 @@ export class ChatSession {
|
|
|
1242
1337
|
* cache the delta depends on.
|
|
1243
1338
|
*
|
|
1244
1339
|
* Tool resolution here is side-effect free because public capacity
|
|
1245
|
-
* preflights use this same merge path
|
|
1246
|
-
* {@link commitActiveTools} after native
|
|
1340
|
+
* preflights use this same merge path with owner establishment disabled.
|
|
1341
|
+
* Successful turn commit sites call {@link commitActiveTools} after native
|
|
1342
|
+
* inference finishes.
|
|
1247
1343
|
*
|
|
1248
1344
|
* MTP auto-default: if neither `defaultConfig` nor `overlay`
|
|
1249
1345
|
* sets `enableMtp` AND the underlying model exposes
|
|
1250
|
-
* `hasMtpWeights()` returning `true
|
|
1251
|
-
*
|
|
1252
|
-
*
|
|
1253
|
-
*
|
|
1254
|
-
*
|
|
1255
|
-
*
|
|
1256
|
-
*
|
|
1346
|
+
* `hasMtpWeights()` returning `true` AND
|
|
1347
|
+
* {@link ChatSession#mtpAutoDefaultAllowed} agrees, set
|
|
1348
|
+
* `enableMtp = true` so the speculative-decode path runs out of the
|
|
1349
|
+
* box on MTP-capable checkpoints. An explicit `false` from either
|
|
1350
|
+
* source wins (the undefined-check below preserves it). This
|
|
1351
|
+
* duck-typed check also covers Gemma4 with an external draft
|
|
1352
|
+
* attached — DSpark or Google assistant (`hasMtpWeights()` reports
|
|
1353
|
+
* the external draft there, not in-checkpoint MTP heads).
|
|
1257
1354
|
*/
|
|
1258
|
-
mergeConfig(overlay) {
|
|
1355
|
+
mergeConfig(overlay, establishCacheOwner = true) {
|
|
1356
|
+
if (this.disposed) {
|
|
1357
|
+
throw new Error('ChatSession: session has been disposed');
|
|
1358
|
+
}
|
|
1259
1359
|
const merged = {
|
|
1260
1360
|
...this.defaultConfig,
|
|
1261
1361
|
...overlay,
|
|
1262
1362
|
reuseCache: true,
|
|
1263
1363
|
};
|
|
1364
|
+
const requestedOwnerId = merged.cacheOwnerId;
|
|
1365
|
+
if (this.cacheOwnerId === null) {
|
|
1366
|
+
if (establishCacheOwner) {
|
|
1367
|
+
this.cacheOwnerId = requestedOwnerId === undefined || requestedOwnerId === '' ? randomUUID() : requestedOwnerId;
|
|
1368
|
+
merged.cacheOwnerId = this.cacheOwnerId;
|
|
1369
|
+
this.nativeCacheOwnerIds.add(this.cacheOwnerId);
|
|
1370
|
+
}
|
|
1371
|
+
}
|
|
1372
|
+
else if (requestedOwnerId !== undefined && requestedOwnerId !== '' && requestedOwnerId !== this.cacheOwnerId) {
|
|
1373
|
+
throw new Error(`ChatSession: cacheOwnerId cannot change after the session owner is established; create a new ChatSession for owner ${requestedOwnerId}`);
|
|
1374
|
+
}
|
|
1375
|
+
else {
|
|
1376
|
+
merged.cacheOwnerId = this.cacheOwnerId;
|
|
1377
|
+
if (establishCacheOwner)
|
|
1378
|
+
this.nativeCacheOwnerIds.add(this.cacheOwnerId);
|
|
1379
|
+
}
|
|
1264
1380
|
// Tools are part of the committed conversation state. Constructor
|
|
1265
1381
|
// defaults seed that state, but must not overwrite a tool set committed by
|
|
1266
1382
|
// a later successful turn. A current-call overlay is the only higher
|
|
@@ -1271,11 +1387,32 @@ export class ChatSession {
|
|
|
1271
1387
|
}
|
|
1272
1388
|
if (merged.enableMtp === undefined &&
|
|
1273
1389
|
typeof this.model.hasMtpWeights === 'function' &&
|
|
1274
|
-
this.model.hasMtpWeights()
|
|
1390
|
+
this.model.hasMtpWeights() &&
|
|
1391
|
+
this.mtpAutoDefaultAllowed()) {
|
|
1275
1392
|
merged.enableMtp = true;
|
|
1276
1393
|
}
|
|
1277
1394
|
return merged;
|
|
1278
1395
|
}
|
|
1396
|
+
/**
|
|
1397
|
+
* Whether an MTP-capable model wants `enableMtp` on for a caller who set
|
|
1398
|
+
* nothing. `mtpAutoEnabled()` is authoritative in both directions when
|
|
1399
|
+
* present; otherwise {@link MTP_AUTO_DEFAULT_SUPPRESSED_MODELS} is matched
|
|
1400
|
+
* along the prototype CHAIN rather than by a bare `constructor.name`, so it
|
|
1401
|
+
* still fires through the `makeStreamingModel` wrapper subclass that
|
|
1402
|
+
* `@mlx-node/lm` actually hands to `ChatSession`.
|
|
1403
|
+
*/
|
|
1404
|
+
mtpAutoDefaultAllowed() {
|
|
1405
|
+
if (typeof this.model.mtpAutoEnabled === 'function') {
|
|
1406
|
+
return this.model.mtpAutoEnabled();
|
|
1407
|
+
}
|
|
1408
|
+
for (let proto = Object.getPrototypeOf(this.model); proto !== null; proto = Object.getPrototypeOf(proto)) {
|
|
1409
|
+
const name = proto.constructor?.name;
|
|
1410
|
+
if (name !== undefined && MTP_AUTO_DEFAULT_SUPPRESSED_MODELS.has(name)) {
|
|
1411
|
+
return false;
|
|
1412
|
+
}
|
|
1413
|
+
}
|
|
1414
|
+
return true;
|
|
1415
|
+
}
|
|
1279
1416
|
/**
|
|
1280
1417
|
* Render the exact full prompt and constrain generation to the physical KV
|
|
1281
1418
|
* window before native code allocates a block. Models that do not expose
|
|
@@ -1336,9 +1473,9 @@ export class ChatSession {
|
|
|
1336
1473
|
* native side gets the full conversation re-rendered with the new
|
|
1337
1474
|
* image set.
|
|
1338
1475
|
*/
|
|
1339
|
-
async runStartPath(userMessage, images, audio, mediaChanged, isFirstTurn, config) {
|
|
1476
|
+
async runStartPath(userMessage, images, audio, mediaChanged, isFirstTurn, config, signal) {
|
|
1340
1477
|
const userMsg = this.buildUserMessage(userMessage, images, audio);
|
|
1341
|
-
return await this.runStartPathWithMessage(userMsg, mediaChanged, isFirstTurn, config);
|
|
1478
|
+
return await this.runStartPathWithMessage(userMsg, mediaChanged, isFirstTurn, config, signal);
|
|
1342
1479
|
}
|
|
1343
1480
|
/**
|
|
1344
1481
|
* Core of {@link runStartPath} that takes a PRE-BUILT pending
|
|
@@ -1348,27 +1485,28 @@ export class ChatSession {
|
|
|
1348
1485
|
* message so the tool-result turn is re-rendered against the full
|
|
1349
1486
|
* history without duplicating the start-path bookkeeping.
|
|
1350
1487
|
*/
|
|
1351
|
-
async runStartPathWithMessage(pendingMessage, mediaChanged, isFirstTurn, config) {
|
|
1488
|
+
async runStartPathWithMessage(pendingMessage, mediaChanged, isFirstTurn, config, signal) {
|
|
1352
1489
|
// Capture pre-state so the restart can be rolled back if the
|
|
1353
|
-
// native call fails. The media-change branch
|
|
1354
|
-
// we know whether the new prefill will succeed, so on failure we
|
|
1355
|
-
// also have to drop turnCount + lastImagesKey/lastAudioKey to force
|
|
1356
|
-
//
|
|
1357
|
-
//
|
|
1490
|
+
// native call fails. The media-change branch releases this owner's caches
|
|
1491
|
+
// BEFORE we know whether the new prefill will succeed, so on failure we
|
|
1492
|
+
// also have to drop turnCount + lastImagesKey/lastAudioKey to force the
|
|
1493
|
+
// next call to re-route through the start path (rather than a delta
|
|
1494
|
+
// continue against released caches).
|
|
1358
1495
|
const wasMediaChangeRestart = mediaChanged && !isFirstTurn;
|
|
1359
1496
|
const historyLenBefore = this.history.length;
|
|
1360
1497
|
// Capacity validation must happen before `prepareStartPath()` because a
|
|
1361
|
-
// media-change restart
|
|
1362
|
-
// is a request error and must leave both JS history and native state
|
|
1498
|
+
// media-change restart releases native owner state. A rejected oversized
|
|
1499
|
+
// prompt is a request error and must leave both JS history and native state
|
|
1500
|
+
// intact.
|
|
1363
1501
|
const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingMessage), config);
|
|
1364
|
-
this.prepareStartPath(mediaChanged, isFirstTurn);
|
|
1502
|
+
await this.prepareStartPath(mediaChanged, isFirstTurn, constrainedConfig.cacheOwnerId);
|
|
1365
1503
|
this.history.push(pendingMessage);
|
|
1366
1504
|
try {
|
|
1367
1505
|
// Pass a shallow snapshot so later pushes to `this.history`
|
|
1368
1506
|
// (e.g. the assistant reply below) don't retroactively mutate
|
|
1369
1507
|
// what the native side / any mock observed as its `messages`
|
|
1370
1508
|
// argument.
|
|
1371
|
-
const result = await this.
|
|
1509
|
+
const result = await this.runNonStreamingNative('start', this.history.slice(), withReplayReasoning(constrainedConfig, this.model), signal);
|
|
1372
1510
|
this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
|
|
1373
1511
|
this.turnCount++;
|
|
1374
1512
|
this.needsFullReplay = false;
|
|
@@ -1390,7 +1528,7 @@ export class ChatSession {
|
|
|
1390
1528
|
// consistent with turnCount.
|
|
1391
1529
|
this.history.length = historyLenBefore;
|
|
1392
1530
|
if (wasMediaChangeRestart) {
|
|
1393
|
-
//
|
|
1531
|
+
// This owner's caches were released by prepareStartPath() but the new prefill
|
|
1394
1532
|
// failed. Force the next call to re-route through the start
|
|
1395
1533
|
// path with the (preserved) prior history.
|
|
1396
1534
|
this.turnCount = 0;
|
|
@@ -1416,10 +1554,10 @@ export class ChatSession {
|
|
|
1416
1554
|
// See `runStartPath` for the full rationale.
|
|
1417
1555
|
const wasMediaChangeRestart = mediaChanged && !isFirstTurn;
|
|
1418
1556
|
const historyLenBefore = this.history.length;
|
|
1419
|
-
// See the sync start path: reject before a media restart can
|
|
1420
|
-
// state, and make the output budget explicit before allocating KV blocks.
|
|
1557
|
+
// See the sync start path: reject before a media restart can release native
|
|
1558
|
+
// owner state, and make the output budget explicit before allocating KV blocks.
|
|
1421
1559
|
const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingMessage), config);
|
|
1422
|
-
this.prepareStartPath(mediaChanged, isFirstTurn);
|
|
1560
|
+
await this.prepareStartPath(mediaChanged, isFirstTurn, constrainedConfig.cacheOwnerId);
|
|
1423
1561
|
// Stage the pending message on the pending history BEFORE the
|
|
1424
1562
|
// stream starts — the native call reads it synchronously via
|
|
1425
1563
|
// `model.chatStreamSessionStart(history, config)`.
|
|
@@ -1490,7 +1628,7 @@ export class ChatSession {
|
|
|
1490
1628
|
// consistent with turnCount.
|
|
1491
1629
|
this.history.length = historyLenBefore;
|
|
1492
1630
|
if (wasMediaChangeRestart) {
|
|
1493
|
-
//
|
|
1631
|
+
// This owner's caches were released by prepareStartPath() but the new
|
|
1494
1632
|
// prefill never reached a successful done:true. Force the
|
|
1495
1633
|
// next call to re-route through the start path with the
|
|
1496
1634
|
// preserved prior history.
|
|
@@ -1504,8 +1642,9 @@ export class ChatSession {
|
|
|
1504
1642
|
/**
|
|
1505
1643
|
* Shared pre-start bookkeeping for both `send()` and `sendStream()`:
|
|
1506
1644
|
*
|
|
1507
|
-
* - On an image-change restart (turn >= 1),
|
|
1508
|
-
*
|
|
1645
|
+
* - On an image-change restart (turn >= 1), release this session's native
|
|
1646
|
+
* cache owner so the new image set gets a fresh prefill without wiping
|
|
1647
|
+
* concurrently live owners. History is
|
|
1509
1648
|
* intentionally preserved — `chatSessionStart` receives the full
|
|
1510
1649
|
* accumulated conversation plus the new user turn so the jinja
|
|
1511
1650
|
* render walks every prior turn and every prior image again
|
|
@@ -1516,9 +1655,19 @@ export class ChatSession {
|
|
|
1516
1655
|
* turn.
|
|
1517
1656
|
* - On a fresh / reset history, re-inject the system prompt.
|
|
1518
1657
|
*/
|
|
1519
|
-
prepareStartPath(mediaChanged, isFirstTurn) {
|
|
1658
|
+
async prepareStartPath(mediaChanged, isFirstTurn, cacheOwnerId) {
|
|
1520
1659
|
if (mediaChanged && !isFirstTurn) {
|
|
1521
|
-
|
|
1660
|
+
// Await owner release so it has drained through the model thread before
|
|
1661
|
+
// the replacement start is enqueued. This is deliberately owner-scoped
|
|
1662
|
+
// on continuously batched models: a model-wide reset would invalidate
|
|
1663
|
+
// unrelated live ChatSessions. Third-party/exclusive models predating
|
|
1664
|
+
// the owner lifecycle retain the full-reset fallback.
|
|
1665
|
+
if (this.model.hasBlockPagedCache?.() === true && this.model.releaseCacheOwner && cacheOwnerId) {
|
|
1666
|
+
await this.model.releaseCacheOwner(cacheOwnerId);
|
|
1667
|
+
}
|
|
1668
|
+
else {
|
|
1669
|
+
await this.model.resetCaches();
|
|
1670
|
+
}
|
|
1522
1671
|
}
|
|
1523
1672
|
if (this.history.length === 0 && this.system != null) {
|
|
1524
1673
|
this.history.push({ role: 'system', content: this.system });
|
|
@@ -1544,7 +1693,7 @@ export class ChatSession {
|
|
|
1544
1693
|
for (let i = this.history.length - 1; i >= 0; i--) {
|
|
1545
1694
|
const msg = this.history[i];
|
|
1546
1695
|
if (msg?.role === 'user' && msg.images && msg.images.length > 0) {
|
|
1547
|
-
return
|
|
1696
|
+
return computeByteListKey(msg.images);
|
|
1548
1697
|
}
|
|
1549
1698
|
}
|
|
1550
1699
|
return null;
|
|
@@ -1558,7 +1707,7 @@ export class ChatSession {
|
|
|
1558
1707
|
for (let i = this.history.length - 1; i >= 0; i--) {
|
|
1559
1708
|
const msg = this.history[i];
|
|
1560
1709
|
if (msg?.role === 'user' && msg.audio && msg.audio.length > 0) {
|
|
1561
|
-
return
|
|
1710
|
+
return computeByteListKey(msg.audio);
|
|
1562
1711
|
}
|
|
1563
1712
|
}
|
|
1564
1713
|
return null;
|