feat(coding-agent): allow thinking renderers to replace blocks

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
 

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

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<T = unknown> = (
 ) => 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<Pick<AssistantMessage, "timestamp" | "responseId" | "api" | "provider" | "model">>;
+	/**
+	 * 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<Pick<ThinkingContent, "itemId" | "thinkingSignature">>;
 	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<T = unknown>(customType: string, renderer: MessageRenderer<T>): 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;
 
 	// =========================================================================

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));