Emil_Shanaty/my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Rework custom tools API with CustomToolContext

Browse Source 
- 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<any, any> 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
This commit is contained in:
Mario Zechner
2025-12-31 12:05:24 +01:00
parent b123df5fab
commit 568150f18b
27 changed files with 336 additions and 289 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/core/custom-tools/index.ts
+8 -4
View File
@@ -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";
packages/coding-agent/src/core/custom-tools/loader.ts
+4 -4
View File
@@ -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<string>(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),
packages/coding-agent/src/core/custom-tools/types.ts
+86 -44
View File
@@ -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<ExecResult>;
/** 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<any> | 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<TDetails = any> = AgentToolResult<TDetails>;
/**
* 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<TParams extends TSchema = TSchema, TDetails = any>
extends AgentTool<TParams, TDetails> {
export interface CustomTool<TParams extends TSchema = TSchema, TDetails = any> {
/** 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<TParams>,
signal: AbortSignal | undefined,
onUpdate: AgentToolUpdateCallback<TDetails> | undefined,
ctx: CustomToolContext,
): Promise<AgentToolResult<TDetails>>;
/** Called on session lifecycle events - use to reconstruct state or cleanup resources */
onSession?: (event: SessionEvent) => void | Promise<void>;
onSession?: (event: CustomToolSessionEvent, ctx: CustomToolContext) => void | Promise<void>;
/** Custom rendering for tool call display - return a Component */
renderCall?: (args: Static<TParams>, theme: Theme) => Component;
/** Custom rendering for tool result display - return a Component */
renderResult?: (result: AgentToolResult<TDetails>, options: RenderResultOptions, theme: Theme) => Component;
renderResult?: (result: CustomToolResult<TDetails>, options: RenderResultOptions, theme: Theme) => Component;
}
/** Factory function that creates a custom tool or array of tools */
export type CustomToolFactory = (
pi: ToolAPI,
) => CustomAgentTool<any> | CustomAgentTool[] | Promise<CustomAgentTool | CustomAgentTool[]>;
pi: CustomToolAPI,
) => CustomTool<any, any> | CustomTool<any, any>[] | Promise<CustomTool<any, any> | CustomTool<any, any>[]>;
/** 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;
}
packages/coding-agent/src/core/custom-tools/wrapper.ts
+28
View File
@@ -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));
}