From 5882cdb43a1f7ff4edffe4efef7a637fb837656b Mon Sep 17 00:00:00 2001 From: shoucandanghehe Date: Fri, 12 Jun 2026 10:40:30 +0800 Subject: [PATCH] feat(coding-agent): allow thinking renderers to replace blocks --- docs/extensions.md | 15 +- packages/coding-agent/CHANGELOG.md | 4 + .../src/extensibility/extensions/types.ts | 27 ++- .../src/modes/components/assistant-message.ts | 185 ++++++++++++++---- .../modes/components/transcript-container.ts | 64 ++++-- .../assistant-message-mermaid.test.ts | 163 +++++++++++++-- .../components/transcript-container.test.ts | 159 +++++++++++++++ 7 files changed, 546 insertions(+), 71 deletions(-) diff --git a/docs/extensions.md b/docs/extensions.md index 337a1d65c83..2fa5d7dfa31 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -389,11 +389,22 @@ import { Container, Text } from "@oh-my-pi/pi-tui"; pi.registerAssistantThinkingRenderer((context, theme) => { const container = new Container(); container.addChild(new Text(theme.fg("dim", `thinking chars: ${context.text.length}`), 1, 0)); - return container; + return container; // legacy shorthand for { type: "append", component: container } }); ``` -Used by interactive rendering to add display-only supplemental UI below each visible assistant thinking block. The renderer receives the already-visible thinking text, content/thinking indexes, theme, and a `requestRender()` callback for async renderers. All registered renderers that return a component are appended in registration order. Renderers must not mutate messages; the original thinking block remains the provider/session source of truth. +Used by interactive rendering to customize each visible assistant thinking block. The renderer receives parent assistant message metadata, provider thinking content metadata, the current full trimmed thinking text snapshot, content/thinking indexes, theme, and a `requestRender()` callback for async renderers. `requestRender()` repaints already-mounted renderer output; if the renderer returned `undefined`, it re-runs this thinking block's renderers and schedules a repaint so async work can later mount `{ type: "append" }` or `{ type: "replace" }` output. Renderers must not mutate messages; the thinking block remains the provider/session source of truth. + +`context.message.timestamp` is stable across streaming partials and is the recommended anchor for renderer-local async state keys. `context.message.responseId` may be absent or appear only after provider finalization; do not key streaming state on `responseId` alone. `context.content.itemId` and `context.content.thinkingSignature` identify provider-native reasoning items when available; fall back to `contentIndex` and `thinkingIndex` otherwise. + +Renderer results: + +- `Component`: legacy shorthand for appending supplemental UI below the original thinking Markdown. +- `{ type: "append", component: Component }`: explicit append behavior, equivalent to returning a bare component. +- `{ type: "replace", component: Component }`: suppresses the default thinking Markdown and renders the replacement component instead. +- `undefined`: renders nothing extra and leaves the default thinking Markdown visible. + +If multiple renderers return `{ type: "replace", component }`, the first replacement wins and later replacements are ignored with a warning. Append renderers still run after a replacement, in registration order. ## Tool call/result renderer diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c5595a8932f..624b6406dc6 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- Added explicit assistant thinking renderer result arms for extensions: `{ type: "append", component: Component }` keeps the existing supplemental rendering behavior, while `{ type: "replace", component: Component }` suppresses the default thinking Markdown so extensions can render the block themselves. Bare `Component` returns remain the legacy append shorthand. Thinking render contexts now include parent assistant message metadata and provider thinking content metadata for async renderer state. + ## [16.3.9] - 2026-07-06 ### Added diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 82d513257a6..a031dc48521 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -17,6 +17,7 @@ import type { import type { CompactionResult } from "@oh-my-pi/pi-agent-core/compaction"; import type { Api, + AssistantMessage, AssistantMessageEvent, AssistantMessageEventStream, Context, @@ -27,6 +28,7 @@ import type { SimpleStreamOptions, Static, TextContent, + ThinkingContent, TSchema, } from "@oh-my-pi/pi-ai"; import type { OAuthCredentials, OAuthLoginCallbacks } from "@oh-my-pi/pi-ai/oauth/types"; @@ -921,16 +923,37 @@ export type MessageRenderer = ( ) => Component | undefined; export interface AssistantThinkingRenderContext { + /** + * Parent assistant message metadata. `timestamp` is stable across streaming partials + * and is the recommended anchor for renderer-local async state keys. `responseId` + * may be absent or appear only after provider finalization; do not key streaming + * state on `responseId` alone. + */ + message: Readonly>; + /** + * Provider thinking content item currently being rendered. `itemId` and + * `thinkingSignature` identify provider-native reasoning items when available; + * fall back to `contentIndex` and `thinkingIndex` otherwise. + */ + content: Readonly>; contentIndex: number; thinkingIndex: number; + /** Current full trimmed thinking snapshot. */ text: string; + /** Repaint mounted output, or re-run this thinking block's renderers when no component was mounted yet. */ requestRender(): void; } +export type AssistantThinkingRenderResult = + | Component + | { type: "append"; component: Component } + | { type: "replace"; component: Component } + | undefined; + export type AssistantThinkingRenderer = ( context: AssistantThinkingRenderContext, theme: Theme, -) => Component | undefined; +) => AssistantThinkingRenderResult; // ============================================================================ // Command Registration @@ -1088,7 +1111,7 @@ export interface ExtensionAPI { /** Register a custom renderer for CustomMessageEntry. */ registerMessageRenderer(customType: string, renderer: MessageRenderer): void; - /** Register a renderer for assistant thinking blocks. Rendered after the original thinking text. */ + /** Register a renderer for assistant thinking blocks. Legacy Component returns append after the original thinking text. */ registerAssistantThinkingRenderer(renderer: AssistantThinkingRenderer): void; // ========================================================================= diff --git a/packages/coding-agent/src/modes/components/assistant-message.ts b/packages/coding-agent/src/modes/components/assistant-message.ts index e09fc8c15d9..5c8d5b0dae4 100644 --- a/packages/coding-agent/src/modes/components/assistant-message.ts +++ b/packages/coding-agent/src/modes/components/assistant-message.ts @@ -1,8 +1,18 @@ -import type { AssistantMessage, ImageContent } from "@oh-my-pi/pi-ai"; -import { Container, Image, type ImageBudget, ImageProtocol, Markdown, Spacer, TERMINAL, Text } from "@oh-my-pi/pi-tui"; -import { formatNumber } from "@oh-my-pi/pi-utils"; +import type { AssistantMessage, ImageContent, ThinkingContent } from "@oh-my-pi/pi-ai"; +import { + type Component, + Container, + Image, + type ImageBudget, + ImageProtocol, + Markdown, + Spacer, + TERMINAL, + Text, +} from "@oh-my-pi/pi-tui"; +import { formatNumber, logger } from "@oh-my-pi/pi-utils"; import chalk from "chalk"; -import type { AssistantThinkingRenderer } from "../../extensibility/extensions/types"; +import type { AssistantThinkingRenderer, AssistantThinkingRenderResult } from "../../extensibility/extensions/types"; import { getMarkdownTheme, theme } from "../../modes/theme/theme"; import { getPreviewLines, resolveImageOptions, TRUNCATE_LENGTHS } from "../../tools/render-utils"; import { canonicalizeMessage, formatThinkingForDisplay, hasDisplayableThinking } from "../../utils/thinking-display"; @@ -164,6 +174,24 @@ function lerpHex(from: string, to: string, t: number): string { return `#${((1 << 24) + (r << 16) + (g << 8) + b).toString(16).slice(1)}`; } +interface ThinkingExtensionComponents { + append: Component[]; + replace?: Component; +} + +function isComponent(result: AssistantThinkingRenderResult): result is Component { + return typeof result === "object" && result !== null && typeof (result as Component).render === "function"; +} + +function isStructuredThinkingResult( + result: AssistantThinkingRenderResult, +): result is { type: "append" | "replace"; component: Component } { + if (typeof result !== "object" || result === null || isComponent(result)) return false; + const candidate = result as { type?: unknown; component?: unknown }; + if (!isComponent(candidate.component as AssistantThinkingRenderResult)) return false; + return candidate.type === "append" || candidate.type === "replace"; +} + /** * Component that renders a complete assistant message */ @@ -233,6 +261,8 @@ export class AssistantMessageComponent extends Container { * session-wide {@link sharedSpeedTracker} can't surface a previous turn's rate * on a fresh block that has no live token throughput of its own. */ #thinkingRateLive = false; + #visibleThinkingUsesRenderers = false; + #thinkingRenderRefreshQueued = false; constructor( message?: AssistantMessage, @@ -451,6 +481,17 @@ export class AssistantMessageComponent extends Container { return this.#blockVersion; } + /** + * Default assistant text/thinking streams are append-only while visible rows + * keep prior content and grow at the bottom. Once a visible thinking block is + * handed to extension renderers, any renderer can later switch from `undefined` + * to `{ type: "replace" }` via `requestRender()`, so keep those rows in the + * repaintable live region instead of marking them safe for native scrollback. + */ + isTranscriptBlockAppendOnly(): boolean { + return !this.#visibleThinkingUsesRenderers; + } + markTranscriptBlockFinalized(): void { this.#transcriptBlockFinalized = true; this.#stopThinkingAnimation(); @@ -577,25 +618,101 @@ export class AssistantMessageComponent extends Container { } } - #appendThinkingExtensions(contentIndex: number, thinkingIndex: number, text: string): void { + #requestThinkingRender(rerunRenderers: boolean): void { + if (!rerunRenderers) { + this.onImageUpdate?.(); + return; + } + if (this.#thinkingRenderRefreshQueued) return; + this.#thinkingRenderRefreshQueued = true; + queueMicrotask(() => { + try { + if (this.#lastMessage) { + this.#fastPathKey = undefined; + this.#fastPathItems = undefined; + this.updateContent(this.#lastMessage, { transient: this.#lastUpdateTransient }); + } else { + super.invalidate(); + } + this.onImageUpdate?.(); + } catch (error) { + logger.warn("Assistant thinking renderer refresh failed", { + error: error instanceof Error ? error.message : String(error), + }); + } finally { + this.#thinkingRenderRefreshQueued = false; + } + }); + } + + #renderThinkingExtensions( + message: AssistantMessage, + content: ThinkingContent, + contentIndex: number, + thinkingIndex: number, + text: string, + ): ThinkingExtensionComponents { + const append: Component[] = []; + let replacement: Component | undefined; for (const renderer of this.thinkingRenderers) { + let rerunRenderersOnRequest = true; try { - const component = renderer( + const result: AssistantThinkingRenderResult = renderer( { + message: { + timestamp: message.timestamp, + responseId: message.responseId, + api: message.api, + provider: message.provider, + model: message.model, + }, + content: { + itemId: content.itemId, + thinkingSignature: content.thinkingSignature, + }, contentIndex, thinkingIndex, text, - requestRender: () => this.onImageUpdate?.(), + requestRender: () => this.#requestThinkingRender(rerunRenderersOnRequest), }, theme, ); - if (component) { - this.#contentContainer.addChild(component); + if (!result) continue; + if (isComponent(result)) { + rerunRenderersOnRequest = false; + append.push(result); + continue; + } + if (!isStructuredThinkingResult(result)) { + logger.warn("Assistant thinking renderer returned an invalid result", { + contentIndex, + thinkingIndex, + }); + continue; + } + if (result.type === "replace") { + if (replacement) { + logger.warn("Assistant thinking replacement renderer ignored because one is already active", { + contentIndex, + thinkingIndex, + }); + continue; + } + rerunRenderersOnRequest = false; + replacement = result.component; + continue; } - } catch { - // Ignore extension renderer failures and keep the original thinking block visible. + rerunRenderersOnRequest = false; + append.push(result.component); + } catch (error) { + logger.warn("Assistant thinking renderer failed", { + contentIndex, + thinkingIndex, + error: error instanceof Error ? error.message : String(error), + }); } } + return { append, replace: replacement }; } #computeShapeKey(message: AssistantMessage): string { @@ -628,19 +745,9 @@ export class AssistantMessageComponent extends Container { if (errorPresentation.kind === "full" && !(message.stopReason === "error" && this.#errorPinned)) { return false; } - // Extension stability: if thinking renderers exist and any tracked thinking - // block's text changed, extensions may produce a different child count. - if (this.thinkingRenderers.length > 0 && this.#fastPathItems) { - for (const item of this.#fastPathItems) { - if (item.blockType === "thinking") { - const content = message.content[item.contentIndex]; - if (content?.type === "thinking") { - const display = resolveThinkingDisplay(content, this.proseOnlyThinking); - if (display.text !== item.lastText) return false; - } - } - } - } + // Thinking renderers are arbitrary extension code: their output can change + // even when the message shape and text are identical. + if (this.thinkingRenderers.length > 0) return false; return true; } @@ -677,6 +784,7 @@ export class AssistantMessageComponent extends Container { this.#fastPathItems = undefined; return false; } + item.md.transientRenderCache = transient; if (newText !== item.lastText) { // Only the last (actively streaming) block may mutate in place: a // delta into an earlier block would invalidate rows the settled @@ -756,6 +864,7 @@ export class AssistantMessageComponent extends Container { // Fast path: reuse Markdown children when shape is stable during streaming if (this.#tryFastPathUpdate(message, opts)) return; + this.#visibleThinkingUsesRenderers = false; // Clear content container this.#contentContainer.clear(); this.#thinkingDots = undefined; @@ -801,15 +910,25 @@ export class AssistantMessageComponent extends Container { (c.type === "thinking" && resolveThinkingDisplay(c, this.proseOnlyThinking).visible), ); - // Thinking traces in thinkingText color, italic - const md = new Markdown(thinkingText, 1, 0, getMarkdownTheme(), { - color: (text: string) => theme.fg("thinkingText", text), - italic: true, - }); - md.transientRenderCache = this.#lastUpdateTransient; - this.#contentContainer.addChild(md); - captureItems?.push({ md, contentIndex: i, blockType: "thinking", lastText: thinkingText }); - this.#appendThinkingExtensions(i, thinkingIndex, thinkingText); + if (this.thinkingRenderers.length > 0) { + this.#visibleThinkingUsesRenderers = true; + } + const thinkingComponents = this.#renderThinkingExtensions(message, content, i, thinkingIndex, thinkingText); + if (thinkingComponents.replace) { + this.#contentContainer.addChild(thinkingComponents.replace); + } else { + // Thinking traces in thinkingText color, italic + const md = new Markdown(thinkingText, 1, 0, getMarkdownTheme(), { + color: (text: string) => theme.fg("thinkingText", text), + italic: true, + }); + md.transientRenderCache = this.#lastUpdateTransient; + this.#contentContainer.addChild(md); + captureItems?.push({ md, contentIndex: i, blockType: "thinking", lastText: thinkingText }); + } + for (const component of thinkingComponents.append) { + this.#contentContainer.addChild(component); + } thinkingIndex += 1; if (hasVisibleContentAfter) { this.#contentContainer.addChild(new Spacer(1)); diff --git a/packages/coding-agent/src/modes/components/transcript-container.ts b/packages/coding-agent/src/modes/components/transcript-container.ts index c66b081c1ee..43f2c035273 100644 --- a/packages/coding-agent/src/modes/components/transcript-container.ts +++ b/packages/coding-agent/src/modes/components/transcript-container.ts @@ -15,6 +15,13 @@ import { */ interface FinalizableBlock { isTranscriptBlockFinalized?(): boolean; + /** + * Blocks may opt out of native scrollback commit promotion even when their + * current frame is otherwise eligible. Replacement thinking renderers own + * previously rendered rows and can re-layout them after async state changes, + * so those rows must stay in the repaintable live region. + */ + isTranscriptBlockAppendOnly?(): boolean; /** * Monotonic content version for blocks that can still mutate *after* * reporting finalized (e.g. `AssistantMessageComponent`: the inline error @@ -51,6 +58,13 @@ function getBlockVersion(child: Component): number | undefined { const fn = (child as Component & FinalizableBlock).getTranscriptBlockVersion; return fn ? fn.call(child) : undefined; } +function canCommitLiveBlock(child: Component): boolean { + const fn = (child as Component & FinalizableBlock).isTranscriptBlockAppendOnly; + return fn ? fn.call(child) : true; +} +function isLiveRegionBoundary(child: Component): boolean { + return !isBlockFinalized(child) || !canCommitLiveBlock(child); +} /** Clamped read of a block's declared settled rows (see {@link FinalizableBlock}). */ function getBlockSettledRows(child: Component): number { @@ -216,12 +230,13 @@ export class TranscriptContainer const index = children.indexOf(component); if (index < 0) return false; for (let i = 0; i <= index; i++) { - if (!isBlockFinalized(children[i]!)) return true; + if (isLiveRegionBoundary(children[i]!)) return true; } - // Every block at/before `index` finalized: the live region starts at the - // first unfinalized block below it, or at the last child when none exists. + // Every block at/before `index` finalized and can commit: the live region + // starts at the first live/non-committable block below it, or at the last + // child when none exists. for (let i = index + 1; i < children.length; i++) { - if (!isBlockFinalized(children[i]!)) return false; + if (isLiveRegionBoundary(children[i]!)) return false; } return index === children.length - 1; } @@ -275,16 +290,19 @@ export class TranscriptContainer const count = this.children.length; - // The commit boundary stops at the earliest still-mutating block. A - // block that has not finalized must gate it: out-of-band inserts - // (TTSR/todo cards) can append a finalized block *below* a tool that is - // still awaiting its result, and committing rows there would strand the - // tool's history rows on a mid-stream preview the late result never - // reaches. + // The commit boundary stops at the earliest still-mutating or explicitly + // non-committable block. A block that has not finalized must gate it: + // out-of-band inserts (TTSR/todo cards) can append a finalized block + // *below* a tool that is still awaiting its result, and committing rows + // there would strand the tool's history rows on a mid-stream preview the + // late result never reaches. A finalized block may also opt out when it can + // still repaint prior rows through async renderer state, e.g. replacement- + // capable assistant thinking. let liveStartIndex = -1; let hasLiveBlock = false; for (let i = 0; i < count; i++) { - if (!isBlockFinalized(this.children[i]!)) { + const child = this.children[i]!; + if (isLiveRegionBoundary(child)) { liveStartIndex = i; hasLiveBlock = true; break; @@ -327,6 +345,9 @@ export class TranscriptContainer // post-finalize re-layouts, and expand toggles remain visible. const previous = previousSegments[i]; const finalized = isBlockFinalized(child); + const commitAllowed = canCommitLiveBlock(child); + const inLiveRegion = i >= liveStartIndex; + const liveBoundary = !finalized || !commitAllowed; const version = getBlockVersion(child); const committedReusable = previous !== undefined && @@ -393,17 +414,20 @@ export class TranscriptContainer const sep = row > 0 && !isPlainBlank(lines[row - 1]!) ? 1 : 0; // The separator before the first live block stays in the committed - // prefix (it is deterministic once the prior block's body is - // settled); the boundary then extends through the live block's - // declared settled rows, mapped from its raw render into the - // stripped contribution. + // prefix (it is deterministic once the prior block's body is settled). + // Append-only live blocks may extend the exactness boundary through + // declared settled rows, mapped from raw render into stripped + // contribution; non-committable blocks keep the boundary at their first + // content row so replacement-capable renderers stay in the live region. if (hasLiveBlock && i === liveStartIndex) { let settled = 0; - const settledRaw = getBlockSettledRows(child); - if (settledRaw > 0) { - let lead = 0; - while (lead < raw.length && isPlainBlank(raw[lead]!)) lead++; - settled = Math.max(0, Math.min(contribution.length, settledRaw - lead)); + if (commitAllowed) { + const settledRaw = getBlockSettledRows(child); + if (settledRaw > 0) { + let lead = 0; + while (lead < raw.length && isPlainBlank(raw[lead]!)) lead++; + settled = Math.max(0, Math.min(contribution.length, settledRaw - lead)); + } } this.#nativeScrollbackLiveRegionStart = row + sep + settled; } diff --git a/packages/coding-agent/test/modes/components/assistant-message-mermaid.test.ts b/packages/coding-agent/test/modes/components/assistant-message-mermaid.test.ts index 33582e8bd42..30da1a894b2 100644 --- a/packages/coding-agent/test/modes/components/assistant-message-mermaid.test.ts +++ b/packages/coding-agent/test/modes/components/assistant-message-mermaid.test.ts @@ -150,7 +150,64 @@ describe("AssistantMessageComponent settled-row commit boundary", () => { describe("AssistantMessageComponent thinking renderers", () => { it("renders all extension outputs below visible thinking blocks in registration order", () => { - const contexts: Array<{ contentIndex: number; thinkingIndex: number; text: string }> = []; + const contexts: Array<{ + message: { timestamp: number; responseId?: string; api: string; provider: string; model: string }; + content: { itemId?: string; thinkingSignature?: string }; + contentIndex: number; + thinkingIndex: number; + text: string; + }> = []; + const message = { + ...createAssistantMessage(""), + responseId: "resp-1", + content: [ + { + type: "thinking" as const, + thinking: "I should inspect the input.", + itemId: "item-1", + thinkingSignature: "sig-1", + }, + ], + }; + const component = new AssistantMessageComponent(message, false, undefined, [ + context => { + contexts.push({ + message: context.message, + content: context.content, + contentIndex: context.contentIndex, + thinkingIndex: context.thinkingIndex, + text: context.text, + }); + return new Text("first note", 1, 0); + }, + () => new Text("second note", 1, 0), + ]); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).toContain("I should inspect the input."); + expect(rendered.indexOf("I should inspect the input.")).toBeLessThan(rendered.indexOf("first note")); + expect(rendered.indexOf("first note")).toBeLessThan(rendered.indexOf("second note")); + expect(contexts).toEqual([ + { + message: { + timestamp: message.timestamp, + responseId: "resp-1", + api: "anthropic-messages", + provider: "anthropic", + model: "claude-sonnet-4-5", + }, + content: { + itemId: "item-1", + thinkingSignature: "sig-1", + }, + contentIndex: 0, + thinkingIndex: 0, + text: "I should inspect the input.", + }, + ]); + }); + + it("preserves legacy component returns that define structured result fields", () => { const component = new AssistantMessageComponent( { ...createAssistantMessage(""), @@ -159,23 +216,62 @@ describe("AssistantMessageComponent thinking renderers", () => { false, undefined, [ - context => { - contexts.push({ - contentIndex: context.contentIndex, - thinkingIndex: context.thinkingIndex, - text: context.text, - }); - return new Text("first note", 1, 0); - }, - () => new Text("second note", 1, 0), + () => + Object.assign(new Text("legacy component with structured fields", 1, 0), { + type: "replace", + component: new Text("wrong structured replacement", 1, 0), + }), ], ); const rendered = Bun.stripANSI(component.render(120).join("\n")); expect(rendered).toContain("I should inspect the input."); - expect(rendered.indexOf("I should inspect the input.")).toBeLessThan(rendered.indexOf("first note")); - expect(rendered.indexOf("first note")).toBeLessThan(rendered.indexOf("second note")); - expect(contexts).toEqual([{ contentIndex: 0, thinkingIndex: 0, text: "I should inspect the input." }]); + expect(rendered).toContain("legacy component with structured fields"); + expect(rendered).not.toContain("wrong structured replacement"); + }); + + it("lets extension renderers replace the default thinking block while preserving appenders", () => { + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "Sensitive chain of thought." }], + }, + false, + undefined, + [ + () => ({ type: "append", component: new Text("early audit note", 1, 0) }), + () => ({ type: "replace", component: new Text("redacted outline", 1, 0) }), + () => new Text("legacy appended note", 1, 0), + ], + ); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).not.toContain("Sensitive chain of thought."); + expect(rendered).toContain("redacted outline"); + expect(rendered).toContain("early audit note"); + expect(rendered).toContain("legacy appended note"); + expect(rendered.indexOf("redacted outline")).toBeLessThan(rendered.indexOf("early audit note")); + expect(rendered.indexOf("early audit note")).toBeLessThan(rendered.indexOf("legacy appended note")); + }); + + it("uses the first replacement renderer when multiple extensions try to replace thinking", () => { + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "Original thinking." }], + }, + false, + undefined, + [ + () => ({ type: "replace", component: new Text("first replacement", 1, 0) }), + () => ({ type: "replace", component: new Text("second replacement", 1, 0) }), + ], + ); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).not.toContain("Original thinking."); + expect(rendered).toContain("first replacement"); + expect(rendered).not.toContain("second replacement"); }); it("keeps original thinking visible when an extension renderer throws", () => { @@ -198,7 +294,7 @@ describe("AssistantMessageComponent thinking renderers", () => { expect(rendered).not.toContain("renderer failed"); }); - it("keeps async renderer components mounted when they request a render", () => { + it("keeps async renderer components mounted when they request a render", async () => { let renderRequests = 0; let rendererCalls = 0; let mountedNote: Text | undefined; @@ -226,6 +322,7 @@ describe("AssistantMessageComponent thinking renderers", () => { expect(Bun.stripANSI(component.render(120).join("\n"))).toContain("translation loading"); mountedNote?.setText("translation ready"); requestRender?.(); + await Promise.resolve(); const rendered = Bun.stripANSI(component.render(120).join("\n")); expect(renderRequests).toBe(1); @@ -234,6 +331,44 @@ describe("AssistantMessageComponent thinking renderers", () => { expect(rendered).not.toContain("translation loading"); }); + it("mounts a replacement when an async renderer becomes ready after streaming stops", async () => { + let ready = false; + let renderRequests = 0; + let rendererCalls = 0; + let requestRender: (() => void) | undefined; + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "Sensitive async thinking." }], + }, + false, + () => { + renderRequests += 1; + }, + [ + context => { + rendererCalls += 1; + requestRender = context.requestRender; + if (!ready) return undefined; + return { type: "replace", component: new Text(`redacted ${context.text.length}`, 1, 0) }; + }, + ], + ); + + expect(Bun.stripANSI(component.render(120).join("\n"))).toContain("Sensitive async thinking."); + expect(rendererCalls).toBe(1); + + ready = true; + requestRender?.(); + await Promise.resolve(); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(renderRequests).toBe(1); + expect(rendererCalls).toBe(2); + expect(rendered).toContain("redacted 25"); + expect(rendered).not.toContain("Sensitive async thinking."); + }); + it("does not invoke extension renderers when thinking is hidden", () => { let rendererCalled = false; const component = new AssistantMessageComponent( diff --git a/packages/coding-agent/test/modes/components/transcript-container.test.ts b/packages/coding-agent/test/modes/components/transcript-container.test.ts index 73f6bbaa894..c829194628c 100644 --- a/packages/coding-agent/test/modes/components/transcript-container.test.ts +++ b/packages/coding-agent/test/modes/components/transcript-container.test.ts @@ -125,6 +125,28 @@ class VersionedFinalizedBlock implements Component { } } +class NonCommittableFinalizedBlock implements Component { + #lines: string[]; + + constructor(lines: string[]) { + this.#lines = lines; + } + + isTranscriptBlockFinalized(): boolean { + return true; + } + + isTranscriptBlockAppendOnly(): boolean { + return false; + } + + invalidate(): void {} + + render(_width: number): string[] { + return [...this.#lines]; + } +} + beforeAll(() => { initTheme(); }); @@ -206,6 +228,143 @@ describe("TranscriptContainer", () => { expect(container.getNativeScrollbackLiveRegionStart()).toBeUndefined(); }); + it("keeps replacement-rendered assistant thinking in the live region", () => { + const defaultContainer = new TranscriptContainer(); + const defaultAssistant = new AssistantMessageComponent(); + defaultAssistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "Final append-only thinking." }], + }), + ); + defaultAssistant.markTranscriptBlockFinalized(); + defaultContainer.addChild(defaultAssistant); + defaultContainer.render(80); + expect(defaultContainer.getNativeScrollbackLiveRegionStart()).toBeUndefined(); + + const replacementContainer = new TranscriptContainer(); + const replacementAssistant = new AssistantMessageComponent(undefined, false, undefined, [ + ({ text }) => ({ type: "replace", component: new Text(`redacted${".".repeat(text.length)}`, 1, 0) }), + ]); + replacementAssistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "x" }], + }), + ); + replacementContainer.addChild(replacementAssistant); + replacementContainer.render(80); + replacementAssistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "xx" }], + }), + ); + + const rendered = plain(replacementContainer.render(80)); + expect(rendered).toContain("redacted.."); + expect(rendered).not.toContain("xx"); + expect(replacementContainer.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + + it("keeps pending replacement-capable assistant thinking in the live region", () => { + let ready = false; + const container = new TranscriptContainer(); + const assistant = new AssistantMessageComponent(undefined, false, undefined, [ + ({ text }) => (ready ? { type: "replace", component: new Text(`redacted ${text.length}`, 1, 0) } : undefined), + ]); + assistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "raw" }], + }), + ); + container.addChild(assistant); + container.render(80); + assistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "raw grows" }], + }), + ); + + expect(plain(container.render(80))).toContain("raw grows"); + expect(container.getNativeScrollbackLiveRegionStart()).toBe(0); + + ready = true; + assistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "raw grows" }], + }), + ); + + const rendered = plain(container.render(80)); + expect(rendered).toContain("redacted 9"); + expect(rendered).not.toContain("raw grows"); + expect(container.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + it("keeps finalized replacement-capable assistant thinking in the live region", () => { + const defaultContainer = new TranscriptContainer(); + const defaultAssistant = new AssistantMessageComponent(); + defaultAssistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "Final append-only thinking." }], + }), + ); + defaultAssistant.markTranscriptBlockFinalized(); + defaultContainer.addChild(defaultAssistant); + defaultContainer.render(80); + expect(defaultContainer.getNativeScrollbackLiveRegionStart()).toBeUndefined(); + + const replacementContainer = new TranscriptContainer(); + const replacementAssistant = new AssistantMessageComponent(undefined, false, undefined, [() => undefined]); + replacementAssistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "raw final thinking" }], + }), + ); + replacementAssistant.markTranscriptBlockFinalized(); + replacementContainer.addChild(replacementAssistant); + + const rendered = plain(replacementContainer.render(80)); + expect(rendered).toContain("raw final thinking"); + expect(replacementContainer.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + + it("keeps finalized replacement-capable assistant thinking in the live region above later finalized blocks", () => { + const container = new TranscriptContainer(); + const assistant = new AssistantMessageComponent(undefined, false, undefined, [() => undefined]); + assistant.updateContent( + makeAssistantMessage({ + content: [{ type: "thinking", thinking: "raw final thinking" }], + }), + ); + assistant.markTranscriptBlockFinalized(); + container.addChild(assistant); + container.addChild(new MutableBlock(["next finalized block"])); + + const rendered = plain(container.render(80)); + expect(rendered).toContain("raw final thinking"); + expect(rendered).toContain("next finalized block"); + expect(container.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + + it("treats blocks below finalized non-committable blocks as live", () => { + const container = new TranscriptContainer(); + const volatile = new NonCommittableFinalizedBlock(["volatile"]); + const card = new MutableBlock(["notification"]); + container.addChild(volatile); + container.addChild(card); + + expect(container.render(40)).toEqual(["volatile", "", "notification"]); + expect(container.isBlockInLiveRegion(card)).toBe(true); + expect(container.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + + it("pins the live region at empty finalized non-committable blocks", () => { + const container = new TranscriptContainer(); + container.addChild(new NonCommittableFinalizedBlock([])); + container.addChild(new MutableBlock(["later finalized block"])); + + expect(container.render(40)).toEqual(["later finalized block"]); + expect(container.getNativeScrollbackLiveRegionStart()).toBe(0); + }); + it("keeps an unfinalized block below the seam when a finalized block is appended below it", () => { const container = new TranscriptContainer(); // A foreground tool whose args are still streaming (no result yet).