Learn Agents from Pi
中文
On this page

14. The Extension System

Once your agent starts being used by different teams, customization requests never stop: add an internal search tool, intercept dangerous commands, change the system prompt, show token cost in the status bar, save a report after a task completes. Merging all of these into the core would bloat the kernel. The goal of an extension system is to productize the points of customization.

An extension system is not a scripts directory

An extension system must provide at least three kinds of capability:

  • Lifecycle hooks: listen to session, turn, model request, tool call, compaction, and shutdown events.
  • Registration: register tools, commands, keyboard shortcuts, CLI flags, status displays, and message renderers.
  • Runtime context: access to cwd, the session, events, the UI, configuration, and safety APIs.

Extensions should influence the agent through controlled APIs, not by mutating internal objects directly. Otherwise a single extension can corrupt the log, the events, and the permission model. A common shape is a factory function: the extension exports a function, the runtime passes in a pi object, and the extension registers its hooks and capabilities on it.

export default function myExtension(pi: ExtensionApi): void {
  pi.on("tool_call", async (event, ctx) => {
    if (isDangerous(event.call)) {
      return { block: "This command is blocked by policy." };
    }
  });
  pi.registerTool(internalSearchTool);
  pi.registerCommand("report", { handler: writeReport });
}

Lifecycle hooks

A practical set of hooks might include:

project_trust
session_start
before_agent_start
before_model_request
after_model_response
tool_call
tool_result
session_before_compact
session_shutdown

Some hooks are purely notifications; others can block or rewrite. The semantics of each hook must be explicit. tool_call can return block to prevent tool execution or rewrite the arguments; before_agent_start can modify messages before the model request; session_before_compact can supply a custom summary; notification hooks like turn_start are not allowed to rewrite state.

Blocking hooks fail closed by default: when an extension errors, it is better to block a dangerous operation than to proceed. This matters most for permission-related hooks — a safety extension that throws an exception, if treated as "allow," is the same as having no safety extension at all. Conversely, a purely observational hook (metrics, logging) that errors should not drag down the whole turn; catch it and log it.

Registering tools

When an extension registers a tool, it must provide a definition of the same quality as a built-in tool: name, description, schema, execute function, output truncation, error semantics, and an optional UI renderer. A custom tool that writes files must participate in the same file-write queue; otherwise it and the built-in edit/write will clobber each other. Tools that declare an execution mode (parallel-safe or must run serially) are scheduled uniformly by the kernel executor, so the extension does not need to add its own locking.

Tool descriptions must also enter the "available tools" section of the system prompt. Otherwise the model won't know when to use the extension tool. A description should name the tool explicitly — don't write "use this tool," because in a flattened prompt the model may not know what "this" refers to.

Declarative extension surfaces: skills, commands, and prompt templates

Not every extension point needs code. A mature system turns several high-frequency customizations into declarative formats, so users who don't write TypeScript can still extend the agent:

  • skills: a Markdown file with YAML frontmatter (a conventional name is SKILL.md), whose frontmatter declares at least a name and a description. The name and description need length limits (for example, name no longer than 64 characters, description no longer than 1024), because they enter the skills catalog in the system prompt — an unbounded description quietly eats the context budget. A skill's body is an instruction the model loads on demand and does not keep resident in the context.
  • slash commands: bind a common prompt or a fixed series of actions to a /name, expanding into the corresponding user message or action when the user types the command.
  • prompt templates: reusable prompts with placeholders, filled in with parameters at runtime and injected as user input.

What these three share is that they are data, not code, which makes them easier to distribute, review, and constrain. But they are still bound by the same boundaries — the skills catalog is only injected when its loader tool is active (Chapter 11), and project-local declarative resources are equally gated by project trust (an untrusted project does not automatically load its SKILL.md). Declarative does not mean harmless: the text in a skill file can be prompt injection just as easily, and the runtime must treat it as project content rather than a system instruction.

Custom message types

Extensions often need to place things in the session that aren't standard conversation: a local command execution, a generated report, a notification. This relies on two capabilities laid down in Chapters 6 and 8 — the custom message entry, plus the convertToLlm projection hook. An extension declares a custom message type, specifying whether it enters the model context and how it renders; at projection time, the kernel filters out the types that shouldn't reach the context and rewrites the types that should be converted into standard messages.

This lets an extension persist its own state (written into the session log, still there after resuming) without polluting the context sent to the model. The UI paints these custom entries into an appropriate shape through the accompanying message renderer.

Extensions and safety

An extension is generally arbitrary code. Installing an extension means trusting it with every resource within the process's permissions. This is exactly where Chapter 12's project trust earns its keep: an untrusted project should not automatically load and execute project-local extensions. The product must tell users this fact and distinguish project extensions, user-global extensions, and built-in extensions — the three carry different trust levels.

The extension API must also prevent bypassing the permission gate. If an extension registers a dangerous_shell tool, that tool must still pass through the tool call permission check. Otherwise a user who has disabled built-in bash gets bypassed by an extension tool. Extensions can extend capabilities, but they cannot extend the permission model with exceptions.

The self-hosting test

The way to judge whether the extension API is good enough is to implement part of the product's own capabilities with it. For example:

  • The status bar is registered through the extension API.
  • Permission confirmation is implemented through the tool_call hook.
  • Custom compaction is implemented through session_before_compact.
  • Dangerous-command interception is implemented through the tool_call hook.
  • Automatic checkpoints are implemented through turn_start and session_shutdown.

If any of these requirements forces a change to core code, the extension surface isn't complete enough yet. Mature systems often do build a fair number of their built-in features as extensions — this is both a stress test of the API and a way for third parties to replace or enhance those features.

Exercise

Implement a minimal extension system.

Acceptance criteria:

  • An extension can register a read-only tool, and the tool description enters the system prompt.
  • An extension can deny a bash command before the tool call runs (fail closed).
  • An extension can append a tool guide to the system prompt.
  • An extension can register a custom message type and decide, through the projection hook, whether it enters the model context.
  • Extension errors become explicit events, and an observational hook that errors does not drag down the turn.
  • In an untrusted project, project extensions do not run automatically.