How to Write a Prime Agent Extension with Custom Tools and Event Handlers

To write a Prime Agent extension, export a default factory function that receives an ExtensionAPI object, then call pi.registerTool() to add custom capabilities and pi.on() to subscribe to lifecycle events like tool_call_started and tool_call_finished.

Prime Agent is a modular AI agent framework developed by PrimeIntellect-ai that allows developers to extend core functionality without modifying the source. By leveraging the ExtensionAPI defined in the core runtime, you can inject custom tools for the LLM to invoke and react to internal state changes through strongly-typed event handlers.

Understanding the Extension Architecture

The extension system centers on two critical files in the repository. The interface contracts live in packages/coding-agent/src/core/extensions/types.ts, which exports the ExtensionAPI type, ExtensionFactory signature, and event definitions. The runtime instantiation occurs in packages/coding-agent/src/core/extensions/loader.ts, where the framework constructs an ExtensionAPI instance for each extension and executes its exported factory function. This architecture ensures that extensions run in a controlled environment with explicit access to agent internals.

Creating Your First Extension File

An extension is a TypeScript or JavaScript module that exports a single default function conforming to the ExtensionFactory type. This function receives the ExtensionAPI object (commonly named pi) and uses it to register functionality.

Create a new file in your project or under packages/coding-agent/examples/extensions/:

import type { ExtensionAPI, ToolDefinition } from "@earendil-works/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // Extension logic executes here during agent startup
  console.log("Custom extension loaded");
}

The factory executes synchronously during the loading phase, so any async setup should be deferred until tool invocations or event handlers trigger.

Registering Custom Tools

Custom tools allow the LLM to execute code outside the base agent capabilities. In packages/coding-agent/examples/extensions/tools.ts, the pattern involves calling pi.registerTool() with a ToolDefinition object containing the tool name, description, JSON schema for arguments, and an async call implementation.

import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "execute_bash",
    description: "Runs a shell command on the host machine and returns stdout/stderr.",
    inputSchema: {
      type: "object",
      properties: {
        cmd: { type: "string", description: "The bash command to execute" }
      },
      required: ["cmd"]
    },
    async call({ args }) {
      const { cmd } = args;
      const { exec } = await import("node:child_process/promises");
      const { stdout, stderr } = await exec(cmd);
      
      return {
        output: stdout.trim(),
        error: stderr.trim() || undefined
      };
    }
  });
}

The call method receives a ToolInvocation object containing args validated against your inputSchema. It must return a ToolResult with output and optional error fields. Once registered, the tool appears in the agent's available function set and can be invoked via LLM tool calls.

Handling Extension Events

Extensions react to agent lifecycle events using pi.on(), which accepts an event name from the ExtensionEvent union and a handler callback. The reference implementation in packages/coding-agent/examples/extensions/working-indicator.ts demonstrates listening for tool_call_started to provide user feedback.

export default function (pi: ExtensionAPI) {
  // Notify when a tool begins execution
  pi.on("tool_call_started", (event) => {
    const { toolName, toolCallId } = event;
    pi.emitMessage({
      type: "status",
      text: `Starting tool: ${toolName} (${toolCallId})`
    });
  });

  // Clean up notification when tool completes
  pi.on("tool_call_finished", (event) => {
    const { toolName, result } = event;
    const status = result.error ? "failed" : "succeeded";
    pi.emitMessage({
      type: "status",
      text: `Tool ${toolName} ${status}`
    });
  });
}

Available events include session_info_changed (for metadata updates), extension_ui_request (for interactive dialogs), and tool_call_started/tool_call_finished (for execution instrumentation). Handlers receive strongly-typed payloads defined in packages/coding-agent/src/core/extensions/types.ts.

Loading and Testing Your Extension

The runtime discovers extensions through the loader mechanism in packages/coding-agent/src/core/extensions/loader.ts. To activate your extension, pass the file path via the --extension CLI flag when starting the agent:

prime-agent --extension ./packages/coding-agent/examples/extensions/tools.ts

Alternatively, add the path to the extensions array in your runtime configuration file. The loader imports your module, instantiates the ExtensionAPI, and executes the factory function, binding all registered tools and event listeners into the active session.

For validation, use the test utilities in packages/coding-agent/test/utilities.ts to programmatically load your extension via loadExtensionFromFactory and assert that session.getActiveToolNames() includes your custom tool names.

Summary

• Export a factory function that accepts ExtensionAPI (the pi object) as its sole argument.
• Register tools using pi.registerTool() with a name, description, inputSchema, and async call method.
• Subscribe to events via pi.on() using event names like tool_call_started or session_info_changed.
• Reference core types in packages/coding-agent/src/core/extensions/types.ts and loading logic in loader.ts for implementation details.
• Load extensions at runtime using the --extension CLI flag or configuration array.

Frequently Asked Questions

What is the ExtensionAPI in Prime Agent?

The ExtensionAPI is the primary interface injected into every extension factory function, defined in packages/coding-agent/src/core/extensions/types.ts. It provides methods like registerTool() for adding capabilities and on() for event subscription, acting as the bridge between third-party code and the agent's internal session management.

How do I register a custom tool in a Prime Agent extension?

Call pi.registerTool() inside your exported factory function, passing a ToolDefinition object that includes the tool's name, description, JSON inputSchema, and an async call method that executes the logic. The tool becomes immediately available to the LLM once the extension loads.

Which events can a Prime Agent extension listen to?

Extensions can subscribe to events defined in the ExtensionEvent union, including tool_call_started, tool_call_finished, session_info_changed, and extension_ui_request. Each event emits a specific payload documented in packages/coding-agent/src/core/extensions/types.ts, allowing handlers to react to execution state or user interactions.

How do I load a Prime Agent extension at runtime?

Pass the absolute or relative path to your extension file using the --extension command-line flag when launching the agent, or include the path in the extensions configuration array. The loader in packages/coding-agent/src/core/extensions/loader.ts will import the module and execute the default factory function with a fresh ExtensionAPI instance.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}