Extension System

Overview

Pi’s extension system allows you to customize and extend the agent’s behavior through event hooks, custom tools, commands, and UI components. Extensions are TypeScript modules that run in the same process as the agent.

Creating an Extension

Basic Structure

Extensions export a factory function that receives the Pi API:

// ~/.pi/agent/extensions/my-extension.ts
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";

export default function (pi: ExtensionAPI) {
  // Extension initialization code here
  
  pi.on("session_start", async (event, ctx) => {
    console.log("Session started");
  });
}

Extension API

The pi object provides the full extension API:

interface ExtensionAPI {
  // Event subscription
  on(event: string, handler: ExtensionHandler): void;
  
  // Tool registration
  registerTool(tool: ToolDefinition): void;
  
  // Command registration
  registerCommand(name: string, options: RegisteredCommand): void;
  
  // UI and messaging
  sendMessage(message: CustomMessage, options?: MessageOptions): void;
  sendUserMessage(content: string | Content[], options?: MessageOptions): void;
  
  // State management
  appendEntry(customType: string, data?: unknown): void;
  setSessionName(name: string): void;
  getSessionName(): string | undefined;
  setLabel(entryId: string, label: string | undefined): void;
  
  // Model and tools
  setModel(model: Model): Promise<boolean>;
  getThinkingLevel(): ThinkingLevel;
  setThinkingLevel(level: ThinkingLevel): void;
  getActiveTools(): string[];
  getAllTools(): ToolInfo[];
  setActiveTools(toolNames: string[]): void;
  
  // Provider registration
  registerProvider(name: string, config: ProviderConfig): void;
  
  // Shared event bus
  events: EventBus;
}

Location: packages/coding-agent/src/core/extensions/types.ts:916

Event Hooks

Session Lifecycle

pi.on("session_start", async (event, ctx) => {
  // Called when session loads
  // Use to restore extension state from session entries
});

pi.on("session_before_switch", async (event, ctx) => {
  // Can cancel session switches
  if (shouldBlock) {
    return { cancel: true };
  }
});

pi.on("session_switch", async (event, ctx) => {
  // Called after switching sessions
  console.log("Switched from:", event.previousSessionFile);
});

pi.on("session_shutdown", async (event, ctx) => {
  // Cleanup on exit
});

Agent Lifecycle

pi.on("before_agent_start", async (event, ctx) => {
  // Inject context before LLM call
  return {
    message: {
      customType: "context",
      content: "Additional context for the LLM",
      display: true,
    },
    systemPrompt: modifiedSystemPrompt,  // Optional
  };
});

pi.on("agent_start", async (event, ctx) => {
  // Agent loop started
});

pi.on("turn_start", async (event, ctx) => {
  // New LLM call started
  console.log("Turn index:", event.turnIndex);
});

pi.on("turn_end", async (event, ctx) => {
  // LLM call completed
  console.log("Message:", event.message);
  console.log("Tool results:", event.toolResults);
});

pi.on("agent_end", async (event, ctx) => {
  // Agent loop finished
});

Message Events

pi.on("message_start", async (event, ctx) => {
  // Message started (user, assistant, or toolResult)
});

pi.on("message_update", async (event, ctx) => {
  // Token streamed (assistant messages only)
  if (event.assistantMessageEvent.type === "text") {
    process.stdout.write(event.assistantMessageEvent.text);
  }
});

pi.on("message_end", async (event, ctx) => {
  // Message completed
});

Tool Events

pi.on("tool_call", async (event, ctx) => {
  // Block dangerous commands
  if (event.toolName === "bash" && event.input.command.includes("rm -rf /")) {
    return {
      block: true,
      reason: "Dangerous command blocked",
    };
  }
});

pi.on("tool_result", async (event, ctx) => {
  // Modify or redact tool results
  if (event.toolName === "read" && event.details?.filePath.endsWith(".env")) {
    return {
      content: [{ type: "text", text: "[REDACTED]" }],
    };
  }
});

pi.on("tool_execution_start", async (event, ctx) => {
  console.log(`Tool ${event.toolName} started`);
});

pi.on("tool_execution_end", async (event, ctx) => {
  console.log(`Tool ${event.toolName} completed`);
});

Context Manipulation

pi.on("context", async (event, ctx) => {
  // Modify messages before sending to LLM
  const messages = event.messages.filter(m => {
    // Remove sensitive messages
    return m.role !== "custom" || !m.content.includes("SECRET");
  });
  
  return { messages };
});

Input Interception

pi.on("input", async (event, ctx) => {
  // Transform user input
  if (event.text.startsWith("@")) {
    return {
      action: "transform",
      text: `Find references to: ${event.text.slice(1)}`,
    };
  }
  
  // Handle input completely
  if (event.text === "exit") {
    ctx.shutdown();
    return { action: "handled" };
  }
  
  return { action: "continue" };
});

Compaction Hooks

pi.on("session_before_compact", async (event, ctx) => {
  // Custom compaction logic
  const { preparation, signal } = event;
  
  // Generate custom summary
  const summary = await generateCustomSummary(
    preparation.messagesToSummarize,
    signal
  );
  
  return {
    compaction: {
      summary,
      firstKeptEntryId: preparation.firstKeptEntryId,
      tokensBefore: preparation.tokensBefore,
      details: { custom: "data" },
    },
  };
});

pi.on("session_compact", async (event, ctx) => {
  // React to compaction
  console.log("Compacted", event.compactionEntry.tokensBefore, "tokens");
});

Extension Context

All event handlers receive a context object:

interface ExtensionContext {
  // UI methods (if available)
  ui: ExtensionUIContext;
  hasUI: boolean;
  
  // Session state
  cwd: string;
  sessionManager: ReadonlySessionManager;
  model: Model | undefined;
  
  // Model registry
  modelRegistry: ModelRegistry;
  
  // Agent state
  isIdle(): boolean;
  abort(): void;
  hasPendingMessages(): boolean;
  
  // System control
  shutdown(): void;
  compact(options?: CompactOptions): void;
  
  // Context info
  getContextUsage(): ContextUsage | undefined;
  getSystemPrompt(): string;
}

Location: packages/coding-agent/src/core/extensions/types.ts:261

Custom Commands

pi.registerCommand("deploy", {
  description: "Deploy the application",
  handler: async (args, ctx) => {
    const environment = args || "production";
    
    ctx.ui.notify(`Deploying to ${environment}...`);
    
    const result = await ctx.ui.confirm(
      "Confirm Deployment",
      `Deploy to ${environment}?`
    );
    
    if (!result) {
      ctx.ui.notify("Deployment cancelled", "warning");
      return;
    }
    
    // Execute deployment
    pi.sendMessage({
      customType: "deployment",
      content: `Deploying to ${environment}`,
      display: true,
    });
  },
});

Command Context

Command handlers receive an extended context:

interface ExtensionCommandContext extends ExtensionContext {
  // Wait for agent to finish
  waitForIdle(): Promise<void>;
  
  // Session control
  newSession(options?: NewSessionOptions): Promise<{ cancelled: boolean }>;
  fork(entryId: string): Promise<{ cancelled: boolean }>;
  navigateTree(targetId: string, options?: NavigateOptions): Promise<{ cancelled: boolean }>;
  switchSession(sessionPath: string): Promise<{ cancelled: boolean }>;
  
  // Reload resources
  reload(): Promise<void>;
}

UI Integration

Dialogs

pi.registerCommand("select-env", {
  handler: async (args, ctx) => {
    const env = await ctx.ui.select(
      "Select Environment",
      ["development", "staging", "production"]
    );
    
    if (env) {
      await pi.sendMessage({
        customType: "env_change",
        content: `Environment set to: ${env}`,
        display: true,
      });
    }
  },
});

Status Bar

pi.on("session_start", async (event, ctx) => {
  // Set status in footer
  ctx.ui.setStatus("my-extension", "Ready");
});

pi.on("agent_start", async (event, ctx) => {
  ctx.ui.setStatus("my-extension", "Processing...");
});

pi.on("agent_end", async (event, ctx) => {
  ctx.ui.setStatus("my-extension", "Ready");
});

Widgets

import { box, text, list } from "@mariozechner/pi-tui";

pi.on("session_start", async (event, ctx) => {
  // Simple text widget
  ctx.ui.setWidget("info", ["Custom info widget", "Line 2"]);
  
  // Custom component widget
  ctx.ui.setWidget("stats", (tui, theme) => {
    return box(
      { border: "single", borderColor: theme.colors.primary },
      list([
        text("Tool calls: 42", { color: theme.colors.text }),
        text("Tokens: 12345", { color: theme.colors.text }),
      ])
    );
  });
});

Custom Messages

Extensions can send custom message types with custom rendering:

pi.registerMessageRenderer("deployment", (message, options, theme) => {
  const { environment, status } = message.details;
  const color = status === "success" ? theme.colors.success : theme.colors.error;
  
  return box(
    { border: "single", borderColor: color },
    text(`Deployment to ${environment}: ${status}`, { color })
  );
});

pi.sendMessage({
  customType: "deployment",
  content: "Deployment completed",
  display: true,
  details: { environment: "production", status: "success" },
});

State Persistence

Extensions can persist state across sessions:

interface MyState {
  counter: number;
  lastAction: string;
}

pi.on("session_start", async (event, ctx) => {
  // Restore state from session entries
  const entries = ctx.sessionManager.getEntries();
  const myEntries = entries.filter(
    e => e.type === "custom" && e.customType === "my-extension"
  );
  
  const state: MyState = myEntries.length > 0
    ? myEntries[myEntries.length - 1].data
    : { counter: 0, lastAction: "none" };
  
  console.log("Restored state:", state);
});

pi.on("turn_end", async (event, ctx) => {
  // Persist state after each turn
  const state: MyState = {
    counter: getCurrentCounter(),
    lastAction: event.message.role,
  };
  
  pi.appendEntry("my-extension", state);
});

Custom Providers

pi.registerProvider("my-proxy", {
  baseUrl: "https://proxy.example.com",
  apiKey: "PROXY_API_KEY",
  api: "anthropic-messages",
  models: [
    {
      id: "claude-sonnet-4-20250514",
      name: "Claude 4 Sonnet (proxy)",
      reasoning: false,
      input: ["text", "image"],
      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
      contextWindow: 200000,
      maxTokens: 16384,
    },
  ],
});

OAuth Providers

pi.registerProvider("corporate-ai", {
  baseUrl: "https://ai.corp.com",
  api: "openai-responses",
  models: [...],
  oauth: {
    name: "Corporate AI (SSO)",
    async login(callbacks) {
      // Implement OAuth login flow
      const authUrl = "https://ai.corp.com/oauth/authorize";
      callbacks.onAuthUrl(authUrl);
      
      // Wait for callback
      const code = await waitForCallback();
      const credentials = await exchangeCode(code);
      
      return credentials;
    },
    async refreshToken(credentials) {
      // Refresh expired token
      return await refreshCredentials(credentials);
    },
    getApiKey(credentials) {
      return credentials.access_token;
    },
  },
});

Resource Discovery

Provide custom skills, prompts, and themes:

import { join } from "path";

pi.on("resources_discover", async (event, ctx) => {
  const extensionDir = "/path/to/extension";
  
  return {
    skillPaths: [
      join(extensionDir, "skills/code-review.md"),
      join(extensionDir, "skills/testing.md"),
    ],
    promptPaths: [
      join(extensionDir, "prompts/refactor.md"),
    ],
    themePaths: [
      join(extensionDir, "themes/cyberpunk.json"),
    ],
  };
});

Best Practices

Handle UI unavailable gracefully

Check ctx.hasUI before using UI methods:

if (ctx.hasUI) {
  ctx.ui.notify("Task complete");
} else {
  console.log("Task complete");
}

Clean up resources on shutdown

Use session_shutdown to close connections:

pi.on("session_shutdown", async (event, ctx) => {
  await database.close();
  server.close();
});

Use custom entries for state persistence

Don’t rely on in-memory state - persist to session:

pi.appendEntry("my-extension", { state: "data" });

Use event bus for extension communication

Extensions can communicate via shared event bus:

// Extension A
pi.events.emit("task:complete", { result: "data" });

// Extension B
pi.events.on("task:complete", (data) => {
  console.log("Task completed:", data);
});

Get Started

Core Concepts

Guides

Extension System

Overview

Creating an Extension

Basic Structure

Extension API

Event Hooks

Session Lifecycle

Agent Lifecycle

Message Events

Tool Events

Context Manipulation

Input Interception

Compaction Hooks

Extension Context

Custom Commands

Command Context

UI Integration

Dialogs

Status Bar

Widgets

Custom Messages

State Persistence

Custom Providers

OAuth Providers

Resource Discovery

Best Practices

Next Steps

Tools

Sessions

Build docs developers (and LLMs) love

Get Started

Core Concepts

Guides

​Overview

​Creating an Extension

​Basic Structure

​Extension API

​Event Hooks

​Session Lifecycle

​Agent Lifecycle

​Message Events

​Tool Events

​Context Manipulation

​Input Interception

​Compaction Hooks

​Extension Context

​Custom Commands

​Command Context

​UI Integration

​Dialogs

​Status Bar

​Widgets

​Custom Messages

​State Persistence

​Custom Providers

​OAuth Providers

​Resource Discovery

​Best Practices

​Next Steps

Tools

Sessions

Build docs developers (and LLMs) love

Overview

Creating an Extension

Basic Structure

Extension API

Event Hooks

Session Lifecycle

Agent Lifecycle

Message Events

Tool Events

Context Manipulation

Input Interception

Compaction Hooks

Extension Context

Custom Commands

Command Context

UI Integration

Dialogs

Status Bar

Widgets

Custom Messages

State Persistence

Custom Providers

OAuth Providers

Resource Discovery

Best Practices

Next Steps