diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d8802825..e4c5c69c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,8 +2,13 @@ ## [Unreleased] +### Added + +- Added optional `promptSnippet` to `ToolDefinition` for one-line entries in the default system prompt's `Available tools` section. Active extension tools now appear there, including tools registered dynamically after startup. + ### Fixed +- Fixed `pi.registerTool()` dynamic registration after session initialization. Tools registered in `session_start` and later handlers now refresh immediately, become active, and are visible to the LLM without `/reload` ([#1720](https://github.com/badlogic/pi-mono/issues/1720)) - Fixed session message persistence ordering by serializing `AgentSession` event processing, preventing `toolResult` entries from being written before their corresponding assistant tool-call messages when extension handlers are asynchronous ([#1717](https://github.com/badlogic/pi-mono/issues/1717)) ## [0.55.3] - 2026-02-27 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 6a6983b2..fbfd35a3 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -880,6 +880,12 @@ Subscribe to events. See [Events](#events) for event types and return values. Register a custom tool callable by the LLM. See [Custom Tools](#custom-tools) for full details. +`pi.registerTool()` works both during extension load and after startup. You can call it inside `session_start`, command handlers, or other event handlers. New tools are refreshed immediately in the same session, so they appear in `pi.getAllTools()` and are callable by the LLM without `/reload`. + +Use `pi.setActiveTools()` to enable or disable tools (including dynamically added tools) at runtime. + +See [dynamic-tools.ts](../examples/extensions/dynamic-tools.ts) for a full example. + ```typescript import { Type } from "@sinclair/typebox"; import { StringEnum } from "@mariozechner/pi-ai"; @@ -888,6 +894,7 @@ pi.registerTool({ name: "my_tool", label: "My Tool", description: "What this tool does", + promptSnippet: "Summarize or transform text according to action", parameters: Type.Object({ action: StringEnum(["list", "add"] as const), text: Type.Optional(Type.String()), @@ -1116,7 +1123,7 @@ const result = await pi.exec("git", ["status"], { signal, timeout: 5000 }); ### pi.getActiveTools() / pi.getAllTools() / pi.setActiveTools(names) -Manage active tools. +Manage active tools. This works for both built-in tools and dynamically registered tools. ```typescript const active = pi.getActiveTools(); // ["read", "bash", "edit", "write"] @@ -1276,6 +1283,8 @@ export default function (pi: ExtensionAPI) { Register tools the LLM can call via `pi.registerTool()`. Tools appear in the system prompt and can have custom rendering. +Use `promptSnippet` for a short one-line entry in the `Available tools` section in the default system prompt. If omitted, pi falls back to `description`. + Note: Some models are idiots and include the @ prefix in tool path arguments. Built-in tools strip a leading @ before resolving paths. If your custom tool accepts a path, normalize a leading @ as well. ### Tool Definition @@ -1289,6 +1298,7 @@ pi.registerTool({ name: "my_tool", label: "My Tool", description: "What this tool does (shown to LLM)", + promptSnippet: "List or add items in the project todo list", parameters: Type.Object({ action: StringEnum(["list", "add"] as const), // Use StringEnum for Google compatibility text: Type.Optional(Type.String()), @@ -1886,6 +1896,7 @@ All examples in [examples/extensions/](../examples/extensions/). | `question.ts` | Tool with user interaction | `registerTool`, `ui.select` | | `questionnaire.ts` | Multi-step wizard tool | `registerTool`, `ui.custom` | | `todo.ts` | Stateful tool with persistence | `registerTool`, `appendEntry`, `renderResult`, session events | +| `dynamic-tools.ts` | Register tools after startup and during commands | `registerTool`, `session_start`, `registerCommand` | | `truncated-tool.ts` | Output truncation example | `registerTool`, `truncateHead` | | `tool-override.ts` | Override built-in read tool | `registerTool` (same name as built-in) | | **Commands** ||| diff --git a/packages/coding-agent/examples/extensions/README.md b/packages/coding-agent/examples/extensions/README.md index a997e6aa..17091392 100644 --- a/packages/coding-agent/examples/extensions/README.md +++ b/packages/coding-agent/examples/extensions/README.md @@ -33,6 +33,7 @@ cp permission-gate.ts ~/.pi/agent/extensions/ | `question.ts` | Demonstrates `ctx.ui.select()` for asking the user questions with custom UI | | `questionnaire.ts` | Multi-question input with tab bar navigation between questions | | `tool-override.ts` | Override built-in tools (e.g., add logging/access control to `read`) | +| `dynamic-tools.ts` | Register tools after startup (`session_start`) and at runtime via command | | `built-in-tool-renderer.ts` | Custom compact rendering for built-in tools (read, bash, edit, write) while keeping original behavior | | `minimal-mode.ts` | Override built-in tool rendering for minimal display (only tool calls, no output in collapsed mode) | | `truncated-tool.ts` | Wraps ripgrep with proper output truncation (50KB/2000 lines) | diff --git a/packages/coding-agent/examples/extensions/dynamic-tools.ts b/packages/coding-agent/examples/extensions/dynamic-tools.ts new file mode 100644 index 00000000..d5e3e404 --- /dev/null +++ b/packages/coding-agent/examples/extensions/dynamic-tools.ts @@ -0,0 +1,73 @@ +/** + * Dynamic Tools Extension + * + * Demonstrates registering tools after session initialization. + * + * - Registers one tool during session_start + * - Registers additional tools at runtime via /add-echo-tool + */ + +import type { ExtensionAPI } from "@mariozechner/pi-coding-agent"; +import { Type } from "@sinclair/typebox"; + +const ECHO_PARAMS = Type.Object({ + message: Type.String({ description: "Message to echo" }), +}); + +function normalizeToolName(input: string): string | undefined { + const trimmed = input.trim().toLowerCase(); + if (!trimmed) return undefined; + if (!/^[a-z0-9_]+$/.test(trimmed)) return undefined; + return trimmed; +} + +export default function dynamicToolsExtension(pi: ExtensionAPI) { + const registeredToolNames = new Set(); + + const registerEchoTool = (name: string, label: string, prefix: string): boolean => { + if (registeredToolNames.has(name)) { + return false; + } + + registeredToolNames.add(name); + pi.registerTool({ + name, + label, + description: `Echo a message with prefix: ${prefix}`, + promptSnippet: `Echo back user-provided text with ${prefix.trim()} prefix`, + parameters: ECHO_PARAMS, + async execute(_toolCallId, params) { + return { + content: [{ type: "text", text: `${prefix}${params.message}` }], + details: { tool: name, prefix }, + }; + }, + }); + + return true; + }; + + pi.on("session_start", (_event, ctx) => { + registerEchoTool("echo_session", "Echo Session", "[session] "); + ctx.ui.notify("Registered dynamic tool: echo_session", "info"); + }); + + pi.registerCommand("add-echo-tool", { + description: "Register a new echo tool dynamically: /add-echo-tool ", + handler: async (args, ctx) => { + const toolName = normalizeToolName(args); + if (!toolName) { + ctx.ui.notify("Usage: /add-echo-tool (lowercase, numbers, underscores)", "warning"); + return; + } + + const created = registerEchoTool(toolName, `Echo ${toolName}`, `[${toolName}] `); + if (!created) { + ctx.ui.notify(`Tool already registered: ${toolName}`, "warning"); + return; + } + + ctx.ui.notify(`Registered dynamic tool: ${toolName}`, "info"); + }, + }); +} diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index d6f7b955..0e84a830 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -268,6 +268,7 @@ export class AgentSession { // Tool registry for extension getTools/setTools private _toolRegistry: Map = new Map(); + private _toolPromptSnippets: Map = new Map(); // Base system prompt (without extension appends) - used to apply fresh appends each turn private _baseSystemPrompt = ""; @@ -682,8 +683,25 @@ export class AgentSession { return this._resourceLoader.getPrompts().prompts; } + private _normalizePromptSnippet(text: string | undefined): string | undefined { + if (!text) return undefined; + const oneLine = text + .replace(/[\r\n]+/g, " ") + .replace(/\s+/g, " ") + .trim(); + return oneLine.length > 0 ? oneLine : undefined; + } + private _rebuildSystemPrompt(toolNames: string[]): string { - const validToolNames = toolNames.filter((name) => this._baseToolRegistry.has(name)); + const validToolNames = toolNames.filter((name) => this._toolRegistry.has(name)); + const toolSnippets: Record = {}; + for (const name of validToolNames) { + const snippet = this._toolPromptSnippets.get(name); + if (snippet) { + toolSnippets[name] = snippet; + } + } + const loaderSystemPrompt = this._resourceLoader.getSystemPrompt(); const loaderAppendSystemPrompt = this._resourceLoader.getAppendSystemPrompt(); const appendSystemPrompt = @@ -698,6 +716,7 @@ export class AgentSession { customPrompt: loaderSystemPrompt, appendSystemPrompt, selectedTools: validToolNames, + toolSnippets, }); } @@ -1934,6 +1953,7 @@ export class AgentSession { getActiveTools: () => this.getActiveToolNames(), getAllTools: () => this.getAllTools(), setActiveTools: (toolNames) => this.setActiveToolsByName(toolNames), + refreshTools: () => this._refreshToolRegistry(), getCommands, setModel: async (model) => { const key = await this.modelRegistry.getApiKey(model); @@ -1969,6 +1989,60 @@ export class AgentSession { ); } + private _refreshToolRegistry(options?: { activeToolNames?: string[]; includeAllExtensionTools?: boolean }): void { + const previousRegistryNames = new Set(this._toolRegistry.keys()); + const previousActiveToolNames = this.getActiveToolNames(); + + const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? []; + const allCustomTools = [ + ...registeredTools, + ...this._customTools.map((def) => ({ definition: def, extensionPath: "" })), + ]; + this._toolPromptSnippets = new Map( + allCustomTools + .map((registeredTool) => { + const snippet = this._normalizePromptSnippet( + registeredTool.definition.promptSnippet ?? registeredTool.definition.description, + ); + return snippet ? ([registeredTool.definition.name, snippet] as const) : undefined; + }) + .filter((entry): entry is readonly [string, string] => entry !== undefined), + ); + const wrappedExtensionTools = this._extensionRunner + ? wrapRegisteredTools(allCustomTools, this._extensionRunner) + : []; + + const toolRegistry = new Map(this._baseToolRegistry); + for (const tool of wrappedExtensionTools as AgentTool[]) { + toolRegistry.set(tool.name, tool); + } + + if (this._extensionRunner) { + const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner); + this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool])); + } else { + this._toolRegistry = toolRegistry; + } + + const nextActiveToolNames = options?.activeToolNames + ? [...options.activeToolNames] + : [...previousActiveToolNames]; + + if (options?.includeAllExtensionTools) { + for (const tool of wrappedExtensionTools) { + nextActiveToolNames.push(tool.name); + } + } else if (!options?.activeToolNames) { + for (const toolName of this._toolRegistry.keys()) { + if (!previousRegistryNames.has(toolName)) { + nextActiveToolNames.push(toolName); + } + } + } + + this.setActiveToolsByName([...new Set(nextActiveToolNames)]); + } + private _buildRuntime(options: { activeToolNames?: string[]; flagValues?: Map; @@ -2012,52 +2086,14 @@ export class AgentSession { this._applyExtensionBindings(this._extensionRunner); } - const registeredTools = this._extensionRunner?.getAllRegisteredTools() ?? []; - const allCustomTools = [ - ...registeredTools, - ...this._customTools.map((def) => ({ definition: def, extensionPath: "" })), - ]; - const wrappedExtensionTools = this._extensionRunner - ? wrapRegisteredTools(allCustomTools, this._extensionRunner) - : []; - - const toolRegistry = new Map(this._baseToolRegistry); - for (const tool of wrappedExtensionTools as AgentTool[]) { - toolRegistry.set(tool.name, tool); - } - const defaultActiveToolNames = this._baseToolsOverride ? Object.keys(this._baseToolsOverride) : ["read", "bash", "edit", "write"]; const baseActiveToolNames = options.activeToolNames ?? defaultActiveToolNames; - const activeToolNameSet = new Set(baseActiveToolNames); - if (options.includeAllExtensionTools) { - for (const tool of wrappedExtensionTools as AgentTool[]) { - activeToolNameSet.add(tool.name); - } - } - - const extensionToolNames = new Set(wrappedExtensionTools.map((tool) => tool.name)); - const activeBaseTools = Array.from(activeToolNameSet) - .filter((name) => this._baseToolRegistry.has(name) && !extensionToolNames.has(name)) - .map((name) => this._baseToolRegistry.get(name) as AgentTool); - const activeExtensionTools = wrappedExtensionTools.filter((tool) => activeToolNameSet.has(tool.name)); - const activeToolsArray: AgentTool[] = [...activeBaseTools, ...activeExtensionTools]; - - if (this._extensionRunner) { - const wrappedActiveTools = wrapToolsWithExtensions(activeToolsArray, this._extensionRunner); - this.agent.setTools(wrappedActiveTools as AgentTool[]); - - const wrappedAllTools = wrapToolsWithExtensions(Array.from(toolRegistry.values()), this._extensionRunner); - this._toolRegistry = new Map(wrappedAllTools.map((tool) => [tool.name, tool])); - } else { - this.agent.setTools(activeToolsArray); - this._toolRegistry = toolRegistry; - } - - const systemPromptToolNames = Array.from(activeToolNameSet).filter((name) => this._baseToolRegistry.has(name)); - this._baseSystemPrompt = this._rebuildSystemPrompt(systemPromptToolNames); - this.agent.setSystemPrompt(this._baseSystemPrompt); + this._refreshToolRegistry({ + activeToolNames: baseActiveToolNames, + includeAllExtensionTools: options.includeAllExtensionTools, + }); } async reload(): Promise { diff --git a/packages/coding-agent/src/core/extensions/loader.ts b/packages/coding-agent/src/core/extensions/loader.ts index 25825a02..ceaaf28b 100644 --- a/packages/coding-agent/src/core/extensions/loader.ts +++ b/packages/coding-agent/src/core/extensions/loader.ts @@ -119,6 +119,8 @@ export function createExtensionRuntime(): ExtensionRuntime { getActiveTools: notInitialized, getAllTools: notInitialized, setActiveTools: notInitialized, + // registerTool() is valid during extension load; refresh is only needed post-bind. + refreshTools: () => {}, getCommands: notInitialized, setModel: () => Promise.reject(new Error("Extension runtime not initialized")), getThinkingLevel: notInitialized, @@ -162,6 +164,7 @@ function createExtensionAPI( definition: tool, extensionPath: extension.path, }); + runtime.refreshTools(); }, registerCommand(name: string, options: Omit): void { diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 819f006d..ede4c04b 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -244,6 +244,7 @@ export class ExtensionRunner { this.runtime.getActiveTools = actions.getActiveTools; this.runtime.getAllTools = actions.getAllTools; this.runtime.setActiveTools = actions.setActiveTools; + this.runtime.refreshTools = actions.refreshTools; this.runtime.getCommands = actions.getCommands; this.runtime.setModel = actions.setModel; this.runtime.getThinkingLevel = actions.getThinkingLevel; diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 97c1a295..aa88e0ee 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -339,6 +339,8 @@ export interface ToolDefinition SlashCommandInfo[]; export type SetActiveToolsHandler = (toolNames: string[]) => void; +export type RefreshToolsHandler = () => void; + export type SetModelHandler = (model: Model) => Promise; export type GetThinkingLevelHandler = () => ThinkingLevel; @@ -1291,6 +1295,7 @@ export interface ExtensionActions { getActiveTools: GetActiveToolsHandler; getAllTools: GetAllToolsHandler; setActiveTools: SetActiveToolsHandler; + refreshTools: RefreshToolsHandler; getCommands: GetCommandsHandler; setModel: SetModelHandler; getThinkingLevel: GetThinkingLevelHandler; diff --git a/packages/coding-agent/src/core/system-prompt.ts b/packages/coding-agent/src/core/system-prompt.ts index 368ac739..51b72f00 100644 --- a/packages/coding-agent/src/core/system-prompt.ts +++ b/packages/coding-agent/src/core/system-prompt.ts @@ -21,6 +21,8 @@ export interface BuildSystemPromptOptions { customPrompt?: string; /** Tools to include in prompt. Default: [read, bash, edit, write] */ selectedTools?: string[]; + /** Optional one-line tool snippets keyed by tool name. */ + toolSnippets?: Record; /** Text to append to system prompt. */ appendSystemPrompt?: string; /** Working directory. Default: process.cwd() */ @@ -36,6 +38,7 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin const { customPrompt, selectedTools, + toolSnippets, appendSystemPrompt, cwd, contextFiles: providedContextFiles, @@ -94,9 +97,18 @@ export function buildSystemPrompt(options: BuildSystemPromptOptions = {}): strin const docsPath = getDocsPath(); const examplesPath = getExamplesPath(); - // Build tools list based on selected tools (only built-in tools with known descriptions) - const tools = (selectedTools || ["read", "bash", "edit", "write"]).filter((t) => t in toolDescriptions); - const toolsList = tools.length > 0 ? tools.map((t) => `- ${t}: ${toolDescriptions[t]}`).join("\n") : "(none)"; + // Build tools list based on selected tools. + // Built-ins use toolDescriptions. Custom tools can provide one-line snippets. + const tools = selectedTools || ["read", "bash", "edit", "write"]; + const toolsList = + tools.length > 0 + ? tools + .map((name) => { + const snippet = toolSnippets?.[name] ?? toolDescriptions[name] ?? name; + return `- ${name}: ${snippet}`; + }) + .join("\n") + : "(none)"; // Build guidelines based on which tools are actually available const guidelinesList: string[] = []; diff --git a/packages/coding-agent/test/agent-session-dynamic-tools.test.ts b/packages/coding-agent/test/agent-session-dynamic-tools.test.ts new file mode 100644 index 00000000..d9a0223e --- /dev/null +++ b/packages/coding-agent/test/agent-session-dynamic-tools.test.ts @@ -0,0 +1,75 @@ +import { existsSync, mkdirSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { getModel } from "@mariozechner/pi-ai"; +import { Type } from "@sinclair/typebox"; +import { afterEach, beforeEach, describe, expect, it } from "vitest"; +import { DefaultResourceLoader } from "../src/core/resource-loader.js"; +import { createAgentSession } from "../src/core/sdk.js"; +import { SessionManager } from "../src/core/session-manager.js"; +import { SettingsManager } from "../src/core/settings-manager.js"; + +describe("AgentSession dynamic tool registration", () => { + let tempDir: string; + let agentDir: string; + + beforeEach(() => { + tempDir = join(tmpdir(), `pi-dynamic-tool-test-${Date.now()}-${Math.random().toString(36).slice(2)}`); + agentDir = join(tempDir, "agent"); + mkdirSync(agentDir, { recursive: true }); + }); + + afterEach(() => { + if (tempDir && existsSync(tempDir)) { + rmSync(tempDir, { recursive: true, force: true }); + } + }); + + it("refreshes tool registry when tools are registered after initialization", async () => { + const settingsManager = SettingsManager.create(tempDir, agentDir); + const sessionManager = SessionManager.inMemory(); + + const resourceLoader = new DefaultResourceLoader({ + cwd: tempDir, + agentDir, + settingsManager, + extensionFactories: [ + (pi) => { + pi.on("session_start", () => { + pi.registerTool({ + name: "dynamic_tool", + label: "Dynamic Tool", + description: "Tool registered from session_start", + promptSnippet: "Run dynamic test behavior", + parameters: Type.Object({}), + execute: async () => ({ + content: [{ type: "text", text: "ok" }], + details: {}, + }), + }); + }); + }, + ], + }); + await resourceLoader.reload(); + + const { session } = await createAgentSession({ + cwd: tempDir, + agentDir, + model: getModel("anthropic", "claude-sonnet-4-5")!, + settingsManager, + sessionManager, + resourceLoader, + }); + + expect(session.getAllTools().map((tool) => tool.name)).not.toContain("dynamic_tool"); + + await session.bindExtensions({}); + + expect(session.getAllTools().map((tool) => tool.name)).toContain("dynamic_tool"); + expect(session.getActiveToolNames()).toContain("dynamic_tool"); + expect(session.systemPrompt).toContain("- dynamic_tool: Run dynamic test behavior"); + + session.dispose(); + }); +}); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 551eee54..ae900447 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -60,6 +60,7 @@ describe("ExtensionRunner", () => { getActiveTools: () => [], getAllTools: () => [], setActiveTools: () => {}, + refreshTools: () => {}, getCommands: () => [], setModel: async () => false, getThinkingLevel: () => "off", diff --git a/packages/coding-agent/test/system-prompt.test.ts b/packages/coding-agent/test/system-prompt.test.ts index af20f155..586bf8f8 100644 --- a/packages/coding-agent/test/system-prompt.test.ts +++ b/packages/coding-agent/test/system-prompt.test.ts @@ -37,4 +37,19 @@ describe("buildSystemPrompt", () => { expect(prompt).toContain("- write:"); }); }); + + describe("custom tool snippets", () => { + test("includes custom tools in available tools section", () => { + const prompt = buildSystemPrompt({ + selectedTools: ["read", "dynamic_tool"], + toolSnippets: { + dynamic_tool: "Run dynamic test behavior", + }, + contextFiles: [], + skills: [], + }); + + expect(prompt).toContain("- dynamic_tool: Run dynamic test behavior"); + }); + }); });