diff --git a/packages/coding-agent/src/core/hooks/index.ts b/packages/coding-agent/src/core/hooks/index.ts index db222989..7524b605 100644 --- a/packages/coding-agent/src/core/hooks/index.ts +++ b/packages/coding-agent/src/core/hooks/index.ts @@ -5,6 +5,8 @@ export type { AgentEndEvent, AgentStartEvent, BashToolResultEvent, + CustomMessageRenderer, + CustomMessageRenderOptions, CustomToolResultEvent, EditToolResultEvent, ExecResult, diff --git a/packages/coding-agent/src/core/hooks/loader.ts b/packages/coding-agent/src/core/hooks/loader.ts index 18e09c19..5f4fef75 100644 --- a/packages/coding-agent/src/core/hooks/loader.ts +++ b/packages/coding-agent/src/core/hooks/loader.ts @@ -10,7 +10,7 @@ import { fileURLToPath } from "node:url"; import type { Attachment } from "@mariozechner/pi-agent-core"; import { createJiti } from "jiti"; import { getAgentDir } from "../../config.js"; -import type { HookAPI, HookFactory } from "./types.js"; +import type { CustomMessageRenderer, HookAPI, HookFactory } from "./types.js"; // Create require function to resolve module paths at runtime const require = createRequire(import.meta.url); @@ -61,6 +61,8 @@ export interface LoadedHook { resolvedPath: string; /** Map of event type to handler functions */ handlers: Map; + /** Map of customType to custom message renderer */ + customMessageRenderers: Map; /** Set the send handler for this hook's pi.send() */ setSendHandler: (handler: SendHandler) => void; } @@ -110,16 +112,18 @@ function resolveHookPath(hookPath: string, cwd: string): string { } /** - * Create a HookAPI instance that collects handlers. - * Returns the API and a function to set the send handler later. + * Create a HookAPI instance that collects handlers and renderers. + * Returns the API, renderers map, and a function to set the send handler later. */ function createHookAPI(handlers: Map): { api: HookAPI; + customMessageRenderers: Map; setSendHandler: (handler: SendHandler) => void; } { let sendHandler: SendHandler = () => { // Default no-op until mode sets the handler }; + const customMessageRenderers = new Map(); const api: HookAPI = { on(event: string, handler: HandlerFn): void { @@ -130,10 +134,14 @@ function createHookAPI(handlers: Map): { send(text: string, attachments?: Attachment[]): void { sendHandler(text, attachments); }, + renderCustomMessage(customType: string, renderer: CustomMessageRenderer): void { + customMessageRenderers.set(customType, renderer); + }, } as HookAPI; return { api, + customMessageRenderers, setSendHandler: (handler: SendHandler) => { sendHandler = handler; }, @@ -164,13 +172,13 @@ async function loadHook(hookPath: string, cwd: string): Promise<{ hook: LoadedHo // Create handlers map and API const handlers = new Map(); - const { api, setSendHandler } = createHookAPI(handlers); + const { api, customMessageRenderers, setSendHandler } = createHookAPI(handlers); // Call factory to register handlers factory(api); return { - hook: { path: hookPath, resolvedPath, handlers, setSendHandler }, + hook: { path: hookPath, resolvedPath, handlers, customMessageRenderers, setSendHandler }, error: null, }; } catch (err) { diff --git a/packages/coding-agent/src/core/hooks/runner.ts b/packages/coding-agent/src/core/hooks/runner.ts index e7646a19..2d4c9d2a 100644 --- a/packages/coding-agent/src/core/hooks/runner.ts +++ b/packages/coding-agent/src/core/hooks/runner.ts @@ -5,6 +5,7 @@ import { spawn } from "node:child_process"; import type { LoadedHook, SendHandler } from "./loader.js"; import type { + CustomMessageRenderer, ExecOptions, ExecResult, HookError, @@ -203,6 +204,20 @@ export class HookRunner { return false; } + /** + * Get a custom message renderer for the given customType. + * Returns the first renderer found across all hooks, or undefined if none. + */ + getCustomMessageRenderer(customType: string): CustomMessageRenderer | undefined { + for (const hook of this.hooks) { + const renderer = hook.customMessageRenderers.get(customType); + if (renderer) { + return renderer; + } + } + return undefined; + } + /** * Create the event context for handlers. */ diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index 8d621ae8..6b46a903 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -7,9 +7,11 @@ import type { AppMessage, Attachment } from "@mariozechner/pi-agent-core"; import type { ImageContent, Model, TextContent, ToolResultMessage } from "@mariozechner/pi-ai"; +import type { Component } from "@mariozechner/pi-tui"; +import type { Theme } from "../../modes/interactive/theme/theme.js"; import type { CompactionPreparation, CompactionResult } from "../compaction.js"; import type { ModelRegistry } from "../model-registry.js"; -import type { CompactionEntry, SessionManager } from "../session-manager.js"; +import type { CompactionEntry, CustomMessageEntry, SessionManager } from "../session-manager.js"; import type { BashToolDetails, FindToolDetails, @@ -368,6 +370,24 @@ export interface SessionEventResult { */ export type HookHandler = (event: E, ctx: HookEventContext) => Promise; +/** + * Options passed to custom message renderers. + */ +export interface CustomMessageRenderOptions { + /** Whether the view is expanded */ + expanded: boolean; +} + +/** + * Renderer for custom message entries. + * Hooks register these to provide custom TUI rendering for their CustomMessageEntry types. + */ +export type CustomMessageRenderer = ( + entry: CustomMessageEntry, + options: CustomMessageRenderOptions, + theme: Theme, +) => Component | null; + /** * HookAPI passed to hook factory functions. * Hooks use pi.on() to subscribe to events and pi.send() to inject messages. @@ -388,6 +408,13 @@ export interface HookAPI { * If the agent is idle, a new agent loop is started. */ send(text: string, attachments?: Attachment[]): void; + + /** + * Register a custom renderer for CustomMessageEntry with a specific customType. + * The renderer is called when rendering the entry in the TUI. + * Return null to use the default renderer. + */ + renderCustomMessage(customType: string, renderer: CustomMessageRenderer): void; } /** diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index e245f233..c78a6bd9 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -340,6 +340,7 @@ function createFactoryFromLoadedHook(loaded: LoadedHook): HookFactory { function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; factory: HookFactory }>): LoadedHook[] { return definitions.map((def) => { const handlers = new Map Promise>>(); + const customMessageRenderers = new Map(); let sendHandler: (text: string, attachments?: any[]) => void = () => {}; const api = { @@ -351,6 +352,9 @@ function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; fa send: (text: string, attachments?: any[]) => { sendHandler(text, attachments); }, + renderCustomMessage: (customType: string, renderer: any) => { + customMessageRenderers.set(customType, renderer); + }, }; def.factory(api as any); @@ -359,6 +363,7 @@ function createLoadedHooksFromDefinitions(definitions: Array<{ path?: string; fa path: def.path ?? "", resolvedPath: def.path ?? "", handlers, + customMessageRenderers, setSendHandler: (handler: (text: string, attachments?: any[]) => void) => { sendHandler = handler; }, diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index ed6fbdfb..8d2323d8 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -131,6 +131,8 @@ export interface SessionTreeNode { export interface SessionContext { messages: AppMessage[]; + /** Entries in the current path (root to leaf). Use to identify custom_message entries for rendering. */ + entries: SessionEntry[]; thinkingLevel: string; model: { provider: string; modelId: string } | null; } @@ -290,7 +292,7 @@ export function buildSessionContext( } if (!leaf) { - return { messages: [], thinkingLevel: "off", model: null }; + return { messages: [], entries: [], thinkingLevel: "off", model: null }; } // Walk from leaf to root, collecting path @@ -318,16 +320,18 @@ export function buildSessionContext( } } - // Build messages - handle compaction ordering correctly + // Build messages and collect corresponding entries // When there's a compaction, we need to: - // 1. Emit summary first + // 1. Emit summary first (entry = compaction) // 2. Emit kept messages (from firstKeptEntryId up to compaction) // 3. Emit messages after compaction const messages: AppMessage[] = []; + const contextEntries: SessionEntry[] = []; if (compaction) { // Emit summary first messages.push(createSummaryMessage(compaction.summary, compaction.timestamp)); + contextEntries.push(compaction); // Find compaction index in path const compactionIdx = path.findIndex((e) => e.type === "compaction" && e.id === compaction.id); @@ -342,8 +346,10 @@ export function buildSessionContext( if (foundFirstKept) { if (entry.type === "message") { messages.push(entry.message); + contextEntries.push(entry); } else if (entry.type === "custom_message") { messages.push(createCustomMessage(entry)); + contextEntries.push(entry); } } } @@ -353,10 +359,13 @@ export function buildSessionContext( const entry = path[i]; if (entry.type === "message") { messages.push(entry.message); + contextEntries.push(entry); } else if (entry.type === "custom_message") { messages.push(createCustomMessage(entry)); + contextEntries.push(entry); } else if (entry.type === "branch_summary") { messages.push(createSummaryMessage(entry.summary, entry.timestamp)); + contextEntries.push(entry); } } } else { @@ -364,15 +373,18 @@ export function buildSessionContext( for (const entry of path) { if (entry.type === "message") { messages.push(entry.message); + contextEntries.push(entry); } else if (entry.type === "custom_message") { messages.push(createCustomMessage(entry)); + contextEntries.push(entry); } else if (entry.type === "branch_summary") { messages.push(createSummaryMessage(entry.summary, entry.timestamp)); + contextEntries.push(entry); } } } - return { messages, thinkingLevel, model }; + return { messages, entries: contextEntries, thinkingLevel, model }; } /** diff --git a/packages/coding-agent/test/compaction-hooks.test.ts b/packages/coding-agent/test/compaction-hooks.test.ts index d056dc37..b24618e8 100644 --- a/packages/coding-agent/test/compaction-hooks.test.ts +++ b/packages/coding-agent/test/compaction-hooks.test.ts @@ -63,6 +63,7 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => { path: "test-hook", resolvedPath: "/test/test-hook.ts", handlers, + customMessageRenderers: new Map(), setSendHandler: () => {}, }; } @@ -238,6 +239,7 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => { ], ], ]), + customMessageRenderers: new Map(), setSendHandler: () => {}, }; @@ -281,6 +283,7 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => { ], ], ]), + customMessageRenderers: new Map(), setSendHandler: () => {}, }; @@ -303,6 +306,7 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => { ], ], ]), + customMessageRenderers: new Map(), setSendHandler: () => {}, };