pi-agent-squad 0.8.4 → 0.8.5

package/index.ts

@@ -36,6 +36,27 @@ import { openSubagentSessionOverlay } from "./session-ui.ts";
 import { getFinalOutput, spawnInteractiveSubagent } from "./spawn.ts";
 import { deadlockMessage, MessageWaitGraph } from "./wait-graph.ts";
 import { ActiveRunRegistry, type ActiveRun } from "./active-runs.ts";
+import {
+	boundTaskText,
+	isTerminalTaskStatus,
+	isRetryableTaskStatus,
+	reconstructTaskStates,
+	TASK_STATE_ENTRY,
+	TASK_STATE_LIMITS,
+	TASK_STATE_VERSION,
+	validateLaunchText,
+	type PersistedTaskState,
+	type PersistedTaskStatus,
+} from "./task-state.ts";
+import { planTaskRecovery } from "./task-recovery.ts";
+import {
+	BACKGROUND_EVENT_TYPE,
+	branchHasResultDelivery,
+	RecoveryOutbox,
+	resultDeliveryId,
+	type BackgroundEventDetails,
+} from "./task-delivery.ts";
+import { formatTaskStatus, type LiveTaskView } from "./task-status.ts";
 
 const MESSAGE_ROOT_BASE = path.join(
 	os.tmpdir(),
@@ -51,7 +72,6 @@ const RUNNING_WIDGET_TICK_MS = 1000;
 const RUNNING_WIDGET_NAV_DEBOUNCE_MS = 150;
 const RUNNING_WIDGET_SPINNER = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
 const INCOMING_MESSAGE_TYPE = "pi_message_request";
-const BACKGROUND_EVENT_TYPE = "pi_subagent_background_event";
 
 interface IncomingMessageDetails {
 	id?: string;
@@ -61,15 +81,6 @@ interface IncomingMessageDetails {
 	createdAt?: number;
 }
 
-interface BackgroundEventDetails {
-	agent?: string;
-	address?: string;
-	status?: "done" | "error";
-	body?: string;
-	elapsedMs?: number;
-	runId?: string;
-}
-
 type TranscriptEventKind = "message" | "done" | "error";
 
 function customMessageText(content: unknown): string {
@@ -351,7 +362,12 @@ interface AsyncTask {
 	runId: string;
 	agent: string;
 	address: string;
+	task: string;
+	cwd: string;
+	readOnly: boolean;
 	startedAt: number;
+	updatedAt: number;
+	timeoutAt: number;
 	sessionGeneration: number;
 }
 
@@ -579,7 +595,14 @@ export class RunningSubagentWidgetController {
 	}
 }
 
-export default function (pi: ExtensionAPI) {
+export interface AgentSquadRuntimeDependencies {
+	spawnInteractiveSubagent: typeof spawnInteractiveSubagent;
+}
+
+export function createAgentSquadExtension(
+	pi: ExtensionAPI,
+	dependencies: AgentSquadRuntimeDependencies = { spawnInteractiveSubagent },
+) {
 	const isChild = process.env[ENV_ROLE] === "child";
 
 	if (isChild) {
@@ -589,6 +612,7 @@ export default function (pi: ExtensionAPI) {
 	}
 
 	// ===== main-agent mode =====
+	const ownerRuntimeId = randomUUID();
 	pi.registerMessageRenderer<IncomingMessageDetails>(
 		INCOMING_MESSAGE_TYPE,
 		(message, options, theme) => {
@@ -650,13 +674,142 @@ export default function (pi: ExtensionAPI) {
 		string,
 		{ release: () => void; timer?: ReturnType<typeof setTimeout> }
 	>();
+	const persistedTasks = new Map<string, PersistedTaskState>();
 	let sessionContext: any;
+	let runtimeContext: any;
+	let currentSessionToken: string | undefined;
+	let acceptingTasks = false;
+	let shutdownStarted = false;
+	let router: ReturnType<typeof createMessageRouter> | undefined;
 	let terminalInputUnsubscribe: (() => void) | undefined;
 	let sessionOverlayOpen = false;
 	let closeSessionOverlay: (() => void) | undefined;
 	let pendingOpenActivityId: string | undefined;
 	let lastNavigation: { direction: -1 | 1; at: number } | undefined;
 
+	function appendTaskSnapshot(state: PersistedTaskState): PersistedTaskState {
+		pi.appendEntry(TASK_STATE_ENTRY, state);
+		persistedTasks.set(state.runId, state);
+		return state;
+	}
+
+	function transitionTask(
+		runId: string,
+		expected: readonly PersistedTaskStatus[],
+		patch: Partial<Omit<PersistedTaskState, "version" | "runId" | "updatedAt">>,
+		now = Date.now(),
+	): PersistedTaskState | undefined {
+		const current = persistedTasks.get(runId);
+		if (!current || !expected.includes(current.status)) return undefined;
+		const next: PersistedTaskState = {
+			...current,
+			...patch,
+			version: TASK_STATE_VERSION,
+			runId,
+			updatedAt: now,
+		};
+		return appendTaskSnapshot(next);
+	}
+
+	function reportPersistenceFailure(error: unknown, prefix: string): void {
+		const message = `${prefix}: ${error instanceof Error ? error.message : String(error)}`;
+		try {
+			runtimeContext?.ui?.notify?.(message, "warning");
+		} catch {
+			/* UI may already be closing */
+		}
+		console.error(`[pi-agent-squad] ${message}`);
+	}
+
+	const recoveryOutbox = new RecoveryOutbox({
+		getSessionToken: () => currentSessionToken,
+		getBranch: () => runtimeContext?.sessionManager?.getBranch?.() ?? [],
+		getTaskState: (runId) => persistedTasks.get(runId),
+		send: (message, options) => pi.sendMessage(message, options),
+		markInjected: (state) => {
+			transitionTask(
+				state.runId,
+				["completed"],
+				{ status: "completed", resultInjected: true },
+			);
+		},
+		onError: (error, state) =>
+			reportPersistenceFailure(
+				error,
+				`Could not finish result delivery bookkeeping for run ${state.runId.slice(0, 8)}`,
+			),
+	});
+
+	function refreshPersistedTasks(entries: readonly unknown[]): void {
+		persistedTasks.clear();
+		for (const [runId, state] of reconstructTaskStates(entries)) {
+			persistedTasks.set(runId, state);
+		}
+	}
+
+	function liveTaskRunIds(): Set<string> {
+		return new Set([
+			...tasks.keys(),
+			...activeRuns.list().map((run) => run.runId),
+		]);
+	}
+
+	function reconcileCurrentBranch(): void {
+		const branch = runtimeContext?.sessionManager?.getBranch?.() ?? [];
+		const currentLiveIds = liveTaskRunIds();
+		const liveSnapshots = [...persistedTasks.values()].filter(
+			(state) =>
+				state.ownerRuntimeId === ownerRuntimeId &&
+				currentLiveIds.has(state.runId) &&
+				(state.status === "starting" || state.status === "running"),
+		);
+		refreshPersistedTasks(branch);
+		for (const state of liveSnapshots) {
+			if (!persistedTasks.has(state.runId)) persistedTasks.set(state.runId, state);
+		}
+		const plan = planTaskRecovery(
+			persistedTasks,
+			ownerRuntimeId,
+			currentLiveIds,
+			(deliveryId) => branchHasResultDelivery(branch, deliveryId),
+		);
+		for (const stale of plan.interrupt) {
+			const now = Date.now();
+			try {
+				transitionTask(
+					stale.runId,
+					["starting", "running"],
+					{
+						status: "interrupted",
+						error: "Owning Pi runtime ended before the task reached a terminal state.",
+						endedAt: now,
+					},
+					now,
+				);
+			} catch (error) {
+				reportPersistenceFailure(
+					error,
+					`Could not mark stale run ${stale.runId.slice(0, 8)} interrupted`,
+				);
+			}
+		}
+		for (const delivered of plan.markInjected) {
+			try {
+				transitionTask(
+					delivered.runId,
+					["completed"],
+					{ status: "completed", resultInjected: true },
+				);
+			} catch (error) {
+				reportPersistenceFailure(
+					error,
+					`Could not repair delivery state for run ${delivered.runId.slice(0, 8)}`,
+				);
+			}
+		}
+		for (const pending of plan.deliver) recoveryOutbox.enqueue(pending);
+	}
+
 	const bindPoolLifecycle = (): void => {
 		pool.setProcessExitHandler((agentName, runId) => {
 			const run = activeRuns.resolveExact(runId);
@@ -715,6 +868,9 @@ export default function (pi: ExtensionAPI) {
 
 	pi.on("session_start", (event, ctx) => {
 		sessionGeneration++;
+		shutdownStarted = false;
+		acceptingTasks = false;
+		recoveryOutbox.cancelAll();
 		if (poolDisposed) {
 			pool = new SubagentPool(messageRoot);
 			bindPoolLifecycle();
@@ -725,6 +881,8 @@ export default function (pi: ExtensionAPI) {
 			ctx.sessionManager?.getSessionId?.() ??
 			(event as any).sessionId ??
 			"ephemeral";
+		currentSessionToken = `${sessionId}:${sessionGeneration}`;
+		runtimeContext = ctx;
 		messageRoot = sessionRoot(sessionId);
 		ensureChannelRoot(MESSAGE_ROOT_BASE);
 		sweepMessageRoots(MESSAGE_ROOT_BASE);
@@ -733,6 +891,8 @@ export default function (pi: ExtensionAPI) {
 		pool.setWorkingDirectory(cwd);
 		runningWidget.attach(ctx);
 		sessionContext = ctx.mode === "tui" ? ctx : undefined;
+		reconcileCurrentBranch();
+		acceptingTasks = true;
 		lastNavigation = undefined;
 		terminalInputUnsubscribe?.();
 		terminalInputUnsubscribe =
@@ -760,7 +920,42 @@ export default function (pi: ExtensionAPI) {
 				: undefined;
 	});
 
-	pi.on("session_shutdown", () => {
+	pi.on("session_shutdown", (event, ctx) => {
+		if (shutdownStarted) return;
+		shutdownStarted = true;
+		acceptingTasks = false;
+		recoveryOutbox.cancelAll();
+		const now = Date.now();
+		for (const state of [...persistedTasks.values()]) {
+			if (
+				state.ownerRuntimeId !== ownerRuntimeId ||
+				(state.status !== "starting" && state.status !== "running")
+			) {
+				continue;
+			}
+			const reason = boundTaskText(
+				[
+					"Owning Pi runtime shut down before the task reached a terminal state.",
+					event.reason ? `Reason: ${event.reason}.` : "",
+				]
+					.filter(Boolean)
+					.join(" "),
+				TASK_STATE_LIMITS.error,
+			);
+			try {
+				transitionTask(
+					state.runId,
+					["starting", "running"],
+					{ status: "interrupted", error: reason, endedAt: now },
+					now,
+				);
+			} catch (error) {
+				reportPersistenceFailure(
+					error,
+					`Could not persist shutdown interruption for run ${state.runId.slice(0, 8)}`,
+				);
+			}
+		}
 		sessionGeneration++;
 		for (const controller of runControllers.values()) controller.abort();
 		runControllers.clear();
@@ -772,16 +967,21 @@ export default function (pi: ExtensionAPI) {
 		sessionContext = undefined;
 		pendingOpenActivityId = undefined;
 		sessionHandles.clear();
-		activeRuns.clear();
 		for (const pending of pendingMainReplyEdges.values()) {
 			if (pending.timer) clearTimeout(pending.timer);
 			pending.release();
 		}
 		pendingMainReplyEdges.clear();
 		waitGraph.clear();
-		runningWidget.shutdown();
 		pool.dispose();
 		poolDisposed = true;
+		router?.dispose();
+		router = undefined;
+		activeRuns.clear();
+		tasks.clear();
+		runningWidget.shutdown();
+		currentSessionToken = undefined;
+		runtimeContext = undefined;
 		sweepMessageRoots(MESSAGE_ROOT_BASE);
 	});
 
@@ -809,7 +1009,6 @@ export default function (pi: ExtensionAPI) {
 	// ---- message router: subagent -> main injects into main session; child
 	// messages resolve through the active-run registry before being delivered
 	// to a direct/background session or the resident pool. ----
-	let router: ReturnType<typeof createMessageRouter> | undefined;
 	const routedTimeout = (msg: MessageRequest): number =>
 		Math.min(
 			MAX_SUBAGENT_TIMEOUT_SECONDS * 1000,
@@ -1009,11 +1208,6 @@ export default function (pi: ExtensionAPI) {
 		router.start();
 	});
 
-	pi.on("session_shutdown", () => {
-		router?.dispose();
-		router = undefined;
-	});
-
 	/** Write a reply back to the message sender */
 	function writeReplyTo(msg: MessageRequest, content: string, root = messageRoot): boolean {
 		try {
@@ -1045,20 +1239,13 @@ export default function (pi: ExtensionAPI) {
 		return `${agent.name}#${runId.slice(0, 8)}`;
 	}
 
-	function createDirectRun(
+	function resolveDirectRunAddress(
 		agent: AgentConfig,
 		runId: string,
-		mode: "background" | "task",
 		requestedAddress?: string,
-		readOnly = agent.readOnly === true,
-		runCwd = cwd,
-	): {
-		address: string;
-		session: Promise<SubagentSessionHandle>;
-		resolveSession: (session: SubagentSessionHandle) => void;
-		rejectSession: (error: Error) => void;
-	} {
+	): string {
 		const address = requestedAddress?.trim() || directRunAddress(agent, runId);
+		validateLaunchText("address", address);
 		if (
 			(requestedAddress &&
 				(!/^[A-Za-z0-9_.-]+$/.test(address) ||
@@ -1073,6 +1260,25 @@ export default function (pi: ExtensionAPI) {
 				`Invalid or unavailable runtime address "${address}". Use a unique name matching [A-Za-z0-9_.-] that is not a logical agent name or "main".`,
 			);
 		}
+		return address;
+	}
+
+	function createDirectRun(
+		agent: AgentConfig,
+		runId: string,
+		mode: "background" | "task",
+		address: string,
+		readOnly = agent.readOnly === true,
+		runCwd = cwd,
+	): {
+		address: string;
+		session: Promise<SubagentSessionHandle>;
+		resolveSession: (session: SubagentSessionHandle) => void;
+		rejectSession: (error: Error) => void;
+	} {
+		if (activeRuns.hasAddress(address)) {
+			throw new Error(`Runtime address "${address}" became unavailable before launch.`);
+		}
 		let resolveSession!: (session: SubagentSessionHandle) => void;
 		let rejectSession!: (error: Error) => void;
 		const session = new Promise<SubagentSessionHandle>((resolve, reject) => {
@@ -1122,29 +1328,77 @@ export default function (pi: ExtensionAPI) {
 		timeoutMs?: number,
 		requestedAddress?: string,
 		readOnly = agent.readOnly === true,
+		retryOf?: string,
 	): { runId: string; address: string } {
+		if (!acceptingTasks) throw new Error("This Pi session is shutting down and cannot accept new subagent tasks.");
+		validateLaunchText("agent", agent.name);
+		validateLaunchText("task", taskText, { allowEmpty: true });
+		const runCwd = normalizeCwd(cwdOverride ?? cwd);
+		validateLaunchText("cwd", runCwd);
 		const runId = randomUUID();
-		const controller = new AbortController();
-		const direct = createDirectRun(
-			agent,
+		const address = resolveDirectRunAddress(agent, runId, requestedAddress);
+		const startedAt = Date.now();
+		const effectiveTimeoutMs = Math.max(
+			1000,
+			timeoutMs ?? DEFAULT_SUBAGENT_TIMEOUT_SECONDS * 1000,
+		);
+		const timeoutAt = Math.min(Number.MAX_SAFE_INTEGER, startedAt + effectiveTimeoutMs);
+		const starting: PersistedTaskState = {
+			version: TASK_STATE_VERSION,
 			runId,
-			"background",
-			requestedAddress,
+			address,
+			agent: agent.name,
+			task: taskText,
+			cwd: runCwd,
 			readOnly,
-			cwdOverride ?? cwd,
-		);
+			mode: "async",
+			status: "starting",
+			startedAt,
+			updatedAt: startedAt,
+			timeoutAt,
+			...(retryOf ? { retryOf } : {}),
+			ownerRuntimeId,
+		};
+		appendTaskSnapshot(starting);
+
+		const controller = new AbortController();
+		let direct: ReturnType<typeof createDirectRun>;
+		try {
+			direct = createDirectRun(agent, runId, "background", address, readOnly, runCwd);
+		} catch (error) {
+			const now = Date.now();
+			transitionTask(
+				runId,
+				["starting"],
+				{
+					status: "failed",
+					error: boundTaskText(
+						error instanceof Error ? error.message : String(error),
+						TASK_STATE_LIMITS.error,
+					),
+					endedAt: now,
+				},
+				now,
+			);
+			throw error;
+		}
 		runControllers.set(runId, controller);
 		tasks.set(runId, {
 			runId,
 			agent: agent.name,
-			address: direct.address,
-			startedAt: Date.now(),
+			address,
+			task: taskText,
+			cwd: runCwd,
+			readOnly,
+			startedAt,
+			updatedAt: startedAt,
+			timeoutAt,
 			sessionGeneration,
 		});
 		runningWidget.start(runId, agent.name, taskText, "background");
 
 		const start = async () => {
-			return await spawnInteractiveSubagent({
+			return await dependencies.spawnInteractiveSubagent({
 				agent,
 				task: taskText,
 				address: direct.address,
@@ -1153,7 +1407,29 @@ export default function (pi: ExtensionAPI) {
 				runId,
 				childIndex: 0,
 				signal: controller.signal,
-				timeoutMs,
+				timeoutMs: effectiveTimeoutMs,
+				persistSession: true,
+				onStarted: (childSessionFile) => {
+					const now = Date.now();
+					const running = transitionTask(
+						runId,
+						["starting"],
+						{
+							status: "running",
+							...(childSessionFile
+								? {
+										childSessionFile: boundTaskText(
+											childSessionFile,
+											TASK_STATE_LIMITS.childSessionFile,
+										),
+									}
+								: {}),
+						},
+						now,
+					);
+					const task = tasks.get(runId);
+					if (task && running) task.updatedAt = running.updatedAt;
+				},
 				onSession: (session) => {
 					direct.resolveSession(session);
 					registerSessionHandle(runId, session);
@@ -1168,47 +1444,82 @@ export default function (pi: ExtensionAPI) {
 				activeRuns.remove(runId);
 				const task = tasks.get(runId);
 				if (!task || task.sessionGeneration !== sessionGeneration) return;
-				const finalStatus: "done" | "error" =
-					result.exitCode === 0 && result.stopReason !== "error" ? "done" : "error";
 				const output = getFinalOutput(result.messages) || "(no text output)";
 				const failureReason =
 					result.exitCode !== 0
-						? result.stderr || result.errorMessage || `process exited with code ${result.exitCode}`
+						? result.errorMessage || result.stderr || `process exited with code ${result.exitCode}`
 						: result.stopReason === "error"
 							? result.errorMessage || "subagent reported an error"
 							: "";
-				const head =
-					finalStatus === "done"
-						? `Subagent ${agent.name} done (run ${runId.slice(0, 8)})`
-						: `Subagent ${agent.name} failed: ${failureReason}`;
-				const displayBody =
-					finalStatus === "done"
-						? output
-						: [
-								failureReason,
+				const cancelled =
+					result.exitCode === 124 ||
+					result.stopReason === "aborted" ||
+					/\b(?:timed out|cancelled|canceled|aborted)\b/i.test(result.errorMessage ?? "");
+				const finalStatus: PersistedTaskStatus =
+					result.exitCode === 0 && result.stopReason !== "error" && result.stopReason !== "aborted"
+						? "completed"
+						: cancelled
+							? "cancelled"
+							: "failed";
+				const now = Date.now();
+				try {
+					if (finalStatus === "completed") {
+						const completed = transitionTask(
+							runId,
+							["starting", "running"],
+							{
+								status: "completed",
+								resultSummary: boundTaskText(output, TASK_STATE_LIMITS.resultSummary),
+								deliveryId: resultDeliveryId(runId),
+								resultInjected: false,
+								endedAt: now,
+							},
+							now,
+						);
+						if (completed) {
+							recoveryOutbox.enqueue(completed, { body: output, content: output });
+						}
+					} else {
+						const diagnostic = boundTaskText(
+							failureReason || result.errorMessage || `Subagent ${finalStatus}.`,
+							TASK_STATE_LIMITS.error,
+						);
+						const terminal = transitionTask(
+							runId,
+							["starting", "running"],
+							{ status: finalStatus, error: diagnostic, endedAt: now },
+							now,
+						);
+						if (terminal) {
+							const displayBody = [
+								diagnostic,
 								output === "(no text output)" ? "" : `**Partial result**\n\n${output}`,
 							]
 								.filter(Boolean)
 								.join("\n\n");
-				try {
-					pi.sendMessage(
-						{
-							customType: BACKGROUND_EVENT_TYPE,
-							content: `[background subagent ${finalStatus}] ${head}\n\n--- result ---\n${output}`,
-							display: true,
-							details: {
-								agent: agent.name,
-								address: direct.address,
-								status: finalStatus,
-								body: displayBody,
-								elapsedMs: Date.now() - task.startedAt,
-								runId,
-							},
-						},
-						{ triggerTurn: true, deliverAs: "steer" },
+							pi.sendMessage(
+								{
+									customType: BACKGROUND_EVENT_TYPE,
+									content: `[background subagent error] ${agent.name} (run ${runId.slice(0, 8)}): ${diagnostic}`,
+									display: true,
+									details: {
+										agent: agent.name,
+										address,
+										status: "error",
+										body: displayBody,
+										elapsedMs: now - task.startedAt,
+										runId,
+									},
+								},
+								{ triggerTurn: true, deliverAs: "steer" },
+							);
+						}
+					}
+				} catch (error) {
+					reportPersistenceFailure(
+						error,
+						`Could not finalize background run ${runId.slice(0, 8)}`,
 					);
-				} catch {
-					/* session is closed */
 				}
 			})
 			.catch((err) => {
@@ -1216,26 +1527,47 @@ export default function (pi: ExtensionAPI) {
 				direct.rejectSession(err instanceof Error ? err : new Error(String(err)));
 				const task = tasks.get(runId);
 				const errorText = err instanceof Error ? err.message : String(err);
+				const current = persistedTasks.get(runId);
+				if (!current || isTerminalTaskStatus(current.status)) return;
+				const cancelled =
+					controller.signal.aborted ||
+					/\b(?:timed out|cancelled|canceled|aborted)\b/i.test(errorText);
+				const now = Date.now();
 				try {
 					if (task && task.sessionGeneration !== sessionGeneration) return;
-					pi.sendMessage(
+					const terminal = transitionTask(
+						runId,
+						["starting", "running"],
 						{
-							customType: BACKGROUND_EVENT_TYPE,
-							content: `[background subagent error] ${agent.name} (run ${runId.slice(0, 8)}): ${errorText}`,
-							display: true,
-							details: {
-								agent: agent.name,
-								address: direct.address,
-								status: "error",
-								body: errorText,
-								elapsedMs: task ? Date.now() - task.startedAt : undefined,
-								runId,
-							},
+							status: cancelled ? "cancelled" : "failed",
+							error: boundTaskText(errorText, TASK_STATE_LIMITS.error),
+							endedAt: now,
 						},
-						{ triggerTurn: true, deliverAs: "steer" },
+						now,
+					);
+					if (terminal) {
+						pi.sendMessage(
+							{
+								customType: BACKGROUND_EVENT_TYPE,
+								content: `[background subagent error] ${agent.name} (run ${runId.slice(0, 8)}): ${errorText}`,
+								display: true,
+								details: {
+									agent: agent.name,
+									address,
+									status: "error",
+									body: errorText,
+									elapsedMs: task ? now - task.startedAt : undefined,
+									runId,
+								},
+							},
+							{ triggerTurn: true, deliverAs: "steer" },
+						);
+					}
+				} catch (error) {
+					reportPersistenceFailure(
+						error,
+						`Could not record background failure for run ${runId.slice(0, 8)}`,
 					);
-				} catch {
-					/* ignore */
 				}
 			})
 			.finally(() => {
@@ -1247,7 +1579,7 @@ export default function (pi: ExtensionAPI) {
 				runningWidget.finish(runId);
 			});
 
-		return { runId, address: direct.address };
+		return { runId, address };
 	}
 
 	// ---- adaptive orchestration prompt; `/orchestrate on` installs the opt-in prompt ----
@@ -1277,11 +1609,12 @@ export default function (pi: ExtensionAPI) {
 			"While running, a subagent may send_message (to=main) to reach you, or contact other subagents — reply promptly with reply_message. Synchronous tasks must use wait=false when contacting main.",
 		].join(" "),
 		parameters: Type.Object({
-			agent: Type.String({ description: "Subagent name" }),
-			task: Type.String({ description: "Task description for the subagent" }),
-			cwd: Type.Optional(Type.String({ description: "Working directory for the subagent" })),
+			agent: Type.String({ maxLength: TASK_STATE_LIMITS.agent, description: "Subagent name" }),
+			task: Type.String({ maxLength: TASK_STATE_LIMITS.task, description: "Task description for the subagent" }),
+			cwd: Type.Optional(Type.String({ maxLength: TASK_STATE_LIMITS.cwd, description: "Working directory for the subagent" })),
 			as: Type.Optional(
 				Type.String({
+					maxLength: TASK_STATE_LIMITS.address,
 					description: "Optional unique runtime address for this run (for example actor-a); identity still comes from agent",
 				}),
 			),
@@ -1301,6 +1634,12 @@ export default function (pi: ExtensionAPI) {
 		}),
 
 		async execute(_toolCallId, params, signal, onUpdate, ctx) {
+			if (!acceptingTasks) {
+				return {
+					content: [{ type: "text", text: "This Pi session is shutting down and cannot accept new subagent tasks." }],
+					details: undefined,
+				};
+			}
 			const agents = discoverAgents(ctx.cwd);
 			const agent = agents.find((a) => a.name === params.agent);
 			if (!agent) {
@@ -1351,7 +1690,9 @@ export default function (pi: ExtensionAPI) {
 			else signal?.addEventListener("abort", forwardAbort, { once: true });
 			let direct: ReturnType<typeof createDirectRun>;
 			try {
-				direct = createDirectRun(agent, runId, "task", params.as, readOnly, params.cwd ?? cwd);
+				validateLaunchText("task", params.task, { allowEmpty: true });
+				const address = resolveDirectRunAddress(agent, runId, params.as);
+				direct = createDirectRun(agent, runId, "task", address, readOnly, params.cwd ?? cwd);
 			} catch (error) {
 				signal?.removeEventListener("abort", forwardAbort);
 				runControllers.delete(runId);
@@ -1373,7 +1714,7 @@ export default function (pi: ExtensionAPI) {
 			runningWidget.start(runId, agent.name, params.task, "task");
 			let result: Awaited<ReturnType<typeof spawnInteractiveSubagent>>;
 			try {
-				result = await spawnInteractiveSubagent({
+				result = await dependencies.spawnInteractiveSubagent({
 					agent,
 					task: params.task,
 					address: direct.address,
@@ -1509,25 +1850,122 @@ export default function (pi: ExtensionAPI) {
 		},
 	});
 
+	pi.registerCommand("subagent-retry", {
+		description: "Retry an interrupted, failed, or cancelled background task",
+		getArgumentCompletions: (prefix) => {
+			const items = [...persistedTasks.values()]
+				.filter((state) => isRetryableTaskStatus(state.status))
+				.sort((a, b) => b.updatedAt - a.updatedAt)
+				.map((state) => ({
+					value: state.runId,
+					label: state.runId,
+					description: `${state.address} · ${state.status}`,
+				}))
+				.filter((item) => item.value.startsWith(prefix.trim()));
+			return items.length > 0 ? items : null;
+		},
+		handler: async (args, ctx) => {
+			runtimeContext = ctx;
+			reconcileCurrentBranch();
+			const raw = Array.isArray(args) ? String(args[0] ?? "") : String(args ?? "");
+			const oldRunId = raw.trim().split(/\s+/)[0] ?? "";
+			if (!oldRunId) {
+				ctx.ui.notify("Usage: /subagent-retry <runId>", "warning");
+				return;
+			}
+			const previous = persistedTasks.get(oldRunId);
+			if (!previous) {
+				ctx.ui.notify(`No task "${oldRunId}" exists on the current branch.`, "error");
+				return;
+			}
+			if (!isRetryableTaskStatus(previous.status)) {
+				ctx.ui.notify(
+					`Task ${oldRunId.slice(0, 8)} is ${previous.status} and cannot be retried. Only interrupted, failed, or cancelled tasks are retryable.`,
+					"warning",
+				);
+				return;
+			}
+			const agent = discoverAgents(ctx.cwd).find((candidate) => candidate.name === previous.agent);
+			if (!agent) {
+				ctx.ui.notify(
+					`Cannot retry run ${oldRunId.slice(0, 8)} because subagent "${previous.agent}" no longer exists.`,
+					"error",
+				);
+				return;
+			}
+			const originalTimeoutMs =
+				previous.timeoutAt !== undefined
+					? Math.max(
+							MIN_SUBAGENT_TIMEOUT_SECONDS * 1000,
+							Math.min(
+								MAX_SUBAGENT_TIMEOUT_SECONDS * 1000,
+								previous.timeoutAt - previous.startedAt,
+							),
+						)
+					: DEFAULT_SUBAGENT_TIMEOUT_SECONDS * 1000;
+			try {
+				const launched = launchBackground(
+					agent,
+					previous.task,
+					previous.cwd,
+					originalTimeoutMs,
+					undefined,
+					previous.readOnly,
+					previous.runId,
+				);
+				ctx.ui.notify(
+					`Retry started at ${launched.address} (new run ${launched.runId.slice(0, 8)}, retry of ${previous.runId.slice(0, 8)}).`,
+					"info",
+				);
+			} catch (error) {
+				ctx.ui.notify(
+					`Failed to retry run ${oldRunId.slice(0, 8)}: ${error instanceof Error ? error.message : String(error)}`,
+					"error",
+				);
+			}
+		},
+	});
+
 	pi.registerCommand("subagent-status", {
-		description: "Show active subagent runs and runtime addresses",
+		description: "Show live subagent tasks and current-branch task history",
 		handler: async (_args, ctx) => {
-			const list = [...tasks.values()];
-			const lines = list.map((t) => {
-					const age = Math.round((Date.now() - t.startedAt) / 1000);
-					return `- ${t.address} run ${t.runId.slice(0, 8)} ${t.agent} running (${age}s ago)`;
+			runtimeContext = ctx;
+			reconcileCurrentBranch();
+			const live: LiveTaskView[] = [...tasks.values()].map((task) => {
+				const persisted = persistedTasks.get(task.runId);
+				return {
+					runId: task.runId,
+					address: task.address,
+					agent: task.agent,
+					task: task.task,
+					status:
+						persisted?.status === "starting" || persisted?.status === "running"
+							? persisted.status
+							: "running",
+					startedAt: task.startedAt,
+					updatedAt: persisted?.updatedAt ?? task.updatedAt,
+				};
 			});
-			const active = activeRuns.list();
-			const activeLines = active.map(
+			const taskRunIds = new Set(live.map((task) => task.runId));
+			const otherActive = activeRuns
+				.list()
+				.filter((run) => !taskRunIds.has(run.runId))
+				.slice(0, 10);
+			const activeLines = otherActive.map(
 				(run) => `- ${run.address} -> ${run.agent} [${run.mode}] (run ${run.runId.slice(0, 8)})`,
 			);
+			const status = formatTaskStatus(live, persistedTasks);
 			ctx.ui.notify(
 				[
-					lines.length ? `Background tasks:\n${lines.join("\n")}` : "Background tasks: none",
-					activeLines.length ? `Active run registry:\n${activeLines.join("\n")}` : "Active run registry: none",
+					status,
+					activeLines.length
+						? `Other live processes:\n${activeLines.join("\n")}${activeRuns.list().length - taskRunIds.size > 10 ? "\n- … more live processes omitted" : ""}`
+						: "",
 				].join("\n"),
 				"info",
 			);
 		},
 	});
 }
+
+export default createAgentSquadExtension;