From ee64be00ed94d898f73378bf4dd06732a59384cc Mon Sep 17 00:00:00 2001 From: Sarah Wooders Date: Tue, 10 Feb 2026 11:30:18 -0800 Subject: [PATCH] Add custom tools support MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add custom tools support - Add AgentTool, AgentToolResult, AnyAgentTool types - Add tools option to SessionOptions - Add tool-helpers.ts with jsonResult, readStringParam, etc. - Implement tool registration via register_external_tools control request - Handle execute_external_tool requests in stream() 🐾 Generated with [Letta Code](https://letta.com) --- src/index.ts | 15 +++++ src/session.ts | 151 ++++++++++++++++++++++++++++++++++++++++++-- src/tool-helpers.ts | 147 ++++++++++++++++++++++++++++++++++++++++++ src/types.ts | 87 +++++++++++++++++++++++++ 4 files changed, 395 insertions(+), 5 deletions(-) create mode 100644 src/tool-helpers.ts diff --git a/src/index.ts b/src/index.ts index e2d9bd8..6486991 100644 --- a/src/index.ts +++ b/src/index.ts @@ -54,10 +54,25 @@ export type { ImageContent, MessageContentItem, SendMessage, + // Tool types + AgentTool, + AgentToolResult, + AgentToolResultContent, + AgentToolUpdateCallback, + AnyAgentTool, } from "./types.js"; export { Session } from "./session.js"; +// Tool helpers +export { + jsonResult, + readStringParam, + readNumberParam, + readBooleanParam, + readStringArrayParam, +} from "./tool-helpers.js"; + /** * Create a new agent with a default conversation. * Returns the agentId which can be used with resumeSession or createSession. diff --git a/src/session.ts b/src/session.ts index f803e2a..3d99f4c 100644 --- a/src/session.ts +++ b/src/session.ts @@ -19,6 +19,8 @@ import type { CanUseToolResponseAllow, CanUseToolResponseDeny, SendMessage, + AnyAgentTool, + ExecuteExternalToolRequest, } from "./types.js"; @@ -33,6 +35,7 @@ export class Session implements AsyncDisposable { private _sessionId: string | null = null; private _conversationId: string | null = null; private initialized = false; + private externalTools: Map = new Map(); constructor( @@ -40,6 +43,13 @@ export class Session implements AsyncDisposable { ) { // Note: Validation happens in public API functions (createSession, createAgent, etc.) this.transport = new SubprocessTransport(options); + + // Store external tools in a map for quick lookup + if (options.tools) { + for (const tool of options.tools) { + this.externalTools.set(tool.name, tool); + } + } } /** @@ -78,7 +88,18 @@ export class Session implements AsyncDisposable { this._conversationId = initMsg.conversation_id; this.initialized = true; - sessionLog("init", `initialized: agent=${initMsg.agent_id} conversation=${initMsg.conversation_id} model=${initMsg.model} tools=${initMsg.tools?.length || 0}`); + // Register external tools with CLI + if (this.externalTools.size > 0) { + await this.registerExternalTools(); + } + + // Include external tool names in the tools list + const allTools = [ + ...initMsg.tools, + ...Array.from(this.externalTools.keys()), + ]; + + sessionLog("init", `initialized: agent=${initMsg.agent_id} conversation=${initMsg.conversation_id} model=${initMsg.model} tools=${allTools.length} (${this.externalTools.size} external)`); return { type: "init", @@ -86,7 +107,7 @@ export class Session implements AsyncDisposable { sessionId: initMsg.session_id, conversationId: initMsg.conversation_id, model: initMsg.model, - tools: initMsg.tools, + tools: allTools, }; } } @@ -140,17 +161,34 @@ export class Session implements AsyncDisposable { sessionLog("stream", `starting stream (agent=${this._agentId}, conversation=${this._conversationId})`); for await (const wireMsg of this.transport.messages()) { - // Handle CLI → SDK control requests (e.g., can_use_tool) + // Handle CLI → SDK control requests (e.g., can_use_tool, execute_external_tool) if (wireMsg.type === "control_request") { const controlReq = wireMsg as ControlRequest; - sessionLog("stream", `control_request: subtype=${controlReq.request.subtype} tool=${(controlReq.request as CanUseToolControlRequest).tool_name || "N/A"}`); - if (controlReq.request.subtype === "can_use_tool") { + // Widen to string to allow SDK-extension subtypes not in the protocol union + const subtype: string = controlReq.request.subtype; + sessionLog("stream", `control_request: subtype=${subtype} tool=${(controlReq.request as CanUseToolControlRequest).tool_name || "N/A"}`); + + if (subtype === "can_use_tool") { await this.handleCanUseTool( controlReq.request_id, controlReq.request as CanUseToolControlRequest ); continue; } + if (subtype === "execute_external_tool") { + // SDK extension: not in protocol ControlRequestBody union, extract fields via Record + const rawReq = controlReq.request as Record; + await this.handleExecuteExternalTool( + controlReq.request_id, + { + subtype: "execute_external_tool", + tool_call_id: rawReq.tool_call_id as string, + tool_name: rawReq.tool_name as string, + input: rawReq.input as Record, + } + ); + continue; + } } const sdkMsg = this.transformMessage(wireMsg); @@ -178,6 +216,109 @@ export class Session implements AsyncDisposable { } } + /** + * Register external tools with the CLI + */ + private async registerExternalTools(): Promise { + const toolDefs = Array.from(this.externalTools.values()).map((tool) => ({ + name: tool.name, + label: tool.label, + description: tool.description, + // Convert TypeBox schema to plain JSON Schema + parameters: this.schemaToJsonSchema(tool.parameters), + })); + + sessionLog("registerTools", `registering ${toolDefs.length} external tools: ${toolDefs.map(t => t.name).join(", ")}`); + + await this.transport.write({ + type: "control_request", + request_id: `register_tools_${Date.now()}`, + request: { + subtype: "register_external_tools", + tools: toolDefs, + }, + }); + } + + /** + * Convert TypeBox schema to JSON Schema + */ + private schemaToJsonSchema(schema: unknown): Record { + // TypeBox schemas are already JSON Schema compatible + // Just need to extract the schema object + if (schema && typeof schema === "object") { + // TypeBox schemas have these JSON Schema properties + const s = schema as Record; + return { + type: s.type, + properties: s.properties, + required: s.required, + additionalProperties: s.additionalProperties, + description: s.description, + }; + } + return { type: "object" }; + } + + /** + * Handle execute_external_tool control request from CLI + */ + private async handleExecuteExternalTool( + requestId: string, + req: ExecuteExternalToolRequest + ): Promise { + const tool = this.externalTools.get(req.tool_name); + + if (!tool) { + // Tool not found - send error result + sessionLog("executeExternalTool", `ERROR: unknown tool ${req.tool_name}`); + await this.transport.write({ + type: "control_response", + response: { + subtype: "external_tool_result", + request_id: requestId, + tool_call_id: req.tool_call_id, + content: [{ type: "text", text: `Unknown external tool: ${req.tool_name}` }], + is_error: true, + }, + }); + return; + } + + try { + sessionLog("executeExternalTool", `executing ${req.tool_name} (call_id=${req.tool_call_id})`); + // Execute the tool + const result = await tool.execute(req.tool_call_id, req.input); + + // Send success result + await this.transport.write({ + type: "control_response", + response: { + subtype: "external_tool_result", + request_id: requestId, + tool_call_id: req.tool_call_id, + content: result.content, + is_error: false, + }, + }); + sessionLog("executeExternalTool", `${req.tool_name} completed successfully`); + } catch (err) { + // Send error result + const errorMessage = err instanceof Error ? err.message : String(err); + sessionLog("executeExternalTool", `${req.tool_name} failed: ${errorMessage}`); + await this.transport.write({ + type: "control_response", + response: { + subtype: "external_tool_result", + request_id: requestId, + tool_call_id: req.tool_call_id, + content: [{ type: "text", text: `Tool execution error: ${errorMessage}` }], + is_error: true, + }, + }); + } + } + /** * Handle can_use_tool control request from CLI (Claude SDK compatible format) */ diff --git a/src/tool-helpers.ts b/src/tool-helpers.ts new file mode 100644 index 0000000..6a24ce1 --- /dev/null +++ b/src/tool-helpers.ts @@ -0,0 +1,147 @@ +/** + * Tool Helpers + * + * Helper functions for creating tool results and parsing parameters. + * Matches the API from pi-coding-agent. + */ + +import type { AgentToolResult } from "./types.js"; + +/** + * Create a JSON tool result + */ +export function jsonResult(payload: unknown): AgentToolResult { + return { + content: [ + { + type: "text", + text: JSON.stringify(payload, null, 2), + }, + ], + details: payload, + }; +} + +/** + * Options for reading string parameters + */ +export interface StringParamOptions { + required?: boolean; + trim?: boolean; + label?: string; + allowEmpty?: boolean; +} + +/** + * Read a string parameter from tool args + */ +export function readStringParam( + params: Record, + key: string, + options?: StringParamOptions & { required: true }, +): string; +export function readStringParam( + params: Record, + key: string, + options?: StringParamOptions, +): string | undefined; +export function readStringParam( + params: Record, + key: string, + options: StringParamOptions = {}, +): string | undefined { + const { required = false, trim = true, label = key, allowEmpty = false } = options; + const raw = params[key]; + if (typeof raw !== "string") { + if (required) throw new Error(`${label} required`); + return undefined; + } + const value = trim ? raw.trim() : raw; + if (!value && !allowEmpty) { + if (required) throw new Error(`${label} required`); + return undefined; + } + return value; +} + +/** + * Read a number parameter from tool args + */ +export function readNumberParam( + params: Record, + key: string, + options: { required?: boolean; label?: string; integer?: boolean } = {}, +): number | undefined { + const { required = false, label = key, integer = false } = options; + const raw = params[key]; + let value: number | undefined; + if (typeof raw === "number" && Number.isFinite(raw)) { + value = raw; + } else if (typeof raw === "string") { + const trimmed = raw.trim(); + if (trimmed) { + const parsed = Number.parseFloat(trimmed); + if (Number.isFinite(parsed)) value = parsed; + } + } + if (value === undefined) { + if (required) throw new Error(`${label} required`); + return undefined; + } + return integer ? Math.trunc(value) : value; +} + +/** + * Read a boolean parameter from tool args + */ +export function readBooleanParam( + params: Record, + key: string, + options: { required?: boolean; label?: string } = {}, +): boolean | undefined { + const { required = false, label = key } = options; + const raw = params[key]; + if (typeof raw === "boolean") { + return raw; + } + if (typeof raw === "string") { + const lower = raw.toLowerCase().trim(); + if (lower === "true" || lower === "1" || lower === "yes") return true; + if (lower === "false" || lower === "0" || lower === "no") return false; + } + if (required) throw new Error(`${label} required`); + return undefined; +} + +/** + * Read a string array parameter from tool args + */ +export function readStringArrayParam( + params: Record, + key: string, + options: StringParamOptions = {}, +): string[] | undefined { + const { required = false, label = key } = options; + const raw = params[key]; + if (Array.isArray(raw)) { + const values = raw + .filter((entry): entry is string => typeof entry === "string") + .map((entry) => entry.trim()) + .filter(Boolean); + if (values.length === 0) { + if (required) throw new Error(`${label} required`); + return undefined; + } + return values; + } + if (typeof raw === "string") { + const value = raw.trim(); + if (!value) { + if (required) throw new Error(`${label} required`); + return undefined; + } + return [value]; + } + if (required) throw new Error(`${label} required`); + return undefined; +} diff --git a/src/types.ts b/src/types.ts index 45517a7..b6471f2 100644 --- a/src/types.ts +++ b/src/types.ts @@ -115,6 +115,64 @@ export type MemoryItem = */ export type MemoryPreset = "persona" | "human" | "skills" | "loaded_skills"; +// ═══════════════════════════════════════════════════════════════ +// TOOL TYPES (matches pi-agent-core) +// ═══════════════════════════════════════════════════════════════ + +/** + * Tool result content block + */ +export interface AgentToolResultContent { + type: "text" | "image"; + text?: string; + data?: string; // base64 for images + mimeType?: string; +} + +/** + * Tool result (matches pi-agent-core) + */ +export interface AgentToolResult { + content: AgentToolResultContent[]; + details?: T; +} + +/** + * Tool update callback (for streaming tool progress) + */ +export type AgentToolUpdateCallback = (update: Partial>) => void; + +/** + * Agent tool definition (matches pi-agent-core) + */ +export interface AgentTool { + /** Display label */ + label: string; + + /** Tool name (used in API calls) */ + name: string; + + /** Description shown to the model */ + description: string; + + /** JSON Schema for parameters (TypeBox or plain object) */ + parameters: TParams; + + /** Execution function */ + execute: ( + toolCallId: string, + args: unknown, + signal?: AbortSignal, + onUpdate?: AgentToolUpdateCallback, + ) => Promise>; +} + +/** + * Convenience type for tools with any params + */ +// eslint-disable-next-line @typescript-eslint/no-explicit-any +export type AnyAgentTool = AgentTool; + // ═══════════════════════════════════════════════════════════════ // SESSION OPTIONS // ═══════════════════════════════════════════════════════════════ @@ -155,6 +213,9 @@ export interface InternalSessionOptions { permissionMode?: PermissionMode; canUseTool?: CanUseToolCallback; + // Custom tools + tools?: AnyAgentTool[]; + // Process settings cwd?: string; } @@ -183,6 +244,12 @@ export interface CreateSessionOptions { /** Custom permission callback - called when tool needs approval */ canUseTool?: CanUseToolCallback; + + /** + * Custom tools that execute locally in the SDK process. + * These tools are registered with the CLI and executed when the LLM calls them. + */ + tools?: AnyAgentTool[]; } /** @@ -228,6 +295,12 @@ export interface CreateAgentOptions { /** Custom permission callback - called when tool needs approval */ canUseTool?: CanUseToolCallback; + + /** + * Custom tools that execute locally in the SDK process. + * These tools are registered with the CLI and executed when the LLM calls them. + */ + tools?: AnyAgentTool[]; } // ═══════════════════════════════════════════════════════════════ @@ -305,3 +378,17 @@ export type SDKMessage = | SDKReasoningMessage | SDKResultMessage | SDKStreamEventMessage; + +// ═══════════════════════════════════════════════════════════════ +// EXTERNAL TOOL PROTOCOL TYPES +// ═══════════════════════════════════════════════════════════════ + +/** + * Request to execute an external tool (CLI → SDK) + */ +export interface ExecuteExternalToolRequest { + subtype: "execute_external_tool"; + tool_call_id: string; + tool_name: string; + input: Record; +}