From ee45472c05ffd8ba86cd59ba8ac0778d7f6b5a86 Mon Sep 17 00:00:00 2001 From: metaphorics <152830360+metaphorics@users.noreply.github.com> Date: Sun, 14 Jun 2026 05:14:41 +0900 Subject: [PATCH] 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 | 5 + .../src/modes/components/tool-execution.ts | 46 ++++- .../coding-agent/src/modes/theme/theme.ts | 46 +++-- .../test/theme-epoch-fallback.test.ts | 42 +++++ .../test/tool-execution-memoization.test.ts | 175 ++++++++++++++++++ 5 files changed, 296 insertions(+), 18 deletions(-) create mode 100644 packages/coding-agent/test/theme-epoch-fallback.test.ts create mode 100644 packages/coding-agent/test/tool-execution-memoization.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 7fc96cb5e..d5c517dfa 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/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 diff --git a/packages/coding-agent/src/modes/components/tool-execution.ts b/packages/coding-agent/src/modes/components/tool-execution.ts index 8d659cb1d..d98e6edee 100644 --- a/packages/coding-agent/src/modes/components/tool-execution.ts +++ b/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 { diff --git a/packages/coding-agent/src/modes/theme/theme.ts b/packages/coding-agent/src/modes/theme/theme.ts index a0e2390e1..72ea5b306 100644 --- a/packages/coding-agent/src/modes/theme/theme.ts +++ b/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 = ""; stopThemeWatcher(); - if (onThemeChangeCallback) { - onThemeChangeCallback(); - } + notifyThemeChange(); } /** @@ -2259,7 +2258,7 @@ export async function setSymbolPreset(preset: SymbolPreset): Promise { theme = await loadTheme("dark", getCurrentThemeOptions()); if (requestId !== themeLoadRequestId) return; } - onThemeChangeCallback?.(); + notifyThemeChange(); } /** @@ -2288,7 +2287,7 @@ export async function setColorBlindMode(enabled: boolean): Promise { 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 { 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) }); diff --git a/packages/coding-agent/test/theme-epoch-fallback.test.ts b/packages/coding-agent/test/theme-epoch-fallback.test.ts new file mode 100644 index 000000000..005d3f251 --- /dev/null +++ b/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); + }); +}); diff --git a/packages/coding-agent/test/tool-execution-memoization.test.ts b/packages/coding-agent/test/tool-execution-memoization.test.ts new file mode 100644 index 000000000..148df9e97 --- /dev/null +++ b/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); + }); +});