@oh-my-pi/pi-coding-agent 16.2.12 → 16.3.0

package/src/modes/controllers/streaming-reveal.ts

@@ -11,6 +11,7 @@ export const CATCHUP_FRAMES = 8;
 type AssistantContentBlock = AssistantMessage["content"][number];
 type DisplayThinkingContentBlock = Extract<AssistantContentBlock, { type: "thinking" }> & { rawThinking?: string };
 type StreamingRevealComponent = Pick<AssistantMessageComponent, "updateContent">;
+type GraphemeSlicer = (index: number, text: string, units: number) => string;
 
 type StreamingRevealControllerOptions = {
 	getSmoothStreaming(): boolean;
@@ -44,12 +45,29 @@ function countGraphemesFrom(text: string, start: number): { count: number; tailS
 	}
 	return { count, tailStart };
 }
+/** Segment `text` from code-unit offset `start`, walking up to `clusters`
+ *  graphemes. Returns the code-unit END of the final cluster walked, its START
+ *  (`lastStart`), and how many clusters were found (`count` may be less than
+ *  `clusters` if the suffix is shorter than requested). */
+function segmentFrom(text: string, start: number, clusters: number): { end: number; lastStart: number; count: number } {
+	let count = 0;
+	let lastStart = start;
+	let end = start;
+	for (const seg of getSegmenter().segment(start === 0 ? text : text.slice(start))) {
+		count += 1;
+		lastStart = start + seg.index;
+		end = start + seg.index + seg.segment.length;
+		if (count >= clusters) break;
+	}
+	return { end, lastStart, count };
+}
 
 /** Memoizes per-block grapheme counts across reveal ticks. Streaming blocks only
  *  grow by appending, and an append can only alter the final grapheme cluster of
  *  the previous text, so only the suffix from that cluster needs re-segmenting. */
-class BlockUnitCounter {
+export class BlockUnitCounter {
 	#entries = new Map<number, { text: string; count: number; tailStart: number }>();
+	#sliceEntries = new Map<number, { text: string; units: number; end: number; lastStart: number }>();
 
 	count(index: number, text: string): number {
 		const entry = this.#entries.get(index);
@@ -69,6 +87,29 @@ class BlockUnitCounter {
 
 	reset(): void {
 		this.#entries.clear();
+		this.#sliceEntries.clear();
+	}
+	/** Slice `text` to its first `units` graphemes. Memoized across reveal ticks:
+	 *  streaming blocks grow only by appending and the reveal target advances
+	 *  monotonically, so a previously sliced prefix is reused and only the suffix
+	 *  from the boundary cluster is re-segmented. Only an exact (text, units) hit
+	 *  skips segmentation entirely — an append can extend the boundary cluster, so
+	 *  the incremental path still re-segments from that cluster's start. */
+	slice(index: number, text: string, units: number): string {
+		if (units <= 0 || text.length === 0) return "";
+		const entry = this.#sliceEntries.get(index);
+		if (entry !== undefined && entry.text === text && entry.units === units) {
+			return entry.end >= text.length ? text : text.slice(0, entry.end);
+		}
+		if (entry !== undefined && (entry.text === text || text.startsWith(entry.text)) && units >= entry.units) {
+			const extra = units - entry.units + 1;
+			const seg = segmentFrom(text, entry.lastStart, extra);
+			this.#sliceEntries.set(index, { text, units, end: seg.end, lastStart: seg.lastStart });
+			return seg.end >= text.length ? text : text.slice(0, seg.end);
+		}
+		const seg = segmentFrom(text, 0, units);
+		this.#sliceEntries.set(index, { text, units, end: seg.end, lastStart: seg.lastStart });
+		return seg.end >= text.length ? text : text.slice(0, seg.end);
 	}
 }
 
@@ -104,20 +145,24 @@ function revealTextBlock(
 	block: Extract<AssistantContentBlock, { type: "text" }>,
 	remaining: number,
 	units: number,
+	index: number,
+	sliceOf: GraphemeSlicer,
 ): AssistantContentBlock {
 	if (remaining <= 0) return block.text.length === 0 ? block : { ...block, text: "" };
 	if (remaining >= units) return block;
-	return { ...block, text: sliceGraphemes(block.text, remaining) };
+	return { ...block, text: sliceOf(index, block.text, remaining) };
 }
 
 function revealThinkingBlock(
 	block: Extract<AssistantContentBlock, { type: "thinking" }>,
 	remaining: number,
 	units: number,
+	index: number,
+	sliceOf: GraphemeSlicer,
 ): AssistantContentBlock {
 	if (remaining <= 0) return block.thinking.length === 0 ? block : { ...block, thinking: "" };
 	if (remaining >= units) return block;
-	return { ...block, thinking: sliceGraphemes(block.thinking, remaining) };
+	return { ...block, thinking: sliceOf(index, block.thinking, remaining) };
 }
 
 export function buildDisplayMessage(
@@ -126,6 +171,7 @@ export function buildDisplayMessage(
 	hideThinking: boolean,
 	proseOnly = true,
 	countOf: (index: number, text: string) => number = (_index, text) => countGraphemes(text),
+	sliceOf: GraphemeSlicer = (_index, text, units) => sliceGraphemes(text, units),
 ): AssistantMessage {
 	let remaining = Math.max(0, Math.floor(revealed));
 	const content: AssistantContentBlock[] = [];
@@ -133,7 +179,7 @@ export function buildDisplayMessage(
 		const block = target.content[i]!;
 		if (block.type === "text") {
 			const units = countOf(i, block.text);
-			content.push(revealTextBlock(block, remaining, units));
+			content.push(revealTextBlock(block, remaining, units, i, sliceOf));
 			remaining = Math.max(0, remaining - units);
 		} else if (block.type === "thinking" && !hideThinking) {
 			const formatted = formatThinkingForDisplay(block.thinking, proseOnly);
@@ -144,7 +190,7 @@ export function buildDisplayMessage(
 					thinking: formatted,
 					rawThinking: block.thinking,
 				};
-				content.push(revealThinkingBlock(displayBlock, remaining, units));
+				content.push(revealThinkingBlock(displayBlock, remaining, units, i, sliceOf));
 				remaining = Math.max(0, remaining - units);
 			} else {
 				content.push(block);
@@ -174,6 +220,8 @@ export class StreamingRevealController {
 	#smoothStreaming = true;
 	readonly #unitCounter = new BlockUnitCounter();
 	readonly #countOf = (index: number, text: string): number => this.#unitCounter.count(index, text);
+	readonly #sliceOf = (index: number, text: string, units: number): string =>
+		this.#unitCounter.slice(index, text, units);
 
 	constructor(options: StreamingRevealControllerOptions) {
 		this.#getSmoothStreaming = options.getSmoothStreaming;
@@ -181,6 +229,16 @@ export class StreamingRevealController {
 		this.#getProseOnlyThinking = options.getProseOnlyThinking;
 		this.#requestRender = options.requestRender;
 	}
+	#build(target: AssistantMessage, revealed: number): AssistantMessage {
+		return buildDisplayMessage(
+			target,
+			revealed,
+			this.#hideThinkingBlock,
+			this.#proseOnlyThinking,
+			this.#countOf,
+			this.#sliceOf,
+		);
+	}
 
 	begin(component: StreamingRevealComponent, message: AssistantMessage): void {
 		this.stop();
@@ -192,10 +250,7 @@ export class StreamingRevealController {
 		this.#smoothStreaming = this.#getSmoothStreaming();
 		if (!this.#smoothStreaming) {
 			const total = this.#visibleUnits(message);
-			component.updateContent(
-				buildDisplayMessage(message, total, this.#hideThinkingBlock, this.#proseOnlyThinking, this.#countOf),
-				{ transient: true },
-			);
+			component.updateContent(this.#build(message, total), { transient: true });
 			return;
 		}
 		const total = this.#visibleUnits(message);
@@ -203,18 +258,9 @@ export class StreamingRevealController {
 			// A tool call is a transcript-order boundary: finish any leading
 			// assistant text before EventController renders the separate tool card.
 			this.#revealed = total;
-			component.updateContent(
-				buildDisplayMessage(
-					message,
-					this.#revealed,
-					this.#hideThinkingBlock,
-					this.#proseOnlyThinking,
-					this.#countOf,
-				),
-				{
-					transient: true,
-				},
-			);
+			component.updateContent(this.#build(message, this.#revealed), {
+				transient: true,
+			});
 			return;
 		}
 		this.#renderCurrent();
@@ -229,10 +275,7 @@ export class StreamingRevealController {
 		if (!this.#component) return;
 		if (!this.#smoothStreaming) {
 			const total = this.#visibleUnits(message);
-			this.#component.updateContent(
-				buildDisplayMessage(message, total, this.#hideThinkingBlock, this.#proseOnlyThinking, this.#countOf),
-				{ transient: true },
-			);
+			this.#component.updateContent(this.#build(message, total), { transient: true });
 			return;
 		}
 		const total = this.#visibleUnits(message);
@@ -241,18 +284,9 @@ export class StreamingRevealController {
 			// assistant text before EventController renders the separate tool card.
 			this.#revealed = total;
 			this.#stopTimer();
-			this.#component.updateContent(
-				buildDisplayMessage(
-					message,
-					this.#revealed,
-					this.#hideThinkingBlock,
-					this.#proseOnlyThinking,
-					this.#countOf,
-				),
-				{
-					transient: true,
-				},
-			);
+			this.#component.updateContent(this.#build(message, this.#revealed), {
+				transient: true,
+			});
 			return;
 		}
 		if (this.#revealed > total) {
@@ -309,16 +343,7 @@ export class StreamingRevealController {
 		// Every controller render is an in-flight streaming snapshot, even when
 		// smooth reveal has temporarily caught up to the current target. The
 		// message_end handler performs the only stable non-transient render.
-		this.#component.updateContent(
-			buildDisplayMessage(
-				this.#target,
-				this.#revealed,
-				this.#hideThinkingBlock,
-				this.#proseOnlyThinking,
-				this.#countOf,
-			),
-			{ transient: true },
-		);
+		this.#component.updateContent(this.#build(this.#target, this.#revealed), { transient: true });
 	}
 
 	#syncTimer(total = this.#target ? this.#visibleUnits(this.#target) : 0): void {
@@ -356,12 +381,9 @@ export class StreamingRevealController {
 			return;
 		}
 		this.#revealed = Math.min(total, this.#revealed + nextStep(total - this.#revealed));
-		component.updateContent(
-			buildDisplayMessage(target, this.#revealed, this.#hideThinkingBlock, this.#proseOnlyThinking, this.#countOf),
-			{
-				transient: true,
-			},
-		);
+		component.updateContent(this.#build(target, this.#revealed), {
+			transient: true,
+		});
 		this.#requestRender();
 		if (this.#revealed >= total) {
 			this.#stopTimer();

package/src/modes/controllers/tool-args-reveal.ts

@@ -6,11 +6,289 @@ type ToolArgsRevealComponent = {
 	updateArgs(args: unknown, toolCallId?: string): void;
 };
 
+// Top-level string args a renderer reads mid-stream. The streamed-args decode
+// reads these fields incrementally between throttled full-JSON parses so a
+// long payload updates preview args at reveal cadence instead of stalling for
+// STREAMING_JSON_PARSE_MIN_GROWTH bytes at a time. Nested-array modes (edit
+// patch/replace `edits[].diff`) still fall through to the throttled parse.
+const STREAMING_STRING_KEYS_BY_TOOL: Record<string, readonly string[]> = {
+	write: ["content"],
+	edit: ["input"],
+	eval: ["code"],
+};
+
+/** String fields the streamed-args decode reads incrementally for `toolName`. */
+export function streamingStringKeysForTool(toolName: string, rawInput: boolean): readonly string[] | undefined {
+	if (rawInput) return undefined;
+	return STREAMING_STRING_KEYS_BY_TOOL[toolName];
+}
+
 type ToolArgsRevealControllerOptions = {
 	getSmoothStreaming(): boolean;
 	requestRender(): void;
 };
 
+type StreamingJsonStringExtractorResult = {
+	values: Record<string, string>;
+	changed: boolean;
+};
+
+function decodeJsonStringEscape(ch: string): string {
+	switch (ch) {
+		case '"':
+		case "\\":
+		case "/":
+			return ch;
+		case "b":
+			return "\b";
+		case "f":
+			return "\f";
+		case "n":
+			return "\n";
+		case "r":
+			return "\r";
+		case "t":
+			return "\t";
+		default:
+			return ch;
+	}
+}
+
+function isHexDigit(ch: string): boolean {
+	return (ch >= "0" && ch <= "9") || (ch >= "a" && ch <= "f") || (ch >= "A" && ch <= "F");
+}
+
+type StreamingJsonStringExtractorState = "scan" | "candidate" | "afterCandidate" | "beforeValue" | "target";
+
+class StreamingJsonStringExtractor {
+	readonly #keys: Set<string>;
+	#source = "";
+	#offset = 0;
+	/** `{`/`[` nesting outside strings. Candidate keys match only at depth 1 —
+	 *  the top level of the args object — so a nested object's key (e.g.
+	 *  `{"meta":{"content":…}}`) is never captured as a streamed top-level arg. */
+	#depth = 0;
+	#state: StreamingJsonStringExtractorState = "scan";
+	#candidate = "";
+	#candidateEscaped = false;
+	#candidateUnicode = "";
+	#matchedKey: string | undefined;
+	#targetKey: string | undefined;
+	#targetEscaped = false;
+	#targetUnicode = "";
+	#values: Record<string, string> = {};
+	#changed = false;
+
+	constructor(keys: readonly string[]) {
+		this.#keys = new Set(keys);
+	}
+
+	reset(): void {
+		this.#source = "";
+		this.#offset = 0;
+		this.#depth = 0;
+		this.#state = "scan";
+		this.#candidate = "";
+		this.#candidateEscaped = false;
+		this.#candidateUnicode = "";
+		this.#matchedKey = undefined;
+		this.#targetKey = undefined;
+		this.#targetEscaped = false;
+		this.#targetUnicode = "";
+		this.#values = {};
+		this.#changed = false;
+	}
+
+	update(prefix: string): StreamingJsonStringExtractorResult {
+		if (!prefix.startsWith(this.#source)) {
+			this.reset();
+		}
+		this.#source = prefix;
+		this.#changed = false;
+		while (this.#offset < prefix.length) {
+			const ch = prefix[this.#offset]!;
+			switch (this.#state) {
+				case "scan":
+					this.#scan(ch);
+					break;
+				case "candidate":
+					this.#readCandidate(ch);
+					break;
+				case "afterCandidate":
+					this.#afterCandidate(ch);
+					break;
+				case "beforeValue":
+					this.#beforeValue(ch);
+					break;
+				case "target":
+					this.#readTarget(ch);
+					break;
+			}
+		}
+		return { values: { ...this.#values }, changed: this.#changed };
+	}
+
+	#scan(ch: string): void {
+		if (ch === '"') {
+			this.#candidate = "";
+			this.#candidateEscaped = false;
+			this.#candidateUnicode = "";
+			this.#state = "candidate";
+		} else if (ch === "{" || ch === "[") {
+			this.#depth++;
+		} else if (ch === "}" || ch === "]") {
+			this.#depth--;
+		}
+		this.#offset++;
+	}
+
+	#readCandidate(ch: string): void {
+		if (this.#candidateUnicode) {
+			this.#readCandidateUnicode(ch);
+			return;
+		}
+		if (this.#candidateEscaped) {
+			if (ch === "u") {
+				this.#candidateUnicode = "u";
+			} else {
+				this.#candidate += decodeJsonStringEscape(ch);
+				this.#candidateEscaped = false;
+			}
+			this.#offset++;
+			return;
+		}
+		if (ch === "\\") {
+			this.#candidateEscaped = true;
+			this.#offset++;
+			return;
+		}
+		if (ch === '"') {
+			this.#matchedKey = this.#depth === 1 && this.#keys.has(this.#candidate) ? this.#candidate : undefined;
+			this.#state = "afterCandidate";
+			this.#offset++;
+			return;
+		}
+		this.#candidate += ch;
+		this.#offset++;
+	}
+
+	#readCandidateUnicode(ch: string): void {
+		if (isHexDigit(ch)) {
+			this.#candidateUnicode += ch;
+			if (this.#candidateUnicode.length === 5) {
+				this.#candidate += String.fromCharCode(Number.parseInt(this.#candidateUnicode.slice(1), 16));
+				this.#candidateUnicode = "";
+				this.#candidateEscaped = false;
+			}
+		} else {
+			this.#candidate += this.#candidateUnicode + ch;
+			this.#candidateUnicode = "";
+			this.#candidateEscaped = false;
+		}
+		this.#offset++;
+	}
+
+	#afterCandidate(ch: string): void {
+		if (/\s/.test(ch)) {
+			this.#offset++;
+			return;
+		}
+		const matchedKey = this.#matchedKey;
+		this.#matchedKey = undefined;
+		if (ch === ":" && matchedKey) {
+			this.#targetKey = matchedKey;
+			this.#state = "beforeValue";
+			this.#offset++;
+			return;
+		}
+		this.#state = "scan";
+	}
+
+	#beforeValue(ch: string): void {
+		if (/\s/.test(ch)) {
+			this.#offset++;
+			return;
+		}
+		if (ch === '"' && this.#targetKey) {
+			if (this.#values[this.#targetKey]) {
+				this.#values[this.#targetKey] = "";
+				this.#changed = true;
+			}
+			this.#targetEscaped = false;
+			this.#targetUnicode = "";
+			this.#state = "target";
+			this.#offset++;
+			return;
+		}
+		this.#targetKey = undefined;
+		this.#state = "scan";
+	}
+
+	#readTarget(ch: string): void {
+		if (this.#targetUnicode) {
+			this.#readTargetUnicode(ch);
+			return;
+		}
+		if (this.#targetEscaped) {
+			if (ch === "u") {
+				this.#targetUnicode = "u";
+			} else {
+				this.#appendTarget(decodeJsonStringEscape(ch));
+				this.#targetEscaped = false;
+			}
+			this.#offset++;
+			return;
+		}
+		if (ch === "\\") {
+			this.#targetEscaped = true;
+			this.#offset++;
+			return;
+		}
+		if (ch === '"') {
+			this.#targetKey = undefined;
+			this.#state = "scan";
+			this.#offset++;
+			return;
+		}
+		this.#appendTarget(ch);
+		this.#offset++;
+	}
+
+	#readTargetUnicode(ch: string): void {
+		if (isHexDigit(ch)) {
+			this.#targetUnicode += ch;
+			if (this.#targetUnicode.length === 5) {
+				this.#appendTarget(String.fromCharCode(Number.parseInt(this.#targetUnicode.slice(1), 16)));
+				this.#targetUnicode = "";
+				this.#targetEscaped = false;
+			}
+		} else {
+			this.#appendTarget(this.#targetUnicode + ch);
+			this.#targetUnicode = "";
+			this.#targetEscaped = false;
+		}
+		this.#offset++;
+	}
+
+	#appendTarget(text: string): void {
+		if (!this.#targetKey || text.length === 0) return;
+		this.#values[this.#targetKey] = `${this.#values[this.#targetKey] ?? ""}${text}`;
+		this.#changed = true;
+	}
+}
+
+function createStringExtractor(keys: readonly string[] | undefined): StreamingJsonStringExtractor | undefined {
+	return keys && keys.length > 0 ? new StreamingJsonStringExtractor(keys) : undefined;
+}
+
+function sameStringKeys(a: readonly string[], b: readonly string[] | undefined): boolean {
+	if (a.length !== (b?.length ?? 0)) return false;
+	for (let i = 0; i < a.length; i++) {
+		if (a[i] !== b?.[i]) return false;
+	}
+	return true;
+}
+
 type RevealEntry = {
 	component: ToolArgsRevealComponent | undefined;
 	/** Latest raw streamed argument text (JSON for function tools, raw text for custom tools). */
@@ -29,6 +307,9 @@ type RevealEntry = {
 	displayArgs: Record<string, unknown>;
 	/** Raw prefix carried by `displayArgs.__partialJson`. */
 	displayPrefix: string;
+	/** JSON string fields decoded incrementally between full JSON parses. */
+	streamingStringKeys: readonly string[];
+	stringExtractor: StreamingJsonStringExtractor | undefined;
 };
 
 /** Clamp a slice end into `text`, never splitting a surrogate pair: a prefix
@@ -45,7 +326,7 @@ function clampSliceEnd(text: string, end: number): number {
 type ToolArgsRevealTarget = {
 	rawInput: boolean;
 	exposeRawPartialJson: boolean;
-	fullArgs: Record<string, unknown>;
+	streamingStringKeys?: readonly string[];
 };
 
 type DisplayArgsStep = {
@@ -62,6 +343,7 @@ function resetDisplayState(entry: RevealEntry): void {
 	entry.parsedLen = 0;
 	entry.displayArgs = initialDisplayArgs();
 	entry.displayPrefix = "";
+	entry.stringExtractor?.reset();
 }
 
 /** Display args for a revealed prefix. Function-tool JSON is parsed at the same
@@ -91,6 +373,11 @@ function displayArgsForPrefix(entry: RevealEntry, prefix: string, forceParse = f
 			parsedChanged = true;
 		}
 	}
+	const extracted = entry.stringExtractor?.update(prefix);
+	if (extracted?.changed) {
+		entry.parsedArgs = { ...entry.parsedArgs, ...extracted.values };
+		parsedChanged = true;
+	}
 
 	const rawPrefixChanged = entry.exposeRawPartialJson && prefix !== entry.displayPrefix;
 	if (!parsedChanged && !rawPrefixChanged) return { args: entry.displayArgs, changed: false };
@@ -102,6 +389,39 @@ function displayArgsForPrefix(entry: RevealEntry, prefix: string, forceParse = f
 	return { args, changed: true };
 }
 
+type StreamedToolArgsSource = {
+	/** Custom-tool raw text stream (`customWireName` tools): never JSON-parsed. */
+	rawInput: boolean;
+	/** Provider-parsed arguments, spread UNDER the fresh decode: a dialect
+	 *  projector may carry keys a raw re-parse cannot recover, but any key the
+	 *  fresh parse does recover wins — provider parses lag the stream by up to
+	 *  STREAMING_JSON_PARSE_MIN_GROWTH bytes mid-stream. */
+	fullArgs?: Record<string, unknown>;
+	/** See {@link streamingStringKeysForTool}. */
+	streamingStringKeys?: readonly string[];
+};
+
+/**
+ * One-shot decode of a streamed tool-call argument buffer into display args —
+ * the same decode the live reveal applies frame-by-frame, for paths that see
+ * the buffer once (transcript rebuilds on theme change, settings, focus
+ * replay). Keeps a rebuilt preview identical to the live preview: parsed
+ * fields come from a fresh parse of the full buffer, `streamingStringKeys`
+ * fields from the incremental string decoder (which also wins ties in the
+ * live path), never from the provider's throttled `arguments`.
+ */
+export function decodeStreamedToolArgs(partialJson: string, source: StreamedToolArgsSource): Record<string, unknown> {
+	if (source.rawInput) {
+		return { input: partialJson, __partialJson: partialJson };
+	}
+	const parsed = parseStreamingJson<Record<string, unknown>>(partialJson);
+	const args: Record<string, unknown> = source.fullArgs ? { ...source.fullArgs, ...parsed } : { ...parsed };
+	const extracted = createStringExtractor(source.streamingStringKeys)?.update(partialJson);
+	if (extracted) Object.assign(args, extracted.values);
+	args.__partialJson = partialJson;
+	return args;
+}
+
 /**
  * Paces streamed tool-call arguments the same way StreamingRevealController
  * paces assistant text: providers that deliver `partialJson` in large batches
@@ -129,16 +449,15 @@ export class ToolArgsRevealController {
 
 	/**
 	 * Record the latest streamed argument text for a tool call and return the
-	 * args to render right now. With smoothing disabled the full target passes
-	 * through in the caller's legacy shape (`{ ...args, __partialJson }`).
+	 * args to render right now. With smoothing disabled nothing is paced — the
+	 * full received buffer decodes in one step — but the entry still runs the
+	 * incremental string decoder + parse throttle, so streamed text fields
+	 * (write `content`, edit bodies, eval `code`) stay fresh between the
+	 * provider's own throttled full-JSON parses instead of lagging up to
+	 * STREAMING_JSON_PARSE_MIN_GROWTH bytes behind.
 	 */
 	setTarget(id: string, partialJson: string, target: ToolArgsRevealTarget): Record<string, unknown> {
-		const { rawInput, exposeRawPartialJson, fullArgs } = target;
-		if (!this.#getSmoothStreaming()) {
-			// Toggle may flip mid-call: drop any live entry so ticks stop.
-			this.#entries.delete(id);
-			return { ...fullArgs, __partialJson: partialJson };
-		}
+		const { rawInput, exposeRawPartialJson, streamingStringKeys } = target;
 		let entry = this.#entries.get(id);
 		if (!entry) {
 			entry = {
@@ -151,13 +470,21 @@ export class ToolArgsRevealController {
 				parsedLen: 0,
 				displayArgs: initialDisplayArgs(),
 				displayPrefix: "",
+				streamingStringKeys: streamingStringKeys ?? [],
+				stringExtractor: createStringExtractor(streamingStringKeys),
 			};
 			this.#entries.set(id, entry);
 		} else {
-			if (entry.rawInput !== rawInput || entry.exposeRawPartialJson !== exposeRawPartialJson) {
+			if (
+				entry.rawInput !== rawInput ||
+				entry.exposeRawPartialJson !== exposeRawPartialJson ||
+				!sameStringKeys(entry.streamingStringKeys, streamingStringKeys)
+			) {
 				entry.rawInput = rawInput;
 				entry.exposeRawPartialJson = exposeRawPartialJson;
 				resetDisplayState(entry);
+				entry.streamingStringKeys = streamingStringKeys ?? [];
+				entry.stringExtractor = createStringExtractor(streamingStringKeys);
 			}
 			// Streams only append; a non-prefix target means a rewind — snap into range.
 			if (!partialJson.startsWith(entry.target)) {
@@ -166,6 +493,9 @@ export class ToolArgsRevealController {
 			}
 			entry.target = partialJson;
 		}
+		// Toggle may flip mid-call: snap the reveal to everything received so
+		// pacing stops (and never restarts while the toggle stays off).
+		if (!this.#getSmoothStreaming()) entry.revealed = entry.target.length;
 		entry.revealed = clampSliceEnd(entry.target, entry.revealed);
 		this.#syncTimer();
 		return displayArgsForPrefix(entry, entry.target.slice(0, entry.revealed)).args;