@bacnh85/pi-subagent 0.1.0

package/index.ts

@@ -0,0 +1,701 @@
+/**
+ * pi-subagent — Minimal-overhead sub-agent extension for pi.
+ *
+ * Provides a `subagent` tool that delegates tasks to specialized agents
+ * running in isolated in-process SDK sessions. Supports three modes:
+ *
+ *   - Single:  { agent: "scout", task: "find auth code" }
+ *   - Parallel: { tasks: [{ agent: "scout", task: "..." }, ...] }
+ *   - Chain:    { chain: [{ agent: "scout", task: "..." }, ...] }
+ *
+ * Compared to process-spawning, this saves ~4-11K tokens per sub-agent
+ * by using the pi SDK directly with a minimal system prompt, no AGENTS.md,
+ * no extensions, no skills, no thinking, and no compaction.
+ */
+
+import * as path from "node:path";
+import type { Model } from "@earendil-works/pi-ai";
+import { getModel } from "@earendil-works/pi-ai/compat";
+import { StringEnum } from "@earendil-works/pi-ai";
+import {
+	AuthStorage,
+	CONFIG_DIR_NAME,
+	type ExtensionAPI,
+	getAgentDir,
+	getMarkdownTheme,
+	ModelRegistry,
+} from "@earendil-works/pi-coding-agent";
+import { Container, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
+import { Type } from "typebox";
+
+import { type AgentConfig, type AgentScope, discoverAgents, invalidateAgentCache } from "./agents.ts";
+import {
+	type SubAgentResult,
+	getFinalOutput,
+	getResultOutput,
+	isFailedResult,
+	mapWithConcurrencyLimit,
+	runSubAgent,
+} from "./runner.ts";
+import {
+	aggregateUsage,
+	formatUsageStats,
+	renderSingleResult,
+} from "./render.ts";
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+const MAX_PARALLEL_TASKS = 8;
+const MAX_CONCURRENCY = 4;
+const PER_TASK_OUTPUT_CAP = 50 * 1024; // 50 KB per parallel task
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function truncateParallelOutput(output: string): string {
+	const byteLength = Buffer.byteLength(output, "utf8");
+	if (byteLength <= PER_TASK_OUTPUT_CAP) return output;
+
+	let truncated = output.slice(0, PER_TASK_OUTPUT_CAP);
+	while (Buffer.byteLength(truncated, "utf8") > PER_TASK_OUTPUT_CAP) {
+		truncated = truncated.slice(0, -1);
+	}
+	return `${truncated}\n\n[Output truncated: ${byteLength - Buffer.byteLength(truncated, "utf8")} bytes omitted.]`;
+}
+
+function resolveModel(
+	modelName: string | undefined,
+	parentModel: Model | undefined,
+): Model | null {
+	if (modelName) {
+		// Try as provider/id first, then fall back to anthropic/id
+		const parts = modelName.split("/");
+		if (parts.length === 2) {
+			return getModel(parts[0], parts[1]) ?? null;
+		}
+		// Assume Anthropic shorthand
+		return getModel("anthropic", modelName) ?? null;
+	}
+	return parentModel ?? null;
+}
+
+// ---------------------------------------------------------------------------
+// Tool parameter schema
+// ---------------------------------------------------------------------------
+
+const TaskItem = Type.Object({
+	agent: Type.String({ description: "Name of the agent to invoke" }),
+	task: Type.String({ description: "Task to delegate to the agent" }),
+	cwd: Type.Optional(Type.String({ description: "Working directory for the agent" })),
+});
+
+const ChainItem = Type.Object({
+	agent: Type.String({ description: "Name of the agent to invoke" }),
+	task: Type.String({ description: "Task with optional {previous} placeholder for prior output" }),
+	cwd: Type.Optional(Type.String({ description: "Working directory for the agent" })),
+});
+
+const AgentScopeSchema = StringEnum(["user", "project", "both"] as const, {
+	description:
+		'Which agent directories to use. Default: "user". Use "both" to include project-local agents.',
+	default: "user",
+});
+
+const SubagentParams = Type.Object({
+	agent: Type.Optional(Type.String({ description: "Name of the agent to invoke (single mode)" })),
+	task: Type.Optional(Type.String({ description: "Task to delegate (single mode)" })),
+	tasks: Type.Optional(
+		Type.Array(TaskItem, { description: "Array of {agent, task} for parallel execution" }),
+	),
+	chain: Type.Optional(
+		Type.Array(ChainItem, {
+			description: "Array of {agent, task} for sequential execution with {previous}",
+		}),
+	),
+	agentScope: Type.Optional(AgentScopeSchema),
+	confirmProjectAgents: Type.Optional(
+		Type.Boolean({
+			description: "Prompt before running project-local agents. Default: true.",
+			default: true,
+		}),
+	),
+	cwd: Type.Optional(Type.String({ description: "Working directory (single mode)" })),
+});
+
+// ---------------------------------------------------------------------------
+// Details type
+// ---------------------------------------------------------------------------
+
+interface SubagentDetails {
+	mode: "single" | "parallel" | "chain";
+	agentScope: AgentScope;
+	projectAgentsDir: string | null;
+	results: SubAgentResult[];
+}
+
+// ---------------------------------------------------------------------------
+// Extension entry point
+// ---------------------------------------------------------------------------
+
+export default function (pi: ExtensionAPI) {
+	// Invalidate agent cache on reload so edited agent files take effect
+	pi.on("session_start", (event) => {
+		if (event.reason === "reload") invalidateAgentCache();
+	});
+
+	// Resolve bundled agents directory relative to this extension file
+	const bundledAgentsDir = path.resolve(__dirname, "agents");
+
+	pi.registerTool({
+		name: "subagent",
+		label: "Subagent",
+		description: [
+			"Delegate tasks to specialized subagents with isolated context (SDK-based, minimal overhead).",
+			"Modes: single (agent + task), parallel (tasks array, max 8, 4 concurrent), chain (sequential with {previous}).",
+			`Default agent scope is "user" (from ${path.join(getAgentDir(), "agents")}).`,
+			`To enable project-local agents in ${CONFIG_DIR_NAME}/agents, set agentScope: "both" or "project".`,
+		].join(" "),
+		parameters: SubagentParams,
+
+		async execute(_toolCallId, params, signal, onUpdate, ctx) {
+			const agentScope: AgentScope = params.agentScope ?? "user";
+			const discovery = discoverAgents(ctx.cwd, agentScope, bundledAgentsDir);
+			const agents = discovery.agents;
+			const confirmProjectAgents = params.confirmProjectAgents ?? true;
+
+			const hasChain = (params.chain?.length ?? 0) > 0;
+			const hasTasks = (params.tasks?.length ?? 0) > 0;
+			const hasSingle = Boolean(params.agent && params.task);
+			const modeCount = Number(hasChain) + Number(hasTasks) + Number(hasSingle);
+
+			const makeDetails =
+				(mode: "single" | "parallel" | "chain") =>
+				(results: SubAgentResult[]): SubagentDetails => ({
+					mode,
+					agentScope,
+					projectAgentsDir: discovery.projectAgentsDir,
+					results,
+				});
+
+			// Validate: exactly one mode
+			if (modeCount !== 1) {
+				const available = agents.map((a) => `${a.name} (${a.source})`).join(", ") || "none";
+				return {
+					content: [
+						{
+							type: "text",
+							text: [
+								"Invalid parameters. Provide exactly one mode:",
+								"  single: { agent, task }",
+								"  parallel: { tasks: [...] }",
+								"  chain: { chain: [...] }",
+								`Available agents: ${available}`,
+							].join("\n"),
+						},
+					],
+					details: makeDetails("single")([]),
+				};
+			}
+
+			// Confirm project-local agents
+			if (
+				(agentScope === "project" || agentScope === "both") &&
+				confirmProjectAgents &&
+				ctx.hasUI
+			) {
+				const requestedAgentNames = new Set<string>();
+				if (params.chain) for (const s of params.chain) requestedAgentNames.add(s.agent);
+				if (params.tasks) for (const t of params.tasks) requestedAgentNames.add(t.agent);
+				if (params.agent) requestedAgentNames.add(params.agent);
+
+				const projectAgentsRequested = Array.from(requestedAgentNames)
+					.map((name) => agents.find((a) => a.name === name))
+					.filter((a): a is AgentConfig => a?.source === "project");
+
+				if (projectAgentsRequested.length > 0) {
+					const names = projectAgentsRequested.map((a) => a.name).join(", ");
+					const dir = discovery.projectAgentsDir ?? "(unknown)";
+					const ok = await ctx.ui.confirm(
+						"Run project-local agents?",
+						`Agents: ${names}\nSource: ${dir}\n\nProject agents are repo-controlled. Only continue for trusted repositories.`,
+					);
+					if (!ok) {
+						return {
+							content: [{ type: "text", text: "Canceled: project-local agents not approved." }],
+							details: makeDetails(hasChain ? "chain" : hasTasks ? "parallel" : "single")([]),
+						};
+					}
+				}
+			}
+
+			// Shared auth/model setup for SDK sessions
+			const authStorage = AuthStorage.create();
+			const modelRegistry = ModelRegistry.create(authStorage);
+
+			// Helper: inject parent's API key into child auth storage
+			async function injectApiKey(model: Model): Promise<void> {
+				const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+				if (auth.ok && auth.apiKey) {
+					authStorage.setRuntimeApiKey(model.provider, auth.apiKey);
+				}
+			}
+
+			// Helper: run a single agent via SDK
+			async function runOne(
+				agentName: string,
+				task: string,
+				cwd: string | undefined,
+			): Promise<SubAgentResult> {
+				const agent = agents.find((a) => a.name === agentName);
+
+				if (!agent) {
+					const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
+					return {
+						agent: agentName,
+						task,
+						exitCode: 1,
+						messages: [],
+						stderr: `Unknown agent: "${agentName}". Available: ${available}.`,
+						usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+						errorMessage: `Unknown agent: "${agentName}"`,
+					};
+				}
+
+				const model = resolveModel(agent.model, ctx.model);
+				if (!model) {
+					return {
+						agent: agentName,
+						task,
+						exitCode: 1,
+						messages: [],
+						stderr: `No model resolved for agent "${agentName}". Configure a model in the agent definition or select one in the parent session.`,
+						usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+						errorMessage: "No model resolved",
+					};
+				}
+
+				// Inject parent's API key so --api-key and other runtime overrides work
+				await injectApiKey(model);
+
+				const tools = agent.tools ?? ["read", "bash", "edit", "write", "grep", "find", "ls"];
+
+				return runSubAgent({
+					cwd: cwd ?? ctx.cwd,
+					systemPrompt: agent.systemPrompt,
+					task,
+					tools,
+					model,
+					authStorage,
+					modelRegistry,
+					signal,
+					agentName,
+				});
+			}
+
+			// --- Chain mode ---
+			if (params.chain && params.chain.length > 0) {
+				const results: SubAgentResult[] = [];
+				let previousOutput = "";
+
+				for (let i = 0; i < params.chain.length; i++) {
+					const step = params.chain[i];
+					const taskWithContext = step.task.replace(/\{previous\}/g, previousOutput);
+
+					const result = await runOne(step.agent, taskWithContext, step.cwd);
+					results.push(result);
+
+					const isError = isFailedResult(result);
+					if (isError) {
+						const errorMsg = getResultOutput(result);
+						if (onUpdate) {
+							onUpdate({
+								content: [{ type: "text", text: errorMsg }],
+								details: makeDetails("chain")(results),
+							});
+						}
+						return {
+							content: [
+								{
+									type: "text",
+									text: `Chain stopped at step ${i + 1} (${step.agent}): ${errorMsg}`,
+								},
+							],
+							details: makeDetails("chain")(results),
+							isError: true,
+						};
+					}
+
+					previousOutput = getFinalOutput(result.messages);
+
+					if (onUpdate) {
+						onUpdate({
+							content: [{ type: "text", text: getFinalOutput(result.messages) || "(no output)" }],
+							details: makeDetails("chain")(results),
+						});
+					}
+				}
+
+				const last = results[results.length - 1];
+				return {
+					content: [
+						{ type: "text", text: getFinalOutput(last.messages) || "(no output)" },
+					],
+					details: makeDetails("chain")(results),
+				};
+			}
+
+			// --- Parallel mode ---
+			if (params.tasks && params.tasks.length > 0) {
+				if (params.tasks.length > MAX_PARALLEL_TASKS) {
+					return {
+						content: [
+							{
+								type: "text",
+								text: `Too many parallel tasks (${params.tasks.length}). Max is ${MAX_PARALLEL_TASKS}.`,
+							},
+						],
+						details: makeDetails("parallel")([]),
+					};
+				}
+
+				const allResults: SubAgentResult[] = new Array(params.tasks.length);
+				// Initialize placeholder results for streaming
+				for (let i = 0; i < params.tasks.length; i++) {
+					allResults[i] = {
+						agent: params.tasks[i].agent,
+						task: params.tasks[i].task,
+						exitCode: -1,
+						messages: [],
+						stderr: "",
+						usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+					};
+				}
+
+				const emitParallelUpdate = () => {
+					if (onUpdate) {
+						const running = allResults.filter((r) => r.exitCode === -1).length;
+						const done = allResults.filter((r) => r.exitCode !== -1).length;
+						onUpdate({
+							content: [
+								{
+									type: "text",
+									text: `Parallel: ${done}/${allResults.length} done, ${running} running...`,
+								},
+							],
+							details: makeDetails("parallel")([...allResults]),
+						});
+					}
+				};
+
+				const results = await mapWithConcurrencyLimit(
+					params.tasks,
+					MAX_CONCURRENCY,
+					async (t, index) => {
+						const result = await runOne(t.agent, t.task, t.cwd);
+						allResults[index] = result;
+						emitParallelUpdate();
+						return result;
+					},
+				);
+
+				const successCount = results.filter((r) => !isFailedResult(r)).length;
+				const summaries = results.map((r) => {
+					const output = truncateParallelOutput(getResultOutput(r));
+					const status = isFailedResult(r)
+						? `failed${r.stopReason ? ` (${r.stopReason})` : ""}`
+						: "completed";
+					return `### [${r.agent}] ${status}\n\n${output}`;
+				});
+
+				return {
+					content: [
+						{
+							type: "text",
+							text: `Parallel: ${successCount}/${results.length} succeeded\n\n${summaries.join("\n\n---\n\n")}`,
+						},
+					],
+					details: makeDetails("parallel")(results),
+				};
+			}
+
+			// --- Single mode ---
+			if (params.agent && params.task) {
+				const result = await runOne(params.agent, params.task, params.cwd);
+				const isError = isFailedResult(result);
+
+				if (onUpdate) {
+					onUpdate({
+						content: [
+							{ type: "text", text: getFinalOutput(result.messages) || "(running...)" },
+						],
+						details: makeDetails("single")([result]),
+					});
+				}
+
+				if (isError) {
+					const errorMsg = getResultOutput(result);
+					return {
+						content: [
+							{
+								type: "text",
+								text: `Agent ${result.stopReason || "failed"}: ${errorMsg}`,
+							},
+						],
+						details: makeDetails("single")([result]),
+						isError: true,
+					};
+				}
+
+				return {
+					content: [
+						{ type: "text", text: getFinalOutput(result.messages) || "(no output)" },
+					],
+					details: makeDetails("single")([result]),
+				};
+			}
+
+			// Should not reach here due to validation above
+			const available = agents.map((a) => `${a.name} (${a.source})`).join(", ") || "none";
+			return {
+				content: [{ type: "text", text: `Invalid parameters. Available agents: ${available}` }],
+				details: makeDetails("single")([]),
+			};
+		},
+
+		// ------------------------------------------------------------------
+		// TUI rendering
+		// ------------------------------------------------------------------
+
+		renderCall(args, theme, _context) {
+			const scope: AgentScope = args.agentScope ?? "user";
+			const fg = theme.fg.bind(theme);
+
+			// Chain
+			if (args.chain && args.chain.length > 0) {
+				let text =
+					fg("toolTitle", theme.bold("subagent ")) +
+					fg("accent", `chain (${args.chain.length} steps)`) +
+					fg("muted", ` [${scope}]`);
+				for (let i = 0; i < Math.min(args.chain.length, 3); i++) {
+					const step = args.chain[i];
+					const cleanTask = step.task.replace(/\{previous\}/g, "").trim();
+					const preview = cleanTask.length > 40 ? `${cleanTask.slice(0, 40)}...` : cleanTask;
+					text +=
+						"\n  " +
+						fg("muted", `${i + 1}.`) +
+						" " +
+						fg("accent", step.agent) +
+						fg("dim", ` ${preview}`);
+				}
+				if (args.chain.length > 3)
+					text += `\n  ${fg("muted", `... +${args.chain.length - 3} more`)}`;
+				return new Text(text, 0, 0);
+			}
+
+			// Parallel
+			if (args.tasks && args.tasks.length > 0) {
+				let text =
+					fg("toolTitle", theme.bold("subagent ")) +
+					fg("accent", `parallel (${args.tasks.length} tasks)`) +
+					fg("muted", ` [${scope}]`);
+				for (const t of args.tasks.slice(0, 3)) {
+					const preview = t.task.length > 40 ? `${t.task.slice(0, 40)}...` : t.task;
+					text += `\n  ${fg("accent", t.agent)}${fg("dim", ` ${preview}`)}`;
+				}
+				if (args.tasks.length > 3)
+					text += `\n  ${fg("muted", `... +${args.tasks.length - 3} more`)}`;
+				return new Text(text, 0, 0);
+			}
+
+			// Single
+			const agentName = args.agent || "...";
+			const preview = args.task
+				? args.task.length > 60
+					? `${args.task.slice(0, 60)}...`
+					: args.task
+				: "...";
+			let text =
+				fg("toolTitle", theme.bold("subagent ")) +
+				fg("accent", agentName) +
+				fg("muted", ` [${scope}]`);
+			text += `\n  ${fg("dim", preview)}`;
+			return new Text(text, 0, 0);
+		},
+
+		renderResult(result, { expanded }, theme, _context) {
+			const details = result.details as SubagentDetails | undefined;
+			if (!details || details.results.length === 0) {
+				const text = result.content[0];
+				return new Text(text?.type === "text" ? text.text : "(no output)", 0, 0);
+			}
+
+			const fg = theme.fg.bind(theme);
+			const mdTheme = getMarkdownTheme();
+
+			// --- Single ---
+			if (details.mode === "single" && details.results.length === 1) {
+				return renderSingleResult(details.results[0], expanded, theme);
+			}
+
+			// --- Chain ---
+			if (details.mode === "chain") {
+				const successCount = details.results.filter((r) => !isFailedResult(r)).length;
+				const icon =
+					successCount === details.results.length
+						? fg("success", "✓")
+						: fg("error", "✗");
+
+				if (expanded) {
+					const container = new Container();
+					container.addChild(
+						new Text(
+							icon +
+								" " +
+								fg("toolTitle", theme.bold("chain ")) +
+								fg("accent", `${successCount}/${details.results.length} steps`),
+							0,
+							0,
+						),
+					);
+					for (const r of details.results) {
+						container.addChild(new Spacer(1));
+						const stepIcon = isFailedResult(r) ? fg("error", "✗") : fg("success", "✓");
+						container.addChild(
+							new Text(
+								fg("muted", `─── Step ${r.exitCode !== -1 ? "" : "?"}: `) +
+									fg("accent", r.agent) +
+									` ${stepIcon}`,
+								0,
+								0,
+							),
+						);
+						if (r.errorMessage) {
+							container.addChild(
+								new Text(fg("error", `Error: ${r.errorMessage}`), 0, 0),
+							);
+						}
+						const finalOutput = getResultOutput(r);
+						if (finalOutput) {
+							container.addChild(new Spacer(1));
+							container.addChild(new Markdown(finalOutput.trim(), 0, 0, mdTheme));
+						}
+						const usageStr = formatUsageStats(r.usage, r.model);
+						if (usageStr)
+							container.addChild(new Text(fg("dim", usageStr), 0, 0));
+					}
+					const totalUsage = formatUsageStats(aggregateUsage(details.results));
+					if (totalUsage) {
+						container.addChild(new Spacer(1));
+						container.addChild(new Text(fg("dim", `Total: ${totalUsage}`), 0, 0));
+					}
+					return container;
+				}
+
+				let text =
+					icon +
+					" " +
+					fg("toolTitle", theme.bold("chain ")) +
+					fg("accent", `${successCount}/${details.results.length} steps`);
+				for (const r of details.results) {
+					const stepIcon = isFailedResult(r) ? fg("error", "✗") : fg("success", "✓");
+					text += `\n  ${stepIcon} ${fg("accent", r.agent)}`;
+				}
+				const totalUsage = formatUsageStats(aggregateUsage(details.results));
+				if (totalUsage) text += `\n${fg("dim", totalUsage)}`;
+				text += `\n${fg("muted", "(Ctrl+O to expand)")}`;
+				return new Text(text, 0, 0);
+			}
+
+			// --- Parallel ---
+			if (details.mode === "parallel") {
+				const running = details.results.filter((r) => r.exitCode === -1).length;
+				const successCount = details.results.filter(
+					(r) => r.exitCode !== -1 && !isFailedResult(r),
+				).length;
+				const failCount = details.results.filter(
+					(r) => r.exitCode !== -1 && isFailedResult(r),
+				).length;
+				const isRunning = running > 0;
+				const icon = isRunning
+					? fg("warning", "⏳")
+					: failCount > 0
+						? fg("warning", "◐")
+						: fg("success", "✓");
+				const status = isRunning
+					? `${successCount + failCount}/${details.results.length} done, ${running} running`
+					: `${successCount}/${details.results.length} tasks`;
+
+				if (expanded && !isRunning) {
+					const container = new Container();
+					container.addChild(
+						new Text(
+							`${icon} ${fg("toolTitle", theme.bold("parallel "))}${fg("accent", status)}`,
+							0,
+							0,
+						),
+					);
+					for (const r of details.results) {
+						container.addChild(new Spacer(1));
+						const taskIcon = isFailedResult(r)
+							? fg("error", "✗")
+							: fg("success", "✓");
+						container.addChild(
+							new Text(
+								fg("muted", "─── ") + fg("accent", r.agent) + ` ${taskIcon}`,
+								0,
+								0,
+							),
+						);
+						container.addChild(
+							new Text(fg("muted", "Task: ") + fg("dim", r.task), 0, 0),
+						);
+						if (r.errorMessage) {
+							container.addChild(
+								new Text(fg("error", `Error: ${r.errorMessage}`), 0, 0),
+							);
+						}
+						const finalOutput = getResultOutput(r);
+						if (finalOutput) {
+							container.addChild(new Spacer(1));
+							container.addChild(
+								new Markdown(finalOutput.trim(), 0, 0, mdTheme),
+							);
+						}
+						const taskUsage = formatUsageStats(r.usage, r.model);
+						if (taskUsage)
+							container.addChild(new Text(fg("dim", taskUsage), 0, 0));
+					}
+					const totalUsage = formatUsageStats(aggregateUsage(details.results));
+					if (totalUsage) {
+						container.addChild(new Spacer(1));
+						container.addChild(new Text(fg("dim", `Total: ${totalUsage}`), 0, 0));
+					}
+					return container;
+				}
+
+				let text = `${icon} ${fg("toolTitle", theme.bold("parallel "))}${fg("accent", status)}`;
+				for (const r of details.results) {
+					const taskIcon =
+						r.exitCode === -1
+							? fg("warning", "⏳")
+							: isFailedResult(r)
+								? fg("error", "✗")
+								: fg("success", "✓");
+					text += `\n  ${taskIcon} ${fg("accent", r.agent)}`;
+				}
+				if (!isRunning) {
+					const totalUsage = formatUsageStats(aggregateUsage(details.results));
+					if (totalUsage) text += `\n${fg("dim", totalUsage)}`;
+				}
+				if (!expanded) text += `\n${fg("muted", "(Ctrl+O to expand)")}`;
+				return new Text(text, 0, 0);
+			}
+
+			const fallback = result.content[0];
+			return new Text(fallback?.type === "text" ? fallback.text : "(no output)", 0, 0);
+		},
+	});
+}