Add CustomMessageEntry rendering infrastructure

- Add renderCustomMessage to HookAPI for registering custom renderers
- Add CustomMessageRenderer type and CustomMessageRenderOptions
- Store customMessageRenderers in LoadedHook
- Add getCustomMessageRenderer(customType) to HookRunner
- SessionContext.entries now aligned with messages (same length, corresponding indices)

TUI can now correlate messages with their source entries to identify
custom_message entries and use hook-provided renderers.
This commit is contained in:
Mario Zechner 2025-12-26 23:02:53 +01:00
parent 3ecaaa5937
commit 11a7845ceb
7 changed files with 83 additions and 10 deletions

View file

@ -5,6 +5,8 @@ export type {
AgentEndEvent,
AgentStartEvent,
BashToolResultEvent,
CustomMessageRenderer,
CustomMessageRenderOptions,
CustomToolResultEvent,
EditToolResultEvent,
ExecResult,

View file

@ -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<string, HandlerFn[]>;
/** Map of customType to custom message renderer */
customMessageRenderers: Map<string, CustomMessageRenderer>;
/** 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<string, HandlerFn[]>): {
api: HookAPI;
customMessageRenderers: Map<string, CustomMessageRenderer>;
setSendHandler: (handler: SendHandler) => void;
} {
let sendHandler: SendHandler = () => {
// Default no-op until mode sets the handler
};
const customMessageRenderers = new Map<string, CustomMessageRenderer>();
const api: HookAPI = {
on(event: string, handler: HandlerFn): void {
@ -130,10 +134,14 @@ function createHookAPI(handlers: Map<string, HandlerFn[]>): {
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<string, HandlerFn[]>();
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) {

View file

@ -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.
*/

View file

@ -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<E, R = void> = (event: E, ctx: HookEventContext) => Promise<R>;
/**
* 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<T = unknown> = (
entry: CustomMessageEntry<T>,
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<T = unknown>(customType: string, renderer: CustomMessageRenderer<T>): void;
}
/**

View file

@ -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<string, Array<(...args: unknown[]) => Promise<unknown>>>();
const customMessageRenderers = new Map<string, any>();
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 ?? "<inline>",
resolvedPath: def.path ?? "<inline>",
handlers,
customMessageRenderers,
setSendHandler: (handler: (text: string, attachments?: any[]) => void) => {
sendHandler = handler;
},

View file

@ -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 };
}
/**