@mlx-node/lm 0.0.7 → 0.0.9

package/dist/chat-session.js

@@ -1,3 +1,146 @@
+/**
+ * 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 and passes the complete structured transcript
+ *     on every turn. Native code renders that transcript with the
+ *     checkpoint-provided chat template, verifies the completed history
+ *     against the committed cache, and appends the template-authored suffix
+ *     to the exact cached token IDs. Incompatible or edited history cold
+ *     replays the complete render.
+ *
+ *   - 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 still gets incremental prefill
+ *     on a token-prefix hit. Prompt structure is never reconstructed
+ *     from Rust string literals.
+ *
+ *   - `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.
+ *
+ *   - `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 native media guards use when a history cannot be continued
+ * from the held image/audio state. The session layer recognizes this exact
+ * prefix and transparently replays the complete structured conversation.
+ *
+ * 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 path as a rejected promise, and on
+ * the streaming 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
@@ -45,23 +188,40 @@ function toAssistantToolCalls(toolCalls) {
 }
 /**
  * Build an assistant `ChatMessage` from a just-completed turn's
- * decoded text + tool-call list. The assistant entry is appended to
- * `this.history` after every successful turn and is later read back
- * by the native `chatSessionStart` cold-replay path (image-change
- * mid-session restart, `startFromHistory*`, server-side
- * `SessionRegistry` cache-miss rebuild). Dropping the `toolCalls`
- * field here would orphan any subsequent `{role: 'tool', ...}`
- * entries on replay — the jinja template would render a
- * `<tool_response>` for a call that was never declared on the
- * preceding assistant turn, corrupting the conversation structure
- * and changing model behavior after a restart.
+ * decoded text, exact raw text, tool-call list, reasoning body, and resolved
+ * thinking mode. LFM2 replays the exact raw content because its checkpoint
+ * template does not consume structured reasoning; Qwen/Gemma retain their
+ * structured fields. The assistant entry is appended to `this.history` after every
+ * successful turn and is later read back by the native
+ * `chatSessionStart` cold-replay path (image-change mid-session
+ * restart, `startFromHistory*`, server-side `SessionRegistry`
+ * cache-miss rebuild). Dropping `toolCalls`, `reasoningContent`, or
+ * `thinkingEnabled` changes the rendered assistant bytes on replay:
+ * tool responses lose their declaring call, reasoning disappears, or
+ * an empty disabled-thinking channel is reinterpreted under the
+ * current turn's mode.
  */
-function buildAssistantMessage(text, toolCalls) {
+function buildAssistantMessage(text, toolCalls, thinking, thinkingEnabled, rawText, replayRawText) {
+    if (replayRawText) {
+        return {
+            role: 'assistant',
+            content: rawText ?? text,
+            thinkingEnabled,
+        };
+    }
     const calls = toAssistantToolCalls(toolCalls);
+    const message = {
+        role: 'assistant',
+        content: text,
+        thinkingEnabled,
+    };
     if (calls) {
-        return { role: 'assistant', content: text, toolCalls: calls };
+        message.toolCalls = calls;
     }
-    return { role: 'assistant', content: text };
+    if (thinking != null) {
+        message.reasoningContent = thinking;
+    }
+    return message;
 }
 /**
  * Count the `ok`-status tool calls in a `ChatResult.toolCalls` /
@@ -83,6 +243,63 @@ function countOkToolCalls(toolCalls) {
     }
     return n;
 }
+/**
+ * Select the assistant text committed after a successful stream.
+ *
+ * Default emitters expose parsed text on the terminal event, so an empty
+ * value is authoritative for a tool-only turn and must not fall back to raw
+ * streamed tool markup. Gemma's channel-aware emitter instead sends visible
+ * text as deltas and an empty terminal value for ordinary no-tool replies, so
+ * retain the accumulated visible text in that distinct shape.
+ */
+function selectCommittedStreamText(finalText, accumulatedVisible, terminalTextAuthoritative) {
+    if (finalText == null)
+        return accumulatedVisible;
+    if (terminalTextAuthoritative === true)
+        return finalText;
+    if (terminalTextAuthoritative === false)
+        return accumulatedVisible;
+    if (finalText !== '')
+        return finalText;
+    return accumulatedVisible;
+}
+/** Mirror the native default used by `resolve_include_reasoning`. */
+function includesReasoning(config) {
+    return config.includeReasoning ?? config.reasoningEffort !== 'none';
+}
+/**
+ * Session history must retain reasoning even when the caller hides it. Ask
+ * native finalization for the full parsed turn without mutating the committed
+ * request config; the public view is redacted again below.
+ */
+function withReplayReasoning(config, model) {
+    if (includesReasoning(config))
+        return config;
+    if (model.supportsReplayReasoningCapture?.() !== true)
+        return config;
+    // A zero-budget "none" turn cannot produce a reasoning body worth
+    // replaying. Preserve the caller's native suppression flag for this common
+    // short-generation path (for example title generation).
+    if (config.reasoningEffort === 'none' && (config.thinkingTokenBudget ?? 0) <= 0) {
+        return config;
+    }
+    return { ...config, includeReasoning: true };
+}
+function publicChatResult(result, config) {
+    if (includesReasoning(config))
+        return result;
+    const safeRaw = result.publicRawText;
+    return { ...result, thinking: undefined, rawText: safeRaw ?? result.text };
+}
+function publicStreamEvent(event, config) {
+    if (includesReasoning(config))
+        return event;
+    if (!event.done) {
+        return event.isReasoning === true ? null : event;
+    }
+    const safeRaw = event.publicRawText;
+    return { ...event, thinking: null, rawText: safeRaw ?? event.text };
+}
 /**
  * Compute a stable hex-encoded identity key for a list of image
  * byte buffers.
@@ -94,75 +311,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.
@@ -176,22 +376,32 @@ export class ChatSession {
     model;
     system;
     defaultConfig;
+    /** Tool definitions are conversation state for deterministic template replay. */
+    activeTools;
     /**
      * Full conversation history tracked on the TS side. Appended to on
-     * every successful turn. Only read back when the image-change path
-     * triggers a restart — normal text continues use the server-side
-     * cache, not this array.
+     * every successful turn and sent on every role-aware native turn so
+     * the model-provided template remains the sole prompt authority.
      */
     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 turn 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
@@ -217,6 +427,7 @@ export class ChatSession {
         this.model = model;
         this.system = options.system;
         this.defaultConfig = options.defaultConfig ?? {};
+        this.activeTools = this.defaultConfig.tools;
     }
     /**
      * Number of completed turns. Increments only after a successful
@@ -229,6 +440,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
@@ -252,9 +512,10 @@ export class ChatSession {
     /**
      * Send a user message and resolve with the assistant reply.
      *
-     * Turn 0 and any turn whose image set has changed dispatch through
-     * `chatSessionStart` with the full history. All other turns use
-     * the cheap `chatSessionContinue` delta path.
+     * Turn 0 and any turn whose image set changed dispatch through
+     * `chatSessionStart`. Later turns pass the same complete structured
+     * history through `chatSessionContinue`; native code renders the
+     * model template and reuses KV on an exact token-prefix match.
      */
     async send(userMessage, opts = {}) {
         if (this.inFlight) {
@@ -265,27 +526,49 @@ 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 forced 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 ask native code to verify/reuse the templated history.
             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 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(buildAssistantMessage(result.text, result.toolCalls));
+            // Role-aware continuation: pass the complete structured transcript.
+            // Native code renders it with the checkpoint template and only reuses
+            // the live cache when the resulting tokens exactly extend that cache.
+            const pendingUser = { role: 'user', content: userMessage };
+            const pendingHistory = this.historyWithPending(pendingUser);
+            const constrainedConfig = await this.constrainToContextCapacity(pendingHistory, mergedConfig);
+            let result;
+            try {
+                result = await this.model.chatSessionContinue(pendingHistory, withReplayReasoning(constrainedConfig, this.model));
+            }
+            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 continuation. 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 continuation path has NOT pushed
+                // `userMessage` yet, so `runStartPath` pushing it adds no duplicate.
+                return await this.runStartPath(userMessage, undefined, undefined, true, false, constrainedConfig);
+            }
+            this.history.push(pendingUser);
+            this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
             this.turnCount++;
+            this.commitActiveTools(constrainedConfig);
             this.recordToolCallFanout(result.toolCalls);
-            return result;
+            return publicChatResult(result, constrainedConfig);
         }
         finally {
             this.inFlight = false;
@@ -311,36 +594,78 @@ 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 pendingHistory = this.historyWithPending(pendingUser);
+            const constrainedConfig = await this.constrainToContextCapacity(pendingHistory, mergedConfig);
             let sawFinal = false;
             let accumulated = '';
+            let accumulatedVisible = '';
             let finalRaw = null;
+            let finalReplayRaw = null;
+            let finalTextAuthoritative;
             let finalToolCalls;
+            let finalThinking = null;
+            let finalThinkingEnabled = false;
+            // 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(pendingHistory, withReplayReasoning(constrainedConfig, this.model), opts.signal)) {
+                        if (event.done) {
+                            if (event.finishReason !== 'error') {
+                                sawFinal = true;
+                                finalRaw = event.text;
+                                finalReplayRaw = event.rawText;
+                                finalTextAuthoritative = event.textAuthoritative;
+                                finalToolCalls = event.toolCalls;
+                                finalThinking = event.thinking;
+                                finalThinkingEnabled = event.thinkingEnabled;
+                            }
+                        }
+                        else {
+                            accumulated += event.text;
+                            if (event.isReasoning !== true) {
+                                accumulatedVisible += event.text;
+                            }
                         }
+                        const publicEvent = publicStreamEvent(event, constrainedConfig);
+                        if (publicEvent !== null)
+                            yield publicEvent;
                     }
-                    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 +675,22 @@ 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(selectCommittedStreamText(finalRaw, accumulatedVisible, finalTextAuthoritative), finalToolCalls, finalThinking, finalThinkingEnabled, finalReplayRaw, this.model.replaysAssistantRawText?.() === true));
                     this.turnCount++;
+                    this.commitActiveTools(constrainedConfig);
                     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 {
@@ -364,9 +698,9 @@ export class ChatSession {
         }
     }
     /**
-     * Send a tool-result turn. Always dispatches
-     * `chatSessionContinueTool` — tool turns never change image state,
-     * so there is no restart path here.
+     * Send a tool-result turn. The declaring assistant tool call and the
+     * pending result are both included in the full history passed to
+     * `chatSessionContinueTool`.
      *
      * Rejects if the prior assistant turn emitted more than one `ok`
      * tool call: the chat-session API only supports exactly one tool
@@ -377,6 +711,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 +732,84 @@ 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 pendingHistory = this.historyWithPending(toolMsg);
+            const constrainedConfig = await this.constrainToContextCapacity(pendingHistory, 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(pendingHistory, withReplayReasoning(constrainedConfig, this.model));
+                this.history.push({ role: 'tool', content, toolCallId, isError });
+                this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
+                this.turnCount++;
+                this.commitActiveTools(constrainedConfig);
+                this.recordToolCallFanout(result.toolCalls);
+                return publicChatResult(result, 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 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 +817,124 @@ 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 pendingHistory = this.historyWithPending(toolMsg);
+            const constrainedConfig = await this.constrainToContextCapacity(pendingHistory, 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 finalReplayRaw = null;
+            let finalTextAuthoritative;
             let finalToolCalls;
+            let finalThinking = null;
+            let finalThinkingEnabled = false;
+            // 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(pendingHistory, withReplayReasoning(constrainedConfig, this.model), signal)) {
+                        if (event.done) {
+                            if (event.finishReason !== 'error') {
+                                sawFinal = true;
+                                finalRaw = event.text;
+                                finalReplayRaw = event.rawText;
+                                finalTextAuthoritative = event.textAuthoritative;
+                                finalToolCalls = event.toolCalls;
+                                finalThinking = event.thinking;
+                                finalThinkingEnabled = event.thinkingEnabled;
+                            }
                         }
+                        else {
+                            accumulated += event.text;
+                            if (event.isReasoning !== true) {
+                                accumulatedVisible += event.text;
+                            }
+                        }
+                        const publicEvent = publicStreamEvent(event, constrainedConfig);
+                        if (publicEvent !== null)
+                            yield publicEvent;
                     }
-                    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(selectCommittedStreamText(finalRaw, accumulatedVisible, finalTextAuthoritative), finalToolCalls, finalThinking, finalThinkingEnabled, finalReplayRaw, this.model.replaysAssistantRawText?.() === true));
                     this.turnCount++;
+                    this.commitActiveTools(constrainedConfig);
                     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 +942,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 +969,11 @@ export class ChatSession {
         this.model.resetCaches();
         this.history = [];
         this.lastImagesKey = null;
+        this.lastAudioKey = null;
         this.turnCount = 0;
         this.unresolvedOkToolCallCount = null;
+        this.needsFullReplay = false;
+        this.activeTools = this.defaultConfig.tools;
     }
     /**
      * Prime the session history without running inference.
@@ -527,12 +1039,17 @@ export class ChatSession {
         this.inFlight = true;
         try {
             const mergedConfig = this.mergeConfig(config);
-            const result = await this.model.chatSessionStart(this.history.slice(), mergedConfig);
-            this.history.push(buildAssistantMessage(result.text, result.toolCalls));
+            const historySnapshot = this.history.slice();
+            const constrainedConfig = await this.constrainToContextCapacity(historySnapshot, mergedConfig);
+            const result = await this.model.chatSessionStart(historySnapshot, withReplayReasoning(constrainedConfig, this.model));
+            this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
             this.turnCount++;
+            this.needsFullReplay = false;
             this.lastImagesKey = this.computeTrailingImagesKey();
+            this.lastAudioKey = this.computeTrailingAudioKey();
+            this.commitActiveTools(constrainedConfig);
             this.recordToolCallFanout(result.toolCalls);
-            return result;
+            return publicChatResult(result, constrainedConfig);
         }
         finally {
             this.inFlight = false;
@@ -562,23 +1079,38 @@ 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 finalReplayRaw = null;
+            let finalTextAuthoritative;
             let finalToolCalls;
+            let finalThinking = null;
+            let finalThinkingEnabled = false;
             try {
-                for await (const event of this.model.chatStreamSessionStart(historySnapshot, mergedConfig, signal)) {
+                for await (const event of this.model.chatStreamSessionStart(historySnapshot, withReplayReasoning(constrainedConfig, this.model), signal)) {
                     if (event.done) {
                         if (event.finishReason !== 'error') {
                             sawFinal = true;
                             finalRaw = event.text;
+                            finalReplayRaw = event.rawText;
+                            finalTextAuthoritative = event.textAuthoritative;
                             finalToolCalls = event.toolCalls;
+                            finalThinking = event.thinking;
+                            finalThinkingEnabled = event.thinkingEnabled;
                         }
                     }
                     else {
                         accumulated += event.text;
+                        if (event.isReasoning !== true) {
+                            accumulatedVisible += event.text;
+                        }
                     }
-                    yield event;
+                    const publicEvent = publicStreamEvent(event, constrainedConfig);
+                    if (publicEvent !== null)
+                        yield publicEvent;
                 }
             }
             finally {
@@ -588,9 +1120,12 @@ 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(selectCommittedStreamText(finalRaw, accumulatedVisible, finalTextAuthoritative), finalToolCalls, finalThinking, finalThinkingEnabled, finalReplayRaw, this.model.replaysAssistantRawText?.() === true));
                     this.turnCount++;
+                    this.needsFullReplay = false;
                     this.lastImagesKey = this.computeTrailingImagesKey();
+                    this.lastAudioKey = this.computeTrailingAudioKey();
+                    this.commitActiveTools(constrainedConfig);
                     this.recordToolCallFanout(finalToolCalls);
                 }
             }
@@ -691,18 +1226,108 @@ export class ChatSession {
         const n = countOkToolCalls(toolCalls);
         this.unresolvedOkToolCallCount = n > 0 ? n : null;
     }
+    /**
+     * Persist the effective tools only when their turn commits successfully.
+     * Preflights and failed/abandoned turns intentionally never call this.
+     */
+    commitActiveTools(config) {
+        if (config.tools !== undefined) {
+            this.activeTools = config.tools;
+        }
+    }
     /**
      * Merge default + per-call config and force `reuseCache: true`.
      * 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.
+     *
+     * Tool resolution here is side-effect free because public capacity
+     * preflights use this same merge path. Successful turn commit sites call
+     * {@link commitActiveTools} after native inference finishes.
+     *
+     * 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,
         };
+        // Tools are part of the committed conversation state. Constructor
+        // defaults seed that state, but must not overwrite a tool set committed by
+        // a later successful turn. A current-call overlay is the only higher
+        // precedence source; it remains provisional until a success site calls
+        // commitActiveTools().
+        if (overlay?.tools === undefined && this.activeTools !== undefined) {
+            merged.tools = this.activeTools;
+        }
+        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,76 +1336,128 @@ 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);
-            this.history.push(buildAssistantMessage(result.text, result.toolCalls));
+            const result = await this.model.chatSessionStart(this.history.slice(), withReplayReasoning(constrainedConfig, this.model));
+            this.history.push(buildAssistantMessage(result.text, result.toolCalls, result.thinking, result.thinkingEnabled, result.rawText, this.model.replaysAssistantRawText?.() === true));
             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.commitActiveTools(constrainedConfig);
             this.recordToolCallFanout(result.toolCalls);
-            return result;
+            return publicChatResult(result, constrainedConfig);
         }
         catch (err) {
             // 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 finalReplayRaw = null;
+        let finalTextAuthoritative;
         let finalToolCalls;
+        let finalThinking = null;
+        let finalThinkingEnabled = false;
         // 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, withReplayReasoning(constrainedConfig, this.model), signal)) {
                 if (event.done) {
                     if (event.finishReason !== 'error') {
                         sawFinal = true;
                         finalRaw = event.text;
+                        finalReplayRaw = event.rawText;
+                        finalTextAuthoritative = event.textAuthoritative;
                         finalToolCalls = event.toolCalls;
+                        finalThinking = event.thinking;
+                        finalThinkingEnabled = event.thinkingEnabled;
                     }
                 }
                 else {
                     accumulated += event.text;
+                    if (event.isReasoning !== true) {
+                        accumulatedVisible += event.text;
+                    }
                 }
-                yield event;
+                const publicEvent = publicStreamEvent(event, constrainedConfig);
+                if (publicEvent !== null)
+                    yield publicEvent;
             }
         }
         finally {
@@ -793,22 +1470,33 @@ 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(selectCommittedStreamText(finalRaw, accumulatedVisible, finalTextAuthoritative), finalToolCalls, finalThinking, finalThinkingEnabled, finalReplayRaw, this.model.replaysAssistantRawText?.() === true));
                 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.commitActiveTools(constrainedConfig);
                 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 +1516,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 +1549,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