initial commit

index.ts

@@ -0,0 +1,559 @@
+// Subagent Extension - Registers a tool for delegating work to prompt-defined
+// subagents with constrained tool permissions.
+import { spawn } from "node:child_process";
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { fileURLToPath } from "node:url";
+import type { ExtensionAPI, Theme } from "@mariozechner/pi-coding-agent";
+import {
+  parseFrontmatter,
+  withFileMutationQueue,
+} from "@mariozechner/pi-coding-agent";
+import { Text } from "@mariozechner/pi-tui";
+import { Type } from "typebox";
+
+interface PromptConfig {
+  name: string;
+  description: string;
+  approvedTools: string[];
+  deniedTools: string[];
+  systemPrompt: string;
+  filePath: string;
+}
+
+interface SubagentToolActivity {
+  id: string;
+  toolName: string;
+  summary: string;
+  status: "running" | "done" | "error";
+}
+
+interface SubagentStatus {
+  state: "starting" | "thinking" | "running" | "done";
+  toolCallCount: number;
+  activeToolCalls: number;
+  recentToolCalls: SubagentToolActivity[];
+}
+
+interface SubagentResult {
+  agent: string;
+  task: string;
+  tools: string[];
+  exitCode: number;
+  output: string;
+  stderr: string;
+  status: SubagentStatus;
+  error?: string;
+}
+
+const EXTENSION_DIR = path.dirname(fileURLToPath(import.meta.url));
+const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");
+
+// Parse Tool List - Frontmatter may use YAML arrays or comma-delimited strings.
+function parseToolList(value: unknown): string[] {
+  if (!value) return [];
+  if (Array.isArray(value)) {
+    return value
+      .map(String)
+      .map((tool) => tool.trim())
+      .filter(Boolean);
+  }
+  if (typeof value === "string") {
+    return value
+      .split(",")
+      .map((tool) => tool.trim())
+      .filter(Boolean);
+  }
+  return [];
+}
+
+// Discover Prompts - Load prompt markdown files from this extension's prompts dir.
+function discoverPrompts(): PromptConfig[] {
+  if (!fs.existsSync(PROMPTS_DIR)) return [];
+
+  const prompts: PromptConfig[] = [];
+  const entries = fs.readdirSync(PROMPTS_DIR, { withFileTypes: true });
+  for (const entry of entries) {
+    if (!entry.name.endsWith(".md")) continue;
+    if (!entry.isFile() && !entry.isSymbolicLink()) continue;
+
+    const filePath = path.join(PROMPTS_DIR, entry.name);
+    const content = fs.readFileSync(filePath, "utf-8");
+    const { frontmatter, body } =
+      parseFrontmatter<Record<string, unknown>>(content);
+
+    if (
+      typeof frontmatter.name !== "string" ||
+      typeof frontmatter.description !== "string"
+    ) {
+      continue;
+    }
+
+    prompts.push({
+      name: frontmatter.name,
+      description: frontmatter.description,
+      approvedTools: parseToolList(
+        frontmatter.approved_tools ?? frontmatter.allowed_tools,
+      ),
+      deniedTools: parseToolList(frontmatter.denied_tools),
+      systemPrompt: body.trim(),
+      filePath,
+    });
+  }
+
+  return prompts.sort((a, b) => a.name.localeCompare(b.name));
+}
+
+// Resolve Tools - Use exactly one permission mode. approved_tools/allowed_tools
+// is a whitelist; denied_tools is a blacklist over the currently active tools.
+function resolveTools(agent: PromptConfig, activeTools: string[]): string[] {
+  if (agent.approvedTools.length > 0) {
+    return [...new Set(agent.approvedTools)].filter(
+      (tool) => tool !== "subagent",
+    );
+  }
+
+  const denied = new Set(agent.deniedTools);
+  return [...new Set(activeTools)].filter(
+    (tool) => tool !== "subagent" && !denied.has(tool),
+  );
+}
+
+// Write Prompt - pi accepts appended system prompts via file path.
+async function writePromptToTempFile(
+  agentName: string,
+  prompt: string,
+): Promise<{ dir: string; filePath: string }> {
+  const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
+  const safeName = agentName.replace(/[^\w.-]+/g, "_");
+  const filePath = path.join(dir, `${safeName}.md`);
+  await withFileMutationQueue(filePath, async () => {
+    await fs.promises.writeFile(filePath, prompt, {
+      encoding: "utf-8",
+      mode: 0o600,
+    });
+  });
+  return { dir, filePath };
+}
+
+// Pi Invocation - Prefer the current pi entrypoint when running inside pi.
+function getPiInvocation(args: string[]): { command: string; args: string[] } {
+  const currentScript = process.argv[1];
+  const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+  if (currentScript && !isBunVirtualScript && fs.existsSync(currentScript)) {
+    return { command: process.execPath, args: [currentScript, ...args] };
+  }
+
+  const execName = path.basename(process.execPath).toLowerCase();
+  const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
+  if (!isGenericRuntime) return { command: process.execPath, args };
+  return { command: "pi", args };
+}
+
+// Truncate One Line - Keep live status compact and stable in the TUI.
+function truncateOneLine(text: string, maxLength: number): string {
+  const firstLine = text.split(/\r?\n/, 1)[0]?.trim() ?? "";
+  return firstLine.length > maxLength
+    ? `${firstLine.slice(0, Math.max(0, maxLength - 3))}...`
+    : firstLine;
+}
+
+// Format Tool Args - Keep status updates compact enough for the tool UI.
+function formatToolArgs(args: unknown): string {
+  if (!args || typeof args !== "object") return "";
+
+  const record = args as Record<string, unknown>;
+  const priorityKeys = ["path", "command", "query", "pattern", "file", "agent"];
+  const key = priorityKeys.find((candidate) => record[candidate] !== undefined);
+  const value = key ? record[key] : undefined;
+  if (value === undefined) return "";
+
+  const text = typeof value === "string" ? value : JSON.stringify(value);
+  const compact = text.replace(/\s+/g, " ").trim();
+  return compact.length > 80 ? `${compact.slice(0, 77)}...` : compact;
+}
+
+// Format Status Text - Plain fallback for non-TUI consumers and JSON logs.
+function formatStatusText(result: SubagentResult): string {
+  const statusText =
+    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
+  const lines = [
+    `Subagent: ${result.agent}`,
+    `Task: ${result.task}`,
+    `Status: ${statusText}`,
+    `Tool Calls: ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
+  ];
+
+  if (result.status.recentToolCalls.length > 0) {
+    lines.push("", "Tool Calls - Last 3:");
+    for (const call of result.status.recentToolCalls.slice(0, 3)) {
+      const mark =
+        call.status === "done" ? "✓" : call.status === "error" ? "✗" : "•";
+      const suffix = call.summary ? `  ${call.summary}` : "";
+      lines.push(`- ${mark} ${call.toolName}${suffix}`);
+    }
+  }
+
+  return lines.join("\n");
+}
+
+// Render Subagent Summary - Compact collapsed view; Ctrl+O expands details.
+function renderSubagentSummary(result: SubagentResult, theme: Theme): Text {
+  const statusText =
+    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
+  const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
+  const latest = result.status.recentToolCalls[0];
+  const latestText = latest
+    ? ` · ${latest.toolName}${latest.summary ? ` ${latest.summary}` : ""}`
+    : "";
+  const text =
+    theme.fg(statusColor, statusText) +
+    theme.fg(
+      "muted",
+      ` · ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active${latestText}`,
+    );
+  return new Text(text, 0, 0);
+}
+
+// Render Subagent Status - Expanded TUI renderer with theme-aware colors.
+function renderSubagentStatus(result: SubagentResult, theme: Theme): Text {
+  const statusText =
+    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
+  const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
+
+  const lines = [
+    `${theme.fg("dim", "Subagent:")} ${theme.fg("accent", result.agent)}`,
+    `${theme.fg("dim", "Task:")} ${theme.fg("muted", result.task)}`,
+    `${theme.fg("dim", "Status:")} ${theme.fg(statusColor, statusText)}`,
+    `${theme.fg("dim", "Tool Calls:")} ${theme.fg(
+      "muted",
+      `${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
+    )}`,
+  ];
+
+  if (result.status.recentToolCalls.length > 0) {
+    lines.push("", theme.fg("dim", "Tool Calls - Last 3:"));
+    for (const call of result.status.recentToolCalls.slice(0, 3)) {
+      const mark =
+        call.status === "done"
+          ? theme.fg("success", "✓")
+          : call.status === "error"
+            ? theme.fg("error", "✗")
+            : theme.fg("warning", "•");
+      const tool = theme.fg("toolTitle", call.toolName.padEnd(6));
+      const summary = call.summary ? ` ${theme.fg("muted", call.summary)}` : "";
+      lines.push(`- ${mark} ${tool}${summary}`);
+    }
+  }
+
+  return new Text(lines.join("\n"), 0, 0);
+}
+
+// Run Agent - Spawn an isolated pi process in JSON mode and collect final text.
+async function runAgent(
+  cwd: string,
+  agent: PromptConfig,
+  task: string,
+  tools: string[],
+  signal?: AbortSignal,
+  onUpdate?: (partial: {
+    content: { type: "text"; text: string }[];
+    details: SubagentResult;
+  }) => void,
+): Promise<SubagentResult> {
+  let tmpDir: string | null = null;
+  let tmpPromptPath: string | null = null;
+
+  const result: SubagentResult = {
+    agent: agent.name,
+    task,
+    tools,
+    exitCode: 0,
+    output: "",
+    stderr: "",
+    status: {
+      state: "starting",
+      toolCallCount: 0,
+      activeToolCalls: 0,
+      recentToolCalls: [],
+    },
+  };
+
+  const activeToolIds = new Set<string>();
+
+  const emitUpdate = () => {
+    onUpdate?.({
+      content: [{ type: "text", text: formatStatusText(result) }],
+      details: {
+        ...result,
+        status: {
+          ...result.status,
+          recentToolCalls: [...result.status.recentToolCalls],
+        },
+      },
+    });
+  };
+
+  const rememberToolCall = (activity: SubagentToolActivity) => {
+    const existing = result.status.recentToolCalls.findIndex(
+      (call) => call.id === activity.id,
+    );
+    if (existing >= 0) result.status.recentToolCalls.splice(existing, 1);
+    result.status.recentToolCalls.unshift(activity);
+    result.status.recentToolCalls = result.status.recentToolCalls.slice(0, 3);
+  };
+
+  try {
+    const args = ["--mode", "json", "-p", "--no-session"];
+    args.push("--tools", tools.join(","));
+
+    if (agent.systemPrompt) {
+      const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt);
+      tmpDir = tmp.dir;
+      tmpPromptPath = tmp.filePath;
+      args.push("--append-system-prompt", tmpPromptPath);
+    }
+
+    args.push(`Task: ${task}`);
+    emitUpdate();
+
+    const exitCode = await new Promise<number>((resolve) => {
+      const invocation = getPiInvocation(args);
+      const proc = spawn(invocation.command, invocation.args, {
+        cwd,
+        shell: false,
+        stdio: ["ignore", "pipe", "pipe"],
+      });
+
+      let buffer = "";
+      let aborted = false;
+
+      const processLine = (line: string) => {
+        if (!line.trim()) return;
+        try {
+          const event = JSON.parse(line);
+          if (
+            event.type === "message_start" &&
+            event.message?.role === "assistant"
+          ) {
+            result.status.state = "thinking";
+            emitUpdate();
+          }
+          if (event.type === "tool_execution_start") {
+            const id = String(
+              event.toolCallId ?? result.status.toolCallCount + 1,
+            );
+            activeToolIds.add(id);
+            result.status.state = "running";
+            result.status.toolCallCount += 1;
+            result.status.activeToolCalls = activeToolIds.size;
+            rememberToolCall({
+              id,
+              toolName: String(event.toolName ?? "tool"),
+              summary: formatToolArgs(event.args),
+              status: "running",
+            });
+            emitUpdate();
+          }
+          if (event.type === "tool_execution_end") {
+            const id = String(event.toolCallId ?? "");
+            if (id) activeToolIds.delete(id);
+            result.status.activeToolCalls = activeToolIds.size;
+            const previous = result.status.recentToolCalls.find(
+              (call) => call.id === id,
+            );
+            rememberToolCall({
+              id,
+              toolName: String(event.toolName ?? previous?.toolName ?? "tool"),
+              summary: previous?.summary ?? "",
+              status: event.isError ? "error" : "done",
+            });
+            emitUpdate();
+          }
+          if (
+            event.type === "message_end" &&
+            event.message?.role === "assistant"
+          ) {
+            for (const part of event.message.content ?? []) {
+              if (part.type === "text") result.output = part.text;
+            }
+            if (event.message.stopReason === "error") {
+              result.error =
+                event.message.errorMessage ?? "Subagent stopped with an error.";
+            }
+            emitUpdate();
+          }
+        } catch {
+          // Ignore non-JSON output.
+        }
+      };
+
+      proc.stdout.on("data", (data) => {
+        buffer += data.toString();
+        const lines = buffer.split("\n");
+        buffer = lines.pop() ?? "";
+        for (const line of lines) processLine(line);
+      });
+
+      proc.stderr.on("data", (data) => {
+        result.stderr += data.toString();
+      });
+
+      proc.on("close", (code) => {
+        if (buffer.trim()) processLine(buffer);
+        resolve(aborted ? 130 : (code ?? 0));
+      });
+
+      proc.on("error", (error) => {
+        result.error = error.message;
+        resolve(1);
+      });
+
+      if (signal) {
+        const killProc = () => {
+          aborted = true;
+          proc.kill("SIGTERM");
+          setTimeout(() => {
+            if (!proc.killed) proc.kill("SIGKILL");
+          }, 5000);
+        };
+        if (signal.aborted) killProc();
+        else signal.addEventListener("abort", killProc, { once: true });
+      }
+    });
+
+    result.exitCode = exitCode;
+    result.status.state = "done";
+    result.status.activeToolCalls = 0;
+    if (exitCode === 130) result.error = "Subagent was aborted.";
+    return result;
+  } finally {
+    if (tmpPromptPath) await fs.promises.rm(tmpPromptPath, { force: true });
+    if (tmpDir) await fs.promises.rm(tmpDir, { force: true, recursive: true });
+  }
+}
+
+// Format Prompt List - Keep available prompt hints compact for the LLM.
+function formatPromptList(prompts: PromptConfig[]): string {
+  if (prompts.length === 0) return "none";
+  return prompts
+    .map((prompt) => `${prompt.name}: ${prompt.description}`)
+    .join("; ");
+}
+
+const SubagentParams = Type.Object({
+  agent: Type.String({
+    description: "Name of the prompt-defined subagent to invoke",
+  }),
+  task: Type.String({ description: "Task to delegate to the subagent" }),
+  cwd: Type.Optional(
+    Type.String({
+      description:
+        "Working directory for the subagent process. Defaults to current cwd.",
+    }),
+  ),
+});
+
+export default function (pi: ExtensionAPI) {
+  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 result = await runAgent(
+        cwd,
+        agent,
+        params.task,
+        tools,
+        signal,
+        _onUpdate,
+      );
+      const failed = result.exitCode !== 0 || Boolean(result.error);
+      const fallback = result.error || result.stderr || "(no output)";
+
+      return {
+        content: [{ type: "text", text: result.output || fallback }],
+        details: result,
+        isError: failed,
+      };
+    },
+
+    renderCall(args, theme) {
+      const task = truncateOneLine(args.task, 90);
+      const text =
+        theme.fg("toolTitle", theme.bold("subagent ")) +
+        theme.fg("accent", args.agent) +
+        (task ? ` ${theme.fg("dim", task)}` : "");
+      return new Text(text, 0, 0);
+    },
+
+    renderResult(result, { expanded, isPartial }, theme) {
+      const details = result.details as Partial<SubagentResult> | undefined;
+      if (isPartial && details?.status && details.agent && details.task) {
+        return expanded
+          ? renderSubagentStatus(details as SubagentResult, theme)
+          : renderSubagentSummary(details as SubagentResult, theme);
+      }
+
+      const text = result.content[0];
+      return new Text(text?.type === "text" ? text.text : "", 0, 0);
+    },
+  });
+}