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 fs from "node:fs";
import * as path from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { FINALIZE_TOOL_NAME, MAX_FINALIZE_RETRIES } from "./src/constants.ts";
import { formatPromptList, formatSubagentContent } from "./src/format.ts";
import { discoverPrompts } from "./src/prompts.ts";
import { renderSubagentCall, renderSubagentResult } from "./src/render.ts";
import { runAgent } from "./src/runner.ts";
import { getSubagentSessionPath } from "./src/session.ts";
import { FinalizeStatus, type SubagentResult } from "./src/types.ts";
import { resolveTools, SubagentParams } from "./src/tools.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,
        };
      },
    });
  }

  pi.registerTool({
    name: "subagent",
    label: "Subagent",
    description:
      "Delegate a task to a prompt-defined subagent from this extension's prompts/ directory. " +
      `Available at startup: ${formatPromptList(discoverPrompts())}`,
    promptSnippet:
      "Delegate tasks to subagents by name. Subagent prompts live in prompts/*.md and define approved_tools/denied_tools.",
    parameters: SubagentParams,

    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      const prompts = discoverPrompts();
      const agent = prompts.find((prompt) => prompt.name === params.agent);
      if (!agent) {
        return {
          content: [
            {
              type: "text",
              text: `Unknown subagent: ${params.agent}. Available: ${formatPromptList(prompts)}`,
            },
          ],
          details: { available: prompts.map((prompt) => prompt.name) },
          isError: true,
        };
      }

      if (agent.approvedTools.length > 0 && agent.deniedTools.length > 0) {
        return {
          content: [
            {
              type: "text",
              text:
                `Invalid subagent config for ${agent.name}: define either approved_tools/allowed_tools ` +
                "or denied_tools, not both.",
            },
          ],
          details: { agent: agent.name, filePath: agent.filePath },
          isError: true,
        };
      }

      const activeTools = pi.getActiveTools();
      const tools = resolveTools(agent, activeTools);
      if (tools.length === 0) {
        return {
          content: [
            {
              type: "text",
              text: `Subagent ${agent.name} has no approved tools after applying denied_tools.`,
            },
          ],
          details: {
            agent: agent.name,
            approvedTools: agent.approvedTools,
            deniedTools: agent.deniedTools,
          },
          isError: true,
        };
      }

      const cwd = path.resolve(ctx.cwd, params.cwd ?? ".");
      const sessionId = params.sessionId ?? randomUUID();
      const sessionPath = getSubagentSessionPath(cwd, agent.name, sessionId);
      await fs.promises.mkdir(path.dirname(sessionPath), { recursive: true });

      let result: SubagentResult | null = null;
      for (
        let retryCount = 0;
        retryCount <= MAX_FINALIZE_RETRIES;
        retryCount += 1
      ) {
        const task =
          retryCount === 0
            ? `Task: ${params.task}`
            : [
                "Your previous response did not finalize correctly.",
                `If you are finished, call ${FINALIZE_TOOL_NAME}.`,
                "If you are not finished, continue the original task using available tools as needed.",
                `Original task: ${params.task}`,
              ].join("\n\n");

        result = await runAgent(
          cwd,
          agent,
          task,
          tools,
          sessionId,
          sessionPath,
          signal,
          _onUpdate,
        );

        if (result.finalized) break;
        if (result.exitCode !== 0 || result.error) break;
      }

      if (!result) {
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                FinalizeStatus.ERROR,
                sessionId,
                undefined,
                "Subagent did not run.",
              ),
            },
          ],
          details: { sessionId },
          isError: true,
        };
      }

      if (!result.finalized) {
        const fallback =
          result.error ||
          result.stderr ||
          `Subagent did not call ${FINALIZE_TOOL_NAME}.`;
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                FinalizeStatus.ERROR,
                sessionId,
                undefined,
                fallback,
              ),
            },
          ],
          details: { sessionId },
          isError: true,
        };
      }

      if (result.finalized.status === FinalizeStatus.ERROR) {
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                FinalizeStatus.ERROR,
                sessionId,
                result.finalized.result,
                result.finalized.error ?? "Subagent failed.",
              ),
            },
          ],
          details: { sessionId, finalized: result.finalized },
          isError: true,
        };
      }

      return {
        content: [
          {
            type: "text",
            text: formatSubagentContent(
              FinalizeStatus.SUCCESS,
              sessionId,
              result.finalized.result,
            ),
          },
        ],
        details: { sessionId, finalized: result.finalized },
        isError: false,
      };
    },

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

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