From 568150f18bfa2d472484eb6d15b1fe8afc999d8c Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 31 Dec 2025 12:05:24 +0100 Subject: [PATCH] Rework custom tools API with CustomToolContext - CustomAgentTool renamed to CustomTool - ToolAPI renamed to CustomToolAPI - ToolContext renamed to CustomToolContext - ToolSessionEvent renamed to CustomToolSessionEvent - Added CustomToolContext parameter to execute() and onSession() - CustomToolFactory now returns CustomTool for type compatibility - dispose() replaced with onSession({ reason: 'shutdown' }) - Added wrapCustomTool() to convert CustomTool to AgentTool - Session exposes setToolUIContext() instead of leaking internals - Fix ToolExecutionComponent to sync with toolOutputExpanded state - Update all custom tool examples for new API --- .../examples/custom-tools/hello/index.ts | 5 +- .../examples/custom-tools/question/index.ts | 4 +- .../examples/custom-tools/subagent/index.ts | 10 +- .../examples/custom-tools/todo/index.ts | 14 +- .../examples/hooks/custom-compaction.ts | 3 +- .../coding-agent/examples/sdk/05-tools.ts | 4 +- .../examples/sdk/12-full-control.ts | 7 +- .../coding-agent/src/core/agent-session.ts | 33 +++-- .../core/compaction/branch-summarization.ts | 4 +- .../src/core/custom-tools/index.ts | 12 +- .../src/core/custom-tools/loader.ts | 8 +- .../src/core/custom-tools/types.ts | 130 ++++++++++++------ .../src/core/custom-tools/wrapper.ts | 28 ++++ .../coding-agent/src/core/hooks/runner.ts | 38 +---- packages/coding-agent/src/core/hooks/types.ts | 22 +-- packages/coding-agent/src/core/index.ts | 6 +- packages/coding-agent/src/core/sdk.ts | 40 ++++-- .../coding-agent/src/core/session-manager.ts | 32 +++-- packages/coding-agent/src/index.ts | 9 +- .../interactive/components/tool-execution.ts | 6 +- .../src/modes/interactive/interactive-mode.ts | 92 ++++++------- packages/coding-agent/src/modes/print-mode.ts | 40 +++--- .../coding-agent/src/modes/rpc/rpc-mode.ts | 43 +++--- .../test/compaction-hooks-example.test.ts | 14 +- .../test/compaction-hooks.test.ts | 11 +- .../test/session-manager/save-entry.test.ts | 2 +- .../session-manager/tree-traversal.test.ts | 8 +- 27 files changed, 336 insertions(+), 289 deletions(-) create mode 100644 packages/coding-agent/src/core/custom-tools/wrapper.ts diff --git a/packages/coding-agent/examples/custom-tools/hello/index.ts b/packages/coding-agent/examples/custom-tools/hello/index.ts index f9057fad..c2bf07b8 100644 --- a/packages/coding-agent/examples/custom-tools/hello/index.ts +++ b/packages/coding-agent/examples/custom-tools/hello/index.ts @@ -10,9 +10,10 @@ const factory: CustomToolFactory = (_pi) => ({ }), async execute(_toolCallId, params) { + const { name } = params as { name: string }; return { - content: [{ type: "text", text: `Hello, ${params.name}!` }], - details: { greeted: params.name }, + content: [{ type: "text", text: `Hello, ${name}!` }], + details: { greeted: name }, }; }, }); diff --git a/packages/coding-agent/examples/custom-tools/question/index.ts b/packages/coding-agent/examples/custom-tools/question/index.ts index 76c068ca..6949efcc 100644 --- a/packages/coding-agent/examples/custom-tools/question/index.ts +++ b/packages/coding-agent/examples/custom-tools/question/index.ts @@ -2,7 +2,7 @@ * Question Tool - Let the LLM ask the user a question with options */ -import type { CustomAgentTool, CustomToolFactory } from "@mariozechner/pi-coding-agent"; +import type { CustomTool, CustomToolFactory } from "@mariozechner/pi-coding-agent"; import { Text } from "@mariozechner/pi-tui"; import { Type } from "@sinclair/typebox"; @@ -18,7 +18,7 @@ const QuestionParams = Type.Object({ }); const factory: CustomToolFactory = (pi) => { - const tool: CustomAgentTool = { + const tool: CustomTool = { name: "question", label: "Question", description: "Ask the user a question and let them pick from options. Use when you need user input to proceed.", diff --git a/packages/coding-agent/examples/custom-tools/subagent/index.ts b/packages/coding-agent/examples/custom-tools/subagent/index.ts index 67a2d526..c9fd89e2 100644 --- a/packages/coding-agent/examples/custom-tools/subagent/index.ts +++ b/packages/coding-agent/examples/custom-tools/subagent/index.ts @@ -20,10 +20,10 @@ import type { AgentToolResult } from "@mariozechner/pi-agent-core"; import type { Message } from "@mariozechner/pi-ai"; import { StringEnum } from "@mariozechner/pi-ai"; import { - type CustomAgentTool, + type CustomTool, + type CustomToolAPI, type CustomToolFactory, getMarkdownTheme, - type ToolAPI, } from "@mariozechner/pi-coding-agent"; import { Container, Markdown, Spacer, Text } from "@mariozechner/pi-tui"; import { Type } from "@sinclair/typebox"; @@ -224,7 +224,7 @@ function writePromptToTempFile(agentName: string, prompt: string): { dir: string type OnUpdateCallback = (partial: AgentToolResult) => void; async function runSingleAgent( - pi: ToolAPI, + pi: CustomToolAPI, agents: AgentConfig[], agentName: string, task: string, @@ -411,7 +411,7 @@ const SubagentParams = Type.Object({ }); const factory: CustomToolFactory = (pi) => { - const tool: CustomAgentTool = { + const tool: CustomTool = { name: "subagent", label: "Subagent", get description() { @@ -433,7 +433,7 @@ const factory: CustomToolFactory = (pi) => { }, parameters: SubagentParams, - async execute(_toolCallId, params, signal, onUpdate) { + async execute(_toolCallId, params, signal, onUpdate, _ctx) { const agentScope: AgentScope = params.agentScope ?? "user"; const discovery = discoverAgents(pi.cwd, agentScope); const agents = discovery.agents; diff --git a/packages/coding-agent/examples/custom-tools/todo/index.ts b/packages/coding-agent/examples/custom-tools/todo/index.ts index d6da1b17..6b4d1feb 100644 --- a/packages/coding-agent/examples/custom-tools/todo/index.ts +++ b/packages/coding-agent/examples/custom-tools/todo/index.ts @@ -9,7 +9,12 @@ */ import { StringEnum } from "@mariozechner/pi-ai"; -import type { CustomAgentTool, CustomToolFactory, ToolSessionEvent } from "@mariozechner/pi-coding-agent"; +import type { + CustomTool, + CustomToolContext, + CustomToolFactory, + CustomToolSessionEvent, +} from "@mariozechner/pi-coding-agent"; import { Text } from "@mariozechner/pi-tui"; import { Type } from "@sinclair/typebox"; @@ -43,11 +48,12 @@ const factory: CustomToolFactory = (_pi) => { * Reconstruct state from session entries. * Scans tool results for this tool and applies them in order. */ - const reconstructState = (event: ToolSessionEvent) => { + const reconstructState = (_event: CustomToolSessionEvent, ctx: CustomToolContext) => { todos = []; nextId = 1; - for (const entry of event.entries) { + // Use getBranch() to get entries on the current branch + for (const entry of ctx.sessionManager.getBranch()) { if (entry.type !== "message") continue; const msg = entry.message; @@ -63,7 +69,7 @@ const factory: CustomToolFactory = (_pi) => { } }; - const tool: CustomAgentTool = { + const tool: CustomTool = { name: "todo", label: "Todo", description: "Manage a todo list. Actions: list, add (text), toggle (id), clear", diff --git a/packages/coding-agent/examples/hooks/custom-compaction.ts b/packages/coding-agent/examples/hooks/custom-compaction.ts index be7795b2..3dd0176b 100644 --- a/packages/coding-agent/examples/hooks/custom-compaction.ts +++ b/packages/coding-agent/examples/hooks/custom-compaction.ts @@ -14,7 +14,6 @@ */ import { complete, getModel } from "@mariozechner/pi-ai"; -import type { CompactionEntry } from "@mariozechner/pi-coding-agent"; import { convertToLlm } from "@mariozechner/pi-coding-agent"; import type { HookAPI } from "@mariozechner/pi-coding-agent/hooks"; @@ -22,7 +21,7 @@ export default function (pi: HookAPI) { pi.on("session_before_compact", async (event, ctx) => { ctx.ui.notify("Custom compaction hook triggered", "info"); - const { preparation, branchEntries, signal } = event; + const { preparation, branchEntries: _, signal } = event; const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId, previousSummary } = preparation; // Use Gemini Flash for summarization (cheaper/faster than most conversation models) diff --git a/packages/coding-agent/examples/sdk/05-tools.ts b/packages/coding-agent/examples/sdk/05-tools.ts index 7772c1dc..f7939688 100644 --- a/packages/coding-agent/examples/sdk/05-tools.ts +++ b/packages/coding-agent/examples/sdk/05-tools.ts @@ -11,7 +11,7 @@ import { Type } from "@sinclair/typebox"; import { bashTool, // read, bash, edit, write - uses process.cwd() - type CustomAgentTool, + type CustomTool, createAgentSession, createBashTool, createCodingTools, // Factory: creates tools for specific cwd @@ -55,7 +55,7 @@ await createAgentSession({ console.log("Specific tools with custom cwd session created"); // Inline custom tool (needs TypeBox schema) -const weatherTool: CustomAgentTool = { +const weatherTool: CustomTool = { name: "get_weather", label: "Get Weather", description: "Get current weather for a city", diff --git a/packages/coding-agent/examples/sdk/12-full-control.ts b/packages/coding-agent/examples/sdk/12-full-control.ts index bdf6a478..5dbe7718 100644 --- a/packages/coding-agent/examples/sdk/12-full-control.ts +++ b/packages/coding-agent/examples/sdk/12-full-control.ts @@ -12,7 +12,7 @@ import { getModel } from "@mariozechner/pi-ai"; import { Type } from "@sinclair/typebox"; import { AuthStorage, - type CustomAgentTool, + type CustomTool, createAgentSession, createBashTool, createReadTool, @@ -42,7 +42,7 @@ const auditHook: HookFactory = (api) => { }; // Inline custom tool -const statusTool: CustomAgentTool = { +const statusTool: CustomTool = { name: "status", label: "Status", description: "Get system status", @@ -68,15 +68,12 @@ const cwd = process.cwd(); const { session } = await createAgentSession({ cwd, agentDir: "/tmp/my-agent", - model, thinkingLevel: "off", authStorage, modelRegistry, - systemPrompt: `You are a minimal assistant. Available: read, bash, status. Be concise.`, - // Use factory functions with the same cwd to ensure path resolution works correctly tools: [createReadTool(cwd), createBashTool(cwd)], customTools: [{ tool: statusTool }], diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 6916d8e0..e2d63e09 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -27,7 +27,7 @@ import { prepareCompaction, shouldCompact, } from "./compaction/index.js"; -import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "./custom-tools/index.js"; +import type { CustomToolContext, CustomToolSessionEvent, LoadedCustomTool } from "./custom-tools/index.js"; import { exportSessionToHtml } from "./export-html.js"; import type { HookContext, @@ -698,7 +698,7 @@ export class AgentSession { } // Emit session event to custom tools - await this.emitToolSessionEvent("new", previousSessionFile); + await this.emitCustomToolSessionEvent("new", previousSessionFile); return true; } @@ -895,7 +895,7 @@ export class AgentSession { throw new Error(`No API key for ${this.model.provider}`); } - const pathEntries = this.sessionManager.getPath(); + const pathEntries = this.sessionManager.getBranch(); const settings = this.settingsManager.getCompactionSettings(); const preparation = prepareCompaction(pathEntries, settings); @@ -1068,7 +1068,7 @@ export class AgentSession { return; } - const pathEntries = this.sessionManager.getPath(); + const pathEntries = this.sessionManager.getBranch(); const preparation = prepareCompaction(pathEntries, settings); if (!preparation) { @@ -1473,7 +1473,7 @@ export class AgentSession { } // Emit session event to custom tools - await this.emitToolSessionEvent("switch", previousSessionFile); + await this.emitCustomToolSessionEvent("switch", previousSessionFile); this.agent.replaceMessages(sessionContext.messages); @@ -1550,7 +1550,7 @@ export class AgentSession { } // Emit session event to custom tools (with reason "branch") - await this.emitToolSessionEvent("branch", previousSessionFile); + await this.emitCustomToolSessionEvent("branch", previousSessionFile); if (!skipConversationRestore) { this.agent.replaceMessages(sessionContext.messages); @@ -1720,7 +1720,7 @@ export class AgentSession { } // Emit to custom tools - await this.emitToolSessionEvent("tree", this.sessionFile); + await this.emitCustomToolSessionEvent("tree", this.sessionFile); this._branchSummaryAbortController = undefined; return { editorText, cancelled: false, summaryEntry }; @@ -1877,20 +1877,23 @@ export class AgentSession { * Emit session event to all custom tools. * Called on session switch, branch, tree navigation, and shutdown. */ - async emitToolSessionEvent( - reason: ToolSessionEvent["reason"], + async emitCustomToolSessionEvent( + reason: CustomToolSessionEvent["reason"], previousSessionFile?: string | undefined, ): Promise { - const event: ToolSessionEvent = { - entries: this.sessionManager.getEntries(), - sessionFile: this.sessionFile, - previousSessionFile, - reason, + if (!this._customTools) return; + + const event: CustomToolSessionEvent = { reason, previousSessionFile }; + const ctx: CustomToolContext = { + sessionManager: this.sessionManager, + modelRegistry: this._modelRegistry, + model: this.agent.state.model, }; + for (const { tool } of this._customTools) { if (tool.onSession) { try { - await tool.onSession(event); + await tool.onSession(event, ctx); } catch (_err) { // Silently ignore tool errors during session events } diff --git a/packages/coding-agent/src/core/compaction/branch-summarization.ts b/packages/coding-agent/src/core/compaction/branch-summarization.ts index 8bca45ff..d3b4f59d 100644 --- a/packages/coding-agent/src/core/compaction/branch-summarization.ts +++ b/packages/coding-agent/src/core/compaction/branch-summarization.ts @@ -102,8 +102,8 @@ export function collectEntriesForBranchSummary( } // Find common ancestor (deepest node that's on both paths) - const oldPath = new Set(session.getPath(oldLeafId).map((e) => e.id)); - const targetPath = session.getPath(targetId); + const oldPath = new Set(session.getBranch(oldLeafId).map((e) => e.id)); + const targetPath = session.getBranch(targetId); // targetPath is root-first, so iterate backwards to find deepest common ancestor let commonAncestorId: string | null = null; diff --git a/packages/coding-agent/src/core/custom-tools/index.ts b/packages/coding-agent/src/core/custom-tools/index.ts index d78b6858..adb0b705 100644 --- a/packages/coding-agent/src/core/custom-tools/index.ts +++ b/packages/coding-agent/src/core/custom-tools/index.ts @@ -4,14 +4,18 @@ export { discoverAndLoadCustomTools, loadCustomTools } from "./loader.js"; export type { + AgentToolResult, AgentToolUpdateCallback, - CustomAgentTool, + CustomTool, + CustomToolAPI, + CustomToolContext, CustomToolFactory, + CustomToolResult, + CustomToolSessionEvent, CustomToolsLoadResult, + CustomToolUIContext, ExecResult, LoadedCustomTool, RenderResultOptions, - SessionEvent, - ToolAPI, - ToolUIContext, } from "./types.js"; +export { wrapCustomTool, wrapCustomTools } from "./wrapper.js"; diff --git a/packages/coding-agent/src/core/custom-tools/loader.ts b/packages/coding-agent/src/core/custom-tools/loader.ts index f4480611..0c51ff95 100644 --- a/packages/coding-agent/src/core/custom-tools/loader.ts +++ b/packages/coding-agent/src/core/custom-tools/loader.ts @@ -17,7 +17,7 @@ import { getAgentDir, isBunBinary } from "../../config.js"; import type { ExecOptions } from "../exec.js"; import { execCommand } from "../exec.js"; import type { HookUIContext } from "../hooks/types.js"; -import type { CustomToolFactory, CustomToolsLoadResult, LoadedCustomTool, ToolAPI } from "./types.js"; +import type { CustomToolAPI, CustomToolFactory, CustomToolsLoadResult, LoadedCustomTool } from "./types.js"; // Create require function to resolve module paths at runtime const require = createRequire(import.meta.url); @@ -104,7 +104,7 @@ function createNoOpUIContext(): HookUIContext { */ async function loadToolWithBun( resolvedPath: string, - sharedApi: ToolAPI, + sharedApi: CustomToolAPI, ): Promise<{ tools: LoadedCustomTool[] | null; error: string | null }> { try { // Try to import directly - will work for tools without @mariozechner/* imports @@ -149,7 +149,7 @@ async function loadToolWithBun( async function loadTool( toolPath: string, cwd: string, - sharedApi: ToolAPI, + sharedApi: CustomToolAPI, ): Promise<{ tools: LoadedCustomTool[] | null; error: string | null }> { const resolvedPath = resolveToolPath(toolPath, cwd); @@ -209,7 +209,7 @@ export async function loadCustomTools( const seenNames = new Set(builtInToolNames); // Shared API object - all tools get the same instance - const sharedApi: ToolAPI = { + const sharedApi: CustomToolAPI = { cwd, exec: (command: string, args: string[], options?: ExecOptions) => execCommand(command, args, options?.cwd ?? cwd, options), diff --git a/packages/coding-agent/src/core/custom-tools/types.ts b/packages/coding-agent/src/core/custom-tools/types.ts index 7bf99407..dc9cfd75 100644 --- a/packages/coding-agent/src/core/custom-tools/types.ts +++ b/packages/coding-agent/src/core/custom-tools/types.ts @@ -5,45 +5,56 @@ * They can provide custom rendering for tool calls and results in the TUI. */ -import type { AgentTool, AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core"; +import type { AgentToolResult, AgentToolUpdateCallback } from "@mariozechner/pi-agent-core"; +import type { Model } from "@mariozechner/pi-ai"; import type { Component } from "@mariozechner/pi-tui"; import type { Static, TSchema } from "@sinclair/typebox"; import type { Theme } from "../../modes/interactive/theme/theme.js"; import type { ExecOptions, ExecResult } from "../exec.js"; import type { HookUIContext } from "../hooks/types.js"; -import type { SessionEntry } from "../session-manager.js"; +import type { ModelRegistry } from "../model-registry.js"; +import type { ReadonlySessionManager } from "../session-manager.js"; /** Alias for clarity */ -export type ToolUIContext = HookUIContext; +export type CustomToolUIContext = HookUIContext; /** Re-export for custom tools to use in execute signature */ -export type { AgentToolUpdateCallback }; +export type { AgentToolResult, AgentToolUpdateCallback }; // Re-export for backward compatibility export type { ExecOptions, ExecResult } from "../exec.js"; /** API passed to custom tool factory (stable across session changes) */ -export interface ToolAPI { +export interface CustomToolAPI { /** Current working directory */ cwd: string; /** Execute a command */ exec(command: string, args: string[], options?: ExecOptions): Promise; - /** UI methods for user interaction (select, confirm, input, notify) */ - ui: ToolUIContext; + /** UI methods for user interaction (select, confirm, input, notify, custom) */ + ui: CustomToolUIContext; /** Whether UI is available (false in print/RPC mode) */ hasUI: boolean; } +/** + * Context passed to tool execute and onSession callbacks. + * Provides access to session state and model information. + */ +export interface CustomToolContext { + /** Session manager (read-only) */ + sessionManager: ReadonlySessionManager; + /** Model registry - use for API key resolution and model retrieval */ + modelRegistry: ModelRegistry; + /** Current model (may be undefined if no model is selected yet) */ + model: Model | undefined; +} + /** Session event passed to onSession callback */ -export interface SessionEvent { - /** All session entries (including pre-compaction history) */ - entries: SessionEntry[]; - /** Current session file path, or undefined in --no-session mode */ - sessionFile: string | undefined; - /** Previous session file path, or undefined for "start", "new", and "shutdown" */ - previousSessionFile: string | undefined; +export interface CustomToolSessionEvent { /** Reason for the session event */ reason: "start" | "switch" | "branch" | "new" | "tree" | "shutdown"; + /** Previous session file path, or undefined for "start", "new", and "shutdown" */ + previousSessionFile: string | undefined; } /** Rendering options passed to renderResult */ @@ -54,58 +65,89 @@ export interface RenderResultOptions { isPartial: boolean; } +export type CustomToolResult = AgentToolResult; + /** - * Custom tool with optional lifecycle and rendering methods. + * Custom tool definition. * - * The execute signature inherited from AgentTool includes an optional onUpdate callback - * for streaming progress updates during long-running operations: - * - The callback emits partial results to subscribers (e.g. TUI/RPC), not to the LLM. - * - Partial updates should use the same TDetails type as the final result (use a union if needed). + * Custom tools are standalone - they don't extend AgentTool directly. + * When loaded, they are wrapped in an AgentTool for the agent to use. + * + * The execute callback receives a ToolContext with access to session state, + * model registry, and current model. * * @example * ```typescript - * type Details = - * | { status: "running"; step: number; total: number } - * | { status: "done"; count: number }; + * const factory: CustomToolFactory = (pi) => ({ + * name: "my_tool", + * label: "My Tool", + * description: "Does something useful", + * parameters: Type.Object({ input: Type.String() }), * - * async execute(toolCallId, params, signal, onUpdate) { - * const items = params.items || []; - * for (let i = 0; i < items.length; i++) { - * onUpdate?.({ - * content: [{ type: "text", text: `Step ${i + 1}/${items.length}...` }], - * details: { status: "running", step: i + 1, total: items.length }, - * }); - * await processItem(items[i], signal); + * async execute(toolCallId, params, signal, onUpdate, ctx) { + * // Access session state via ctx.sessionManager + * // Access model registry via ctx.modelRegistry + * // Current model via ctx.model + * return { content: [{ type: "text", text: "Done" }] }; + * }, + * + * onSession(event, ctx) { + * if (event.reason === "shutdown") { + * // Cleanup + * } + * // Reconstruct state from ctx.sessionManager.getEntries() * } - * return { content: [{ type: "text", text: "Done" }], details: { status: "done", count: items.length } }; - * } + * }); * ``` - * - * Progress updates are rendered via renderResult with isPartial: true. */ -export interface CustomAgentTool - extends AgentTool { +export interface CustomTool { + /** Tool name (used in LLM tool calls) */ + name: string; + /** Human-readable label for UI */ + label: string; + /** Description for LLM */ + description: string; + /** Parameter schema (TypeBox) */ + parameters: TParams; + + /** + * Execute the tool. + * @param toolCallId - Unique ID for this tool call + * @param params - Parsed parameters matching the schema + * @param signal - AbortSignal for cancellation + * @param onUpdate - Callback for streaming partial results (for UI, not LLM) + * @param ctx - Context with session manager, model registry, and current model + */ + execute( + toolCallId: string, + params: Static, + signal: AbortSignal | undefined, + onUpdate: AgentToolUpdateCallback | undefined, + ctx: CustomToolContext, + ): Promise>; + /** Called on session lifecycle events - use to reconstruct state or cleanup resources */ - onSession?: (event: SessionEvent) => void | Promise; + onSession?: (event: CustomToolSessionEvent, ctx: CustomToolContext) => void | Promise; /** Custom rendering for tool call display - return a Component */ renderCall?: (args: Static, theme: Theme) => Component; + /** Custom rendering for tool result display - return a Component */ - renderResult?: (result: AgentToolResult, options: RenderResultOptions, theme: Theme) => Component; + renderResult?: (result: CustomToolResult, options: RenderResultOptions, theme: Theme) => Component; } /** Factory function that creates a custom tool or array of tools */ export type CustomToolFactory = ( - pi: ToolAPI, -) => CustomAgentTool | CustomAgentTool[] | Promise; + pi: CustomToolAPI, +) => CustomTool | CustomTool[] | Promise | CustomTool[]>; -/** Loaded custom tool with metadata */ +/** Loaded custom tool with metadata and wrapped AgentTool */ export interface LoadedCustomTool { /** Original path (as specified) */ path: string; /** Resolved absolute path */ resolvedPath: string; - /** The tool instance */ - tool: CustomAgentTool; + /** The original custom tool instance */ + tool: CustomTool; } /** Result from loading custom tools */ @@ -113,5 +155,5 @@ export interface CustomToolsLoadResult { tools: LoadedCustomTool[]; errors: Array<{ path: string; error: string }>; /** Update the UI context for all loaded tools. Call when mode initializes. */ - setUIContext(uiContext: ToolUIContext, hasUI: boolean): void; + setUIContext(uiContext: CustomToolUIContext, hasUI: boolean): void; } diff --git a/packages/coding-agent/src/core/custom-tools/wrapper.ts b/packages/coding-agent/src/core/custom-tools/wrapper.ts new file mode 100644 index 00000000..253f6092 --- /dev/null +++ b/packages/coding-agent/src/core/custom-tools/wrapper.ts @@ -0,0 +1,28 @@ +/** + * Wraps CustomTool instances into AgentTool for use with the agent. + */ + +import type { AgentTool } from "@mariozechner/pi-agent-core"; +import type { CustomTool, CustomToolContext, LoadedCustomTool } from "./types.js"; + +/** + * Wrap a CustomTool into an AgentTool. + * The wrapper injects the ToolContext into execute calls. + */ +export function wrapCustomTool(tool: CustomTool, getContext: () => CustomToolContext): AgentTool { + return { + name: tool.name, + label: tool.label, + description: tool.description, + parameters: tool.parameters, + execute: (toolCallId, params, signal, onUpdate) => + tool.execute(toolCallId, params, signal, onUpdate, getContext()), + }; +} + +/** + * Wrap all loaded custom tools into AgentTools. + */ +export function wrapCustomTools(loadedTools: LoadedCustomTool[], getContext: () => CustomToolContext): AgentTool[] { + return loadedTools.map((lt) => wrapCustomTool(lt.tool, getContext)); +} diff --git a/packages/coding-agent/src/core/hooks/runner.ts b/packages/coding-agent/src/core/hooks/runner.ts index 04a7eae3..998e39b7 100644 --- a/packages/coding-agent/src/core/hooks/runner.ts +++ b/packages/coding-agent/src/core/hooks/runner.ts @@ -108,20 +108,12 @@ export class HookRunner { hasUI?: boolean; }): void { this.getModel = options.getModel; - this.setSendMessageHandler(options.sendMessageHandler); - this.setAppendEntryHandler(options.appendEntryHandler); - if (options.uiContext) { - this.setUIContext(options.uiContext, options.hasUI ?? false); + for (const hook of this.hooks) { + hook.setSendMessageHandler(options.sendMessageHandler); + hook.setAppendEntryHandler(options.appendEntryHandler); } - } - - /** - * Set the UI context for hooks. - * Call this when the mode initializes and UI is available. - */ - setUIContext(uiContext: HookUIContext, hasUI: boolean): void { - this.uiContext = uiContext; - this.hasUI = hasUI; + this.uiContext = options.uiContext ?? noOpUIContext; + this.hasUI = options.hasUI ?? false; } /** @@ -145,26 +137,6 @@ export class HookRunner { return this.hooks.map((h) => h.path); } - /** - * Set the send message handler for all hooks' pi.sendMessage(). - * Call this when the mode initializes. - */ - setSendMessageHandler(handler: SendMessageHandler): void { - for (const hook of this.hooks) { - hook.setSendMessageHandler(handler); - } - } - - /** - * Set the append entry handler for all hooks' pi.appendEntry(). - * Call this when the mode initializes. - */ - setAppendEntryHandler(handler: AppendEntryHandler): void { - for (const hook of this.hooks) { - hook.setAppendEntryHandler(handler); - } - } - /** * Subscribe to hook errors. * @returns Unsubscribe function diff --git a/packages/coding-agent/src/core/hooks/types.ts b/packages/coding-agent/src/core/hooks/types.ts index 2d3cd141..879b111c 100644 --- a/packages/coding-agent/src/core/hooks/types.ts +++ b/packages/coding-agent/src/core/hooks/types.ts @@ -13,27 +13,7 @@ import type { CompactionPreparation, CompactionResult } from "../compaction/inde import type { ExecOptions, ExecResult } from "../exec.js"; import type { HookMessage } from "../messages.js"; import type { ModelRegistry } from "../model-registry.js"; -import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "../session-manager.js"; - -/** - * Read-only view of SessionManager for hooks. - * Hooks should use pi.sendMessage() and pi.appendEntry() for writes. - */ -export type ReadonlySessionManager = Pick< - SessionManager, - | "getCwd" - | "getSessionDir" - | "getSessionId" - | "getSessionFile" - | "getLeafId" - | "getLeafEntry" - | "getEntry" - | "getLabel" - | "getPath" - | "getHeader" - | "getEntries" - | "getTree" ->; +import type { BranchSummaryEntry, CompactionEntry, ReadonlySessionManager, SessionEntry } from "../session-manager.js"; import type { EditToolDetails } from "../tools/edit.js"; import type { diff --git a/packages/coding-agent/src/core/index.ts b/packages/coding-agent/src/core/index.ts index 55e51d99..4b15f6fe 100644 --- a/packages/coding-agent/src/core/index.ts +++ b/packages/coding-agent/src/core/index.ts @@ -14,16 +14,16 @@ export { export { type BashExecutorOptions, type BashResult, executeBash } from "./bash-executor.js"; export type { CompactionResult } from "./compaction/index.js"; export { - type CustomAgentTool, + type CustomTool, + type CustomToolAPI, type CustomToolFactory, type CustomToolsLoadResult, + type CustomToolUIContext, discoverAndLoadCustomTools, type ExecResult, type LoadedCustomTool, loadCustomTools, type RenderResultOptions, - type ToolAPI, - type ToolUIContext, } from "./custom-tools/index.js"; export { type HookAPI, diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index 089691d3..d2779ced 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -35,8 +35,13 @@ import { join } from "path"; import { getAgentDir } from "../config.js"; import { AgentSession } from "./agent-session.js"; import { AuthStorage } from "./auth-storage.js"; -import { discoverAndLoadCustomTools, type LoadedCustomTool } from "./custom-tools/index.js"; -import type { CustomAgentTool } from "./custom-tools/types.js"; +import { + type CustomToolsLoadResult, + discoverAndLoadCustomTools, + type LoadedCustomTool, + wrapCustomTools, +} from "./custom-tools/index.js"; +import type { CustomTool } from "./custom-tools/types.js"; import { discoverAndLoadHooks, HookRunner, type LoadedHook, wrapToolsWithHooks } from "./hooks/index.js"; import type { HookFactory } from "./hooks/types.js"; import { convertToLlm } from "./messages.js"; @@ -99,7 +104,7 @@ export interface CreateAgentSessionOptions { /** Built-in tools to use. Default: codingTools [read, bash, edit, write] */ tools?: Tool[]; /** Custom tools (replaces discovery). */ - customTools?: Array<{ path?: string; tool: CustomAgentTool }>; + customTools?: Array<{ path?: string; tool: CustomTool }>; /** Additional custom tool paths to load (merged with discovery). */ additionalCustomToolPaths?: string[]; @@ -127,17 +132,14 @@ export interface CreateAgentSessionResult { /** The created session */ session: AgentSession; /** Custom tools result (for UI context setup in interactive mode) */ - customToolsResult: { - tools: LoadedCustomTool[]; - setUIContext: (uiContext: any, hasUI: boolean) => void; - }; + customToolsResult: CustomToolsLoadResult; /** Warning if session was restored with a different model than saved */ modelFallbackMessage?: string; } // Re-exports -export type { CustomAgentTool } from "./custom-tools/types.js"; +export type { CustomTool } from "./custom-tools/types.js"; export type { HookAPI, HookFactory } from "./hooks/types.js"; export type { Settings, SkillsSettings } from "./settings-manager.js"; export type { Skill } from "./skills.js"; @@ -219,7 +221,7 @@ export async function discoverHooks( export async function discoverCustomTools( cwd?: string, agentDir?: string, -): Promise> { +): Promise> { const resolvedCwd = cwd ?? process.cwd(); const resolvedAgentDir = agentDir ?? getDefaultAgentDir(); @@ -507,7 +509,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const builtInTools = options.tools ?? createCodingTools(cwd); time("createCodingTools"); - let customToolsResult: { tools: LoadedCustomTool[]; setUIContext: (ctx: any, hasUI: boolean) => void }; + let customToolsResult: CustomToolsLoadResult; if (options.customTools !== undefined) { // Use provided custom tools const loadedTools: LoadedCustomTool[] = options.customTools.map((ct) => ({ @@ -517,17 +519,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} })); customToolsResult = { tools: loadedTools, + errors: [], setUIContext: () => {}, }; } else { // Discover custom tools, merging with additional paths const configuredPaths = [...settingsManager.getCustomToolPaths(), ...(options.additionalCustomToolPaths ?? [])]; - const result = await discoverAndLoadCustomTools(configuredPaths, cwd, Object.keys(allTools), agentDir); + customToolsResult = await discoverAndLoadCustomTools(configuredPaths, cwd, Object.keys(allTools), agentDir); time("discoverAndLoadCustomTools"); - for (const { path, error } of result.errors) { + for (const { path, error } of customToolsResult.errors) { console.error(`Failed to load custom tool "${path}": ${error}`); } - customToolsResult = result; } let hookRunner: HookRunner | undefined; @@ -549,7 +551,15 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} } } - let allToolsArray: Tool[] = [...builtInTools, ...customToolsResult.tools.map((lt) => lt.tool as unknown as Tool)]; + // Wrap custom tools with context getter (agent is assigned below, accessed at execute time) + let agent: Agent; + const wrappedCustomTools = wrapCustomTools(customToolsResult.tools, () => ({ + sessionManager, + modelRegistry, + model: agent.state.model, + })); + + let allToolsArray: Tool[] = [...builtInTools, ...wrappedCustomTools]; time("combineTools"); if (hookRunner) { allToolsArray = wrapToolsWithHooks(allToolsArray, hookRunner) as Tool[]; @@ -581,7 +591,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} const slashCommands = options.slashCommands ?? discoverSlashCommands(cwd, agentDir); time("discoverSlashCommands"); - const agent = new Agent({ + agent = new Agent({ initialState: { systemPrompt, model, diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index c3e17714..dc47f749 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -159,19 +159,21 @@ export interface SessionInfo { allMessagesText: string; } -/** - * Read-only interface for SessionManager. - * Used by compaction/summarization utilities that only need to read session data. - */ -export interface ReadonlySessionManager { - getLeafId(): string | null; - getEntry(id: string): SessionEntry | undefined; - getPath(fromId?: string): SessionEntry[]; - getEntries(): SessionEntry[]; - getChildren(parentId: string): SessionEntry[]; - getTree(): SessionTreeNode[]; - getLabel(id: string): string | undefined; -} +export type ReadonlySessionManager = Pick< + SessionManager, + | "getCwd" + | "getSessionDir" + | "getSessionId" + | "getSessionFile" + | "getLeafId" + | "getLeafEntry" + | "getEntry" + | "getLabel" + | "getBranch" + | "getHeader" + | "getEntries" + | "getTree" +>; /** Generate a unique short ID (8 hex chars, collision-checked) */ function generateId(byId: { has(id: string): boolean }): string { @@ -772,7 +774,7 @@ export class SessionManager { * Includes all entry types (messages, compaction, model changes, etc.). * Use buildSessionContext() to get the resolved messages for the LLM. */ - getPath(fromId?: string): SessionEntry[] { + getBranch(fromId?: string): SessionEntry[] { const path: SessionEntry[] = []; const startId = fromId ?? this.leafId; let current = startId ? this.byId.get(startId) : undefined; @@ -908,7 +910,7 @@ export class SessionManager { * Returns the new session file path, or undefined if not persisting. */ createBranchedSession(leafId: string): string | undefined { - const path = this.getPath(leafId); + const path = this.getBranch(leafId); if (path.length === 0) { throw new Error(`Entry ${leafId} not found`); } diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 1dae4688..4efaacd4 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -35,15 +35,16 @@ export { // Custom tools export type { AgentToolUpdateCallback, - CustomAgentTool, + CustomTool, + CustomToolAPI, + CustomToolContext, CustomToolFactory, + CustomToolSessionEvent, CustomToolsLoadResult, + CustomToolUIContext, ExecResult, LoadedCustomTool, RenderResultOptions, - SessionEvent as ToolSessionEvent, - ToolAPI, - ToolUIContext, } from "./core/custom-tools/index.js"; export { discoverAndLoadCustomTools, loadCustomTools } from "./core/custom-tools/index.js"; export type * from "./core/hooks/index.js"; diff --git a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts index 7124c84b..4f6bfac7 100644 --- a/packages/coding-agent/src/modes/interactive/components/tool-execution.ts +++ b/packages/coding-agent/src/modes/interactive/components/tool-execution.ts @@ -11,7 +11,7 @@ import { type TUI, } from "@mariozechner/pi-tui"; import stripAnsi from "strip-ansi"; -import type { CustomAgentTool } from "../../../core/custom-tools/types.js"; +import type { CustomTool } from "../../../core/custom-tools/types.js"; import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize } from "../../../core/tools/truncate.js"; import { getLanguageFromPath, highlightCode, theme } from "../theme/theme.js"; import { renderDiff } from "./diff.js"; @@ -55,7 +55,7 @@ export class ToolExecutionComponent extends Container { private expanded = false; private showImages: boolean; private isPartial = true; - private customTool?: CustomAgentTool; + private customTool?: CustomTool; private ui: TUI; private result?: { content: Array<{ type: string; text?: string; data?: string; mimeType?: string }>; @@ -67,7 +67,7 @@ export class ToolExecutionComponent extends Container { toolName: string, args: any, options: ToolExecutionOptions = {}, - customTool: CustomAgentTool | undefined, + customTool: CustomTool | undefined, ui: TUI, ) { super(); diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index b6ea952d..809eca88 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -26,7 +26,7 @@ import { import { exec, spawnSync } from "child_process"; import { APP_NAME, getAuthPath, getDebugLogPath } from "../../config.js"; import type { AgentSession, AgentSessionEvent } from "../../core/agent-session.js"; -import type { LoadedCustomTool, SessionEvent as ToolSessionEvent } from "../../core/custom-tools/index.js"; +import type { CustomToolSessionEvent, LoadedCustomTool } from "../../core/custom-tools/index.js"; import type { HookUIContext } from "../../core/hooks/index.js"; import { createCompactionSummaryMessage } from "../../core/messages.js"; import { type SessionContext, SessionManager } from "../../core/session-manager.js"; @@ -350,19 +350,20 @@ export class InteractiveMode { this.chatContainer.addChild(new Spacer(1)); } - // Load session entries if any - const entries = this.session.sessionManager.getEntries(); - - // Set TUI-based UI context for custom tools - const uiContext = this.createHookUIContext(); + // Create and set hook & tool UI context + const uiContext: HookUIContext = { + select: (title, options) => this.showHookSelector(title, options), + confirm: (title, message) => this.showHookConfirm(title, message), + input: (title, placeholder) => this.showHookInput(title, placeholder), + notify: (message, type) => this.showHookNotify(message, type), + custom: (component) => this.showHookCustom(component), + }; this.setToolUIContext(uiContext, true); // Notify custom tools of session start - await this.emitToolSessionEvent({ - entries, - sessionFile: this.session.sessionFile, - previousSessionFile: undefined, + await this.emitCustomToolSessionEvent({ reason: "start", + previousSessionFile: undefined, }); const hookRunner = this.session.hookRunner; @@ -370,34 +371,35 @@ export class InteractiveMode { return; // No hooks loaded } - // Set UI context on hook runner - hookRunner.setUIContext(uiContext, true); + hookRunner.initialize({ + getModel: () => this.session.model, + sendMessageHandler: (message, triggerTurn) => { + const wasStreaming = this.session.isStreaming; + this.session + .sendHookMessage(message, triggerTurn) + .then(() => { + // For non-streaming cases with display=true, update UI + // (streaming cases update via message_end event) + if (!wasStreaming && message.display) { + this.rebuildChatFromMessages(); + } + }) + .catch((err) => { + this.showError(`Hook sendMessage failed: ${err instanceof Error ? err.message : String(err)}`); + }); + }, + appendEntryHandler: (customType, data) => { + this.sessionManager.appendCustomEntry(customType, data); + }, + uiContext, + hasUI: true, + }); // Subscribe to hook errors hookRunner.onError((error) => { this.showHookError(error.hookPath, error.error); }); - // Set up handlers for pi.sendMessage() and pi.appendEntry() - hookRunner.setSendMessageHandler((message, triggerTurn) => { - const wasStreaming = this.session.isStreaming; - this.session - .sendHookMessage(message, triggerTurn) - .then(() => { - // For non-streaming cases with display=true, update UI - // (streaming cases update via message_end event) - if (!wasStreaming && message.display) { - this.rebuildChatFromMessages(); - } - }) - .catch((err) => { - this.showError(`Hook sendMessage failed: ${err instanceof Error ? err.message : String(err)}`); - }); - }); - hookRunner.setAppendEntryHandler((customType, data) => { - this.sessionManager.appendCustomEntry(customType, data); - }); - // Show loaded hooks const hookPaths = hookRunner.getHookPaths(); if (hookPaths.length > 0) { @@ -415,11 +417,15 @@ export class InteractiveMode { /** * Emit session event to all custom tools. */ - private async emitToolSessionEvent(event: ToolSessionEvent): Promise { + private async emitCustomToolSessionEvent(event: CustomToolSessionEvent): Promise { for (const { tool } of this.customTools.values()) { if (tool.onSession) { try { - await tool.onSession(event); + await tool.onSession(event, { + sessionManager: this.session.sessionManager, + modelRegistry: this.session.modelRegistry, + model: this.session.model, + }); } catch (err) { this.showToolError(tool.name, err instanceof Error ? err.message : String(err)); } @@ -436,19 +442,6 @@ export class InteractiveMode { this.ui.requestRender(); } - /** - * Create the UI context for hooks. - */ - private createHookUIContext(): HookUIContext { - return { - select: (title, options) => this.showHookSelector(title, options), - confirm: (title, message) => this.showHookConfirm(title, message), - input: (title, placeholder) => this.showHookInput(title, placeholder), - notify: (message, type) => this.showHookNotify(message, type), - custom: (component) => this.showHookCustom(component), - }; - } - /** * Show a selector for hooks. */ @@ -861,6 +854,7 @@ export class InteractiveMode { this.customTools.get(content.name)?.tool, this.ui, ); + component.setExpanded(this.toolOutputExpanded); this.chatContainer.addChild(component); this.pendingTools.set(content.id, component); } else { @@ -909,6 +903,7 @@ export class InteractiveMode { this.customTools.get(event.toolName)?.tool, this.ui, ); + component.setExpanded(this.toolOutputExpanded); this.chatContainer.addChild(component); this.pendingTools.set(event.toolCallId, component); this.ui.requestRender(); @@ -1158,6 +1153,7 @@ export class InteractiveMode { this.customTools.get(content.name)?.tool, this.ui, ); + component.setExpanded(this.toolOutputExpanded); this.chatContainer.addChild(component); if (message.stopReason === "aborted" || message.stopReason === "error") { @@ -1251,7 +1247,7 @@ export class InteractiveMode { } // Emit shutdown event to custom tools - await this.session.emitToolSessionEvent("shutdown"); + await this.session.emitCustomToolSessionEvent("shutdown"); this.stop(); process.exit(0); diff --git a/packages/coding-agent/src/modes/print-mode.ts b/packages/coding-agent/src/modes/print-mode.ts index fbf3037a..b0aec7fa 100644 --- a/packages/coding-agent/src/modes/print-mode.ts +++ b/packages/coding-agent/src/modes/print-mode.ts @@ -26,25 +26,24 @@ export async function runPrintMode( initialMessage?: string, initialImages?: ImageContent[], ): Promise { - // Load entries once for session start events - const entries = session.sessionManager.getEntries(); - // Hook runner already has no-op UI context by default (set in main.ts) // Set up hooks for print mode (no UI) const hookRunner = session.hookRunner; if (hookRunner) { + hookRunner.initialize({ + getModel: () => session.model, + sendMessageHandler: (message, triggerTurn) => { + session.sendHookMessage(message, triggerTurn).catch((e) => { + console.error(`Hook sendMessage failed: ${e instanceof Error ? e.message : String(e)}`); + }); + }, + appendEntryHandler: (customType, data) => { + session.sessionManager.appendCustomEntry(customType, data); + }, + }); hookRunner.onError((err) => { console.error(`Hook error (${err.hookPath}): ${err.error}`); }); - // Set up handlers - sendHookMessage handles queuing/direct append as needed - hookRunner.setSendMessageHandler((message, triggerTurn) => { - session.sendHookMessage(message, triggerTurn).catch((e) => { - console.error(`Hook sendMessage failed: ${e instanceof Error ? e.message : String(e)}`); - }); - }); - hookRunner.setAppendEntryHandler((customType, data) => { - session.sessionManager.appendCustomEntry(customType, data); - }); // Emit session_start event await hookRunner.emit({ type: "session_start", @@ -55,12 +54,17 @@ export async function runPrintMode( for (const { tool } of session.customTools) { if (tool.onSession) { try { - await tool.onSession({ - entries, - sessionFile: session.sessionFile, - previousSessionFile: undefined, - reason: "start", - }); + await tool.onSession( + { + reason: "start", + previousSessionFile: undefined, + }, + { + sessionManager: session.sessionManager, + modelRegistry: session.modelRegistry, + model: session.model, + }, + ); } catch (_err) { // Silently ignore tool errors } diff --git a/packages/coding-agent/src/modes/rpc/rpc-mode.ts b/packages/coding-agent/src/modes/rpc/rpc-mode.ts index de378612..e9fabf2a 100644 --- a/packages/coding-agent/src/modes/rpc/rpc-mode.ts +++ b/packages/coding-agent/src/modes/rpc/rpc-mode.ts @@ -125,25 +125,25 @@ export async function runRpcMode(session: AgentSession): Promise { }, }); - // Load entries once for session start events - const entries = session.sessionManager.getEntries(); - // Set up hooks with RPC-based UI context const hookRunner = session.hookRunner; if (hookRunner) { - hookRunner.setUIContext(createHookUIContext(), false); + hookRunner.initialize({ + getModel: () => session.agent.state.model, + sendMessageHandler: (message, triggerTurn) => { + session.sendHookMessage(message, triggerTurn).catch((e) => { + output(error(undefined, "hook_send", e.message)); + }); + }, + appendEntryHandler: (customType, data) => { + session.sessionManager.appendCustomEntry(customType, data); + }, + uiContext: createHookUIContext(), + hasUI: false, + }); hookRunner.onError((err) => { output({ type: "hook_error", hookPath: err.hookPath, event: err.event, error: err.error }); }); - // Set up handlers for pi.sendMessage() and pi.appendEntry() - hookRunner.setSendMessageHandler((message, triggerTurn) => { - session.sendHookMessage(message, triggerTurn).catch((e) => { - output(error(undefined, "hook_send", e.message)); - }); - }); - hookRunner.setAppendEntryHandler((customType, data) => { - session.sessionManager.appendCustomEntry(customType, data); - }); // Emit session_start event await hookRunner.emit({ type: "session_start", @@ -155,12 +155,17 @@ export async function runRpcMode(session: AgentSession): Promise { for (const { tool } of session.customTools) { if (tool.onSession) { try { - await tool.onSession({ - entries, - sessionFile: session.sessionFile, - previousSessionFile: undefined, - reason: "start", - }); + await tool.onSession( + { + previousSessionFile: undefined, + reason: "start", + }, + { + sessionManager: session.sessionManager, + modelRegistry: session.modelRegistry, + model: session.model, + }, + ); } catch (_err) { // Silently ignore tool errors } diff --git a/packages/coding-agent/test/compaction-hooks-example.test.ts b/packages/coding-agent/test/compaction-hooks-example.test.ts index deceaeb5..64e6cec6 100644 --- a/packages/coding-agent/test/compaction-hooks-example.test.ts +++ b/packages/coding-agent/test/compaction-hooks-example.test.ts @@ -11,17 +11,11 @@ describe("Documentation example", () => { const exampleHook = (pi: HookAPI) => { pi.on("session_before_compact", async (event: SessionBeforeCompactEvent, ctx) => { // All these should be accessible on the event - const { preparation, branchEntries, signal } = event; + const { preparation, branchEntries } = event; // sessionManager, modelRegistry, and model come from ctx - const { sessionManager, modelRegistry, model } = ctx; - const { - messagesToSummarize, - turnPrefixMessages, - tokensBefore, - firstKeptEntryId, - isSplitTurn, - previousSummary, - } = preparation; + const { sessionManager, modelRegistry } = ctx; + const { messagesToSummarize, turnPrefixMessages, tokensBefore, firstKeptEntryId, isSplitTurn } = + preparation; // Verify types expect(Array.isArray(messagesToSummarize)).toBe(true); diff --git a/packages/coding-agent/test/compaction-hooks.test.ts b/packages/coding-agent/test/compaction-hooks.test.ts index b0b65511..d1fce41e 100644 --- a/packages/coding-agent/test/compaction-hooks.test.ts +++ b/packages/coding-agent/test/compaction-hooks.test.ts @@ -99,16 +99,19 @@ describe.skipIf(!API_KEY)("Compaction hooks", () => { const modelRegistry = new ModelRegistry(authStorage); hookRunner = new HookRunner(hooks, tempDir, sessionManager, modelRegistry); - hookRunner.setUIContext( - { + hookRunner.initialize({ + getModel: () => session.model, + sendMessageHandler: async () => {}, + appendEntryHandler: async () => {}, + uiContext: { select: async () => undefined, confirm: async () => false, input: async () => undefined, notify: () => {}, custom: () => ({ close: () => {}, requestRender: () => {} }), }, - false, - ); + hasUI: false, + }); session = new AgentSession({ agent, diff --git a/packages/coding-agent/test/session-manager/save-entry.test.ts b/packages/coding-agent/test/session-manager/save-entry.test.ts index 2a618986..45015321 100644 --- a/packages/coding-agent/test/session-manager/save-entry.test.ts +++ b/packages/coding-agent/test/session-manager/save-entry.test.ts @@ -42,7 +42,7 @@ describe("SessionManager.saveCustomEntry", () => { expect(customEntry.parentId).toBe(msgId); // Tree structure should be correct - const path = session.getPath(); + const path = session.getBranch(); expect(path).toHaveLength(3); expect(path[0].id).toBe(msgId); expect(path[1].id).toBe(customId); diff --git a/packages/coding-agent/test/session-manager/tree-traversal.test.ts b/packages/coding-agent/test/session-manager/tree-traversal.test.ts index 5fe7610a..fe244710 100644 --- a/packages/coding-agent/test/session-manager/tree-traversal.test.ts +++ b/packages/coding-agent/test/session-manager/tree-traversal.test.ts @@ -122,14 +122,14 @@ describe("SessionManager append and tree traversal", () => { describe("getPath", () => { it("returns empty array for empty session", () => { const session = SessionManager.inMemory(); - expect(session.getPath()).toEqual([]); + expect(session.getBranch()).toEqual([]); }); it("returns single entry path", () => { const session = SessionManager.inMemory(); const id = session.appendMessage(userMsg("hello")); - const path = session.getPath(); + const path = session.getBranch(); expect(path).toHaveLength(1); expect(path[0].id).toBe(id); }); @@ -142,7 +142,7 @@ describe("SessionManager append and tree traversal", () => { const id3 = session.appendThinkingLevelChange("high"); const id4 = session.appendMessage(userMsg("3")); - const path = session.getPath(); + const path = session.getBranch(); expect(path).toHaveLength(4); expect(path.map((e) => e.id)).toEqual([id1, id2, id3, id4]); }); @@ -155,7 +155,7 @@ describe("SessionManager append and tree traversal", () => { const _id3 = session.appendMessage(userMsg("3")); const _id4 = session.appendMessage(assistantMsg("4")); - const path = session.getPath(id2); + const path = session.getBranch(id2); expect(path).toHaveLength(2); expect(path.map((e) => e.id)).toEqual([id1, id2]); });