@mlx-node/lm 0.0.7 → 0.0.8

package/dist/chat-session.js

@@ -1,3 +1,150 @@
+/**
+ * Generic server-side chat session wrapper.
+ *
+ * `ChatSession<M>` is the cross-model chat-session wrapper. It works
+ * against any model that exposes the uniform chat-session NAPI
+ * surface — `chatSessionStart`,
+ * `chatSessionContinue`, `chatSessionContinueTool`, and their
+ * streaming variants plus `resetCaches`. See `SessionCapableModel`
+ * below.
+ *
+ * Design notes:
+ *
+ *   - The session tracks its own `ChatMessage[]` history on the
+ *     TypeScript side. In the common text-continue case the history
+ *     is only appended to and never read back — each `send()` on
+ *     turn >= 1 issues a cheap `chatSessionContinue` delta against
+ *     the live KV cache. The history is kept purely so the
+ *     image-change mid-session path can call `chatSessionStart` with
+ *     the full rebuilt history for a clean re-prefill.
+ *
+ *   - An image hash (`lastImagesKey`) tracks the images bound to the
+ *     current cache. A `send()` call whose image set has changed
+ *     (different bytes or different ordering) triggers a full
+ *     restart: `resetCaches()` → push the new user message (with
+ *     images) to history → `chatSessionStart(history)`.
+ *
+ *   - Text-only `send()` on turn >= 1 takes the cheap delta path.
+ *
+ *   - `sendToolResult` always dispatches `chatSessionContinueTool`,
+ *     since tool turns never change image state. The session enforces
+ *     a strict unresolved-ok-tool-call contract at runtime, driven by
+ *     `unresolvedOkToolCallCount` (derived from `ChatResult.toolCalls`
+ *     after each turn via `countOkToolCalls` /
+ *     `computeTrailingAssistantUnresolvedToolCallCount`):
+ *
+ *       * `null` — the trailing assistant turn has no outstanding ok
+ *         tool call. Plain `send()` / `sendStream()` are the only
+ *         valid entry points; `sendToolResult*()` throws because
+ *         there is nothing for the result to resolve.
+ *       * `1` — exactly one outstanding ok tool call. Plain `send()` /
+ *         `sendStream()` throw (they would orphan the call);
+ *         `sendToolResult*()` is the sole valid forward step and
+ *         dispatches the tool result through the native session.
+ *       * `>1` — a multi-tool-call fan-out that the chat-session API
+ *         cannot progress incrementally (each `sendToolResult*` would
+ *         re-open the assistant turn and weave new replies between
+ *         the sibling results). Both `send()` / `sendStream()` and
+ *         `sendToolResult*()` throw. The only valid recovery is
+ *         `reset()` or `primeHistory()` + `startFromHistory*()` with
+ *         a fully-resolved conversation — there is no "advance past
+ *         the broken turn" path. This mirrors the native ChatML delta
+ *         format which would otherwise silently corrupt multi-call
+ *         conversations.
+ *
+ *   - `sawFinal` gates `turnCount` advance on the streaming path, so
+ *     the session refuses to advance when the stream throws
+ *     mid-decode or yields a final chunk with
+ *     `finishReason: 'error'`.
+ *
+ *   - The `inFlight` guard rejects concurrent `send()` /
+ *     `sendStream()` calls at the class level. The native side
+ *     serializes cache mutation on a single worker thread, so a
+ *     second in-flight call would race the first's cache-save step.
+ *
+ *   - **Cold-restart primitives.** `primeHistory()` plus
+ *     `startFromHistory()` / `startFromHistoryStream()` let a caller
+ *     seed a fresh session with an externally-reconstructed history
+ *     (e.g. a server `ResponseStore` chain) and replay it through the
+ *     native `chatSessionStart` path without going through `send()`.
+ *     These are intended for server-side `SessionRegistry` cache-miss
+ *     cold-start; normal usage stays on `send` / `sendStream` /
+ *     `sendToolResult` / `reset`.
+ *
+ * ## Typical usage
+ *
+ * ```typescript
+ * import { Qwen35Model, ChatSession } from '@mlx-node/lm';
+ *
+ * const model = await Qwen35Model.load('./models/qwen3.5-0.8b');
+ * const session = new ChatSession(model, { system: 'Be concise.' });
+ * const r1 = await session.send('Say hi in one word.');
+ * const r2 = await session.send('Another word?');
+ * await session.reset();
+ * ```
+ */
+import { createHash } from 'node:crypto';
+/**
+ * Typed prefix the native delta path uses to reject a text-only
+ * continuation while the session still holds image/audio KV state
+ * (gemma4 raises this after a media turn). The native session refuses
+ * to advance the cheap delta on top of media KV, so the session layer
+ * recognizes this exact prefix and transparently replays the whole
+ * conversation through the cold start path instead of surfacing the
+ * raw error to the caller.
+ *
+ * MUST stay byte-for-byte identical to the Rust constant
+ * `IMAGE_CHANGE_RESTART_PREFIX` in
+ * `crates/mlx-core/src/engine/cache.rs` — it is not exported across the
+ * NAPI boundary, so the two literals are kept in sync by hand. The
+ * native message starts with this prefix and is delivered as the
+ * `Error.message`: on the sync delta path as a rejected promise, and on
+ * the streaming delta path as a thrown error on the generator's first
+ * iteration (the native worker-thread sink error is re-thrown by the
+ * `packages/lm/src/stream.ts` bridge before any chunk is yielded).
+ */
+const IMAGE_CHANGE_RESTART_PREFIX = 'IMAGE_CHANGE_REQUIRES_SESSION_RESTART:';
+/**
+ * Default resolved by the native shared chat engine when `maxNewTokens` is
+ * absent. Keep this in sync with `extract_chat_params()` in
+ * `crates/mlx-core/src/engine/params.rs`.
+ */
+const NATIVE_DEFAULT_MAX_NEW_TOKENS = 2048;
+/**
+ * Stable, provider-neutral error raised before native inference when a
+ * rendered prompt cannot fit in the model's physically available hot KV
+ * window. The marker is intentionally the canonical string recognized by
+ * pi's overflow recovery, so managed agent sessions compact and retry while
+ * stateless HTTP callers receive a clean request error instead of a native
+ * `BlockAllocator exhausted` failure.
+ */
+export class ContextCapacityError extends Error {
+    promptTokens;
+    effectiveWindowTokens;
+    code = 'context_length_exceeded';
+    constructor(promptTokens, effectiveWindowTokens) {
+        super(`context_length_exceeded: rendered prompt uses ${promptTokens} tokens, ` +
+            `but this model currently has capacity for ${effectiveWindowTokens} tokens`);
+        this.promptTokens = promptTokens;
+        this.effectiveWindowTokens = effectiveWindowTokens;
+        this.name = 'ContextCapacityError';
+    }
+}
+/** Recognize both the typed JS preflight and the native hard backstop. */
+export function isContextCapacityError(error) {
+    return (error instanceof ContextCapacityError ||
+        (error instanceof Error && error.message.startsWith('context_length_exceeded:')));
+}
+/**
+ * Whether `err` is the native media-held delta rejection (see
+ * {@link IMAGE_CHANGE_RESTART_PREFIX}). The native message begins with
+ * the literal prefix and reaches both the sync and streaming bridges
+ * unwrapped (NAPI surfaces `Error.from_reason` as `Error.message`
+ * verbatim), so a `startsWith` match is exact.
+ */
+function isMediaHeldRestartError(err) {
+    return err instanceof Error && err.message.startsWith(IMAGE_CHANGE_RESTART_PREFIX);
+}
 /**
  * Convert the parsed `ToolCallResult[]` emitted by the native chat
  * pipeline into the `ToolCall[]` shape expected by
@@ -94,75 +241,58 @@ function countOkToolCalls(toolCalls) {
  * template.
  *
  * This is a byte-identity check — callers use the key solely to
- * decide whether to restart the server-side session, so a
- * non-cryptographic hash is sufficient. We use FNV-1a 64-bit with a
- * length-prefixed framing so different image counts and different
- * byte lengths cannot collide by accident.
+ * decide whether to restart the server-side session, so any
+ * collision-resistant digest is sufficient. We use SHA-256 (native
+ * `node:crypto`) with a length-prefixed framing so different image
+ * counts and different byte lengths cannot collide by accident.
  *
  * Implementation note: kept fully sync + self-contained so
- * `send()` can stay synchronous in its routing decision and so the
- * module has no external runtime dependencies beyond `@mlx-node/core`
- * and the existing stream bridge.
+ * `send()` can stay synchronous in its routing decision. `node:crypto`
+ * is a Node built-in, so this adds no external runtime dependency
+ * beyond `@mlx-node/core` and the existing stream bridge.
  */
 function computeImagesKey(images) {
-    if (!images || images.length === 0)
+    return computeByteListKey(images);
+}
+/**
+ * Audio counterpart of {@link computeImagesKey}: a stable, order-sensitive
+ * byte-identity key for a list of encoded audio buffers. Used by `send()` /
+ * `sendStream()` to decide whether a new audio set must cold-restart the
+ * server-side session. Shares the exact SHA-256 framing as the image key.
+ */
+function computeAudioKey(audio) {
+    return computeByteListKey(audio);
+}
+/**
+ * SHA-256 byte-identity key for a length-framed list of byte buffers.
+ * Returns `null` for an empty/absent list so callers can distinguish
+ * "no media" from "media changed". Shared by the image and audio keys.
+ *
+ * Uses `node:crypto`'s native SHA-256 rather than a hand-rolled JS hash
+ * loop: hashing large image/audio buffers byte-at-a-time in JS is
+ * 25-60x slower than the native digest (measured: 5MB ~105ms JS loop
+ * vs ~1.8ms native) and this runs synchronously on the event loop
+ * before any `await` in `send()`/`sendStream()`, so the JS loop's cost
+ * was a real head-of-line-blocking stall for every other request
+ * handled by the same process.
+ */
+function computeByteListKey(buffers) {
+    if (!buffers || buffers.length === 0)
         return null;
-    // FNV-1a 64-bit. Split into two 32-bit halves because JavaScript
-    // doesn't have a native 64-bit integer type and BigInt ops are
-    // slow on large byte streams. This emulates 64-bit FNV-1a using
-    // paired 32-bit lo/hi halves — the standard JS idiom.
-    const FNV_OFFSET_LO = 0x84222325 >>> 0;
-    const FNV_OFFSET_HI = 0xcbf29ce4 >>> 0;
-    const FNV_PRIME_LO = 0x000001b3 >>> 0;
-    const FNV_PRIME_HI = 0x00000100 >>> 0;
-    let lo = FNV_OFFSET_LO;
-    let hi = FNV_OFFSET_HI;
-    function mix(byte) {
-        lo = (lo ^ byte) >>> 0;
-        // Multiply (hi:lo) by (FNV_PRIME_HI:FNV_PRIME_LO) mod 2^64.
-        // Break 32-bit halves into 16-bit quarters to keep intermediate
-        // products inside the safe-integer range.
-        const loLo = lo & 0xffff;
-        const loHi = lo >>> 16;
-        const hiLo = hi & 0xffff;
-        const hiHi = hi >>> 16;
-        const pLo = FNV_PRIME_LO & 0xffff;
-        const pLoH = FNV_PRIME_LO >>> 16;
-        const pHi = FNV_PRIME_HI & 0xffff;
-        const pHiH = FNV_PRIME_HI >>> 16;
-        const r0 = loLo * pLo;
-        const r1 = loLo * pLoH + loHi * pLo;
-        const r2 = loLo * pHi + loHi * pLoH + hiLo * pLo;
-        const r3 = loLo * pHiH + loHi * pHi + hiLo * pLoH + hiHi * pLo;
-        const newLo0 = r0 & 0xffff;
-        const carry1 = r0 >>> 16;
-        const sum1 = r1 + carry1;
-        const newLo1 = sum1 & 0xffff;
-        const carry2 = Math.floor(sum1 / 0x10000);
-        const sum2 = r2 + carry2;
-        const newHi0 = sum2 & 0xffff;
-        const carry3 = Math.floor(sum2 / 0x10000);
-        const sum3 = r3 + carry3;
-        const newHi1 = sum3 & 0xffff;
-        lo = ((newLo1 << 16) | newLo0) >>> 0;
-        hi = ((newHi1 << 16) | newHi0) >>> 0;
-    }
-    // Frame each image with a 4-byte little-endian length prefix so
-    // `[ab, c]` and `[a, bc]` hash to distinct values.
-    mix(images.length & 0xff);
-    mix((images.length >>> 8) & 0xff);
-    mix((images.length >>> 16) & 0xff);
-    mix((images.length >>> 24) & 0xff);
-    for (const img of images) {
-        mix(img.byteLength & 0xff);
-        mix((img.byteLength >>> 8) & 0xff);
-        mix((img.byteLength >>> 16) & 0xff);
-        mix((img.byteLength >>> 24) & 0xff);
-        for (let i = 0; i < img.byteLength; i++) {
-            mix(img[i]);
-        }
-    }
-    return hi.toString(16).padStart(8, '0') + lo.toString(16).padStart(8, '0');
+    const hash = createHash('sha256');
+    // Frame each buffer with a 4-byte little-endian length prefix (and a
+    // leading count prefix) so `[ab, c]` and `[a, bc]` — and different
+    // buffer counts — hash to distinct values.
+    const prefix = new Uint8Array(4);
+    const prefixView = new DataView(prefix.buffer);
+    prefixView.setUint32(0, buffers.length, true);
+    hash.update(prefix);
+    for (const buf of buffers) {
+        prefixView.setUint32(0, buf.byteLength, true);
+        hash.update(prefix);
+        hash.update(buf);
+    }
+    return hash.digest('hex');
 }
 /**
  * Cross-model chat session. See module docstring for design notes.
@@ -185,13 +315,22 @@ export class ChatSession {
     history = [];
     /**
      * Hex-encoded byte-identity key of the image set currently bound
-     * to the server's KV cache (FNV-1a 64-bit; see `computeImagesKey`).
+     * to the server's KV cache (SHA-256; see `computeImagesKey`).
      * `null` when no images are cached. A `send()` whose new key
      * differs triggers a full `chatSessionStart` restart.
      */
     lastImagesKey = null;
+    /**
+     * Hex-encoded byte-identity key of the audio set currently bound to the
+     * server's KV cache (see {@link computeAudioKey}). `null` when no audio is
+     * cached. A `send()` whose new key differs triggers a full
+     * `chatSessionStart` restart — the audio counterpart of `lastImagesKey`.
+     */
+    lastAudioKey = null;
     turnCount = 0;
     inFlight = false;
+    /** A failed/abandoned native delta must be followed by a full replay. */
+    needsFullReplay = false;
     /**
      * Count of `ok` tool calls emitted by the prior assistant turn, or
      * `null` when the prior turn produced none. Gates every continuation
@@ -229,6 +368,55 @@ export class ChatSession {
     get hasImages() {
         return this.lastImagesKey !== null;
     }
+    /** Load-time physical context snapshot, when exposed by the model. */
+    contextLimits() {
+        return this.model.contextLimits?.();
+    }
+    /** Authoritative image-input capability of the loaded native model. */
+    supportsImages() {
+        return this.model.supportsImages?.() === true;
+    }
+    /**
+     * Render and validate a complete message list against the model's physical
+     * context window without starting inference or mutating session/native cache
+     * state.
+     *
+     * HTTP streaming callers use this before committing SSE headers so an
+     * oversized prompt can still receive a protocol-shaped 400 response without
+     * delaying those headers until image processing, prefill, or the first
+     * generated token. The returned config carries the same output-budget clamp
+     * applied by the send entry points; those entry points intentionally repeat
+     * the check against their own authoritative history before native dispatch.
+     */
+    async preflightContextCapacity(messages, config) {
+        if (this.inFlight) {
+            throw new Error('ChatSession: cannot preflight context capacity while a send() is in flight');
+        }
+        return await this.constrainToContextCapacity(messages.slice(), this.mergeConfig(config));
+    }
+    /**
+     * Capacity-preflight one pending user/tool message against this session's
+     * preserved history without starting inference or mutating cache state.
+     *
+     * This is the exact counterpart of the delta `send*` paths. It matters for
+     * server-side prompt-cache hits where the HTTP request contains only the new
+     * message while the leased ChatSession owns the earlier conversation.
+     */
+    async preflightPendingContextCapacity(pending, config) {
+        if (this.inFlight) {
+            throw new Error('ChatSession: cannot preflight pending context capacity while a send() is in flight');
+        }
+        if (pending.role === 'user') {
+            this.assertCanSendPlain('sendStream');
+        }
+        else if (pending.role === 'tool') {
+            this.assertCanSendToolResult('sendToolResultStream');
+        }
+        else {
+            throw new Error('ChatSession: pending context capacity preflight requires a user or tool message');
+        }
+        return await this.constrainToContextCapacity(this.historyWithPending(pending), this.mergeConfig(config));
+    }
     /**
      * Count of `ok` tool calls from the most recent assistant turn, or
      * `null` when the trailing turn produced none. Non-null means the
@@ -265,23 +453,43 @@ export class ChatSession {
         try {
             const mergedConfig = this.mergeConfig(opts.config);
             const newImagesKey = computeImagesKey(opts.images);
-            // Only an explicit NEW image set can trigger a restart. Omitting
-            // `images` (newImagesKey === null) is interpreted as "keep the
-            // current image cache state" — the server-side cache already
-            // holds any prior image context, so a text-only follow-up like
-            // "what about the top-right?" can stay on the cheap delta path
-            // even after an image turn.
+            const newAudioKey = computeAudioKey(opts.audio);
+            // Only an explicit NEW image/audio set can trigger a restart. Omitting
+            // `images`/`audio` (key === null) is interpreted as "keep the current
+            // media cache state" — the server-side cache already holds any prior
+            // media context, so a text-only follow-up like "what about the
+            // top-right?" can stay on the cheap delta path even after a media turn.
             const imageChanged = newImagesKey !== null && newImagesKey !== this.lastImagesKey;
+            const audioChanged = newAudioKey !== null && newAudioKey !== this.lastAudioKey;
             const isFirstTurn = this.turnCount === 0;
-            if (isFirstTurn || imageChanged) {
-                return await this.runStartPath(userMessage, opts.images, newImagesKey, imageChanged, isFirstTurn, mergedConfig);
+            const replayRequired = this.needsFullReplay;
+            if (isFirstTurn || imageChanged || audioChanged || replayRequired) {
+                return await this.runStartPath(userMessage, opts.images, opts.audio, imageChanged || audioChanged || replayRequired, isFirstTurn, mergedConfig);
+            }
+            // Delta continue: text-only, images/audio always null. The server
+            // cache already holds all prior turns (including any media from an
+            // earlier restart), so we only need to ship the new user string.
+            const pendingUser = { role: 'user', content: userMessage };
+            const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingUser), mergedConfig);
+            let result;
+            try {
+                result = await this.model.chatSessionContinue(userMessage, null, null, constrainedConfig);
+            }
+            catch (err) {
+                if (!isMediaHeldRestartError(err)) {
+                    this.needsFullReplay = true;
+                    throw err;
+                }
+                // The native session holds media KV (gemma4 after an image/audio
+                // turn) and refused the text delta. Transparently replay the full
+                // conversation through the cold start path. The earlier media turn
+                // already lives in `this.history`, so the start path re-renders it;
+                // the trailing-media keys keep `lastImagesKey`/`lastAudioKey`
+                // consistent across the replay. The delta path has NOT pushed
+                // `userMessage` yet, so `runStartPath` pushing it adds no duplicate.
+                return await this.runStartPath(userMessage, undefined, undefined, true, false, constrainedConfig);
             }
-            // Delta continue: text-only, images always null. The server
-            // cache already holds all prior turns (including any images
-            // from an earlier restart), so we only need to ship the new
-            // user string.
-            const result = await this.model.chatSessionContinue(userMessage, null, mergedConfig);
-            this.history.push({ role: 'user', content: userMessage });
+            this.history.push(pendingUser);
             this.history.push(buildAssistantMessage(result.text, result.toolCalls));
             this.turnCount++;
             this.recordToolCallFanout(result.toolCalls);
@@ -311,36 +519,67 @@ export class ChatSession {
         try {
             const mergedConfig = this.mergeConfig(opts.config);
             const newImagesKey = computeImagesKey(opts.images);
-            // Only an explicit NEW image set can trigger a restart. Omitting
-            // `images` (newImagesKey === null) is interpreted as "keep the
-            // current image cache state" — the server-side cache already
-            // holds any prior image context, so a text-only follow-up like
-            // "what about the top-right?" can stay on the cheap delta path
-            // even after an image turn.
+            const newAudioKey = computeAudioKey(opts.audio);
+            // Only an explicit NEW image/audio set can trigger a restart. Omitting
+            // `images`/`audio` (key === null) is interpreted as "keep the current
+            // media cache state" — the server-side cache already holds any prior
+            // media context, so a text-only follow-up like "what about the
+            // top-right?" can stay on the cheap delta path even after a media turn.
             const imageChanged = newImagesKey !== null && newImagesKey !== this.lastImagesKey;
+            const audioChanged = newAudioKey !== null && newAudioKey !== this.lastAudioKey;
             const isFirstTurn = this.turnCount === 0;
-            if (isFirstTurn || imageChanged) {
-                yield* this.runStartStreamPath(userMessage, opts.images, newImagesKey, imageChanged, isFirstTurn, mergedConfig, opts.signal);
+            const replayRequired = this.needsFullReplay;
+            if (isFirstTurn || imageChanged || audioChanged || replayRequired) {
+                yield* this.runStartStreamPath(userMessage, opts.images, opts.audio, imageChanged || audioChanged || replayRequired, isFirstTurn, mergedConfig, opts.signal);
                 return;
             }
             // Delta continue stream: text-only.
+            const pendingUser = { role: 'user', content: userMessage };
+            const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingUser), mergedConfig);
             let sawFinal = false;
             let accumulated = '';
+            let accumulatedVisible = '';
             let finalRaw = null;
             let finalToolCalls;
+            // Set when the media-held rejection re-routes this turn through the
+            // cold start stream. The replay path owns the history push, turnCount
+            // increment, and media-key rehydration, so the commit `finally` below
+            // must NOT also fire.
+            let delegated = false;
             try {
-                for await (const event of this.model.chatStreamSessionContinue(userMessage, null, mergedConfig, opts.signal)) {
-                    if (event.done) {
-                        if (event.finishReason !== 'error') {
-                            sawFinal = true;
-                            finalRaw = event.text;
-                            finalToolCalls = event.toolCalls;
+                try {
+                    for await (const event of this.model.chatStreamSessionContinue(userMessage, null, null, constrainedConfig, opts.signal)) {
+                        if (event.done) {
+                            if (event.finishReason !== 'error') {
+                                sawFinal = true;
+                                finalRaw = event.text;
+                                finalToolCalls = event.toolCalls;
+                            }
+                        }
+                        else {
+                            accumulated += event.text;
+                            if (event.isReasoning !== true) {
+                                accumulatedVisible += event.text;
+                            }
                         }
+                        yield event;
                     }
-                    else {
-                        accumulated += event.text;
+                }
+                catch (err) {
+                    // The native session holds media KV (gemma4 after an image/audio
+                    // turn) and refused the text delta. The streaming bridge re-throws
+                    // that rejection on the first iteration, BEFORE any chunk is
+                    // emitted — the native guard fires ahead of any prefill, so
+                    // `!sawFinal && accumulated === ''` is guaranteed here. Replay the
+                    // full conversation through the cold start stream. Any non-prefix
+                    // error, or any error after tokens were already emitted, must
+                    // propagate unchanged.
+                    if (!isMediaHeldRestartError(err) || sawFinal || accumulated !== '') {
+                        throw err;
                     }
-                    yield event;
+                    delegated = true;
+                    yield* this.runStartStreamPath(userMessage, undefined, undefined, true, false, constrainedConfig, opts.signal);
+                    return;
                 }
             }
             finally {
@@ -350,13 +589,21 @@ export class ChatSession {
                 // chunks alike. The delta path doesn't push to history until
                 // commit, so the rollback branch is a no-op: nothing to
                 // undo, and the native cache state is managed by the Rust
-                // save_cache_state path on its own.
-                if (sawFinal) {
-                    this.history.push({ role: 'user', content: userMessage });
-                    this.history.push(buildAssistantMessage(finalRaw ?? accumulated, finalToolCalls));
+                // save_cache_state path on its own. When the media-held
+                // rejection delegated to the replay stream, that path already
+                // committed (or rolled back) — so this commit must stay off.
+                if (sawFinal && !delegated) {
+                    this.history.push(pendingUser);
+                    this.history.push(buildAssistantMessage(finalRaw || accumulatedVisible, finalToolCalls));
                     this.turnCount++;
                     this.recordToolCallFanout(finalToolCalls);
                 }
+                else if (!delegated) {
+                    // Qwen commits cancelled/failed delta tokens to its native cached
+                    // history even though this JS turn is intentionally uncommitted.
+                    // The next plain turn must reset and replay the preserved history.
+                    this.needsFullReplay = true;
+                }
             }
         }
         finally {
@@ -377,6 +624,18 @@ export class ChatSession {
      * hit this must tighten the prompt / tool spec or reset the
      * session.
      *
+     * `isError` is the structured tool-error signal. When `true`, the
+     * native renderer prepends a short, model-facing error marker to
+     * `content` inside the wire-format tool block so the model
+     * receives a clear text-level cue that the tool result represents
+     * a failure. The structured field is stored verbatim on the
+     * appended `{ role: 'tool', ... }` history entry so cold-replay
+     * (image-change restart, `startFromHistory*`, server-side
+     * `SessionRegistry` cache-miss rebuild) re-renders the marker
+     * consistently with the live turn. Defaults to `undefined` (no
+     * marker). Pass through verbatim — we do NOT infer error from
+     * `content`.
+     *
      * Appends a `{ role: 'tool', ... }` message to history on success.
      */
     async sendToolResult(toolCallId, content, opts = {}) {
@@ -386,19 +645,77 @@ export class ChatSession {
         this.assertCanSendToolResult('sendToolResult');
         this.inFlight = true;
         try {
-            const mergedConfig = this.mergeConfig(opts.config);
-            const result = await this.model.chatSessionContinueTool(toolCallId, content, mergedConfig);
-            this.history.push({ role: 'tool', content, toolCallId });
-            this.history.push(buildAssistantMessage(result.text, result.toolCalls));
-            this.turnCount++;
-            this.recordToolCallFanout(result.toolCalls);
-            return result;
+            const { isError, config } = opts;
+            const mergedConfig = this.mergeConfig(config);
+            const toolMsg = { role: 'tool', content, toolCallId, isError };
+            const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(toolMsg), mergedConfig);
+            // A cold native session (turnCount===0) has no live KV to delta
+            // against — the typical cause is an interrupted media-held replay
+            // whose rollback wiped the cache and reset the counter while
+            // leaving the unresolved tool-call flag set. Mirror `send()`'s
+            // turn-0 routing: replay the preserved history through the cold
+            // start path instead of dispatching a delta that the native side
+            // would reject with an un-prefixed "requires an initialized
+            // session" error. A normal tool result always follows a prior
+            // tool-call turn (turnCount>=1), so this never fires on the happy
+            // path.
+            if (this.turnCount === 0 || this.needsFullReplay) {
+                return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig);
+            }
+            try {
+                const result = await this.model.chatSessionContinueTool(toolCallId, content, constrainedConfig, isError ?? null);
+                this.history.push({ role: 'tool', content, toolCallId, isError });
+                this.history.push(buildAssistantMessage(result.text, result.toolCalls));
+                this.turnCount++;
+                this.recordToolCallFanout(result.toolCalls);
+                return result;
+            }
+            catch (err) {
+                if (!isMediaHeldRestartError(err)) {
+                    this.needsFullReplay = true;
+                    throw err;
+                }
+                // The native session holds media KV (gemma4 after an image/audio
+                // turn) and refused the tool-result delta. Transparently replay
+                // the full conversation through the cold start path. The prior
+                // media turn already lives in `this.history`, so the start path
+                // re-renders it; the trailing-media keys keep
+                // `lastImagesKey`/`lastAudioKey` consistent across the replay.
+                // The delta path threw before pushing the tool message, so the
+                // restart core pushes it — `isError` rides on that message so the
+                // wire-format error marker is re-rendered, and a tool result
+                // always follows >=1 prior turn so `isFirstTurn` is false.
+                return await this.replayToolResultThroughStartPath(toolMsg, constrainedConfig);
+            }
         }
         finally {
             this.inFlight = false;
         }
     }
-    /** Streaming variant of {@link ChatSession#sendToolResult}. */
+    /**
+     * Cold-replay a tool result through the start path: re-render the
+     * full preserved history (including the prior media turn and the
+     * unresolved tool-call assistant turn) plus this tool message. Used
+     * when the native session is cold (turnCount===0 — e.g. after an
+     * interrupted media-held replay rolled the cache back) and by the
+     * media-held rejection catch. `mediaChanged=true` forces a
+     * resetCaches so the prefill always starts from a guaranteed-clean
+     * cache; `isFirstTurn=false` because a tool result always follows a
+     * prior tool-call turn.
+     */
+    async replayToolResultThroughStartPath(toolMsg, config) {
+        return await this.runStartPathWithMessage(toolMsg, true, false, config);
+    }
+    /**
+     * Streaming variant of {@link ChatSession#sendToolResult}.
+     *
+     * `isError` mirrors the non-streaming entry point — when `true`,
+     * the native renderer prepends a short, model-facing error marker
+     * to `content` inside the wire-format tool block. The structured
+     * field is stored verbatim on the appended `{ role: 'tool', ... }`
+     * history entry so cold-replay re-renders the marker consistently
+     * with the live streaming turn.
+     */
     async *sendToolResultStream(toolCallId, content, opts = {}) {
         if (this.inFlight) {
             throw new Error('ChatSession: concurrent send() not allowed; await the previous call first');
@@ -406,43 +723,107 @@ export class ChatSession {
         this.assertCanSendToolResult('sendToolResultStream');
         this.inFlight = true;
         try {
-            const mergedConfig = this.mergeConfig(opts.config);
+            const { isError, config, signal } = opts;
+            const mergedConfig = this.mergeConfig(config);
+            const toolMsg = { role: 'tool', content, toolCallId, isError };
+            const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(toolMsg), mergedConfig);
+            // A cold native session (turnCount===0) has no live KV to delta
+            // against — typically the residue of an interrupted media-held
+            // replay whose rollback wiped the cache and reset the counter
+            // while leaving the unresolved tool-call flag set. Mirror
+            // `sendStream()`'s turn-0 routing: replay the preserved history
+            // through the cold start stream and return before the
+            // delta/commit machinery so the start path owns the history push,
+            // turnCount increment, and media-key rehydration. A normal tool
+            // result always follows a prior tool-call turn (turnCount>=1), so
+            // this never fires on the happy path.
+            if (this.turnCount === 0 || this.needsFullReplay) {
+                yield* this.replayToolResultThroughStartStreamPath(toolMsg, constrainedConfig, signal);
+                return;
+            }
             let sawFinal = false;
             let accumulated = '';
+            let accumulatedVisible = '';
             let finalRaw = null;
             let finalToolCalls;
+            // Set when the media-held rejection re-routes this tool turn
+            // through the cold start stream. The replay path owns the history
+            // push, turnCount increment, and media-key rehydration, so the
+            // commit `finally` below must NOT also fire.
+            let delegated = false;
             try {
-                for await (const event of this.model.chatStreamSessionContinueTool(toolCallId, content, mergedConfig, opts.signal)) {
-                    if (event.done) {
-                        if (event.finishReason !== 'error') {
-                            sawFinal = true;
-                            finalRaw = event.text;
-                            finalToolCalls = event.toolCalls;
+                try {
+                    for await (const event of this.model.chatStreamSessionContinueTool(toolCallId, content, constrainedConfig, signal, isError ?? null)) {
+                        if (event.done) {
+                            if (event.finishReason !== 'error') {
+                                sawFinal = true;
+                                finalRaw = event.text;
+                                finalToolCalls = event.toolCalls;
+                            }
                         }
+                        else {
+                            accumulated += event.text;
+                            if (event.isReasoning !== true) {
+                                accumulatedVisible += event.text;
+                            }
+                        }
+                        yield event;
                     }
-                    else {
-                        accumulated += event.text;
+                }
+                catch (err) {
+                    // The native session holds media KV (gemma4 after an image/audio
+                    // turn) and refused the tool-result delta. The streaming bridge
+                    // re-throws that rejection on the first iteration, BEFORE any
+                    // chunk is emitted — the native guard fires ahead of any
+                    // prefill, so `!sawFinal && accumulated === ''` is guaranteed
+                    // here. Replay the full conversation through the cold start
+                    // stream with the pending tool message; `isError` rides on it so
+                    // the wire-format error marker is re-rendered. Any non-prefix
+                    // error, or any error after tokens were already emitted, must
+                    // propagate unchanged.
+                    if (!isMediaHeldRestartError(err) || sawFinal || accumulated !== '') {
+                        throw err;
                     }
-                    yield event;
+                    delegated = true;
+                    yield* this.replayToolResultThroughStartStreamPath(toolMsg, constrainedConfig, signal);
+                    return;
                 }
             }
             finally {
                 // finally runs for normal completion, mid-stream throw,
                 // caller `break` (iterator.return() short-circuits the yield),
                 // and error-finish chunks alike. Tool turns never touch
-                // history until commit, so the rollback branch is a no-op.
-                if (sawFinal) {
-                    this.history.push({ role: 'tool', content, toolCallId });
-                    this.history.push(buildAssistantMessage(finalRaw ?? accumulated, finalToolCalls));
+                // history until commit, so the rollback branch is a no-op. When
+                // the media-held rejection delegated to the replay stream, that
+                // path already committed (or rolled back), so this commit stays
+                // off.
+                if (sawFinal && !delegated) {
+                    this.history.push({ role: 'tool', content, toolCallId, isError });
+                    this.history.push(buildAssistantMessage(finalRaw || accumulatedVisible, finalToolCalls));
                     this.turnCount++;
                     this.recordToolCallFanout(finalToolCalls);
                 }
+                else if (!delegated) {
+                    this.needsFullReplay = true;
+                }
             }
         }
         finally {
             this.inFlight = false;
         }
     }
+    /**
+     * Streaming counterpart of {@link replayToolResultThroughStartPath}:
+     * cold-replay a tool result through the start stream. Used by the
+     * turn-0 precheck and the media-held rejection catch in
+     * {@link sendToolResultStream}. The start stream owns the history
+     * push, turnCount increment, and media-key rehydration; callers keep
+     * `delegated`/early-return semantics so the commit `finally` stays
+     * off.
+     */
+    async *replayToolResultThroughStartStreamPath(toolMsg, config, signal) {
+        yield* this.runStartStreamPathWithMessage(toolMsg, true, false, config, signal);
+    }
     /**
      * Reset the session state.
      *
@@ -450,6 +831,23 @@ export class ChatSession {
      * image key, and turn counter so the next `send()` goes through
      * `chatSessionStart` again.
      *
+     * This is a full wipe — safe default for public callers. It always
+     * calls `model.resetCaches()`, which is the ONLY behavior exposed
+     * on the public API because the underlying `SessionCapableModel`
+     * is shared across every `ChatSession` lifetime via the native
+     * `ModelRegistry`: a partial wipe that leaves the shared native
+     * KV cache intact would leak a previous (unrelated) request's
+     * cached prefix into the next `chat_session_start_sync` call. The
+     * server-side warm-lease replay path (where preserving the native
+     * cache is correct) uses its own server-private helper gated by
+     * the `SessionRegistry` HIT signal — the only authoritative proof
+     * that the native cache genuinely belongs to this chain. That
+     * helper lives inside `@mlx-node/server`, never touches the
+     * `@mlx-node/lm` export map, and is not reachable from downstream
+     * consumers. Public consumers of `@mlx-node/lm` have no such HIT
+     * signal, so the public API intentionally offers only the full-wipe
+     * option.
+     *
      * Returns `Promise<void>` for an async-friendly signature even
      * though `resetCaches()` is currently synchronous.
      */
@@ -460,8 +858,10 @@ export class ChatSession {
         this.model.resetCaches();
         this.history = [];
         this.lastImagesKey = null;
+        this.lastAudioKey = null;
         this.turnCount = 0;
         this.unresolvedOkToolCallCount = null;
+        this.needsFullReplay = false;
     }
     /**
      * Prime the session history without running inference.
@@ -527,10 +927,14 @@ export class ChatSession {
         this.inFlight = true;
         try {
             const mergedConfig = this.mergeConfig(config);
-            const result = await this.model.chatSessionStart(this.history.slice(), mergedConfig);
+            const historySnapshot = this.history.slice();
+            const constrainedConfig = await this.constrainToContextCapacity(historySnapshot, mergedConfig);
+            const result = await this.model.chatSessionStart(historySnapshot, constrainedConfig);
             this.history.push(buildAssistantMessage(result.text, result.toolCalls));
             this.turnCount++;
+            this.needsFullReplay = false;
             this.lastImagesKey = this.computeTrailingImagesKey();
+            this.lastAudioKey = this.computeTrailingAudioKey();
             this.recordToolCallFanout(result.toolCalls);
             return result;
         }
@@ -562,12 +966,14 @@ export class ChatSession {
         try {
             const mergedConfig = this.mergeConfig(config);
             const historySnapshot = this.history.slice();
+            const constrainedConfig = await this.constrainToContextCapacity(historySnapshot, mergedConfig);
             let sawFinal = false;
             let accumulated = '';
+            let accumulatedVisible = '';
             let finalRaw = null;
             let finalToolCalls;
             try {
-                for await (const event of this.model.chatStreamSessionStart(historySnapshot, mergedConfig, signal)) {
+                for await (const event of this.model.chatStreamSessionStart(historySnapshot, constrainedConfig, signal)) {
                     if (event.done) {
                         if (event.finishReason !== 'error') {
                             sawFinal = true;
@@ -577,6 +983,9 @@ export class ChatSession {
                     }
                     else {
                         accumulated += event.text;
+                        if (event.isReasoning !== true) {
+                            accumulatedVisible += event.text;
+                        }
                     }
                     yield event;
                 }
@@ -588,9 +997,11 @@ export class ChatSession {
                 // mutated on a successful commit — on any non-success exit,
                 // the primed state is left intact so the caller can retry.
                 if (sawFinal) {
-                    this.history.push(buildAssistantMessage(finalRaw ?? accumulated, finalToolCalls));
+                    this.history.push(buildAssistantMessage(finalRaw || accumulatedVisible, finalToolCalls));
                     this.turnCount++;
+                    this.needsFullReplay = false;
                     this.lastImagesKey = this.computeTrailingImagesKey();
+                    this.lastAudioKey = this.computeTrailingAudioKey();
                     this.recordToolCallFanout(finalToolCalls);
                 }
             }
@@ -696,13 +1107,82 @@ export class ChatSession {
      * The session path is a session-reuse operation by construction —
      * `reuseCache: false` on the continue path would wipe the very
      * cache the delta depends on.
+     *
+     * MTP auto-default: if neither `defaultConfig` nor `overlay`
+     * sets `enableMtp` AND the underlying model exposes
+     * `hasMtpWeights()` returning `true`, set `enableMtp = true` so the
+     * speculative-decode path runs out of the box on MTP-capable
+     * checkpoints. An explicit `false` from either source wins (the
+     * undefined-check below preserves it). This duck-typed check also
+     * covers Gemma4 with an external draft attached — DSpark or Google
+     * assistant (`hasMtpWeights()` reports the external draft there,
+     * not in-checkpoint MTP heads).
      */
     mergeConfig(overlay) {
-        return {
+        const merged = {
             ...this.defaultConfig,
             ...overlay,
             reuseCache: true,
         };
+        if (merged.enableMtp === undefined &&
+            typeof this.model.hasMtpWeights === 'function' &&
+            this.model.hasMtpWeights()) {
+            merged.enableMtp = true;
+        }
+        return merged;
+    }
+    /**
+     * Render the exact full prompt and constrain generation to the physical KV
+     * window before native code allocates a block. Models that do not expose
+     * both the tokenizer seam and a load-time context snapshot retain their
+     * existing behavior.
+     *
+     * The returned config is a copy only when `maxNewTokens` needs clamping.
+     * An omitted output budget stays omitted when the native default fits; it is
+     * made explicit only when the remaining window is smaller than that default.
+     */
+    async constrainToContextCapacity(messages, config) {
+        if (typeof this.model.applyChatTemplate !== 'function' || typeof this.model.contextLimits !== 'function') {
+            return config;
+        }
+        const limits = this.model.contextLimits();
+        const capacity = Math.floor(limits.effectiveWindowTokens);
+        if (!Number.isSafeInteger(capacity) || capacity <= 0) {
+            return config;
+        }
+        const effort = config.reasoningEffort;
+        const enableThinking = effort === 'none' || effort === 'low' ? false : effort === 'medium' || effort === 'high' ? true : null;
+        const tokens = await this.model.applyChatTemplate(messages, true, config.tools ?? null, enableThinking);
+        const hasImages = messages.some((message) => (message.images?.length ?? 0) > 0);
+        let promptTokens = tokens.length;
+        if (hasImages && typeof this.model.expandedPromptTokenCount === 'function') {
+            promptTokens = await this.model.expandedPromptTokenCount(tokens, messages);
+            if (!Number.isSafeInteger(promptTokens) || promptTokens < tokens.length) {
+                throw new Error(`ChatSession: expandedPromptTokenCount returned invalid length ${promptTokens} ` +
+                    `(rendered template length is ${tokens.length})`);
+            }
+        }
+        if (promptTokens > capacity) {
+            throw new ContextCapacityError(promptTokens, capacity);
+        }
+        // The final sampled token is returned without another model forward, so N
+        // generated tokens consume only N-1 additional KV positions.
+        const maxOutput = capacity - promptTokens + 1;
+        const requested = config.maxNewTokens;
+        if (requested === undefined) {
+            return maxOutput >= NATIVE_DEFAULT_MAX_NEW_TOKENS ? config : { ...config, maxNewTokens: maxOutput };
+        }
+        const maxNewTokens = Math.min(requested, maxOutput);
+        return maxNewTokens === requested ? config : { ...config, maxNewTokens };
+    }
+    /** Full history that a pending user/tool turn would render, without mutation. */
+    historyWithPending(pending) {
+        const messages = this.history.slice();
+        if (messages.length === 0 && this.system != null) {
+            messages.push({ role: 'system', content: this.system });
+        }
+        messages.push(pending);
+        return messages;
     }
     /**
      * Shared start-path logic for `send()`. Handles both the turn-0
@@ -711,27 +1191,51 @@ export class ChatSession {
      * native side gets the full conversation re-rendered with the new
      * image set.
      */
-    async runStartPath(userMessage, images, newImagesKey, imageChanged, isFirstTurn, config) {
+    async runStartPath(userMessage, images, audio, mediaChanged, isFirstTurn, config) {
+        const userMsg = this.buildUserMessage(userMessage, images, audio);
+        return await this.runStartPathWithMessage(userMsg, mediaChanged, isFirstTurn, config);
+    }
+    /**
+     * Core of {@link runStartPath} that takes a PRE-BUILT pending
+     * `ChatMessage` (user or tool) instead of building a user message
+     * itself. The cold-restart catch in `sendToolResult()` replays the
+     * conversation through this core with a pending `{ role: 'tool', ... }`
+     * message so the tool-result turn is re-rendered against the full
+     * history without duplicating the start-path bookkeeping.
+     */
+    async runStartPathWithMessage(pendingMessage, mediaChanged, isFirstTurn, config) {
         // Capture pre-state so the restart can be rolled back if the
-        // native call fails. The image-change branch resets caches BEFORE
+        // native call fails. The media-change branch resets caches BEFORE
         // we know whether the new prefill will succeed, so on failure we
-        // also have to drop turnCount + lastImagesKey to force the next
-        // call to re-route through the start path (rather than a delta
-        // continue against wiped caches).
-        const wasImageChangeRestart = imageChanged && !isFirstTurn;
+        // also have to drop turnCount + lastImagesKey/lastAudioKey to force
+        // the next call to re-route through the start path (rather than a
+        // delta continue against wiped caches).
+        const wasMediaChangeRestart = mediaChanged && !isFirstTurn;
         const historyLenBefore = this.history.length;
-        this.prepareStartPath(imageChanged, isFirstTurn);
-        const userMsg = this.buildUserMessage(userMessage, images);
-        this.history.push(userMsg);
+        // Capacity validation must happen before `prepareStartPath()` because a
+        // media-change restart clears native caches. A rejected oversized prompt
+        // is a request error and must leave both JS history and native state intact.
+        const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingMessage), config);
+        this.prepareStartPath(mediaChanged, isFirstTurn);
+        this.history.push(pendingMessage);
         try {
             // Pass a shallow snapshot so later pushes to `this.history`
             // (e.g. the assistant reply below) don't retroactively mutate
             // what the native side / any mock observed as its `messages`
             // argument.
-            const result = await this.model.chatSessionStart(this.history.slice(), config);
+            const result = await this.model.chatSessionStart(this.history.slice(), constrainedConfig);
             this.history.push(buildAssistantMessage(result.text, result.toolCalls));
             this.turnCount++;
-            this.lastImagesKey = newImagesKey;
+            this.needsFullReplay = false;
+            // The start path always re-renders the FULL preserved history, so the
+            // post-restart sticky keys are the trailing media keys of that history,
+            // not the single-turn literal args. A restart driven by a change in only
+            // one modality (e.g. an audio-only turn after an earlier image turn)
+            // would otherwise null the untouched modality's key even though that
+            // media is still live in the native cache, causing a later same-media
+            // turn to be mis-detected as a change and replayed twice.
+            this.lastImagesKey = this.computeTrailingImagesKey();
+            this.lastAudioKey = this.computeTrailingAudioKey();
             this.recordToolCallFanout(result.toolCalls);
             return result;
         }
@@ -739,37 +1243,51 @@ export class ChatSession {
             // Roll back: drop the tentative user push so history stays
             // consistent with turnCount.
             this.history.length = historyLenBefore;
-            if (wasImageChangeRestart) {
+            if (wasMediaChangeRestart) {
                 // Caches were wiped by prepareStartPath() but the new prefill
                 // failed. Force the next call to re-route through the start
                 // path with the (preserved) prior history.
                 this.turnCount = 0;
                 this.lastImagesKey = null;
+                this.lastAudioKey = null;
             }
             throw err;
         }
     }
     /** Streaming counterpart to {@link runStartPath}. */
-    async *runStartStreamPath(userMessage, images, newImagesKey, imageChanged, isFirstTurn, config, signal) {
+    async *runStartStreamPath(userMessage, images, audio, mediaChanged, isFirstTurn, config, signal) {
+        const userMsg = this.buildUserMessage(userMessage, images, audio);
+        yield* this.runStartStreamPathWithMessage(userMsg, mediaChanged, isFirstTurn, config, signal);
+    }
+    /**
+     * Streaming counterpart to {@link runStartPathWithMessage}: replays
+     * through the cold start stream from a PRE-BUILT pending
+     * `ChatMessage`. The cold-restart catch in `sendToolResultStream()`
+     * delegates here with a pending `{ role: 'tool', ... }` message.
+     */
+    async *runStartStreamPathWithMessage(pendingMessage, mediaChanged, isFirstTurn, config, signal) {
         // Capture pre-state so any non-successful exit can roll back.
         // See `runStartPath` for the full rationale.
-        const wasImageChangeRestart = imageChanged && !isFirstTurn;
+        const wasMediaChangeRestart = mediaChanged && !isFirstTurn;
         const historyLenBefore = this.history.length;
-        this.prepareStartPath(imageChanged, isFirstTurn);
-        const userMsg = this.buildUserMessage(userMessage, images);
-        // Stage the user message on the pending history BEFORE the
+        // See the sync start path: reject before a media restart can clear native
+        // state, and make the output budget explicit before allocating KV blocks.
+        const constrainedConfig = await this.constrainToContextCapacity(this.historyWithPending(pendingMessage), config);
+        this.prepareStartPath(mediaChanged, isFirstTurn);
+        // Stage the pending message on the pending history BEFORE the
         // stream starts — the native call reads it synchronously via
         // `model.chatStreamSessionStart(history, config)`.
-        this.history.push(userMsg);
+        this.history.push(pendingMessage);
         let sawFinal = false;
         let accumulated = '';
+        let accumulatedVisible = '';
         let finalRaw = null;
         let finalToolCalls;
         // Snapshot the history before dispatch — see `runStartPath` for
         // the rationale.
         const historySnapshot = this.history.slice();
         try {
-            for await (const event of this.model.chatStreamSessionStart(historySnapshot, config, signal)) {
+            for await (const event of this.model.chatStreamSessionStart(historySnapshot, constrainedConfig, signal)) {
                 if (event.done) {
                     if (event.finishReason !== 'error') {
                         sawFinal = true;
@@ -779,6 +1297,9 @@ export class ChatSession {
                 }
                 else {
                     accumulated += event.text;
+                    if (event.isReasoning !== true) {
+                        accumulatedVisible += event.text;
+                    }
                 }
                 yield event;
             }
@@ -793,22 +1314,32 @@ export class ChatSession {
             // generator was wound down. Mid-stream throws still propagate
             // naturally — finally runs first, then the error continues up.
             if (sawFinal) {
-                this.history.push(buildAssistantMessage(finalRaw ?? accumulated, finalToolCalls));
+                this.history.push(buildAssistantMessage(finalRaw || accumulatedVisible, finalToolCalls));
                 this.turnCount++;
-                this.lastImagesKey = newImagesKey;
+                this.needsFullReplay = false;
+                // The start path always re-renders the FULL preserved history, so the
+                // post-restart sticky keys are the trailing media keys of that history,
+                // not the single-turn literal args. A restart driven by a change in only
+                // one modality (e.g. an audio-only turn after an earlier image turn)
+                // would otherwise null the untouched modality's key even though that
+                // media is still live in the native cache, causing a later same-media
+                // turn to be mis-detected as a change and replayed twice.
+                this.lastImagesKey = this.computeTrailingImagesKey();
+                this.lastAudioKey = this.computeTrailingAudioKey();
                 this.recordToolCallFanout(finalToolCalls);
             }
             else {
                 // Roll back: drop the tentative user push so history stays
                 // consistent with turnCount.
                 this.history.length = historyLenBefore;
-                if (wasImageChangeRestart) {
+                if (wasMediaChangeRestart) {
                     // Caches were wiped by prepareStartPath() but the new
                     // prefill never reached a successful done:true. Force the
                     // next call to re-route through the start path with the
                     // preserved prior history.
                     this.turnCount = 0;
                     this.lastImagesKey = null;
+                    this.lastAudioKey = null;
                 }
             }
         }
@@ -828,24 +1359,26 @@ export class ChatSession {
      *     turn.
      *   - On a fresh / reset history, re-inject the system prompt.
      */
-    prepareStartPath(imageChanged, isFirstTurn) {
-        if (imageChanged && !isFirstTurn) {
+    prepareStartPath(mediaChanged, isFirstTurn) {
+        if (mediaChanged && !isFirstTurn) {
             this.model.resetCaches();
         }
         if (this.history.length === 0 && this.system != null) {
             this.history.push({ role: 'system', content: this.system });
         }
     }
-    /** Build a user `ChatMessage` with or without attached images. */
-    buildUserMessage(userMessage, images) {
-        if (images && images.length > 0) {
-            return { role: 'user', content: userMessage, images };
-        }
-        return { role: 'user', content: userMessage };
+    /** Build a user `ChatMessage` with or without attached images/audio. */
+    buildUserMessage(userMessage, images, audio) {
+        const msg = { role: 'user', content: userMessage };
+        if (images && images.length > 0)
+            msg.images = images;
+        if (audio && audio.length > 0)
+            msg.audio = audio;
+        return msg;
     }
     /**
      * Walk the history backward to find the most recent user message
-     * with images and return its FNV-1a key. Used by
+     * with images and return its SHA-256 key. Used by
      * {@link startFromHistory} and {@link startFromHistoryStream} to
      * hydrate `lastImagesKey` after a cold replay, so subsequent delta
      * continues correctly detect image changes.
@@ -859,6 +1392,20 @@ export class ChatSession {
         }
         return null;
     }
+    /**
+     * Audio counterpart of {@link computeTrailingImagesKey}: walk history
+     * backward to the most recent user message carrying audio and return its
+     * SHA-256 key, so a cold replay hydrates `lastAudioKey` correctly.
+     */
+    computeTrailingAudioKey() {
+        for (let i = this.history.length - 1; i >= 0; i--) {
+            const msg = this.history[i];
+            if (msg?.role === 'user' && msg.audio && msg.audio.length > 0) {
+                return computeAudioKey(msg.audio);
+            }
+        }
+        return null;
+    }
     /**
      * Derive the post-prime value of `unresolvedOkToolCallCount` from
      * the primed history. Walks backward to the most recent assistant