@ferris1225/pi-subagents 0.1.0

package/src/setup.ts

@@ -0,0 +1,222 @@
+/**
+ * Interactive configuration wizard for /subagents-setup.
+ *
+ * Everything is selection-driven (no free-text answers): a multi-select for which
+ * agents to enable, a per-agent single-select for model overrides (fuzzy filter +
+ * paging), and simple menus for the injection toggle and agent scope. Config is
+ * written to <agentDir>/pi-subagents.json.
+ */
+
+import { stat } from "node:fs/promises";
+import type { ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
+import {
+	AGENT_SCOPE_VALUES,
+	BUILTIN_AGENT_NAMES,
+	DEFAULT_CONFIG,
+	DEFAULT_ENABLED_AGENTS,
+	type AgentScope,
+	type SubagentsConfig,
+	errorMessage,
+	getConfigPath,
+	loadConfig,
+	saveConfig,
+} from "./config.ts";
+import { promptSelectMany, promptSelectOne } from "./ui.ts";
+
+const INHERIT = "__inherit__";
+
+/** Short, selection-friendly descriptions for the built-in agents. */
+const MODULE_HINTS: Record<string, string> = {
+	explore: "read-only codebase recon (fast model)",
+	plan: "implementation plan before code (opt-in)",
+	worker: "implement / fix / refactor / test (full tools)",
+	reviewer: "adversarial pre-commit review (read-only)",
+};
+
+function moduleLabel(name: string): string {
+	const hint = MODULE_HINTS[name];
+	return hint ? `${name} — ${hint}` : name;
+}
+
+async function configExists(configPath: string): Promise<boolean> {
+	try {
+		await stat(configPath);
+		return true;
+	} catch {
+		return false;
+	}
+}
+
+/** Ordered "provider/model-id" references, with the active model first. */
+function availableModelRefs(ctx: ExtensionCommandContext): string[] {
+	const refs = new Set<string>();
+	if (ctx.model) refs.add(`${ctx.model.provider}/${ctx.model.id}`);
+	const models =
+		ctx.scopedModels.length > 0 ? ctx.scopedModels.map((entry) => entry.model) : ctx.modelRegistry.getAvailable();
+	for (const model of models) refs.add(`${model.provider}/${model.id}`);
+	return [...refs];
+}
+
+async function pickEnabledAgents(
+	ctx: ExtensionCommandContext,
+	current: readonly string[],
+): Promise<string[] | undefined> {
+	const items = BUILTIN_AGENT_NAMES.map((name) => ({ value: name, label: moduleLabel(name) }));
+	return promptSelectMany(
+		ctx,
+		"Enable which sub-agents?",
+		"Space toggles • Enter confirms • Esc cancels",
+		items,
+		current,
+	);
+}
+
+async function pickAgentModels(
+	ctx: ExtensionCommandContext,
+	enabledAgents: readonly string[],
+	current: Record<string, string>,
+): Promise<Record<string, string> | undefined> {
+	const refs = availableModelRefs(ctx);
+	if (refs.length === 0) {
+		ctx.ui.notify("No Pi models are currently available; model overrides left unchanged.", "warning");
+		return { ...current };
+	}
+
+	const result: Record<string, string> = {};
+	for (const name of enabledAgents) {
+		const currentRef = current[name];
+		const items = [
+			{
+				value: INHERIT,
+				label: currentRef
+					? `(use main session's model — drop override "${currentRef}")`
+					: "(use main session's current model — no override)",
+			},
+			...refs.map((ref) => ({ value: ref, label: ref === currentRef ? `${ref} (current)` : ref })),
+		];
+		const choice = await promptSelectOne(
+			ctx,
+			`Model for "${name}"`,
+			"Type to filter • ↑/↓ • PgUp/PgDn • Enter selects • Esc cancels setup",
+			items,
+		);
+		if (choice === undefined) return undefined; // Esc aborts the whole wizard
+		if (choice !== INHERIT) result[name] = choice;
+	}
+	return result;
+}
+
+async function pickInjection(ctx: ExtensionCommandContext, current: boolean): Promise<boolean | undefined> {
+	const on = "On — inject the delegation directive into the system prompt (recommended)";
+	const off = "Off — do not inject (rely on the tool description alone)";
+	const choice = await ctx.ui.select("Proactive dispatch injection?", [current ? `${on} (current)` : on, current ? off : `${off} (current)`]);
+	if (choice === undefined) return undefined;
+	return choice.startsWith("On");
+}
+
+async function pickScope(ctx: ExtensionCommandContext, current: AgentScope): Promise<AgentScope | undefined> {
+	const labels: Record<AgentScope, string> = {
+		user: "user — built-in + ~/.pi/agent/agents (default)",
+		project: "project — built-in + nearest .pi/agents only",
+		both: "both — user agents, overridden by project agents",
+	};
+	const options = AGENT_SCOPE_VALUES.map((scope) =>
+		scope === current ? `${labels[scope]} (current)` : labels[scope],
+	);
+	const choice = await ctx.ui.select("Which agent directories to discover from?", options);
+	if (choice === undefined) return undefined;
+	const scope = AGENT_SCOPE_VALUES.find((s) => choice.startsWith(s));
+	return scope;
+}
+
+/** Validate that every configured model override still resolves; drop stale ones. */
+function pruneStaleModels(ctx: ExtensionCommandContext, agentModels: Record<string, string>): Record<string, string> {
+	const clean: Record<string, string> = {};
+	for (const [name, ref] of Object.entries(agentModels)) {
+		const slash = ref.indexOf("/");
+		if (slash <= 0) continue;
+		const model = ctx.modelRegistry.find(ref.slice(0, slash), ref.slice(slash + 1));
+		if (model) clean[name] = ref;
+		else ctx.ui.notify(`Dropped stale model override for "${name}": ${ref}`, "warning");
+	}
+	return clean;
+}
+
+async function runFullSetup(ctx: ExtensionCommandContext, configPath: string, base: SubagentsConfig): Promise<void> {
+	const enabled = await pickEnabledAgents(ctx, base.enabledAgents);
+	if (enabled === undefined) return notifyCancelled(ctx);
+
+	const models = await pickAgentModels(ctx, enabled, base.agentModels);
+	if (models === undefined) return notifyCancelled(ctx);
+
+	const injection = await pickInjection(ctx, base.proactiveInjection);
+	if (injection === undefined) return notifyCancelled(ctx);
+
+	const scope = await pickScope(ctx, base.agentScope);
+	if (scope === undefined) return notifyCancelled(ctx);
+
+	const next: SubagentsConfig = {
+		enabledAgents: enabled,
+		agentModels: pruneStaleModels(ctx, models),
+		proactiveInjection: injection,
+		agentScope: scope,
+	};
+	await saveConfig(next, configPath);
+	ctx.ui.notify(`pi-subagents configured. Saved to ${configPath}`, "info");
+}
+
+async function runMenu(ctx: ExtensionCommandContext, configPath: string, config: SubagentsConfig): Promise<void> {
+	const choice = await ctx.ui.select("pi-subagents is already configured. What would you like to change?", [
+		"Enable/disable agents",
+		"Change agent models",
+		"Toggle proactive injection",
+		"Change agent scope",
+		"Full re-setup",
+	]);
+	if (choice === undefined) return notifyCancelled(ctx);
+
+	if (choice.startsWith("Full")) return runFullSetup(ctx, configPath, config);
+
+	let next: SubagentsConfig = { ...config, agentModels: { ...config.agentModels } };
+
+	if (choice.startsWith("Enable")) {
+		const enabled = await pickEnabledAgents(ctx, config.enabledAgents);
+		if (enabled === undefined) return notifyCancelled(ctx);
+		next.enabledAgents = enabled;
+	} else if (choice.startsWith("Change agent models")) {
+		const models = await pickAgentModels(ctx, config.enabledAgents, config.agentModels);
+		if (models === undefined) return notifyCancelled(ctx);
+		next.agentModels = pruneStaleModels(ctx, models);
+	} else if (choice.startsWith("Toggle")) {
+		const injection = await pickInjection(ctx, config.proactiveInjection);
+		if (injection === undefined) return notifyCancelled(ctx);
+		next.proactiveInjection = injection;
+	} else if (choice.startsWith("Change agent scope")) {
+		const scope = await pickScope(ctx, config.agentScope);
+		if (scope === undefined) return notifyCancelled(ctx);
+		next.agentScope = scope;
+	}
+
+	await saveConfig(next, configPath);
+	ctx.ui.notify(`pi-subagents updated. Saved to ${configPath}`, "info");
+}
+
+function notifyCancelled(ctx: ExtensionCommandContext): void {
+	ctx.ui.notify("pi-subagents setup cancelled.", "info");
+}
+
+/** Entry point for the /subagents-setup command. */
+export async function runSetup(ctx: ExtensionCommandContext, configPath: string = getConfigPath()): Promise<void> {
+	if (ctx.mode !== "tui") {
+		ctx.ui.notify("/subagents-setup requires Pi's interactive TUI.", "error");
+		return;
+	}
+	try {
+		const exists = await configExists(configPath);
+		const config = await loadConfig(configPath);
+		if (exists) await runMenu(ctx, configPath, config);
+		else await runFullSetup(ctx, configPath, { ...DEFAULT_CONFIG, enabledAgents: [...DEFAULT_ENABLED_AGENTS] });
+	} catch (error) {
+		ctx.ui.notify(`pi-subagents setup failed: ${errorMessage(error)}`, "error");
+	}
+}

package/src/spawn.ts

@@ -0,0 +1,308 @@
+/**
+ * Sub-agent dispatch: each agent runs as an isolated `pi` child process
+ * (`--mode json -p --no-session`). The agent's system prompt (the .md body) is
+ * written to a temp file and passed via `--append-system-prompt` (which accepts a
+ * file path). Child stdout is a JSON-lines event stream; we accumulate assistant
+ * messages from `message_end` events and stream partial output back via onUpdate.
+ *
+ * Adapted from the official pi example `examples/extensions/subagent`.
+ */
+
+import { spawn } from "node:child_process";
+import { mkdtemp, rm, writeFile } from "node:fs/promises";
+import { existsSync, unlinkSync, rmdirSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { basename, join } from "node:path";
+import type { AgentToolResult } from "@earendil-works/pi-agent-core";
+import type { Message } from "@earendil-works/pi-ai";
+import { withFileMutationQueue } from "@earendil-works/pi-coding-agent";
+import type { AgentConfig, AgentSource } from "./agents.ts";
+
+export const MAX_PARALLEL_TASKS = 8;
+export const MAX_CONCURRENCY = 4;
+export const PER_TASK_OUTPUT_CAP = 50 * 1024;
+/** Max nesting depth for sub-agent -> sub-agent spawning (recursion guard). */
+export const MAX_SUBAGENT_DEPTH = 2;
+export const DEPTH_ENV_VAR = "PI_SUBAGENT_DEPTH";
+
+export interface UsageStats {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cost: number;
+	contextTokens: number;
+	turns: number;
+}
+
+export interface SingleResult {
+	agent: string;
+	agentSource: AgentSource | "unknown";
+	task: string;
+	exitCode: number; // -1 = still running
+	messages: Message[];
+	stderr: string;
+	usage: UsageStats;
+	model?: string;
+	stopReason?: string;
+	errorMessage?: string;
+}
+
+export interface SubagentDetails {
+	mode: "single" | "parallel";
+	results: SingleResult[];
+}
+
+export type OnUpdateCallback = (partial: AgentToolResult<SubagentDetails>) => void;
+
+function emptyUsage(): UsageStats {
+	return { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
+}
+
+export function getFinalOutput(messages: Message[]): string {
+	for (let i = messages.length - 1; i >= 0; i--) {
+		const msg = messages[i];
+		if (msg.role === "assistant") {
+			for (const part of msg.content) {
+				if (part.type === "text") return part.text;
+			}
+		}
+	}
+	return "";
+}
+
+export function isFailedResult(result: SingleResult): boolean {
+	return result.exitCode !== 0 || result.stopReason === "error" || result.stopReason === "aborted";
+}
+
+export function getResultOutput(result: SingleResult): string {
+	if (isFailedResult(result)) {
+		return result.errorMessage || result.stderr || getFinalOutput(result.messages) || "(no output)";
+	}
+	return getFinalOutput(result.messages) || "(no output)";
+}
+
+export 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);
+	const omitted = byteLength - Buffer.byteLength(truncated, "utf8");
+	return `${truncated}\n\n[Output truncated: ${omitted} bytes omitted. Full output preserved in tool details.]`;
+}
+
+export async function mapWithConcurrencyLimit<TIn, TOut>(
+	items: TIn[],
+	concurrency: number,
+	fn: (item: TIn, index: number) => Promise<TOut>,
+): Promise<TOut[]> {
+	if (items.length === 0) return [];
+	const limit = Math.max(1, Math.min(concurrency, items.length));
+	const results: TOut[] = new Array(items.length);
+	let nextIndex = 0;
+	const workers = new Array(limit).fill(null).map(async () => {
+		while (true) {
+			const current = nextIndex++;
+			if (current >= items.length) return;
+			results[current] = await fn(items[current], current);
+		}
+	});
+	await Promise.all(workers);
+	return results;
+}
+
+async function writePromptToTempFile(agentName: string, prompt: string): Promise<{ dir: string; filePath: string }> {
+	const dir = await mkdtemp(join(tmpdir(), "pi-subagents-"));
+	const safeName = agentName.replace(/[^\w.-]+/g, "_");
+	const filePath = join(dir, `prompt-${safeName}.md`);
+	await withFileMutationQueue(filePath, async () => {
+		await writeFile(filePath, prompt, "utf8");
+	});
+	return { dir, filePath };
+}
+
+/** Resolve how to invoke the SAME pi build as the current process. */
+export function getPiInvocation(args: string[]): { command: string; args: string[] } {
+	const currentScript = process.argv[1];
+	const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
+	if (currentScript && !isBunVirtualScript && existsSync(currentScript)) {
+		return { command: process.execPath, args: [currentScript, ...args] };
+	}
+	const execName = basename(process.execPath).toLowerCase();
+	const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
+	if (!isGenericRuntime) return { command: process.execPath, args };
+	return { command: "pi", args };
+}
+
+export function currentSubagentDepth(env: NodeJS.ProcessEnv = process.env): number {
+	const raw = env[DEPTH_ENV_VAR];
+	const parsed = raw === undefined ? 0 : Number.parseInt(raw, 10);
+	return Number.isFinite(parsed) && parsed > 0 ? parsed : 0;
+}
+
+export interface RunSingleOptions {
+	defaultCwd: string;
+	agent: AgentConfig | undefined;
+	agentName: string;
+	task: string;
+	cwd?: string;
+	signal?: AbortSignal;
+	onUpdate?: OnUpdateCallback;
+	makeDetails: (results: SingleResult[]) => SubagentDetails;
+	env?: NodeJS.ProcessEnv;
+}
+
+/** Spawn one agent as an isolated pi child process and collect its output. */
+export async function runSingleAgent(options: RunSingleOptions): Promise<SingleResult> {
+	const { agent, agentName, task, cwd, signal, onUpdate, makeDetails } = options;
+
+	if (!agent) {
+		return {
+			agent: agentName,
+			agentSource: "unknown",
+			task,
+			exitCode: 1,
+			messages: [],
+			stderr: `Unknown agent: "${agentName}".`,
+			usage: emptyUsage(),
+		};
+	}
+
+	const args: string[] = ["--mode", "json", "-p", "--no-session"];
+	if (agent.model) args.push("--model", agent.model);
+	if (agent.tools && agent.tools.length > 0) args.push("--tools", agent.tools.join(","));
+
+	let tmpPromptDir: string | null = null;
+	let tmpPromptPath: string | null = null;
+
+	const currentResult: SingleResult = {
+		agent: agentName,
+		agentSource: agent.source,
+		task,
+		exitCode: 0,
+		messages: [],
+		stderr: "",
+		usage: emptyUsage(),
+		model: agent.model,
+	};
+
+	const emitUpdate = (): void => {
+		onUpdate?.({
+			content: [{ type: "text", text: getFinalOutput(currentResult.messages) || "(running...)" }],
+			details: makeDetails([currentResult]),
+		});
+	};
+
+	try {
+		if (agent.systemPrompt.trim()) {
+			const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt);
+			tmpPromptDir = tmp.dir;
+			tmpPromptPath = tmp.filePath;
+			args.push("--append-system-prompt", tmpPromptPath);
+		}
+
+		args.push(`Task: ${task}`);
+		let wasAborted = false;
+
+		// Increment depth so nested sub-agents can be guarded against runaway recursion.
+		const childDepth = currentSubagentDepth(options.env) + 1;
+		const childEnv: NodeJS.ProcessEnv = {
+			...(options.env ?? process.env),
+			[DEPTH_ENV_VAR]: String(childDepth),
+		};
+
+		const exitCode = await new Promise<number>((resolve) => {
+			const invocation = getPiInvocation(args);
+			const proc = spawn(invocation.command, invocation.args, {
+				cwd: cwd ?? options.defaultCwd,
+				shell: false,
+				stdio: ["ignore", "pipe", "pipe"],
+				env: childEnv,
+			});
+			let buffer = "";
+
+			const processLine = (line: string): void => {
+				if (!line.trim()) return;
+				let event: any;
+				try {
+					event = JSON.parse(line);
+				} catch {
+					return;
+				}
+
+				if (event.type === "message_end" && event.message) {
+					const msg = event.message as Message;
+					currentResult.messages.push(msg);
+					if (msg.role === "assistant") {
+						currentResult.usage.turns++;
+						const usage = (msg as any).usage;
+						if (usage) {
+							currentResult.usage.input += usage.input || 0;
+							currentResult.usage.output += usage.output || 0;
+							currentResult.usage.cacheRead += usage.cacheRead || 0;
+							currentResult.usage.cacheWrite += usage.cacheWrite || 0;
+							currentResult.usage.cost += usage.cost?.total || 0;
+							currentResult.usage.contextTokens = usage.totalTokens || 0;
+						}
+						if (!currentResult.model && (msg as any).model) currentResult.model = (msg as any).model;
+						if ((msg as any).stopReason) currentResult.stopReason = (msg as any).stopReason;
+						if ((msg as any).errorMessage) currentResult.errorMessage = (msg as any).errorMessage;
+					}
+					emitUpdate();
+				}
+
+				if (event.type === "tool_result_end" && event.message) {
+					currentResult.messages.push(event.message as Message);
+					emitUpdate();
+				}
+			};
+
+			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) => {
+				currentResult.stderr += data.toString();
+			});
+
+			proc.on("close", (code) => {
+				if (buffer.trim()) processLine(buffer);
+				resolve(code ?? 0);
+			});
+
+			proc.on("error", () => resolve(1));
+
+			if (signal) {
+				const killProc = (): void => {
+					wasAborted = true;
+					proc.kill("SIGTERM");
+					setTimeout(() => {
+						if (!proc.killed) proc.kill("SIGKILL");
+					}, 5000);
+				};
+				if (signal.aborted) killProc();
+				else signal.addEventListener("abort", killProc, { once: true });
+			}
+		});
+
+		currentResult.exitCode = exitCode;
+		if (wasAborted) throw new Error("Subagent was aborted");
+		return currentResult;
+	} finally {
+		if (tmpPromptPath)
+			try {
+				unlinkSync(tmpPromptPath);
+			} catch {
+				/* ignore */
+			}
+		if (tmpPromptDir)
+			try {
+				rmdirSync(tmpPromptDir);
+			} catch {
+				/* ignore */
+			}
+	}
+}