@d3ara1n/pi-subagent 0.3.0 → 0.4.0

package/src/spawn.ts

@@ -6,7 +6,7 @@
  * 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";
@@ -16,8 +16,6 @@ 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;
 
-/** 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";
 
@@ -128,6 +126,9 @@ export async function spawnSubagent(
 		systemPrompt?: string;
 		subagentRoles?: string[];
 		timeoutMs?: number;
+		depth?: number;
+		maxTurns?: number;
+		maxCost?: number;
 		signal?: AbortSignal;
 		onProgress?: (update: Partial<SubagentResult>) => void;
 	},
@@ -175,6 +176,8 @@ export async function spawnSubagent(
 		// Spawn process
 		const invocation = getPiInvocation(args);
 		let wasAborted = false;
+		let budgetExceeded = false;
+		let wasTimeout = false;
 		let buffer = "";
 
 		const emitProgress = () => {
@@ -189,6 +192,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 +229,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 +242,8 @@ export async function spawnSubagent(
 							result.output = part.text;
 						}
 					}
+
+					checkBudget();
 				}
 
 				emitProgress();
@@ -232,6 +252,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 +262,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 +300,96 @@ 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) => {
 				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);
+				// Budget stops are intentional (success); timeouts are failures (exit 124, Unix convention);
+				// otherwise use the real exit code (signal kills yield null → 0).
+				resolve(budgetExceeded ? 0 : (wasTimeout ? 124 : (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

@@ -5,6 +5,16 @@
 /** Configuration for the subagent extension. */
 export interface SubagentConfig {
   timeoutMs: 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 +24,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,
+  timeoutMs: 600_000,
+  maxConcurrency: 4,
+  maxDepth: 3,
+  maxTurns: 0,
+  maxCost: 0,
+  history: { enabled: true },
   summary: { role: "utility", enabled: true },
   agentOverrides: {},
 };
@@ -41,6 +60,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 (ms). Falls back to config.timeoutMs when unset. */
+  timeoutMs?: 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 +127,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 */

package/src/utils.test.ts

@@ -0,0 +1,252 @@
+/**
+ * Unit tests for pi-subagent pure helpers.
+ *
+ * Zero-dependency: runs on node's built-in test runner.
+ *   node --test packages/pi-subagent/src/utils.test.ts
+ *
+ * These guard the bug fixes introduced during the improvement rounds:
+ * path-injection (sanitizeFilename), concurrency/abort/negative-active
+ * (AsyncSemaphore), provider-error word list (isProviderError), unknown-tool
+ * formatting (previewArgs), output truncation fallback (truncateOutput).
+ */
+
+import { test, describe } from "node:test";
+import assert from "node:assert/strict";
+import {
+	sanitizeFilename,
+	isProviderError,
+	AsyncSemaphore,
+	previewArgs,
+	truncateOutput,
+	formatTokens,
+	effectiveTimeoutMs,
+} from "./utils.ts";
+import type { SubagentResult, SubagentRole } from "./types.ts";
+
+// ── sanitizeFilename: guards the path-injection fix ──
+describe("sanitizeFilename", () => {
+	test("never yields a path separator (no directory traversal)", () => {
+		// Core security contract: result contains no / or \, so it can't escape the dir via path.join.
+		for (const input of ["../../etc", "../passwd", "/etc/passwd", "a/b/c", "a\\b", "..", "///"]) {
+			const out = sanitizeFilename(input);
+			assert.ok(!out.includes("/"), `${input} -> "${out}" still contains /`);
+			assert.ok(!out.includes("\\"), `${input} -> "${out}" still contains \\`);
+		}
+	});
+	test("empty string falls back to unknown", () => {
+		assert.equal(sanitizeFilename(""), "unknown");
+	});
+	test("pure-dots collapses to unknown (leading dots stripped, rest empty)", () => {
+		assert.equal(sanitizeFilename(".."), "unknown");
+		assert.equal(sanitizeFilename("..."), "unknown");
+	});
+	test("special chars become underscores", () => {
+		assert.equal(sanitizeFilename("!!!"), "___");
+		assert.equal(sanitizeFilename("   "), "___");
+		assert.equal(sanitizeFilename("///"), "___");
+		assert.equal(sanitizeFilename("a/b/c"), "a_b_c");
+	});
+	test("keeps normal uuid/alnum/dots/dashes as-is", () => {
+		const id = "019eff4f-b603-7623-9eaa-17d32eb623d9";
+		assert.equal(sanitizeFilename(id), id);
+		assert.equal(sanitizeFilename("call_abc123.json"), "call_abc123.json");
+	});
+});
+
+// ── isProviderError: guards the #9 expanded word list ──
+describe("isProviderError", () => {
+	const mk = (stderr: string, errorMessage = ""): SubagentResult =>
+		({
+			stderr,
+			errorMessage,
+			role: "",
+			task: "",
+			exitCode: 0,
+			messages: [],
+			output: "",
+			usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 },
+			activityLog: [],
+		}) as unknown as SubagentResult;
+
+	test("matches provider error keywords", () => {
+		const cases = [
+			"429 Too Many Requests",
+			"quota exceeded",
+			"rate limit exceeded",
+			"authentication error",
+			"request timeout",
+			"quota exhausted",
+			"service unavailable",
+			"503 Service Unavailable",
+			"internal server error",
+			"temporary failure",
+			"request declined",
+			"server overloaded",
+			"ECONNRESET",
+			"socket hang up",
+			"EPIPE",
+			"network error",
+			"connection refused",
+		];
+		for (const c of cases) {
+			assert.equal(isProviderError(mk(c)), true, `should match: ${c}`);
+		}
+	});
+	test("does not match business/programming errors", () => {
+		assert.equal(isProviderError(mk("TypeError: Cannot read properties of undefined")), false);
+		assert.equal(isProviderError(mk("Error: test failed, expected 5 got 3")), false);
+		assert.equal(isProviderError(mk("AssertionError: values differ")), false);
+		assert.equal(isProviderError(mk("")), false);
+	});
+	test("checks errorMessage too, not just stderr", () => {
+		assert.equal(isProviderError(mk("", "rate limited")), true);
+	});
+});
+
+// ── AsyncSemaphore: guards concurrency cap, negative-active, abort cleanup ──
+describe("AsyncSemaphore", () => {
+	test("never goes negative on extra release", async () => {
+		const s = new AsyncSemaphore(1);
+		await s.acquire();
+		s.release();
+		s.release();
+		s.release();
+		assert.equal((s as any).active, 0);
+	});
+	test("respects concurrency cap (queues beyond max)", async () => {
+		const s = new AsyncSemaphore(2);
+		await s.acquire();
+		await s.acquire();
+		let entered = false;
+		const p = s.acquire().then(() => {
+			entered = true;
+		});
+		await Promise.resolve();
+		await Promise.resolve();
+		assert.equal(entered, false); // still queued
+		s.release();
+		await p;
+		assert.equal(entered, true);
+	});
+	test("abort removes waiter from queue and rejects", async () => {
+		const s = new AsyncSemaphore(1);
+		await s.acquire();
+		const c = new AbortController();
+		const p = s.acquire(c.signal);
+		c.abort();
+		await assert.rejects(p);
+		assert.equal((s as any).waiters.length, 0);
+	});
+	test("releases queued waiters in FIFO order", async () => {
+		const s = new AsyncSemaphore(1);
+		await s.acquire();
+		const order: number[] = [];
+		const p1 = s.acquire().then(() => order.push(1));
+		const p2 = s.acquire().then(() => order.push(2));
+		const p3 = s.acquire().then(() => order.push(3));
+		s.release();
+		await p1;
+		s.release();
+		await p2;
+		s.release();
+		await p3;
+		assert.deepEqual(order, [1, 2, 3]);
+	});
+	test("acquires immediately when under cap", async () => {
+		const s = new AsyncSemaphore(3);
+		await s.acquire();
+		await s.acquire();
+		assert.equal((s as any).active, 2);
+	});
+});
+
+// ── previewArgs: guards the #10 shape-based formatting ──
+describe("previewArgs", () => {
+	test("command -> $ prefix", () => {
+		assert.equal(previewArgs({ command: "ls -la" }), "$ ls -la");
+	});
+	test("command truncated at 60 chars", () => {
+		const long = "x".repeat(70);
+		const r = previewArgs({ command: long });
+		assert.ok(r.startsWith("$ "));
+		assert.ok(r.endsWith("..."));
+		assert.ok(r.length < long.length);
+	});
+	test("file_path is shortened (home -> ~)", () => {
+		const r = previewArgs({ file_path: "/home/user/foo.ts" });
+		assert.ok(r.includes("foo.ts"));
+	});
+	test("url passthrough (truncated when long)", () => {
+		assert.equal(previewArgs({ url: "https://example.com" }), "https://example.com");
+		const longUrl = "https://" + "x".repeat(70);
+		assert.ok(previewArgs({ url: longUrl }).endsWith("..."));
+	});
+	test("query/pattern/regex/search -> /.../  form", () => {
+		assert.equal(previewArgs({ query: "foo" }), "/foo/");
+		assert.equal(previewArgs({ pattern: "bar" }), "/bar/");
+		assert.equal(previewArgs({ regex: "baz" }), "/baz/");
+		assert.equal(previewArgs({ search: "qux" }), "/qux/");
+	});
+	test("empty object falls back to JSON {}", () => {
+		assert.equal(previewArgs({}), "{}");
+	});
+});
+
+// ── effectiveTimeoutMs: guards delegate-role auto-widening ──
+describe("effectiveTimeoutMs", () => {
+	const role = (tools: string[], timeoutMs?: number): SubagentRole =>
+		({ role: "default", description: "", examples: [], decisionTrigger: "", tools, systemPrompt: "", timeoutMs }) as unknown as SubagentRole;
+
+	test("non-delegate role uses base timeout", () => {
+		assert.equal(effectiveTimeoutMs(role(["read", "grep"]), 600000), 600000);
+	});
+	test("delegate role doubles base when no explicit timeout", () => {
+		assert.equal(effectiveTimeoutMs(role(["read", "delegate"]), 600000), 1200000);
+	});
+	test("explicit roleDef.timeoutMs is always honored (no widening)", () => {
+		assert.equal(effectiveTimeoutMs(role(["read", "delegate"], 300000), 600000), 300000);
+	});
+	test("explicit timeout on non-delegate also honored", () => {
+		assert.equal(effectiveTimeoutMs(role(["read"]), 600000), 600000);
+	});
+});
+
+// ── truncateOutput: guards the #2 head+tail fallback ──
+describe("truncateOutput", () => {
+	test("adds truncation header with original length", () => {
+		const big = "x".repeat(60000);
+		const r = truncateOutput(big);
+		assert.ok(r.startsWith("[Output truncated"));
+		assert.ok(r.includes("60000 chars total"));
+		assert.ok(r.includes("[truncated]"));
+	});
+	test("keeps head and tail, drops the middle", () => {
+		// 120000 chars: 40k H + 40k M + 40k T
+		const content = "H".repeat(40000) + "M".repeat(40000) + "T".repeat(40000);
+		const r = truncateOutput(content);
+		assert.ok(r.includes("H"), "head preserved");
+		assert.ok(r.includes("T"), "tail preserved");
+		assert.ok(!r.includes("M"), "middle dropped");
+	});
+});
+
+// ── formatTokens: boundary correctness ──
+describe("formatTokens", () => {
+	test("under 1000 stays raw", () => {
+		assert.equal(formatTokens(0), "0");
+		assert.equal(formatTokens(999), "999");
+	});
+	test("1000-9999 with one decimal place", () => {
+		assert.equal(formatTokens(1000), "1.0k");
+		assert.equal(formatTokens(9500), "9.5k");
+		// 9999/1000 = 9.999, toFixed(1) rounds up to 10.0
+		assert.equal(formatTokens(9999), "10.0k");
+	});
+	test("10000-999999 rounded to integer k", () => {
+		assert.equal(formatTokens(10000), "10k");
+		assert.equal(formatTokens(999999), "1000k");
+	});
+	test(">= 1000000 in M", () => {
+		assert.equal(formatTokens(1000000), "1.0M");
+	});
+});