Add custom tools support

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)
2026-02-10 11:30:18 -08:00
parent b0b82a2a3d
commit ee64be00ed
4 changed files with 395 additions and 5 deletions
--- 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.
--- 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<string, AnyAgentTool> = 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<string, unknown>;
+          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<string, unknown>,
+            }
+          );
+          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<void> {
+    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<string, unknown> {
+    // 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<string, unknown>;
+      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<void> {
+    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)
   */
--- a/src/tool-helpers.ts
+++ 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<unknown> {
+  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<string, unknown>,
+  key: string,
+  options?: StringParamOptions & { required: true },
+): string;
+export function readStringParam(
+  params: Record<string, unknown>,
+  key: string,
+  options?: StringParamOptions,
+): string | undefined;
+export function readStringParam(
+  params: Record<string, unknown>,
+  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<string, unknown>,
+  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<string, unknown>,
+  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<string, unknown>,
+  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;
+}
--- 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<T> {
+  content: AgentToolResultContent[];
+  details?: T;
+}
+
+/**
+ * Tool update callback (for streaming tool progress)
+ */
+export type AgentToolUpdateCallback<T> = (update: Partial<AgentToolResult<T>>) => void;
+
+/**
+ * Agent tool definition (matches pi-agent-core)
+ */
+export interface AgentTool<TParams, TResult> {
+  /** 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<TResult>,
+  ) => Promise<AgentToolResult<TResult>>;
+}
+
+/**
+ * Convenience type for tools with any params
+ */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export type AnyAgentTool = AgentTool<any, unknown>;
+
 // ═══════════════════════════════════════════════════════════════
 // 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<string, unknown>;
+}