pi-agent-squad 0.8.0 → 0.8.1

package/session-ui.ts

@@ -185,7 +185,19 @@ class SubagentSessionComponent implements Component, Focusable {
 
 	private pushItem(item: TranscriptItem): void {
 		this.items.push(item);
-		if (this.items.length > 160) this.items.splice(0, this.items.length - 160);
+		if (this.items.length <= 160) return;
+		this.items.splice(0, this.items.length - 160);
+		const retainedTools = new Set<string>();
+		for (const retained of this.items) {
+			if (retained.kind !== "compact") continue;
+			for (const tool of retained.state.tools) retainedTools.add(tool.id);
+		}
+		for (const id of this.tools.keys()) {
+			if (!retainedTools.has(id)) {
+				this.tools.delete(id);
+				this.toolGroups.delete(id);
+			}
+		}
 	}
 
 	private addUser(text: string): UserTranscriptItem | undefined {
@@ -602,9 +614,15 @@ class SubagentSessionComponent implements Component, Focusable {
 	}
 }
 
-export async function openSubagentSessionOverlay(ctx: any, session: SubagentSessionHandle): Promise<void> {
+export async function openSubagentSessionOverlay(
+	ctx: any,
+	session: SubagentSessionHandle,
+	registerClose?: (close: (() => void) | undefined) => void,
+): Promise<void> {
 	const bufferedEvents: any[] = [];
 	let component: SubagentSessionComponent | undefined;
+	let finishOverlay: (() => void) | undefined;
+	let visibilityTimer: ReturnType<typeof setInterval> | undefined;
 	const unsubscribe = session.subscribe((event) => {
 		if (component) component.pushEvent(event);
 		else bufferedEvents.push(event);
@@ -616,8 +634,9 @@ export async function openSubagentSessionOverlay(ctx: any, session: SubagentSess
 		/* a just-starting process may not have history yet */
 	}
 	try {
-		await ctx.ui.custom<void>(
+		await ctx.ui.custom(
 			(tui: any, theme: any, keybindings: any, done: () => void) => {
+				finishOverlay = () => done();
 				component = new SubagentSessionComponent(
 					tui,
 					theme,
@@ -639,10 +658,32 @@ export async function openSubagentSessionOverlay(ctx: any, session: SubagentSess
 					anchor: "center",
 					margin: 0,
 				},
+				onHandle: (handle: any) => {
+					// The host's hideOverlay() removes the entry without calling
+					// done() or component.dispose(). Detect that external hide and
+					// complete the custom UI promise ourselves.
+					visibilityTimer = setInterval(() => {
+						if (!handle.isHidden() && !handle.isFocused()) {
+							component?.dispose();
+							finishOverlay?.();
+						}
+					}, 500);
+					visibilityTimer.unref?.();
+					registerClose?.(() => {
+						// Pi's host may remove an overlay without calling done()
+						// or disposing its component. Complete the custom promise
+						// explicitly so callers can tear down the session safely.
+						component?.dispose();
+						finishOverlay?.();
+					});
+				},
 			},
 		);
 	} catch (error) {
 		if (!component) unsubscribe();
 		throw error;
+	} finally {
+		if (visibilityTimer) clearInterval(visibilityTimer);
+		registerClose?.(undefined);
 	}
 }

package/session.ts

@@ -3,6 +3,13 @@ export interface SubagentSessionHandle {
 	getMessages(): Promise<any[]>;
 	isStreaming(): Promise<boolean>;
 	send(message: string): Promise<void>;
+	/**
+	 * Deliver a routed message to this exact session and wait until that
+	 * message's turn settles.  The UI only needs `send`; the message router
+	 * uses this stronger operation so it never accidentally falls back to a
+	 * different process with the same logical agent name.
+	 */
+	sendAndWait?(message: string, timeoutMs: number, signal?: AbortSignal): Promise<string>;
 	abort(): Promise<void>;
 	subscribe(listener: (event: any) => void): () => void;
 }

package/spawn.ts

@@ -1,4 +1,4 @@
-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";
@@ -43,6 +43,8 @@ export interface SingleResult {
 export interface SpawnOptions {
 	agent: AgentConfig;
 	task: string;
+	/** Runtime address. Defaults to the logical agent name for compatibility. */
+	address?: string;
 	cwd?: string;
 	messageRoot: string;
 	runId: string;
@@ -58,6 +60,33 @@ export interface InteractiveSpawnOptions extends SpawnOptions {
 
 const DEFAULT_SPAWN_TIMEOUT_MS = 6 * 60 * 60 * 1000;
 const FORCE_KILL_DELAY_MS = 5000;
+const TEARDOWN_ABORT_CAP_MS = 3000;
+const CRASH_STDERR_MAX_CHARS = 2000;
+
+/** Bound a best-effort teardown step so a stuck child can never stall cleanup. */
+async function raceWithCap(promise: Promise<void>, capMs: number): Promise<void> {
+	let capTimer: ReturnType<typeof setTimeout> | undefined;
+	try {
+		await Promise.race([
+			promise,
+			new Promise<void>((resolve) => {
+				capTimer = setTimeout(resolve, capMs);
+				capTimer.unref?.();
+			}),
+		]);
+	} finally {
+		if (capTimer) clearTimeout(capTimer);
+	}
+}
+
+/**
+ * Best-effort access to the child process behind an RpcClient. The field is
+ * private upstream; observing it lets a crash settle the run immediately
+ * instead of waiting out the run timeout. Returns null when unavailable.
+ */
+function getChildProcess(client: RpcClient): ChildProcess | null {
+	return (client as unknown as { process?: ChildProcess | null }).process ?? null;
+}
 
 function getPiInvocation(args: string[]): { command: string; args: string[] } {
 	const currentScript = process.argv[1];
@@ -75,8 +104,13 @@ async function writePromptToTempFile(agentName: string, prompt: string): Promise
 	const tmpDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
 	const safeName = agentName.replace(/[^\w.-]+/g, "_");
 	const filePath = path.join(tmpDir, `prompt-${safeName}.md`);
-	await fs.promises.writeFile(filePath, prompt, { encoding: "utf-8", mode: 0o600 });
-	return filePath;
+	try {
+		await fs.promises.writeFile(filePath, prompt, { encoding: "utf-8", mode: 0o600 });
+		return filePath;
+	} catch (error) {
+		await fs.promises.rm(tmpDir, { recursive: true, force: true }).catch(() => {});
+		throw error;
+	}
 }
 
 export function getFinalOutput(messages: SingleResult["messages"]): string {
@@ -100,10 +134,19 @@ export function getFinalOutput(messages: SingleResult["messages"]): string {
  * one-shot process, this keeps one in-memory session alive for the duration of
  * the task so the main TUI can attach, stream events, steer it, and return to
  * main without losing the subagent's context.
+ *
+ * Lifecycle: the initial prompt, run settle, timeout, abort signal, and child
+ * process exit all feed one terminal-cause promise; whichever fires first
+ * decides the outcome. `finally` then tears down the client, timers, and
+ * listeners exactly once, so a timeout, cancel, or crash can never leave a
+ * waiter or a child process behind.
  */
 export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): Promise<SingleResult> {
 	const { agent, task, cwd, messageRoot, runId, childIndex, signal } = opts;
+	if (signal?.aborted) throw new Error("Subagent was aborted");
+	const address = opts.address ?? agent.name;
 	const timeoutMs = Math.max(1000, opts.timeoutMs ?? DEFAULT_SPAWN_TIMEOUT_MS);
+	const timeoutMessage = `Subagent timed out after ${Math.round(timeoutMs / 1000)}s`;
 	const result: SingleResult = {
 		agent: agent.name,
 		task,
@@ -114,13 +157,52 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 		model: agent.model,
 	};
 
+	// ---- unified cancellable lifecycle -------------------------------------
+	type TerminalCause =
+		| { kind: "settled" }
+		| { kind: "timeout" }
+		| { kind: "aborted" }
+		| { kind: "crashed"; error: Error };
+	let terminalCause: TerminalCause | undefined;
+	let settleLifecycle!: (cause: TerminalCause) => void;
+	const lifecycle = new Promise<TerminalCause>((resolve) => {
+		settleLifecycle = (cause) => {
+			if (terminalCause) return; // first cause wins, later ones are no-ops
+			terminalCause = cause;
+			resolve(cause);
+		};
+	});
+	// Rejects when the child process dies so in-flight session waiters fail
+	// promptly instead of hanging until their own timeout.
+	let rejectCrashed!: (error: Error) => void;
+	const crashed = new Promise<never>((_resolve, reject) => {
+		rejectCrashed = reject;
+	});
+	crashed.catch(() => {}); // a crash after every waiter is gone is not an error
+
+	// Resolved once the initial prompt was accepted; session sends gate on it.
+	let initialPromptAccepted = false;
+	let acceptInitialPrompt!: () => void;
+	let rejectInitialPrompt!: (error: unknown) => void;
+	const initialPromptAcceptedPromise = new Promise<void>((resolve, reject) => {
+		acceptInitialPrompt = () => {
+			initialPromptAccepted = true;
+			resolve();
+		};
+		rejectInitialPrompt = (error: unknown) => {
+			initialPromptAccepted = true;
+			reject(error);
+		};
+	});
+	void initialPromptAcceptedPromise.catch(() => {});
+
 	let tmpPromptPath: string | null = null;
 	let client: RpcClient | undefined;
-	let unsubscribe: (() => void) | undefined;
-	let abortHandler: (() => void) | undefined;
+	let childProcess: ChildProcess | null = null;
+	let unsubscribeEvents: (() => void) | undefined;
+	let signalAbortHandler: (() => void) | undefined;
+	let childExitListener: ((code: number | null, exitSignal: string | null) => void) | undefined;
 	let timeoutTimer: ReturnType<typeof setTimeout> | undefined;
-	let timedOut = false;
-	let wasAborted = false;
 	const sessionListeners = new Set<(event: any) => void>();
 
 	const emitSessionEvent = (event: any) => {
@@ -132,6 +214,12 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 			}
 		}
 	};
+	const sessionClosedError = (): string | undefined => {
+		if (terminalCause?.kind === "aborted") return "Subagent was aborted";
+		if (terminalCause?.kind === "timeout") return timeoutMessage;
+		if (terminalCause?.kind === "crashed") return terminalCause.error.message;
+		return undefined;
+	};
 
 	try {
 		const args: string[] = ["--no-session"];
@@ -153,9 +241,9 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 		env[ENV_ROLE] = ROLE_CHILD;
 		env[ENV_CHANNEL_ROOT] = messageRoot;
 		env[ENV_RUN_ID] = runId;
-		env[ENV_AGENT] = agent.name;
+		env[ENV_AGENT] = address;
 		env[ENV_CHILD_INDEX] = String(childIndex);
-		ensureDir(channelDir(messageRoot, runId, agent.name, childIndex));
+		ensureDir(channelDir(messageRoot, runId, address, childIndex));
 
 		const currentScript = process.argv[1];
 		const cliPath =
@@ -169,8 +257,9 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 			model: agent.model,
 			args,
 		});
+		const activeClient = client;
 
-		unsubscribe = client.onEvent((event: any) => {
+		unsubscribeEvents = client.onEvent((event: any) => {
 			opts.onEvent?.(event);
 			emitSessionEvent(event);
 			if (event?.type === "message_end" && event.message) {
@@ -192,27 +281,93 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 					if (message.errorMessage) result.errorMessage = message.errorMessage;
 				}
 			}
+			// Run settle is owned by the lifecycle (this replaces
+			// RpcClient.waitForIdle, whose internal timer and listener could not
+			// be cancelled during teardown). The subscription above is live
+			// before the prompt is sent, so a very fast run cannot settle in the
+			// gap between prompt acceptance and waiter setup.
+			if (event?.type === "agent_settled") settleLifecycle({ kind: "settled" });
 		});
 
+		// Timeout and abort feed the same terminal-cause promise as settle/crash.
+		timeoutTimer = setTimeout(() => settleLifecycle({ kind: "timeout" }), timeoutMs);
+		timeoutTimer.unref?.();
+		signalAbortHandler = () => settleLifecycle({ kind: "aborted" });
+		if (signal?.aborted) signalAbortHandler();
+		else signal?.addEventListener("abort", signalAbortHandler, { once: true });
+
 		await client.start();
-		const activeClient = client;
-		let acceptInitialPrompt!: () => void;
-		let rejectInitialPrompt!: (error: unknown) => void;
-		const initialPromptAccepted = new Promise<void>((resolve, reject) => {
-			acceptInitialPrompt = resolve;
-			rejectInitialPrompt = reject;
-		});
-		void initialPromptAccepted.catch(() => {});
+		// RpcClient creates its child during start(), so inspect it only now.
+		// Looking before start() always returned null and missed child crashes.
+		childProcess = getChildProcess(activeClient);
+		if (childProcess) {
+			childExitListener = (code, exitSignal) => {
+				const stderr = activeClient.getStderr().trim();
+				const error = new Error(
+					`Subagent process exited unexpectedly (code=${code} signal=${exitSignal})` +
+						(stderr ? `: ${stderr.slice(-CRASH_STDERR_MAX_CHARS)}` : "."),
+				);
+				rejectCrashed(error);
+				settleLifecycle({ kind: "crashed", error });
+			};
+			childProcess.once("exit", childExitListener);
+			if (childProcess.exitCode !== null || childProcess.signalCode !== null) {
+				childExitListener(childProcess.exitCode, childProcess.signalCode);
+			}
+		}
 		const session: SubagentSessionHandle = {
-			agent: agent.name,
+			agent: address,
 			getMessages: () => activeClient.getMessages(),
 			isStreaming: async () => (await activeClient.getState()).isStreaming,
 			send: async (message: string) => {
-				await initialPromptAccepted;
+				await initialPromptAcceptedPromise;
 				const state = await activeClient.getState();
 				if (state.isStreaming) await activeClient.steer(message);
 				else await activeClient.prompt(message);
 			},
+			sendAndWait: async (message: string, messageTimeoutMs: number, messageSignal?: AbortSignal) => {
+				await initialPromptAcceptedPromise;
+				if (messageSignal?.aborted) throw new Error("Message delivery cancelled.");
+				const before = await activeClient.getMessages();
+				// Own settle detection so the timer and listener are always
+				// cancelled below (RpcClient.waitForIdle would leak both).
+				let resolveSettled!: () => void;
+				let rejectSettled!: (error: unknown) => void;
+				const settled = new Promise<void>((resolve, reject) => {
+					resolveSettled = resolve;
+					rejectSettled = reject;
+				});
+				let unsubscribeSettle: (() => void) | undefined;
+				let settleTimer: ReturnType<typeof setTimeout> | undefined;
+				let abortHandler: (() => void) | undefined;
+				const aborted = new Promise<never>((_resolve, reject) => {
+					abortHandler = () => reject(new Error("Message delivery cancelled."));
+					if (messageSignal?.aborted) abortHandler();
+					else messageSignal?.addEventListener("abort", abortHandler, { once: true });
+				});
+				try {
+					unsubscribeSettle = activeClient.onEvent((event: any) => {
+						if (event?.type === "agent_settled") resolveSettled();
+					});
+					settleTimer = setTimeout(() => {
+						rejectSettled(
+							new Error(`Timeout waiting for agent to become idle. Stderr: ${activeClient.getStderr()}`),
+						);
+					}, messageTimeoutMs);
+					settleTimer.unref?.();
+					const state = await activeClient.getState();
+					if (state.isStreaming) await activeClient.steer(message);
+					else await activeClient.prompt(message);
+					// `crashed` fails fast if the child dies mid-message.
+					await Promise.race([settled, aborted, crashed]);
+				} finally {
+					if (settleTimer) clearTimeout(settleTimer);
+					unsubscribeSettle?.();
+					if (messageSignal && abortHandler) messageSignal.removeEventListener("abort", abortHandler);
+				}
+				const after = await activeClient.getMessages();
+				return getFinalOutput(after.slice(before.length) as unknown as SingleResult["messages"]);
+			},
 			abort: () => activeClient.abort(),
 			subscribe: (listener) => {
 				sessionListeners.add(listener);
@@ -221,61 +376,66 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 		};
 		opts.onSession?.(session);
 
-		const timeout = new Promise<never>((_resolve, reject) => {
-			timeoutTimer = setTimeout(() => {
-				timedOut = true;
-				reject(new Error(`Subagent timed out after ${Math.round(timeoutMs / 1000)}s`));
-			}, timeoutMs);
-			timeoutTimer.unref?.();
-		});
-		const aborted = new Promise<never>((_resolve, reject) => {
-			abortHandler = () => {
-				wasAborted = true;
-				reject(new Error("Subagent was aborted"));
-			};
-			if (signal?.aborted) abortHandler();
-			else signal?.addEventListener("abort", abortHandler, { once: true });
-		});
-
-		// Register the settled waiter before prompting so an extremely fast run
-		// cannot finish in the gap between prompt acceptance and listener setup.
-		const settled = activeClient.waitForIdle(timeoutMs + 1000);
-		try {
-			await activeClient.prompt(`Task: ${task}`);
-			acceptInitialPrompt();
-		} catch (error) {
-			rejectInitialPrompt(error);
-			throw error;
-		}
-		try {
-			await Promise.race([settled, timeout, aborted]);
-		} catch (error) {
-			await activeClient.abort().catch(() => {});
-			if (wasAborted) throw error;
-			if (timedOut) {
-				result.exitCode = 124;
-				result.stopReason = "error";
-				result.errorMessage = `Subagent timed out after ${Math.round(timeoutMs / 1000)}s`;
-			} else {
-				result.exitCode = 1;
-				result.stopReason = "error";
-				result.errorMessage = error instanceof Error ? error.message : String(error);
+		// Never launch the initial task for a run that is already terminal.
+		if (!terminalCause) {
+			try {
+				await Promise.race([
+					activeClient.prompt(`Task: ${task}`),
+					lifecycle.then((cause) => {
+						if (cause.kind === "timeout") throw new Error(timeoutMessage);
+						if (cause.kind === "aborted") throw new Error("Subagent was aborted");
+						if (cause.kind === "crashed") throw cause.error;
+					}),
+				]);
+				acceptInitialPrompt();
+			} catch (error) {
+				rejectInitialPrompt(error);
+				throw error;
 			}
 		}
+
+		const cause = terminalCause ?? (await lifecycle);
+		if (cause.kind !== "settled") {
+			// Ask the child to finish its turn, but never let an unresponsive
+			// child stall teardown — stop() below escalates to SIGKILL.
+			await raceWithCap(activeClient.abort().catch(() => {}), TEARDOWN_ABORT_CAP_MS);
+			if (cause.kind === "aborted") throw new Error("Subagent was aborted");
+			if (cause.kind === "crashed") throw cause.error;
+			// timeout: report as a timed-out result (previous behavior)
+			result.exitCode = 124;
+			result.stopReason = "error";
+			result.errorMessage = timeoutMessage;
+		}
 		result.stderr = activeClient.getStderr();
 		return result;
 	} finally {
 		if (timeoutTimer) clearTimeout(timeoutTimer);
-		if (signal && abortHandler) signal.removeEventListener("abort", abortHandler);
-		unsubscribe?.();
-		if (client) await client.stop().catch(() => {});
+		if (signal && signalAbortHandler) signal.removeEventListener("abort", signalAbortHandler);
+		if (childProcess && childExitListener) {
+			try {
+				childProcess.off("exit", childExitListener);
+			} catch {
+				/* process may already be gone */
+			}
+		}
+		unsubscribeEvents?.();
+		// Release anything still gated on the initial prompt so no session
+		// caller can hang on a run that ended before it started.
+		if (!initialPromptAccepted) {
+			rejectInitialPrompt(new Error("Subagent session closed before the initial prompt was accepted."));
+		}
+		if (client) {
+			const childDied =
+				childProcess !== null && (childProcess.exitCode !== null || childProcess.signalCode !== null);
+			if (!childDied) {
+				// stop() waits up to its internal 1s kill timer; skip it when the
+				// child already exited so a crash tears down promptly.
+				await client.stop().catch(() => {});
+			}
+		}
 		emitSessionEvent({
 			type: "session_closed",
-			error: timedOut
-				? `Subagent timed out after ${Math.round(timeoutMs / 1000)}s`
-				: wasAborted
-					? "Subagent was aborted"
-					: undefined,
+			error: sessionClosedError(),
 		});
 		sessionListeners.clear();
 		if (tmpPromptPath)
@@ -291,6 +451,7 @@ export async function spawnInteractiveSubagent(opts: InteractiveSpawnOptions): P
 /** Spawn a subagent (isolated context, message channel + system prompt), parse the JSON event stream */
 export async function spawnSubagent(opts: SpawnOptions): Promise<SingleResult> {
 	const { agent, task, cwd, messageRoot, runId, childIndex, signal } = opts;
+	const address = opts.address ?? agent.name;
 	const timeoutMs = Math.max(1000, opts.timeoutMs ?? DEFAULT_SPAWN_TIMEOUT_MS);
 	const emptyUsage = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, contextTokens: 0, turns: 0 };
 	const result: SingleResult = {
@@ -305,6 +466,8 @@ export async function spawnSubagent(opts: SpawnOptions): Promise<SingleResult> {
 
 	let tmpPromptPath: string | null = null;
 	try {
+		// A cancel that arrives before spawn must not launch the process at all.
+		if (signal?.aborted) throw new Error("Subagent was aborted");
 		const args: string[] = ["--mode", "json", "-p", "--no-session"];
 		if (agent.model) args.push("--model", agent.model);
 		if (agent.thinking) args.push("--thinking", agent.thinking);
@@ -325,9 +488,9 @@ export async function spawnSubagent(opts: SpawnOptions): Promise<SingleResult> {
 		env[ENV_ROLE] = ROLE_CHILD;
 		env[ENV_CHANNEL_ROOT] = messageRoot;
 		env[ENV_RUN_ID] = runId;
-		env[ENV_AGENT] = agent.name;
+		env[ENV_AGENT] = address;
 		env[ENV_CHILD_INDEX] = String(childIndex);
-		ensureDir(channelDir(messageRoot, runId, agent.name, childIndex));
+		ensureDir(channelDir(messageRoot, runId, address, childIndex));
 
 		let wasAborted = false;
 		let timedOut = false;