@d3ara1n/pi-subagent 0.3.0 → 0.5.0

package/src/spawn.ts

@@ -6,18 +6,16 @@
  * Fires onProgress on each event for streaming updates.
  */
 
-import { spawn } from "node:child_process";
+import { spawn, type ChildProcess } 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 { SubagentMessage, SubagentResult } from "./types.ts";
 
-/** Maximum task length before writing to a temp file (avoids CLI arg limits). */
-const TASK_CHAR_LIMIT = 8000;
+/** Max chars for an inline channel block (context or task) before it spills to a temp @file. */
+const INLINE_LIMIT = 8000;
 
-/** Maximum output characters returned to the main model. Larger outputs are truncated. */
-const MAX_OUTPUT_CHARS = 50_000;
 
 const PI_CODING_AGENT_PACKAGE = "@earendil-works/pi-coding-agent";
 
@@ -126,8 +124,15 @@ export async function spawnSubagent(
 		cwd?: string;
 		tools?: string[];
 		systemPrompt?: string;
+		/** Extra context delivered as a separate channel from the task. */
+		context?: string;
+		/** Reference file paths injected as independent @file args (child reads them directly). */
+		contextFiles?: string[];
 		subagentRoles?: string[];
 		timeoutMs?: number;
+		depth?: number;
+		maxTurns?: number;
+		maxCost?: number;
 		signal?: AbortSignal;
 		onProgress?: (update: Partial<SubagentResult>) => void;
 	},
@@ -144,7 +149,6 @@ export async function spawnSubagent(
 	};
 
 	let tmpDir: string | null = null;
-	let tmpFile: string | null = null;
 
 	try {
 		// Build CLI args
@@ -154,27 +158,77 @@ export async function spawnSubagent(
 			args.push("--tools", options.tools.join(","));
 		}
 
-		// Always create temp dir — used for prompt file, long task file, and as PI_SUBAGENT_TMPDIR for subagent work (e.g. git clone)
+		// Temp dir for: large-context/task spill files, and as PI_SUBAGENT_TMPDIR
+		// for subagent bash work (e.g. git clone). The system prompt no longer uses it.
 		tmpDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
 
-		const promptContent = options.systemPrompt?.trim()
-			? options.systemPrompt + `\n\nPI_SUBAGENT_TMPDIR=${tmpDir}`
-			: `PI_SUBAGENT_TMPDIR=${tmpDir}`;
-		tmpFile = path.join(tmpDir, "prompt.md");
-		await fs.promises.writeFile(tmpFile, promptContent, { encoding: "utf-8", mode: 0o600 });
-		args.push("--append-system-prompt", tmpFile);
+		// ── System prompt channel: inline text via --append-system-prompt ──
+		// pi's resolvePromptInput treats an existing path as a file to read and any
+		// non-path string as literal text, so we pass structured blocks directly —
+		// no temp file, zero disk I/O. Multiple flags are joined with "\n\n".
+		if (options.systemPrompt?.trim()) {
+			args.push(
+				"--append-system-prompt",
+				`<subagent_role>\n${options.systemPrompt.trim()}\n</subagent_role>`,
+			);
+		}
+		args.push(
+			"--append-system-prompt",
+			`<subagent_env>\nPI_SUBAGENT_TMPDIR=${tmpDir}\nAvailable as $PI_SUBAGENT_TMPDIR in bash. Use for git clone and scratch files.\n</subagent_env>`,
+		);
+
+		// ── Context channel: independent size gate ──
+		// Large context spills to @ctx.md (pi auto-wraps in <file>); small context
+		// inlines as a structured <context> tag. Decoupled from the task gate so a
+		// large context never drags a short task into a spill file.
+		let contextInline = false;
+		if (options.context && options.context.trim()) {
+			if (options.context.length > INLINE_LIMIT) {
+				const ctxPath = path.join(tmpDir, "context.md");
+				await fs.promises.writeFile(ctxPath, options.context, { encoding: "utf-8", mode: 0o600 });
+				args.push(`@${ctxPath}`);
+			} else {
+				contextInline = true;
+			}
+		}
 
-		if (task.length > TASK_CHAR_LIMIT) {
+		// ── Reference files channel: each as an independent @ argument ──
+		// pi reads each and wraps in <file name="...">. Content never enters the
+		// parent model's context — the child reads it directly.
+		if (options.contextFiles) {
+			for (const f of options.contextFiles) {
+				args.push(`@${f}`);
+			}
+		}
+
+		// ── Task channel: always inline, always the final block ──
+		// The task is an instruction, not reference material — it stays inline so the
+		// child sees it as the primary directive. A pathologically long task spills.
+		let taskInline = true;
+		if (task.length > INLINE_LIMIT) {
 			const taskPath = path.join(tmpDir, "task.md");
 			await fs.promises.writeFile(taskPath, task, { encoding: "utf-8", mode: 0o600 });
 			args.push(`@${taskPath}`);
-		} else {
-			args.push(`Task: ${task}`);
+			taskInline = false;
 		}
 
+		// ── Compose the message body (inline context + task) ──
+		// @file args are injected by pi BEFORE this message (buildInitialMessage),
+		// so the final shape the child sees is:
+		//   [<file>...spilled context / reference files...</file>]
+		//   [<context>...inline context...</context>]
+		//   [<task>...task...</task>]
+		const messageParts: string[] = [];
+		if (contextInline) messageParts.push(`<context>\n${options.context}\n</context>`);
+		if (taskInline) messageParts.push(`<task>\n${task}\n</task>`);
+		const message = messageParts.join("\n\n");
+		if (message) args.push(message);
+
 		// Spawn process
 		const invocation = getPiInvocation(args);
 		let wasAborted = false;
+		let budgetExceeded = false;
+		let wasTimeout = false;
 		let buffer = "";
 
 		const emitProgress = () => {
@@ -189,6 +243,20 @@ export async function spawnSubagent(
 		};
 
 		let thinkingCounter = 0;
+		// O(1) lookup from toolCallId → activityLog index (was linear find → O(n²) on busy runs)
+		const toolCallIndex = new Map<string, number>();
+
+		// Kill the child when the configured turn/cost budget is exceeded.
+		// Called after each assistant message_end (usage already accumulated).
+		const checkBudget = () => {
+			const mt = options.maxTurns ?? 0;
+			const mc = options.maxCost ?? 0;
+			if (budgetExceeded || wasTimeout) return;
+			if ((mt > 0 && result.usage.turns >= mt) || (mc > 0 && result.usage.cost >= mc)) {
+				budgetExceeded = true;
+				killProc("budget");
+			}
+		};
 
 		const processLine = (line: string) => {
 			if (!line.trim()) return;
@@ -212,7 +280,8 @@ export async function spawnSubagent(
 						result.usage.cacheRead += usage.cacheRead || 0;
 						result.usage.cacheWrite += usage.cacheWrite || 0;
 						result.usage.cost += usage.cost?.total || 0;
-						result.usage.contextTokens = usage.totalTokens || 0;
+						// Peak context size, not last-turn size (accumulating is meaningless; max tells how close to the limit)
+						result.usage.contextTokens = Math.max(result.usage.contextTokens, usage.totalTokens || 0);
 					}
 					if (!result.model && msg.model) result.model = msg.model;
 					if (msg.stopReason) result.stopReason = msg.stopReason;
@@ -224,6 +293,8 @@ export async function spawnSubagent(
 							result.output = part.text;
 						}
 					}
+
+					checkBudget();
 				}
 
 				emitProgress();
@@ -232,6 +303,7 @@ export async function spawnSubagent(
 			// Activity log: track thinking blocks and tool calls in arrival order.
 			// Both update in place so the TUI reflects real-time state.
 			if (event.type === "tool_execution_start" && event.toolCallId) {
+				toolCallIndex.set(event.toolCallId, result.activityLog.length);
 				result.activityLog.push({
 					kind: "toolCall",
 					id: event.toolCallId,
@@ -241,8 +313,8 @@ export async function spawnSubagent(
 				});
 				emitProgress();
 			} else if (event.type === "tool_execution_end" && event.toolCallId) {
-				const entry = result.activityLog.find((a) => a.id === event.toolCallId);
-				if (entry) entry.status = event.isError ? "failed" : "done";
+				const idx = toolCallIndex.get(event.toolCallId);
+				if (idx !== undefined) result.activityLog[idx].status = event.isError ? "failed" : "done";
 				emitProgress();
 			}
 
@@ -279,73 +351,106 @@ export async function spawnSubagent(
 		}
 		// Expose tmpdir as env var so subagent bash commands (e.g. git clone) can use it
 		childEnv.PI_SUBAGENT_TMPDIR = tmpDir;
+		// Propagate nesting depth so child delegate calls can bound recursion
+		childEnv.PI_SUBAGENT_DEPTH = String(options.depth ?? 0);
 
 		let timeoutHandle: ReturnType<typeof setTimeout> | undefined;
+		let proc: ChildProcess | undefined;
+
+		// Shared kill helper used by abort, budget, and timeout paths.
+		// Centralizes reason → stopReason mapping and the SIGTERM → 5s → SIGKILL escalation.
+		const escalationTimers: ReturnType<typeof setTimeout>[] = [];
+		const killProc = (reason: "abort" | "budget" | "timeout") => {
+			if (reason === "abort") wasAborted = true;
+			else if (reason === "budget") {
+				result.stopReason = "budget_exceeded";
+				// Human-readable so the caller/TUI never falls back to raw stderr noise.
+				const mt = options.maxTurns ?? 0;
+				const mc = options.maxCost ?? 0;
+				const why = mt > 0 && result.usage.turns >= mt ? `${result.usage.turns} turns` : `$${result.usage.cost.toFixed(4)}`;
+				result.errorMessage = `Budget exceeded (${why}; partial output returned)`;
+			}
+			else if (reason === "timeout") {
+				result.stopReason = "timeout";
+				wasTimeout = true;
+				// Human-readable message so the caller/TUI never falls back to the
+				// raw stderr (which is full of TUI teardown escape sequences).
+				const secs = Math.round((options.timeoutMs ?? 0) / 1000);
+				result.errorMessage = `Timed out after ${secs}s (completed ${result.usage.turns} turn${result.usage.turns === 1 ? "" : "s"})`;
+			}
+			try { proc?.kill("SIGTERM"); } catch { /* ignore */ }
+			escalationTimers.push(setTimeout(() => {
+				try { if (proc && !proc.killed) proc.kill("SIGKILL"); } catch { /* ignore */ }
+			}, 5000));
+		};
 
 		const exitCode = await new Promise<number>((resolve) => {
-			const proc = spawn(invocation.command, invocation.args, {
+			// Register abort BEFORE spawning to close the (tiny) registration window
+			let onAbort: (() => void) | undefined;
+			if (options.signal) {
+				if (options.signal.aborted) { wasAborted = true; resolve(0); return; }
+				onAbort = () => killProc("abort");
+				options.signal.addEventListener("abort", onAbort, { once: true });
+			}
+
+			const p = spawn(invocation.command, invocation.args, {
 				cwd: options.cwd,
 				env: childEnv,
 				shell: false,
 				stdio: ["ignore", "pipe", "pipe"],
 			});
+			proc = p;
 
-			proc.stdout.on("data", (data: Buffer) => {
+			p.stdout.on("data", (data: Buffer) => {
 				buffer += data.toString();
 				const lines = buffer.split("\n");
 				buffer = lines.pop() || "";
 				for (const line of lines) processLine(line);
 			});
 
-			proc.stderr.on("data", (data: Buffer) => {
+			p.stderr.on("data", (data: Buffer) => {
 				result.stderr += data.toString();
 			});
 
-			proc.on("close", (code) => {
+			p.on("close", (code, signal) => {
 				if (timeoutHandle) clearTimeout(timeoutHandle);
+				for (const t of escalationTimers) clearTimeout(t);
+				if (onAbort && options.signal) options.signal.removeEventListener("abort", onAbort);
 				if (buffer.trim()) processLine(buffer);
-				resolve(code ?? 0);
+
+				// External signal death (OOM killer, segfault, kill -9 from elsewhere)
+				// that we didn't trigger. Distinguish from our own budget/timeout/abort kills
+				// which set the flags before we send the signal.
+				const externalKill = signal !== null && !budgetExceeded && !wasTimeout && !wasAborted;
+				if (externalKill) {
+					result.errorMessage = result.errorMessage || `Subagent killed by signal ${signal}`;
+					result.stopReason = "error";
+				}
+
+				// Budget stops are intentional (success); timeouts and external kills
+				// are failures (non-zero); otherwise use the real exit code.
+				resolve(budgetExceeded ? 0 : (wasTimeout || externalKill ? (code ?? 128) : (code ?? 0)));
 			});
 
-			proc.on("error", () => {
+			p.on("error", (err) => {
+				if (timeoutHandle) clearTimeout(timeoutHandle);
+				for (const t of escalationTimers) clearTimeout(t);
+				if (onAbort && options.signal) options.signal.removeEventListener("abort", onAbort);
+				// Surface the real cause (e.g. ENOENT when pi is not in PATH) instead of "unknown error".
+				result.errorMessage = err?.message || String(err);
 				resolve(1);
 			});
 
-			// Handle abort signal
-			if (options.signal) {
-				const killProc = () => {
-					wasAborted = true;
-					proc.kill("SIGTERM");
-					setTimeout(() => {
-						if (!proc.killed) proc.kill("SIGKILL");
-					}, 5000);
-				};
-				if (options.signal.aborted) killProc();
-				else options.signal.addEventListener("abort", killProc, { once: true });
-			}
-
 			// Handle timeout
 			if (options.timeoutMs && options.timeoutMs > 0) {
-				timeoutHandle = setTimeout(() => {
-					if (!proc.killed) {
-						proc.kill("SIGTERM");
-						setTimeout(() => {
-							if (!proc.killed) proc.kill("SIGKILL");
-						}, 5000);
-					}
-				}, options.timeoutMs);
+				timeoutHandle = setTimeout(() => killProc("timeout"), options.timeoutMs);
 			}
 		});
 
 		result.exitCode = exitCode;
 		if (wasAborted) throw new Error("Subagent was aborted");
-
-		// Truncate large outputs: keep head (findings) + tail (summary), drop middle
-		if (result.output.length > MAX_OUTPUT_CHARS) {
-			const head = result.output.slice(0, 30_000);
-			const tail = result.output.slice(-(MAX_OUTPUT_CHARS - 30_050));
-			result.output = `[Output truncated — ${result.output.length} chars total]\n\n${head}\n\n... [truncated] ...\n\n${tail}`;
-		}
+		// NOTE: large outputs are kept raw here — compression/truncation happens in
+		// the extension layer (index.ts) so the summary model can compress first.
 	} finally {
 		// Cleanup temp directory and all contents
 		if (tmpDir) try { fs.rmSync(tmpDir, { recursive: true, force: true }); } catch { /* ignore */ }

package/src/types.ts

@@ -4,7 +4,18 @@
 
 /** Configuration for the subagent extension. */
 export interface SubagentConfig {
-  timeoutMs: number;
+  /** Per-subagent timeout in seconds. Roles that can `delegate` get 2× automatically when no per-role timeout is set. */
+  timeout: number;
+  /** Max number of subagents allowed to run concurrently. Extras queue with a TUI hint. */
+  maxConcurrency: number;
+  /** Max subagent nesting depth (the top-level session is depth 0). */
+  maxDepth: number;
+  /** Default turn budget (0 = unlimited). Per-role maxTurns overrides this. */
+  maxTurns: number;
+  /** Default cost budget in USD (0 = unlimited). Per-role maxCost overrides this. */
+  maxCost: number;
+  /** Persist each delegate run to .pi/subagent/history/{sessionId}/{id}.json for auditing. */
+  history: SubagentHistoryConfig;
   summary: SubagentSummaryConfig;
   /**
    * Per-role overrides from settings.json. Keyed by role name.
@@ -14,13 +25,22 @@ export interface SubagentConfig {
   agentOverrides: Record<string, Partial<SubagentRole> & { disabled?: boolean }>;
 }
 
+export interface SubagentHistoryConfig {
+  enabled: boolean;
+}
+
 export interface SubagentSummaryConfig {
   role: string;
   enabled: boolean;
 }
 
 export const DEFAULT_CONFIG: SubagentConfig = {
-  timeoutMs: 300_000,
+  timeout: 600,
+  maxConcurrency: 4,
+  maxDepth: 3,
+  maxTurns: 0,
+  maxCost: 0,
+  history: { enabled: true },
   summary: { role: "utility", enabled: true },
   agentOverrides: {},
 };
@@ -41,6 +61,12 @@ export interface SubagentRole {
   tools: string[];
   /** If this role has `delegate`, restrict which roles it may spawn. undefined = no restriction. */
   subagentRoles?: string[];
+  /** Per-role timeout override in seconds. Falls back to config.timeout when unset. */
+  timeout?: number;
+  /** Max assistant turns before the run is killed (0 = use config default; unset = unlimited). */
+  maxTurns?: number;
+  /** Max cumulative cost (USD) before the run is killed (0 = use config default; unset = unlimited). */
+  maxCost?: number;
   /** Fallback pi-model-roles role name when this role's model is unavailable (provider error). Defaults to "default". */
   fallbackRole?: string;
 }
@@ -102,6 +128,10 @@ export interface SubagentResult {
   task: string;
   /** Process exit code (-1 = still running for streaming) */
   exitCode: number;
+  /** True while waiting for a concurrency slot (TUI hint only). */
+  queued?: boolean;
+  /** How `output` was prepared for display: raw, compressed by summary model, or mechanically truncated. */
+  outputMethod?: "raw" | "compressed" | "truncated";
   /** All messages from the event stream (assistant + tool results) */
   messages: SubagentMessage[];
   /** Last assistant text output */
@@ -120,6 +150,16 @@ export interface SubagentResult {
   errorMessage?: string;
   /** Real-time activity log: thinking blocks and tool calls in arrival order. */
   activityLog: ActivityEntry[];
+
+  // ── TUI 渲染辅助字段（非子进程产出，由 execute 层填入）──────────
+  /** 运行开始墙钟时间；仅 queued/running 帧存在，供 TUI 实时算耗时。终态帧无此字段。 */
+  startTime?: number;
+  /** 终态总耗时(ms)，由 execute 在结束时写入；覆盖整个 delegate 区间（含 fallback 重试）。 */
+  elapsedMs?: number;
+  /** delegate 传入的引用文件路径（params.files），展开视图渲染用。 */
+  files?: string[];
+  /** delegate 传入的额外上下文（params.context），展开视图渲染用。 */
+  context?: string;
 }
 
 /** TUI details structure passed via tool result details. */