From 569cc3442bf82cb3fc421f969a571e8e8f877b7b Mon Sep 17 00:00:00 2001 From: shoucandanghehe Date: Sun, 31 May 2026 15:52:26 +0800 Subject: [PATCH 1/2] feat(coding-agent): add assistant thinking renderers --- docs/extensions.md | 16 ++++++- .../examples/extensions/README.md | 1 + .../examples/extensions/thinking-note.ts | 13 +++++ .../src/extensibility/extensions/loader.ts | 34 +++++++------ .../src/extensibility/extensions/runner.ts | 39 ++++++++------- .../src/extensibility/extensions/types.ts | 27 +++++++++-- .../src/modes/components/assistant-message.ts | 48 ++++++++++++++++++- .../src/modes/controllers/event-controller.ts | 7 ++- .../src/modes/utils/ui-helpers.ts | 7 ++- .../test/extensions-runner.test.ts | 22 ++++++++- .../assistant-message-mermaid.test.ts | 45 +++++++++++++++-- .../sdk-credential-disabled-bridge.test.ts | 1 + 12 files changed, 212 insertions(+), 48 deletions(-) create mode 100644 packages/coding-agent/examples/extensions/thinking-note.ts diff --git a/docs/extensions.md b/docs/extensions.md index 83cdcfc0d..abb22ebf6 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -112,7 +112,7 @@ Core methods: - `on(event, handler)` - `registerTool`, `registerCommand`, `registerShortcut`, `registerFlag` -- `registerMessageRenderer` +- `registerMessageRenderer`, `registerAssistantThinkingRenderer` - `setLabel`, `getFlag` - `sendMessage`, `sendUserMessage`, `appendEntry`, `exec` - `getActiveTools`, `getAllTools`, `setActiveTools` @@ -359,6 +359,20 @@ pi.registerMessageRenderer("my-type", (message, { expanded }, theme) => { Used by interactive rendering when custom messages are displayed. +## Assistant thinking renderer + +```ts +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; +}); +``` + +Used by interactive rendering to add display-only supplemental UI below each visible assistant thinking block. The renderer receives the original thinking text plus message/content indexes and a `requestRender()` callback for async renderers. It must not mutate messages; the original thinking block remains the provider/session source of truth. + ## Tool call/result renderer Provide `renderCall` / `renderResult` on `registerTool` definitions for custom tool visualization in TUI. diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index f0624d9e0..2cd0e7e2e 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -41,6 +41,7 @@ cp permission-gate.ts ~/.omp/agent/extensions/ | `handoff.ts` | Transfer context to a new focused session via `/handoff ` | | `qna.ts` | Extracts questions from last response into editor via `ctx.ui.setEditorText()` | | `status-line.ts` | Shows turn progress in footer via `ctx.ui.setStatus()` with themed colors | +| `thinking-note.ts` | Adds display-only supplemental UI below assistant thinking blocks | | `snake.ts` | Snake game with custom UI, keyboard handling, and session persistence | ### Git Integration diff --git a/packages/coding-agent/examples/extensions/thinking-note.ts b/packages/coding-agent/examples/extensions/thinking-note.ts new file mode 100644 index 000000000..5871af7f6 --- /dev/null +++ b/packages/coding-agent/examples/extensions/thinking-note.ts @@ -0,0 +1,13 @@ +import type { ExtensionFactory } from "@oh-my-pi/pi-coding-agent"; +import { Container, Text } from "@oh-my-pi/pi-tui"; + +const extension: ExtensionFactory = pi => { + pi.setLabel("Thinking note"); + pi.registerAssistantThinkingRenderer((context, theme) => { + const container = new Container(); + container.addChild(new Text(theme.fg("dim", `thinking chars: ${context.text.length}`), 1, 0)); + return container; + }); +}; + +export default extension; diff --git a/packages/coding-agent/src/extensibility/extensions/loader.ts b/packages/coding-agent/src/extensibility/extensions/loader.ts index 4bb767eb7..40f1c4da9 100644 --- a/packages/coding-agent/src/extensibility/extensions/loader.ts +++ b/packages/coding-agent/src/extensibility/extensions/loader.ts @@ -5,7 +5,8 @@ import type * as fs1 from "node:fs"; import * as fs from "node:fs/promises"; import * as path from "node:path"; import type { ThinkingLevel } from "@oh-my-pi/pi-agent-core"; -import type { ImageContent, Model, TextContent } from "@oh-my-pi/pi-ai"; +import type { ImageContent, Model, TextContent, TSchema } from "@oh-my-pi/pi-ai"; +import * as PiCodingAgent from "@oh-my-pi/pi-coding-agent"; import type { KeyId } from "@oh-my-pi/pi-tui"; import { hasFsCode, isEacces, isEnoent, logger } from "@oh-my-pi/pi-utils"; import * as Zod from "zod/v4"; @@ -22,6 +23,7 @@ import * as TypeBox from "../typebox"; import { resolvePath } from "../utils"; import type { + AssistantThinkingRenderer, Extension, ExtensionAPI, ExtensionContext, @@ -29,6 +31,7 @@ import type { ExtensionRuntime as IExtensionRuntime, LoadExtensionsResult, MessageRenderer, + ProviderConfig, RegisteredCommand, ToolDefinition, } from "./types"; @@ -55,8 +58,7 @@ export class ExtensionRuntimeNotInitializedError extends Error { */ export class ExtensionRuntime implements IExtensionRuntime { flagValues = new Map(); - pendingProviderRegistrations: Array<{ name: string; config: import("./types").ProviderConfig; sourceId: string }> = - []; + pendingProviderRegistrations: Array<{ name: string; config: ProviderConfig; sourceId: string }> = []; sendMessage(): void { throw new ExtensionRuntimeNotInitializedError(); @@ -123,12 +125,12 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { readonly flagValues = new Map(); readonly pendingProviderRegistrations: Array<{ name: string; - config: import("./types").ProviderConfig; + config: ProviderConfig; sourceId: string; }> = []; constructor( - public readonly pi: typeof import("@oh-my-pi/pi-coding-agent"), + public readonly pi: typeof PiCodingAgent, private readonly extension: Extension, private readonly runtime: IExtensionRuntime, private readonly cwd: string, @@ -141,10 +143,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { this.extension.handlers.set(event, list); } - registerTool< - TParams extends import("@oh-my-pi/pi-ai").TSchema = import("@oh-my-pi/pi-ai").TSchema, - TDetails = unknown, - >(tool: ToolDefinition): void { + registerTool(tool: ToolDefinition): void { this.extension.tools.set(tool.name, { definition: tool, extensionPath: this.extension.path, @@ -190,6 +189,10 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { this.extension.messageRenderers.set(customType, renderer as MessageRenderer); } + registerAssistantThinkingRenderer(renderer: AssistantThinkingRenderer): void { + this.extension.assistantThinkingRenderers.push(renderer); + } + getFlag(name: string): boolean | string | undefined { if (!this.extension.flags.has(name)) return undefined; return this.runtime.flagValues.get(name); @@ -253,7 +256,7 @@ class ConcreteExtensionAPI implements ExtensionAPI, IExtensionRuntime { return this.runtime.setSessionName(name); } - registerProvider(name: string, config: import("./types").ProviderConfig): void { + registerProvider(name: string, config: ProviderConfig): void { this.runtime.pendingProviderRegistrations.push({ name, config, sourceId: this.extension.path }); } } @@ -267,6 +270,7 @@ function createExtension(extensionPath: string, resolvedPath: string): Extension resolvedPath, handlers: new Map(), tools: new Map(), + assistantThinkingRenderers: [], messageRenderers: new Map(), commands: new Map(), flags: new Map(), @@ -293,13 +297,7 @@ async function loadExtension( } const extension = createExtension(extensionPath, resolvedPath); - const api = new ConcreteExtensionAPI( - await import("@oh-my-pi/pi-coding-agent"), - extension, - runtime, - cwd, - eventBus, - ); + const api = new ConcreteExtensionAPI(PiCodingAgent, extension, runtime, cwd, eventBus); await factory(api); return { extension, error: null }; @@ -320,7 +318,7 @@ export async function loadExtensionFromFactory( name = "", ): Promise { const extension = createExtension(name, name); - const api = new ConcreteExtensionAPI(await import("@oh-my-pi/pi-coding-agent"), extension, runtime, cwd, eventBus); + const api = new ConcreteExtensionAPI(PiCodingAgent, extension, runtime, cwd, eventBus); await factory(api); return extension; } diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index fbd785914..3b23d3c73 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -10,6 +10,7 @@ import { type Theme, theme } from "../../modes/theme/theme"; import type { SessionManager } from "../../session/session-manager"; import type { AfterProviderResponseEvent, + AssistantThinkingRenderer, BeforeAgentStartEvent, BeforeAgentStartEventResult, BeforeProviderRequestEvent, @@ -343,22 +344,22 @@ export class ExtensionRunner { this.runtime.flagValues.set(name, value); } - static readonly #RESERVED_SHORTCUTS = new Set([ - "ctrl+c", - "ctrl+d", - "ctrl+z", - "ctrl+k", - "ctrl+p", - "ctrl+l", - "ctrl+o", - "ctrl+t", - "ctrl+g", - "shift+tab", - "shift+ctrl+p", - "alt+enter", - "escape", - "enter", - ]); + static readonly #RESERVED_SHORTCUTS: Record = { + "ctrl+c": true, + "ctrl+d": true, + "ctrl+z": true, + "ctrl+k": true, + "ctrl+p": true, + "ctrl+l": true, + "ctrl+o": true, + "ctrl+t": true, + "ctrl+g": true, + "shift+tab": true, + "shift+ctrl+p": true, + "alt+enter": true, + escape: true, + enter: true, + }; getShortcuts(): Map { const allShortcuts = new Map(); @@ -366,7 +367,7 @@ export class ExtensionRunner { for (const [key, shortcut] of ext.shortcuts) { const normalizedKey = key.toLowerCase() as KeyId; - if (ExtensionRunner.#RESERVED_SHORTCUTS.has(normalizedKey)) { + if (ExtensionRunner.#RESERVED_SHORTCUTS[normalizedKey]) { logger.warn("Extension shortcut conflicts with built-in shortcut", { key, extensionPath: shortcut.extensionPath, @@ -419,6 +420,10 @@ export class ExtensionRunner { return undefined; } + getAssistantThinkingRenderers(): AssistantThinkingRenderer[] { + return this.extensions.flatMap(ext => ext.assistantThinkingRenderers); + } + getRegisteredCommands(reserved?: Set): RegisteredCommand[] { this.#commandDiagnostics = []; diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 332427184..77607e8a7 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -11,6 +11,7 @@ import type { AgentMessage, AgentToolResult, AgentToolUpdateCallback, ThinkingLe import type { CompactionResult } from "@oh-my-pi/pi-agent-core/compaction"; import type { Api, + AssistantMessage, AssistantMessageEvent, AssistantMessageEventStream, Context, @@ -25,6 +26,8 @@ import type { import type { OAuthCredentials, OAuthLoginCallbacks } from "@oh-my-pi/pi-ai/utils/oauth/types"; import type * as piCodingAgent from "@oh-my-pi/pi-coding-agent"; import type { AutocompleteItem, Component, EditorTheme, KeyId, TUI } from "@oh-my-pi/pi-tui"; +import type { logger as PiLogger } from "@oh-my-pi/pi-utils"; +import type * as Zod from "zod/v4"; import type { KeybindingsManager } from "../../config/keybindings"; import type { ModelRegistry } from "../../config/model-registry"; import type { EditToolDetails } from "../../edit"; @@ -81,6 +84,7 @@ import type { TurnStartEvent, } from "../shared-events"; import type { SlashCommandInfo } from "../slash-commands"; +import type * as TypeBox from "../typebox"; export type { AppKeybinding, KeybindingsManager } from "../../config/keybindings"; export type { ExecOptions, ExecResult } from "../../exec/exec"; @@ -799,6 +803,19 @@ export type MessageRenderer = ( theme: Theme, ) => Component | undefined; +export interface AssistantThinkingRenderContext { + message: AssistantMessage; + contentIndex: number; + thinkingIndex: number; + text: string; + requestRender(): void; +} + +export type AssistantThinkingRenderer = ( + context: AssistantThinkingRenderContext, + theme: Theme, +) => Component | undefined; + // ============================================================================ // Command Registration // ============================================================================ @@ -830,13 +847,13 @@ export interface ExtensionAPI { // ========================================================================= /** File logger for error/warning/debug messages */ - logger: typeof import("@oh-my-pi/pi-utils").logger; + logger: typeof PiLogger; /** Injected zod-backed typebox shim for legacy `Type.Object(...)` parameter authoring. */ - typebox: typeof import("../typebox"); + typebox: typeof TypeBox; /** Injected zod module for Zod-authored extension tools (canonical going forward). */ - zod: typeof import("zod/v4"); + zod: typeof Zod; /** Injected pi-coding-agent exports for accessing SDK utilities */ pi: typeof piCodingAgent; @@ -950,6 +967,9 @@ 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. */ + registerAssistantThinkingRenderer(renderer: AssistantThinkingRenderer): void; + // ========================================================================= // Actions // ========================================================================= @@ -1232,6 +1252,7 @@ export interface Extension { label?: string; handlers: Map; tools: Map>; + assistantThinkingRenderers: AssistantThinkingRenderer[]; messageRenderers: Map; commands: Map; flags: Map; diff --git a/packages/coding-agent/src/modes/components/assistant-message.ts b/packages/coding-agent/src/modes/components/assistant-message.ts index f64bc6a19..d5338685b 100644 --- a/packages/coding-agent/src/modes/components/assistant-message.ts +++ b/packages/coding-agent/src/modes/components/assistant-message.ts @@ -1,7 +1,9 @@ import type { AssistantMessage, ImageContent, Usage } from "@oh-my-pi/pi-ai"; +import type { Component } from "@oh-my-pi/pi-tui"; import { Container, Image, ImageProtocol, Markdown, Spacer, TERMINAL, Text } from "@oh-my-pi/pi-tui"; import { formatNumber } from "@oh-my-pi/pi-utils"; import { settings } from "../../config/settings"; +import type { AssistantThinkingRenderer } from "../../extensibility/extensions/types"; import { getMarkdownTheme, theme } from "../../modes/theme/theme"; import { isSilentAbort } from "../../session/messages"; import { resolveImageOptions } from "../../tools/render-utils"; @@ -21,6 +23,7 @@ export class AssistantMessageComponent extends Container { message?: AssistantMessage, private hideThinkingBlock = false, private readonly onImageUpdate?: () => void, + private readonly thinkingRenderers: readonly AssistantThinkingRenderer[] = [], ) { super(); @@ -131,6 +134,37 @@ export class AssistantMessageComponent extends Container { } } + #renderThinkingExtensions( + message: AssistantMessage, + contentIndex: number, + thinkingIndex: number, + text: string, + ): Component | undefined { + for (const renderer of this.thinkingRenderers) { + try { + const component = renderer( + { + message, + contentIndex, + thinkingIndex, + text, + requestRender: () => { + if (this.#lastMessage) { + this.updateContent(this.#lastMessage); + } + this.onImageUpdate?.(); + }, + }, + theme, + ); + if (component) return component; + } catch { + // Ignore extension renderer failures and keep the original thinking block visible. + } + } + return undefined; + } + updateContent(message: AssistantMessage): void { this.#lastMessage = message; @@ -146,6 +180,7 @@ export class AssistantMessageComponent extends Container { } // Render content in order + let thinkingIndex = 0; for (let i = 0; i < message.content.length; i++) { const content = message.content[i]; if (content.type === "text" && content.text.trim()) { @@ -166,13 +201,24 @@ export class AssistantMessageComponent extends Container { this.#contentContainer.addChild(new Spacer(1)); } } else { + const thinkingText = content.thinking.trim(); // Thinking traces in thinkingText color, italic this.#contentContainer.addChild( - new Markdown(content.thinking.trim(), 1, 0, getMarkdownTheme(), { + new Markdown(thinkingText, 1, 0, getMarkdownTheme(), { color: (text: string) => theme.fg("thinkingText", text), italic: true, }), ); + const renderedThinkingExtension = this.#renderThinkingExtensions( + message, + i, + thinkingIndex, + thinkingText, + ); + thinkingIndex += 1; + if (renderedThinkingExtension) { + this.#contentContainer.addChild(renderedThinkingExtension); + } if (hasVisibleContentAfter) { this.#contentContainer.addChild(new Spacer(1)); } diff --git a/packages/coding-agent/src/modes/controllers/event-controller.ts b/packages/coding-agent/src/modes/controllers/event-controller.ts index 999d36df4..07a031338 100644 --- a/packages/coding-agent/src/modes/controllers/event-controller.ts +++ b/packages/coding-agent/src/modes/controllers/event-controller.ts @@ -269,8 +269,11 @@ export class EventController { } else if (event.message.role === "assistant") { this.#lastThinkingCount = 0; this.#resetReadGroup(); - this.ctx.streamingComponent = new AssistantMessageComponent(undefined, this.ctx.hideThinkingBlock, () => - this.ctx.ui.requestRender(), + this.ctx.streamingComponent = new AssistantMessageComponent( + undefined, + this.ctx.hideThinkingBlock, + () => this.ctx.ui.requestRender(), + this.ctx.session.extensionRunner?.getAssistantThinkingRenderers(), ); this.ctx.streamingMessage = event.message; this.ctx.chatContainer.addChild(this.ctx.streamingComponent); diff --git a/packages/coding-agent/src/modes/utils/ui-helpers.ts b/packages/coding-agent/src/modes/utils/ui-helpers.ts index 7dd7b47d7..de472c4b2 100644 --- a/packages/coding-agent/src/modes/utils/ui-helpers.ts +++ b/packages/coding-agent/src/modes/utils/ui-helpers.ts @@ -262,8 +262,11 @@ export class UiHelpers { break; } case "assistant": { - const assistantComponent = new AssistantMessageComponent(message, this.ctx.hideThinkingBlock, () => - this.ctx.ui.requestRender(), + const assistantComponent = new AssistantMessageComponent( + message, + this.ctx.hideThinkingBlock, + () => this.ctx.ui.requestRender(), + this.ctx.session.extensionRunner?.getAssistantThinkingRenderers(), ); this.ctx.chatContainer.addChild(assistantComponent); break; diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 5b2749a60..a366e504b 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -305,6 +305,26 @@ describe("ExtensionRunner", () => { const missing = runner.getMessageRenderer("not-exists"); expect(missing).toBeUndefined(); }); + + it("collects assistant thinking renderers", async () => { + const extCode = ` + export default function(pi) { + pi.registerAssistantThinkingRenderer((context, theme) => null); + } + `; + fs.writeFileSync(path.join(extensionsDir, "thinking-renderer.ts"), extCode); + + const result = await loadTestExtensions(); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + + expect(runner.getAssistantThinkingRenderers().length).toBe(1); + }); }); describe("flags", () => { @@ -621,7 +641,7 @@ describe("ExtensionRunner", () => { ` export default function(pi) { pi.on("session_start", async () => { - await new Promise(() => {}); + await Promise.withResolvers().promise; }); } `, 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 a099e69a0..6b62ef97f 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 @@ -2,10 +2,11 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "bun:test import * as path from "node:path"; import type { AssistantMessage } from "@oh-my-pi/pi-ai"; import { resetSettingsForTest, Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; +import type { AssistantThinkingRenderer } from "@oh-my-pi/pi-coding-agent/extensibility/extensions"; import { AssistantMessageComponent } from "@oh-my-pi/pi-coding-agent/modes/components/assistant-message"; import { clearMermaidCache } from "@oh-my-pi/pi-coding-agent/modes/theme/mermaid-cache"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; -import { ImageProtocol, setTerminalImageProtocol, TERMINAL } from "@oh-my-pi/pi-tui"; +import { ImageProtocol, setTerminalImageProtocol, TERMINAL, Text } from "@oh-my-pi/pi-tui"; const originalImageProtocol = TERMINAL.imageProtocol; @@ -29,8 +30,8 @@ function createAssistantMessage(markdown: string): AssistantMessage { }; } -function renderAssistantMessage(markdown: string): string { - const component = new AssistantMessageComponent(createAssistantMessage(markdown)); +function renderAssistantMessage(markdown: string, renderers: readonly AssistantThinkingRenderer[] = []): string { + const component = new AssistantMessageComponent(createAssistantMessage(markdown), false, undefined, renderers); return Bun.stripANSI(component.render(120).join("\n")) .split("\n") .map(line => line.trimEnd()) @@ -74,6 +75,44 @@ describe("AssistantMessageComponent mermaid markdown", () => { }); }); +describe("AssistantMessageComponent thinking renderers", () => { + it("renders extension output below visible thinking blocks", () => { + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "I should inspect the input." }], + }, + false, + undefined, + [() => new Text("translated note", 1, 0)], + ); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).toContain("I should inspect the input."); + expect(rendered).toContain("translated note"); + }); + + it("keeps original thinking visible when an extension renderer throws", () => { + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "I should inspect the input." }], + }, + false, + undefined, + [ + () => { + throw new Error("renderer failed"); + }, + ], + ); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).toContain("I should inspect the input."); + expect(rendered).not.toContain("renderer failed"); + }); +}); + describe("AssistantMessageComponent tool images", () => { it("converts WebP tool images for Kitty terminal rendering", async () => { const webpBase64 = Buffer.from( diff --git a/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts b/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts index ade13ea49..0001a780a 100644 --- a/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts +++ b/packages/coding-agent/test/sdk-credential-disabled-bridge.test.ts @@ -493,6 +493,7 @@ describe("createAgentSession credential_disabled subscription", () => { ], ]), tools: new Map(), + assistantThinkingRenderers: [], messageRenderers: new Map(), commands: new Map(), flags: new Map(), From fbb24fb7dcbb3c28283fb1fea27c3f1f6dd5939e Mon Sep 17 00:00:00 2001 From: shoucandanghehe Date: Sun, 31 May 2026 17:18:29 +0800 Subject: [PATCH 2/2] fix(coding-agent): tighten thinking renderer semantics --- docs/extensions.md | 2 +- packages/coding-agent/CHANGELOG.md | 3 + .../src/extensibility/extensions/types.ts | 2 - .../src/modes/components/assistant-message.ts | 31 ++------ .../assistant-message-mermaid.test.ts | 79 ++++++++++++++++++- 5 files changed, 86 insertions(+), 31 deletions(-) diff --git a/docs/extensions.md b/docs/extensions.md index abb22ebf6..119d0f2cb 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -371,7 +371,7 @@ pi.registerAssistantThinkingRenderer((context, theme) => { }); ``` -Used by interactive rendering to add display-only supplemental UI below each visible assistant thinking block. The renderer receives the original thinking text plus message/content indexes and a `requestRender()` callback for async renderers. It must not mutate messages; the original thinking block remains the provider/session source of truth. +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. ## Tool call/result renderer diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 52a951793..d58846b23 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -1,6 +1,9 @@ # Changelog ## [Unreleased] +### Added + +- Added an extension API for rendering supplemental UI below visible assistant thinking blocks. ## [15.7.3] - 2026-05-31 ### Added diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 77607e8a7..ce8f28147 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -11,7 +11,6 @@ import type { AgentMessage, AgentToolResult, AgentToolUpdateCallback, ThinkingLe import type { CompactionResult } from "@oh-my-pi/pi-agent-core/compaction"; import type { Api, - AssistantMessage, AssistantMessageEvent, AssistantMessageEventStream, Context, @@ -804,7 +803,6 @@ export type MessageRenderer = ( ) => Component | undefined; export interface AssistantThinkingRenderContext { - message: AssistantMessage; contentIndex: number; thinkingIndex: number; text: string; diff --git a/packages/coding-agent/src/modes/components/assistant-message.ts b/packages/coding-agent/src/modes/components/assistant-message.ts index d5338685b..1788e6099 100644 --- a/packages/coding-agent/src/modes/components/assistant-message.ts +++ b/packages/coding-agent/src/modes/components/assistant-message.ts @@ -1,5 +1,4 @@ import type { AssistantMessage, ImageContent, Usage } from "@oh-my-pi/pi-ai"; -import type { Component } from "@oh-my-pi/pi-tui"; import { Container, Image, ImageProtocol, Markdown, Spacer, TERMINAL, Text } from "@oh-my-pi/pi-tui"; import { formatNumber } from "@oh-my-pi/pi-utils"; import { settings } from "../../config/settings"; @@ -134,35 +133,25 @@ export class AssistantMessageComponent extends Container { } } - #renderThinkingExtensions( - message: AssistantMessage, - contentIndex: number, - thinkingIndex: number, - text: string, - ): Component | undefined { + #appendThinkingExtensions(contentIndex: number, thinkingIndex: number, text: string): void { for (const renderer of this.thinkingRenderers) { try { const component = renderer( { - message, contentIndex, thinkingIndex, text, - requestRender: () => { - if (this.#lastMessage) { - this.updateContent(this.#lastMessage); - } - this.onImageUpdate?.(); - }, + requestRender: () => this.onImageUpdate?.(), }, theme, ); - if (component) return component; + if (component) { + this.#contentContainer.addChild(component); + } } catch { // Ignore extension renderer failures and keep the original thinking block visible. } } - return undefined; } updateContent(message: AssistantMessage): void { @@ -209,16 +198,8 @@ export class AssistantMessageComponent extends Container { italic: true, }), ); - const renderedThinkingExtension = this.#renderThinkingExtensions( - message, - i, - thinkingIndex, - thinkingText, - ); + this.#appendThinkingExtensions(i, thinkingIndex, thinkingText); thinkingIndex += 1; - if (renderedThinkingExtension) { - this.#contentContainer.addChild(renderedThinkingExtension); - } if (hasVisibleContentAfter) { this.#contentContainer.addChild(new Spacer(1)); } 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 6b62ef97f..8994face5 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 @@ -76,7 +76,8 @@ describe("AssistantMessageComponent mermaid markdown", () => { }); describe("AssistantMessageComponent thinking renderers", () => { - it("renders extension output below visible thinking blocks", () => { + it("renders all extension outputs below visible thinking blocks in registration order", () => { + const contexts: Array<{ contentIndex: number; thinkingIndex: number; text: string }> = []; const component = new AssistantMessageComponent( { ...createAssistantMessage(""), @@ -84,12 +85,24 @@ describe("AssistantMessageComponent thinking renderers", () => { }, false, undefined, - [() => new Text("translated note", 1, 0)], + [ + 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), + ], ); const rendered = Bun.stripANSI(component.render(120).join("\n")); expect(rendered).toContain("I should inspect the input."); - expect(rendered).toContain("translated note"); + 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." }]); }); it("keeps original thinking visible when an extension renderer throws", () => { @@ -111,6 +124,66 @@ describe("AssistantMessageComponent thinking renderers", () => { expect(rendered).toContain("I should inspect the input."); expect(rendered).not.toContain("renderer failed"); }); + + it("keeps async renderer components mounted when they request a render", () => { + let renderRequests = 0; + let rendererCalls = 0; + let mountedNote: Text | undefined; + let requestRender: (() => void) | undefined; + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "I should inspect the input." }], + }, + false, + () => { + renderRequests += 1; + }, + [ + context => { + rendererCalls += 1; + requestRender = context.requestRender; + const note = new Text("translation loading", 1, 0); + mountedNote ??= note; + return note; + }, + ], + ); + + expect(Bun.stripANSI(component.render(120).join("\n"))).toContain("translation loading"); + mountedNote?.setText("translation ready"); + requestRender?.(); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(renderRequests).toBe(1); + expect(rendererCalls).toBe(1); + expect(rendered).toContain("translation ready"); + expect(rendered).not.toContain("translation loading"); + }); + + it("does not invoke extension renderers when thinking is hidden", () => { + let rendererCalled = false; + const component = new AssistantMessageComponent( + { + ...createAssistantMessage(""), + content: [{ type: "thinking", thinking: "I should inspect the input." }], + }, + true, + undefined, + [ + () => { + rendererCalled = true; + return new Text("hidden note", 1, 0); + }, + ], + ); + + const rendered = Bun.stripANSI(component.render(120).join("\n")); + expect(rendered).toContain("Thinking..."); + expect(rendered).not.toContain("I should inspect the input."); + expect(rendered).not.toContain("hidden note"); + expect(rendererCalled).toBe(false); + }); }); describe("AssistantMessageComponent tool images", () => {