Files
pi-subagents/index.ts

767 lines
24 KiB
TypeScript

// Subagent Extension - Registers a tool for delegating work to prompt-defined
// subagents with constrained tool permissions.
import { spawn } from "node:child_process";
import { createHash, randomUUID } from "node:crypto";
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 SubagentFinalizePayload {
status: "SUCCESS" | "ERROR";
result?: string;
error?: string;
}
interface SubagentResult {
agent: string;
task: string;
tools: string[];
sessionId: string;
exitCode: number;
output: string;
stderr: string;
status: SubagentStatus;
finalized?: SubagentFinalizePayload;
error?: string;
}
const EXTENSION_ENTRY = fileURLToPath(import.meta.url);
const EXTENSION_DIR = path.dirname(EXTENSION_ENTRY);
const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");
const FINALIZE_TOOL_NAME = "subagent_finalize";
const MAX_FINALIZE_RETRIES = 2;
// Format Tool Content - Some clients hide structured details from the model.
function formatSubagentContent(
status: "SUCCESS" | "ERROR",
sessionId: string,
result?: string,
error?: string,
): string {
const header = [`**Status:** ${status}`, `**Session ID:** \`${sessionId}\``];
if (error?.trim()) header.push(`**Error:** ${error.trim()}`);
const body = result?.trimEnd();
return body ? `${header.join(" \n")}\n\n---\n\n${body}` : header.join(" \n");
}
// 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[] {
const withoutDelegationTools = (tools: string[]) =>
tools.filter((tool) => tool !== "subagent" && tool !== FINALIZE_TOOL_NAME);
const resolved =
agent.approvedTools.length > 0
? withoutDelegationTools([...new Set(agent.approvedTools)])
: withoutDelegationTools(
[...new Set(activeTools)].filter(
(tool) => !new Set(agent.deniedTools).has(tool),
),
);
return [...new Set([...resolved, FINALIZE_TOOL_NAME])];
}
// Build Finalize Prompt - Child agents must terminate by calling this tool.
function buildSubagentPrompt(agent: PromptConfig): string {
const finalizePrompt = [
"You are running as a subagent.",
`When the task is complete, call ${FINALIZE_TOOL_NAME} as your final action.`,
"Do not provide the final answer as normal assistant text.",
`${FINALIZE_TOOL_NAME} requires status SUCCESS with result, or status ERROR with error and optional result.`,
].join("\n");
return [agent.systemPrompt, finalizePrompt].filter(Boolean).join("\n\n");
}
// Session Path - Persist child sessions as <agent>_<uuid>.jsonl under a cwd hash.
function getSubagentSessionPath(
cwd: string,
agentName: string,
sessionId: string,
): string {
const cwdHash = createHash("sha256").update(cwd).digest("hex").slice(0, 16);
const safeAgent = agentName.replace(/[^\w.-]+/g, "_");
const safeSessionId = sessionId.replace(/[^\w.-]+/g, "_");
return path.join(
os.homedir(),
".pi",
"subagent-sessions",
cwdHash,
`${safeAgent}_${safeSessionId}.jsonl`,
);
}
// Validate Finalize Payload - Keep the parent contract strict and small.
function validateFinalizePayload(value: unknown): SubagentFinalizePayload | null {
if (!value || typeof value !== "object") return null;
const payload = value as Record<string, unknown>;
if (payload.status === "SUCCESS") {
return typeof payload.result === "string" && payload.result.trim()
? { status: "SUCCESS", result: payload.result }
: null;
}
if (payload.status === "ERROR") {
const result =
typeof payload.result === "string" ? payload.result : undefined;
return typeof payload.error === "string" && payload.error.trim()
? { status: "ERROR", error: payload.error, result }
: null;
}
return null;
}
// 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[],
sessionId: string,
sessionPath: 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,
sessionId,
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",
"--session",
sessionPath,
"--extension",
EXTENSION_ENTRY,
];
args.push("--tools", tools.join(","));
const prompt = buildSubagentPrompt(agent);
if (prompt) {
const tmp = await writePromptToTempFile(agent.name, prompt);
tmpDir = tmp.dir;
tmpPromptPath = tmp.filePath;
args.push("--append-system-prompt", tmpPromptPath);
}
args.push(task);
emitUpdate();
const exitCode = await new Promise<number>((resolve) => {
const invocation = getPiInvocation(args);
const proc = spawn(invocation.command, invocation.args, {
cwd,
env: { ...process.env, PI_SUBAGENT_CHILD: "1" },
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,
);
const toolName = String(event.toolName ?? "tool");
if (toolName === FINALIZE_TOOL_NAME) {
result.finalized = validateFinalizePayload(event.args) ?? undefined;
}
activeToolIds.add(id);
result.status.state = "running";
result.status.toolCallCount += 1;
result.status.activeToolCalls = activeToolIds.size;
rememberToolCall({
id,
toolName,
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" }),
sessionId: Type.Optional(
Type.String({
description:
"Optional sticky subagent session id. Reuse to continue a previous subagent context.",
}),
),
cwd: Type.Optional(
Type.String({
description:
"Working directory for the subagent process. Defaults to current cwd.",
}),
),
});
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.Union([Type.Literal("SUCCESS"), Type.Literal("ERROR")]),
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(
"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("ERROR", sessionId, undefined, fallback),
},
],
details: { sessionId },
isError: true,
};
}
if (result.finalized.status === "ERROR") {
return {
content: [
{
type: "text",
text: formatSubagentContent(
"ERROR",
sessionId,
result.finalized.result,
result.finalized.error ?? "Subagent failed.",
),
},
],
details: { sessionId, finalized: result.finalized },
isError: true,
};
}
return {
content: [
{
type: "text",
text: formatSubagentContent(
"SUCCESS",
sessionId,
result.finalized.result,
),
},
],
details: { sessionId, finalized: result.finalized },
isError: false,
};
},
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);
},
});
}