pi-subagents/index.ts

// Subagent Extension - Registers a tool for delegating work to prompt-defined
// subagents with constrained tool permissions.
import { randomUUID } from "node:crypto";
import * as path from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { FINALIZE_TOOL_NAME } from "./src/constants.ts";
import { formatPromptList, toToolResult } from "./src/format.ts";
import { discoverPrompts } from "./src/prompts.ts";
import { renderSubagentCall, renderSubagentResult } from "./src/render.ts";
import { runAgentWithRetries } from "./src/runner.ts";
import { SubagentParams } from "./src/tools.ts";
import { FinalizeStatus } from "./src/types.ts";
import { validateAgent } from "./src/validate.ts";

export default function (pi: ExtensionAPI) {
  if (process.env.PI_SUBAGENT_CHILD === "1") {
    pi.registerTool({
      name: FINALIZE_TOOL_NAME,
      label: "Subagent Finalize",
      description:
        "Internal subagent-only tool. Call this as your final action when delegated work is complete.",
      promptSnippet:
        "Call subagent_finalize as your final action when subagent work is complete.",
      parameters: Type.Object({
        status: Type.Enum(FinalizeStatus),
        result: Type.Optional(Type.String()),
        error: Type.Optional(Type.String()),
      }),

      async execute(_toolCallId, params) {
        return {
          content: [{ type: "text", text: "Subagent finalized." }],
          details: params,
          terminate: true,
        };
      },
    });
  }

  const prompts = discoverPrompts();

  // Skip Registration - Only register the tool if there's at least one subagent.
  if (prompts.length === 0) return;

  pi.registerTool({
    name: "subagent",
    label: "Subagent",
    description:
      "Delegate a task to a prompt-defined subagent. " +
      `Available at startup: ${formatPromptList(prompts)}`,
    promptSnippet:
      "Delegate tasks to subagents by name. Subagent prompts define approved_tools/denied_tools.",
    parameters: SubagentParams,

    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      // Validate Agent
      const validation = validateAgent(
        params.agent,
        prompts,
        pi.getActiveTools(),
      );
      if (!validation.ok) {
        return {
          content: [{ type: "text", text: validation.text }],
          details: validation.details,
          isError: true,
        };
      }

      // Run Subagent
      const result = await runAgentWithRetries({
        cwd: path.resolve(ctx.cwd, params.cwd ?? "."),
        agent: validation.agent,
        task: params.task,
        tools: validation.tools,
        sessionId: params.sessionId ?? randomUUID(),
        signal,
        onUpdate,
      });

      return toToolResult(result);
    },

    renderCall(args, theme) {
      return renderSubagentCall(args, theme);
    },

    renderResult(result, options, theme) {
      return renderSubagentResult(result, options, theme);
    },
  });
}