fix(coding-agent): memoize tool-result rendering to stop per-frame re-shape stalls

ToolExecutionComponent.#updateDisplay() re-ran renderResult (O(result-size)) on every invalidate - spinner ticks, stream chunks, resizes, keystrokes - so large results blocked the loop for seconds and typing lagged. It now early-returns on an unchanged dirty key (result version, expanded, partial, spinner frame, image visibility, theme epoch), and theme.ts exposes getThemeEpoch() bumped on every theme swap so cached blocks re-render on theme change.

packages/coding-agent/CHANGELOG.md

@@ -2,6 +2,11 @@
 
 ## [Unreleased]
 
+### Fixed
+
+- Fixed the tool-result renderer re-shaping on every `invalidate()` (spinner tick, stream chunk, resize, keystroke), which made large grep/find/read results block the main thread for seconds and made typing sluggish. `ToolExecutionComponent.#updateDisplay()` now memoizes on a dirty key (result version, expand state, partial flag, spinner frame, image visibility, theme epoch, background-task freeze state, the resolved terminal image protocol, and a display-input version that covers streamed call args, the async edit-diff preview, and Kitty image conversions) and a `#displayBuilt` guard that also fast-paths the `#contentText` fallback, so the O(result-size) shaping runs once per change instead of every frame without freezing streamed args, previews, converted images, a backgrounded task settling to its static form, or images that arrive before the async image-protocol probe resolves. Image-bearing results also re-shape on terminal resize (keyed on the resolved image dimensions only when images are present) so inline images rescale, while image-free results never re-shape on resize ([#2484](https://github.com/can1357/oh-my-pi/issues/2484))
+- Fixed `setTheme()` not bumping the theme epoch on its invalid-theme fallback path: a failed theme load swaps the active theme to the dark fallback, so memoized renderers (the tool-result renderer above) must re-shape — previously they kept the failed theme's stale colors until some other state changed ([#2484](https://github.com/can1357/oh-my-pi/issues/2484))
+
 ## [15.12.5] - 2026-06-13
 ### Changed
 

packages/coding-agent/src/modes/components/tool-execution.ts

@@ -16,7 +16,7 @@ import {
 import { getProjectDir, logger, sanitizeText } from "@oh-my-pi/pi-utils";
 import { EDIT_MODE_STRATEGIES, type EditMode, type PerFileDiffPreview } from "../../edit";
 import type { Theme } from "../../modes/theme/theme";
-import { theme } from "../../modes/theme/theme";
+import { getThemeEpoch, theme } from "../../modes/theme/theme";
 import { BASH_DEFAULT_PREVIEW_LINES } from "../../tools/bash";
 import { EVAL_DEFAULT_PREVIEW_LINES } from "../../tools/eval";
 import { isWaitingPollDetails } from "../../tools/job";
@@ -176,6 +176,22 @@ export class ToolExecutionComponent extends Container {
 	#editAllowFuzzy: boolean | undefined;
 	#snapshots?: SnapshotStore;
 	#isPartial = true;
+	#resultVersion = 0;
+	#lastDisplayKey: string | undefined;
+	// Bumped whenever a render input that #rebuildDisplay consumes but the memo
+	// key cannot cheaply hash changes: streamed call args, the async edit-diff
+	// preview, and Kitty PNG conversions. Folded into the dirty key so those
+	// updates are not swallowed by the memo (see #updateDisplay).
+	#displayInputVersion = 0;
+	// Set once #rebuildDisplay has populated the display. Replaces a
+	// #contentBox.children.length probe so the memo fast-path also covers the
+	// #contentText fallback path (which leaves #contentBox empty).
+	#displayBuilt = false;
+	// Number of Image children the last rebuild emitted. Only when this is > 0 does
+	// the memo key fold in viewport-dependent image sizing (resolveImageOptions),
+	// so a terminal resize re-shapes image-bearing results to rescale them without
+	// forcing the common image-free result to re-shape on every resize tick.
+	#renderedImageCount = 0;
 	#tool?: AgentTool;
 	#ui: TUI;
 	#cwd: string;
@@ -281,6 +297,7 @@ export class ToolExecutionComponent extends Container {
 		// signals "nothing meaningful changed" and the renderer can skip.
 		if (args === this.#args) return;
 		this.#args = args;
+		this.#displayInputVersion++;
 		this.#updateSpinnerAnimation();
 		this.#editDiffInFlight = this.#runPreviewDiff();
 		this.#updateDisplay();
@@ -365,6 +382,7 @@ export class ToolExecutionComponent extends Container {
 			if (controller.signal.aborted) return;
 			if (previews) {
 				this.#editDiffPreview = isStreaming ? stabilizeStreamingPreviews(previews) : previews;
+				this.#displayInputVersion++;
 				this.#updateDisplay();
 				this.#ui.requestRender();
 			}
@@ -393,6 +411,7 @@ export class ToolExecutionComponent extends Container {
 			return;
 		}
 		this.#result = result;
+		this.#resultVersion++;
 		this.#isPartial = isPartial;
 		// A `job` poll that found every watched job still running is transient
 		// "still waiting" chrome; keep the block displaceable so the next `job`
@@ -446,6 +465,7 @@ export class ToolExecutionComponent extends Container {
 				.toBase64()
 				.then(data => {
 					this.#convertedImages.set(index, { data, mimeType: "image/png" });
+					this.#displayInputVersion++;
 					this.#updateDisplay();
 					this.#ui.requestRender();
 				})
@@ -674,6 +694,29 @@ export class ToolExecutionComponent extends Container {
 	}
 
 	#updateDisplay(): void {
+		// `TERMINAL.imageProtocol` is resolved by an async capability probe during
+		// TUI startup, so a result rendered before it lands must re-shape once it
+		// does (it gates Image children vs text fallback in #rebuildDisplay); keyed
+		// here for the same reason markdown.ts keys its render cache on it.
+		const key = `${this.#resultVersion}|${this.#expanded}|${this.#isPartial}|${this.#spinnerFrame ?? "-"}|${this.#showImages}|${getThemeEpoch()}|${this.#displayInputVersion}|${this.#backgroundTaskFrozen}|${TERMINAL.imageProtocol ?? "-"}|${this.#imageSizeKey()}`;
+		if (key === this.#lastDisplayKey && this.#displayBuilt) return;
+		this.#lastDisplayKey = key;
+
+		this.#rebuildDisplay();
+		this.#displayBuilt = true;
+	}
+
+	// Viewport-/settings-dependent image sizing folded into the memo key only when
+	// the last rebuild actually emitted images, so a terminal resize re-shapes an
+	// image-bearing result (to rescale it) without re-shaping every image-free
+	// result on each resize tick.
+	#imageSizeKey(): string {
+		if (this.#renderedImageCount === 0) return "-";
+		const o = resolveImageOptions();
+		return `${o.maxWidthCells}:${o.maxHeightCells ?? "-"}`;
+	}
+
+	#rebuildDisplay(): void {
 		// Sync shared mutable render state for component closures
 		this.#renderState.expanded = this.#expanded;
 		this.#renderState.isPartial = this.#isPartial;
@@ -917,6 +960,7 @@ export class ToolExecutionComponent extends Container {
 				}
 			}
 		}
+		this.#renderedImageCount = this.#imageComponents.length;
 	}
 
 	#getCallArgsForRender(): any {

packages/coding-agent/src/modes/theme/theme.ts

@@ -2108,6 +2108,7 @@ var autoDarkTheme: string = "dark";
 var autoLightTheme: string = "light";
 var onThemeChangeCallback: (() => void) | undefined;
 var themeLoadRequestId: number = 0;
+let themeEpoch = 0;
 
 function getCurrentThemeOptions(): CreateThemeOptions {
 	return {
@@ -2160,9 +2161,7 @@ export async function setTheme(
 		if (enableWatcher) {
 			await startThemeWatcher();
 		}
-		if (onThemeChangeCallback) {
-			onThemeChangeCallback();
-		}
+		notifyThemeChange();
 		return { success: true };
 	} catch (error) {
 		if (requestId !== themeLoadRequestId) {
@@ -2171,6 +2170,10 @@ export async function setTheme(
 		// Theme is invalid - fall back to dark theme
 		currentThemeName = "dark";
 		theme = await loadTheme("dark", getCurrentThemeOptions());
+		// The active theme just changed to the fallback — bump the epoch so memoized
+		// renderers (e.g. ToolExecutionComponent) re-shape with the fallback colors
+		// instead of holding the failed theme's stale styling.
+		notifyThemeChange();
 		// Don't start watcher for fallback theme
 		return {
 			success: false,
@@ -2187,9 +2190,7 @@ export async function previewTheme(name: string): Promise<{ success: boolean; er
 			return { success: false, error: "Theme preview superseded by a newer request" };
 		}
 		theme = loadedTheme;
-		if (onThemeChangeCallback) {
-			onThemeChangeCallback();
-		}
+		notifyThemeChange();
 		return { success: true };
 	} catch (error) {
 		if (requestId !== themeLoadRequestId) {
@@ -2236,9 +2237,7 @@ export function setThemeInstance(themeInstance: Theme): void {
 	theme = themeInstance;
 	currentThemeName = "<in-memory>";
 	stopThemeWatcher();
-	if (onThemeChangeCallback) {
-		onThemeChangeCallback();
-	}
+	notifyThemeChange();
 }
 
 /**
@@ -2259,7 +2258,7 @@ export async function setSymbolPreset(preset: SymbolPreset): Promise<void> {
 		theme = await loadTheme("dark", getCurrentThemeOptions());
 		if (requestId !== themeLoadRequestId) return;
 	}
-	onThemeChangeCallback?.();
+	notifyThemeChange();
 }
 
 /**
@@ -2288,7 +2287,7 @@ export async function setColorBlindMode(enabled: boolean): Promise<void> {
 		theme = await loadTheme("dark", getCurrentThemeOptions());
 		if (requestId !== themeLoadRequestId) return;
 	}
-	onThemeChangeCallback?.();
+	notifyThemeChange();
 }
 
 /**
@@ -2302,6 +2301,23 @@ export function onThemeChange(callback: () => void): void {
 	onThemeChangeCallback = callback;
 }
 
+/**
+ * Monotonic counter bumped on any theme-affecting change that should invalidate
+ * cached renders: theme swaps and reloads (including the invalid-theme dark
+ * fallback), theme previews, symbol-preset changes, and color-blind-mode
+ * changes — everything that routes through {@link notifyThemeChange}. Consumers
+ * key cached renders on it so the next render re-shapes their output.
+ */
+export function getThemeEpoch(): number {
+	return themeEpoch;
+}
+
+/** Bump the theme epoch and notify the registered theme-change listener. */
+function notifyThemeChange(): void {
+	themeEpoch++;
+	onThemeChangeCallback?.();
+}
+
 /**
  * Get available symbol presets.
  */
@@ -2354,9 +2370,7 @@ async function startThemeWatcher(): Promise<void> {
 			loadTheme(watchedThemeName, getCurrentThemeOptions())
 				.then(loadedTheme => {
 					theme = loadedTheme;
-					if (onThemeChangeCallback) {
-						onThemeChangeCallback();
-					}
+					notifyThemeChange();
 				})
 				.catch(() => {
 					// Ignore errors (file might be in invalid state while being edited)
@@ -2396,9 +2410,7 @@ function reevaluateAutoTheme(debugLabel: string): void {
 	loadTheme(resolved, getCurrentThemeOptions())
 		.then(loadedTheme => {
 			theme = loadedTheme;
-			if (onThemeChangeCallback) {
-				onThemeChangeCallback();
-			}
+			notifyThemeChange();
 		})
 		.catch(err => {
 			logger.debug(`Theme switch on ${debugLabel} failed`, { error: String(err) });

packages/coding-agent/test/theme-epoch-fallback.test.ts

@@ -0,0 +1,42 @@
+import { afterEach, beforeAll, describe, expect, it } from "bun:test";
+import {
+	getThemeByName,
+	getThemeEpoch,
+	setTheme,
+	setThemeInstance,
+	type Theme,
+} from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
+
+/**
+ * Contract: every change to the *active* theme bumps the theme epoch — including
+ * setTheme()'s fallback path. ToolExecutionComponent (and other memoized
+ * renderers) fold getThemeEpoch() into their dirty key, so a failed theme load
+ * that swapped to the dark fallback without bumping the epoch would leave those
+ * renderers holding the failed theme's stale colors until some other state moved.
+ */
+describe("theme epoch — setTheme fallback", () => {
+	let dark: Theme;
+
+	beforeAll(async () => {
+		const t = await getThemeByName("dark");
+		if (!t) throw new Error("Expected dark theme to exist");
+		dark = t;
+	});
+
+	afterEach(() => {
+		// Leave a deterministic active theme for any later case in this process.
+		setThemeInstance(dark);
+	});
+
+	it("bumps the epoch when an invalid theme name falls back to dark", async () => {
+		setThemeInstance(dark);
+		const before = getThemeEpoch();
+
+		const result = await setTheme("__definitely_not_a_real_theme__");
+
+		// Invalid theme → load throws → dark fallback applied.
+		expect(result.success).toBe(false);
+		// The active theme changed, so the epoch must advance for memoized renderers.
+		expect(getThemeEpoch()).toBeGreaterThan(before);
+	});
+});

packages/coding-agent/test/tool-execution-memoization.test.ts

@@ -0,0 +1,175 @@
+import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test";
+import { stripVTControlCharacters } from "node:util";
+import type { AgentTool } from "@oh-my-pi/pi-agent-core";
+import { ToolExecutionComponent } from "@oh-my-pi/pi-coding-agent/modes/components/tool-execution";
+import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
+import { Text, type TUI } from "@oh-my-pi/pi-tui";
+
+/**
+ * Contract under test (tool-result render memoization):
+ *
+ * `ToolExecutionComponent` shapes a tool result into UI components by calling
+ * the tool's `renderResult` — an O(result-size) pass. A dirty-key guard at the
+ * top of `#updateDisplay()` must collapse the result version, expand state,
+ * partial flag, spinner frame, show-images flag, and theme epoch into one key
+ * and skip `#rebuildDisplay()` when nothing meaningful changed. So:
+ *
+ *   - A flood of `invalidate()` calls (one per render frame) after a final
+ *     result must re-shape EXACTLY ONCE, not once per frame — this is the
+ *     regression guard against the per-frame re-shape stall.
+ *   - A state change that actually alters output (`setExpanded(true)`) must
+ *     force exactly one additional shaping pass, and the new output must be
+ *     observable; a redundant no-op set of the same state must not re-shape.
+ *   - Bumping the result version (a NEW result) must force exactly one more
+ *     shaping pass, and the rendered output must reflect the new result.
+ */
+describe("ToolExecutionComponent tool-result render memoization", () => {
+	beforeAll(async () => {
+		await initTheme();
+	});
+
+	afterEach(() => {
+		vi.restoreAllMocks();
+	});
+
+	// A custom tool whose `renderResult` is the single O(result-size) shaping
+	// function. It echoes the result text so the rendered frame reflects which
+	// result was last shaped — letting us assert the memo never suppresses a
+	// real change, only redundant repaints.
+	function makeShapingTool() {
+		return {
+			name: "custom_render",
+			label: "Custom",
+			renderResult(result: { content: Array<{ type: string; text?: string }> }): Text {
+				const joined = result.content.map(c => c.text ?? "").join("");
+				return new Text(`shaped:${joined}`, 0, 0);
+			},
+		};
+	}
+
+	function finalResult(text: string) {
+		return { content: [{ type: "text", text }] };
+	}
+
+	it("re-shapes once per meaningful change, never per invalidate() frame", () => {
+		const tool = makeShapingTool();
+		const shapeSpy = vi.spyOn(tool, "renderResult");
+		const ui = { requestRender() {} } as unknown as TUI;
+
+		const component = new ToolExecutionComponent(
+			"custom_render",
+			{},
+			{},
+			tool as unknown as AgentTool,
+			ui,
+			process.cwd(),
+		);
+
+		// No result yet: the shaping pass has not run.
+		expect(shapeSpy).toHaveBeenCalledTimes(0);
+
+		// Phase 1 — a final (non-partial) result shapes exactly once, and a
+		// flood of per-frame invalidate()s must NOT re-shape (the regression).
+		component.updateResult(finalResult("ALPHA"), false);
+		expect(shapeSpy).toHaveBeenCalledTimes(1);
+		for (let i = 0; i < 12; i++) component.invalidate();
+		expect(shapeSpy).toHaveBeenCalledTimes(1);
+		expect(stripVTControlCharacters(component.render(80).join("\n"))).toContain("shaped:ALPHA");
+
+		// Phase 2 — a state change that alters output forces exactly one more
+		// shaping pass; further invalidate()s and a redundant same-value set do
+		// not, and the expanded frame is observable.
+		component.setExpanded(true);
+		expect(shapeSpy).toHaveBeenCalledTimes(2);
+		for (let i = 0; i < 12; i++) component.invalidate();
+		component.setExpanded(true);
+		expect(shapeSpy).toHaveBeenCalledTimes(2);
+
+		// Phase 3 — a NEW result (bumped version) forces exactly one more pass,
+		// and the rendered output reflects the new result, not the stale one.
+		component.updateResult(finalResult("BRAVO"), false);
+		expect(shapeSpy).toHaveBeenCalledTimes(3);
+		const frame = stripVTControlCharacters(component.render(80).join("\n"));
+		expect(frame).toContain("shaped:BRAVO");
+		expect(frame).not.toContain("shaped:ALPHA");
+	});
+
+	// Regression: the memo key must also cover streamed call-arg changes. The
+	// dirty key folds in a display-input version bumped by updateArgs(), so a
+	// new args object re-shapes the call preview instead of freezing it at the
+	// first render (the bug: key omitted #args, so once the display was built
+	// every streamed delta was swallowed by the guard).
+	it("re-shapes the call preview when streamed args change, not only on key fields", () => {
+		const tool = {
+			name: "custom_render",
+			label: "Custom",
+			renderCall(args: { cmd?: string }): Text {
+				return new Text(`call:${args?.cmd ?? ""}`, 0, 0);
+			},
+		};
+		const callSpy = vi.spyOn(tool, "renderCall");
+		const ui = { requestRender() {} } as unknown as TUI;
+
+		const component = new ToolExecutionComponent(
+			"custom_render",
+			{ cmd: "A" },
+			{},
+			tool as unknown as AgentTool,
+			ui,
+			process.cwd(),
+		);
+
+		// Constructor shaped the call preview once with the initial args.
+		expect(stripVTControlCharacters(component.render(80).join("\n"))).toContain("call:A");
+		const afterCtor = callSpy.mock.calls.length;
+
+		// A flood of per-frame invalidate()s must NOT re-shape (memo still holds).
+		for (let i = 0; i < 12; i++) component.invalidate();
+		expect(callSpy.mock.calls.length).toBe(afterCtor);
+
+		// A NEW args object (streamed delta) MUST re-shape and reflect the change,
+		// even though no key field (result version, expanded, …) moved.
+		component.updateArgs({ cmd: "B" });
+		expect(callSpy.mock.calls.length).toBe(afterCtor + 1);
+		const frame = stripVTControlCharacters(component.render(80).join("\n"));
+		expect(frame).toContain("call:B");
+		expect(frame).not.toContain("call:A");
+
+		// A same-reference updateArgs is the documented no-op and must not re-shape.
+		const sameArgs = { cmd: "C" };
+		component.updateArgs(sameArgs);
+		const afterReal = callSpy.mock.calls.length;
+		component.updateArgs(sameArgs);
+		expect(callSpy.mock.calls.length).toBe(afterReal);
+	});
+
+	// Regression: freezing a backgrounded task (seal()) flips #backgroundTaskFrozen,
+	// which the render context consumes (context.frozen) — so it must be in the memo
+	// key. The bug: the key omitted it, so once the display was built seal()'s
+	// #updateDisplay() early-returned and the row stayed styled as live progress.
+	it("re-shapes when a background task freezes via seal(), not only on key fields", () => {
+		const tool = makeShapingTool();
+		const shapeSpy = vi.spyOn(tool, "renderResult");
+		const ui = { requestRender() {} } as unknown as TUI;
+
+		const component = new ToolExecutionComponent(
+			"custom_render",
+			{},
+			{},
+			tool as unknown as AgentTool,
+			ui,
+			process.cwd(),
+		);
+
+		// A partial result shapes once; a flood of invalidate()s must not re-shape.
+		component.updateResult(finalResult("RUNNING"), true);
+		const afterResult = shapeSpy.mock.calls.length;
+		for (let i = 0; i < 12; i++) component.invalidate();
+		expect(shapeSpy.mock.calls.length).toBe(afterResult);
+
+		// seal() only flips #backgroundTaskFrozen (no result version / spinner / expand
+		// change), so the memo must still re-shape to settle the row to its frozen form.
+		component.seal();
+		expect(shapeSpy.mock.calls.length).toBe(afterResult + 1);
+	});
+});