@ryan_nookpi/pi-extension-subagent 0.5.7 → 0.5.9

package/index.ts

@@ -28,6 +28,7 @@
  */
 
 import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { SubagentAsyncTasks } from "./async-task-lifecycle.js";
 import { HANG_CHECK_INTERVAL_MS } from "./constants.js";
 import { SUBAGENT_COMMANDS, SUBAGENT_SHORTCUTS, type SubagentCommandName } from "./registration-manifest.js";
 
@@ -40,6 +41,7 @@ interface SubagentCore {
 }
 
 export default function (pi: ExtensionAPI) {
+	const asyncTasks = new SubagentAsyncTasks(pi);
 	let core: SubagentCore | null = null;
 	let corePromise: Promise<SubagentCore> | null = null;
 	/** Serializes session lifecycle work so events apply in dispatch order. */
@@ -54,6 +56,7 @@ export default function (pi: ExtensionAPI) {
 				import("./store.js"),
 			]);
 			const store = storeMod.createStore();
+			store.asyncTasks = asyncTasks;
 			const registrations = commands.registerAll(pi, store);
 			core = { store, commands, registrations, escalation, lifecycle };
 			return core;
@@ -116,11 +119,24 @@ export default function (pi: ExtensionAPI) {
 	}, 0);
 
 	let hangCheckTimer: ReturnType<typeof setInterval> | undefined;
+	let sessionGeneration = 0;
 
 	pi.on("session_start", (_event, ctx) => {
-		enqueue((c) => {
+		const generation = ++sessionGeneration;
+		const sessionId = ctx.sessionManager.getSessionId();
+		// Fence tool admission synchronously, before lazy lifecycle work can yield.
+		try {
+			if (pi.events) asyncTasks.bind(sessionId);
+		} catch (error) {
+			process.stderr.write(`[subagent] session binding failed: ${error instanceof Error ? error.message : error}\n`);
+			return;
+		}
+		enqueue(async (c) => {
+			if (generation !== sessionGeneration) return;
+			if (pi.events) await asyncTasks.provider.whenDiscovered();
+			if (generation !== sessionGeneration || ctx.sessionManager.getSessionId() !== sessionId) return;
 			c.store.disposed = false;
-			c.commands.handleSessionStart(pi, c.store, ctx as unknown as ExtensionContext);
+			c.commands.handleSessionStart(asyncTasks.wrap(pi), c.store, ctx as unknown as ExtensionContext);
 			c.escalation.maybeRegisterAskMaster(pi, ctx);
 		});
 		if (!hangCheckTimer) {
@@ -131,6 +147,8 @@ export default function (pi: ExtensionAPI) {
 	});
 
 	pi.on("session_shutdown", async (event) => {
+		++sessionGeneration;
+		asyncTasks.shutdown();
 		if (hangCheckTimer) {
 			clearInterval(hangCheckTimer);
 			hangCheckTimer = undefined;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ryan_nookpi/pi-extension-subagent",
-  "version": "0.5.7",
+  "version": "0.5.9",
   "description": "Asynchronous subagent delegation for pi with run, batch, chain, and continuation workflows.",
   "license": "MIT",
   "repository": {
@@ -39,6 +39,8 @@
     "format.ts",
     "group-pending.ts",
     "index.ts",
+    "async-task-provider.ts",
+    "async-task-lifecycle.ts",
     "lifecycle.ts",
     "invocation-queue.ts",
     "live-preview.ts",

package/runner.ts

@@ -7,6 +7,7 @@ import { spawn } from "node:child_process";
 import * as fs from "node:fs";
 import type { Message } from "@earendil-works/pi-ai";
 import type { AgentConfig } from "./agents.js";
+import { currentResourceObserver, trackSubagentRunner } from "./async-task-lifecycle.js";
 import { buildClaudeArgs } from "./claude-args.js";
 import { createSidecarWriter } from "./claude-sidecar-writer.js";
 import {
@@ -43,6 +44,7 @@ import {
 } from "./utils/agent-utils.js";
 
 export interface RunSingleAgentSessionConfig {
+	asyncRunId?: number;
 	sessionFile?: string;
 	resumeSessionId?: string;
 	sidecarSessionFile?: string;
@@ -198,7 +200,21 @@ export function computeAgentAliasHints(agents: AgentConfig[]): string {
 
 // ─── Single Agent Execution ──────────────────────────────────────────────────
 
-export async function runSingleAgent(
+export async function runSingleAgent(...args: Parameters<typeof runSingleAgentUntracked>): Promise<SingleResult> {
+	const config = typeof args[8] === "object" ? args[8] : undefined;
+	return trackSubagentRunner(
+		args[2],
+		args[5],
+		(signal) => {
+			const invocation: Parameters<typeof runSingleAgentUntracked> = [...args];
+			invocation[5] = signal;
+			return runSingleAgentUntracked(...invocation);
+		},
+		config?.asyncRunId,
+	);
+}
+
+async function runSingleAgentUntracked(
 	defaultCwd: string,
 	agents: AgentConfig[],
 	agentName: string,
@@ -209,6 +225,7 @@ export async function runSingleAgent(
 	makeDetails: (results: SingleResult[]) => SubagentDetails,
 	sessionConfig?: string | RunSingleAgentSessionConfig,
 ): Promise<SingleResult> {
+	currentResourceObserver()("settled");
 	const normalizedSessionConfig = normalizeRunSessionConfig(sessionConfig);
 	const sessionFile = normalizedSessionConfig.sessionFile;
 	const resumeSessionId = normalizedSessionConfig.resumeSessionId;
@@ -246,7 +263,10 @@ export async function runSingleAgent(
 
 	if (agent.runtime === "claude") {
 		if (resolveClaudeRuntimeMode(defaultCwd) === "sdk") {
+			// SDK query.close() has no process-exit acknowledgement.
+			currentResourceObserver()("unknown");
 			const { runClaudeAgentViaSdk } = await import("./claude-sdk-runner.js");
+			if (signal?.aborted) throw new Error("Subagent was aborted");
 			return runClaudeAgentViaSdk(
 				defaultCwd,
 				agent,
@@ -300,6 +320,7 @@ async function runClaudeAgent(
 	sidecarSessionFile?: string,
 	onDiagnostic?: RunnerDiagnosticSink,
 ): Promise<SingleResult> {
+	const observeResource = currentResourceObserver();
 	try {
 		validateClaudeRuntimeModel(agent.model);
 	} catch (err: any) {
@@ -392,7 +413,9 @@ async function runClaudeAgent(
 		});
 
 		const exitCode = await new Promise<number>((resolve) => {
+			if (signal?.aborted) throw new Error("Subagent was aborted");
 			const proc = spawn("claude", args, { cwd: defaultCwd, shell: false, stdio: ["ignore", "pipe", "pipe"] });
+			observeResource("active");
 			let buffer = "";
 			let procExited = false;
 			let settled = false;
@@ -564,6 +587,7 @@ async function runClaudeAgent(
 			});
 
 			proc.on("exit", (code, terminationSignal) => {
+				observeResource("settled");
 				procExited = true;
 				lastExitCode = normalizeChildExitCode(code, terminationSignal, lastExitCode);
 				lastProcessSignal = terminationSignal;
@@ -583,6 +607,7 @@ async function runClaudeAgent(
 			});
 
 			proc.on("close", (code, terminationSignal) => {
+				observeResource("settled");
 				procExited = true;
 				lastProcessSignal = terminationSignal ?? lastProcessSignal;
 				diagnose({ event: "close", childPid: proc.pid, code, signal: terminationSignal });
@@ -593,6 +618,7 @@ async function runClaudeAgent(
 			});
 
 			proc.on("error", (error) => {
+				observeResource(proc.pid ? "unknown" : "settled");
 				procExited = true;
 				stderrBuf += `[runner] process error: ${error?.message || String(error)}\n`;
 				diagnose({
@@ -658,6 +684,7 @@ async function runPiAgent(
 	persistedSessionBaseOffset = 0,
 	onDiagnostic?: RunnerDiagnosticSink,
 ): Promise<SingleResult> {
+	const observeResource = currentResourceObserver();
 	const args: string[] = ["--mode", "json", "-p"];
 	if (sessionFile) args.push("--session", sessionFile);
 	else args.push("--no-session");
@@ -721,7 +748,9 @@ async function runPiAgent(
 		args.push(normalizeTaskForSubagentPrompt(task));
 
 		const exitCode = await new Promise<number>((resolve) => {
+			if (signal?.aborted) throw new Error("Subagent was aborted");
 			const proc = spawn("pi", args, { cwd: defaultCwd, shell: false, stdio: ["ignore", "pipe", "pipe"] });
+			observeResource("active");
 			let buffer = "";
 			let procExited = false;
 			let settled = false;
@@ -1089,6 +1118,7 @@ async function runPiAgent(
 			});
 
 			proc.on("exit", (code, terminationSignal) => {
+				observeResource("settled");
 				procExited = true;
 				lastExitCode = normalizeChildExitCode(code, terminationSignal, lastExitCode);
 				lastProcessSignal = terminationSignal;
@@ -1105,6 +1135,7 @@ async function runPiAgent(
 			});
 
 			proc.on("close", (code, terminationSignal) => {
+				observeResource("settled");
 				procExited = true;
 				lastProcessSignal = terminationSignal ?? lastProcessSignal;
 				diagnose({ event: "close", childPid: proc.pid, code, signal: terminationSignal });
@@ -1113,6 +1144,7 @@ async function runPiAgent(
 			});
 
 			proc.on("error", (error) => {
+				observeResource(proc.pid ? "unknown" : "settled");
 				procExited = true;
 				appendStderrDiagnostic(currentResult, `process error: ${error?.message || String(error)}`);
 				diagnose({

package/store.ts

@@ -18,6 +18,7 @@ import type { WidgetRenderCtx } from "./widget.js";
 export const COLLAPSED_ITEM_COUNT = 10;
 
 export interface SubagentStore {
+	asyncTasks?: import("./async-task-lifecycle.js").SubagentAsyncTasks;
 	/** True after session_shutdown; async callbacks must not use captured runtime APIs. */
 	disposed: boolean;
 	commandRuns: Map<number, CommandRunState>;

package/tool-execute.ts

@@ -10,6 +10,14 @@ import * as fs from "node:fs";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { createRunActivityRecorder } from "./activity.js";
 import { type AgentConfig, discoverAgents } from "./agents.js";
+import {
+	assertAsyncAdmission,
+	asyncInvocationDetails,
+	completeAsyncInvocation,
+	discardRemovedAsyncInvocation,
+	isExpectedTrackedCancellation,
+	retainAsyncCompletion,
+} from "./async-task-lifecycle.js";
 import { parseSubagentToolCommand, SUBAGENT_CLI_HELP_TEXT } from "./cli.js";
 import {
 	DEFAULT_TURN_COUNT,
@@ -453,6 +461,7 @@ function buildRunStartMessage(runState: CommandRunState, status: "started" | "re
 			`\nContext: ${contextLabel} · turn ${runState.turnCount}`,
 		display: false,
 		details: {
+			...asyncInvocationDetails(),
 			runId: runState.id,
 			agent: runState.agent,
 			task: runState.task,
@@ -515,7 +524,8 @@ function buildRunCompletionMessage(finalized: FinalizedRun, options?: { display?
 	const { runState, result, rawOutput } = finalized;
 	const usage = result ? formatUsageStats(result.usage, result.model) : "";
 	const terminalLabel = getFinalizedRunLabel(finalized);
-	return {
+	completeAsyncInvocation(terminalLabel);
+	return retainAsyncCompletion({
 		customType: "subagent-tool" as const,
 		content:
 			`[subagent:${runState.agent}#${runState.id}] ${terminalLabel}` +
@@ -525,6 +535,7 @@ function buildRunCompletionMessage(finalized: FinalizedRun, options?: { display?
 			`\n\n${rawOutput}`,
 		display: options?.display ?? true,
 		details: {
+			...asyncInvocationDetails(),
 			runId: runState.id,
 			agent: runState.agent,
 			task: runState.task,
@@ -555,12 +566,13 @@ function buildRunCompletionMessage(finalized: FinalizedRun, options?: { display?
 			claudeSessionId: runState.claudeSessionId,
 			claudeProjectDir: runState.claudeProjectDir,
 		},
-	};
+	});
 }
 
 function buildEscalationMessage(runState: CommandRunState, escalationMessage: string, result: SingleResult) {
+	completeAsyncInvocation("completed");
 	const usage = formatUsageStats(result.usage, result.model);
-	return {
+	return retainAsyncCompletion({
 		customType: "subagent-tool" as const,
 		content:
 			`[subagent:${runState.agent}#${runState.id}] escalated` +
@@ -569,6 +581,7 @@ function buildEscalationMessage(runState: CommandRunState, escalationMessage: st
 			`\n\n[ESCALATION] ${escalationMessage}`,
 		display: true,
 		details: {
+			...asyncInvocationDetails(),
 			runId: runState.id,
 			agent: runState.agent,
 			task: runState.task,
@@ -592,7 +605,7 @@ function buildEscalationMessage(runState: CommandRunState, escalationMessage: st
 			claudeSessionId: runState.claudeSessionId,
 			claudeProjectDir: runState.claudeProjectDir,
 		},
-	};
+	});
 }
 
 function buildStrongWaitMessage(runId: number): string {
@@ -660,6 +673,7 @@ function formatBatchSummary(
 	runs: CommandRunState[],
 	terminalStatus: "completed" | "error" | "aborted",
 ): string {
+	completeAsyncInvocation(terminalStatus);
 	const statuses = runs.map((run) => getSubagentRunStatusLabel(run.status, run.errorClass));
 	const headerStatus = runs.map((run, index) => `#${run.id} ${statuses[index]}`).join(", ");
 	const outcomeCounts = formatOutcomeCounts(
@@ -680,6 +694,7 @@ function formatPipelineSummary(
 	stepResults: PipelineStepResult[],
 	terminalStatus: "completed" | "stopped" | "error" | "aborted",
 ): string {
+	completeAsyncInvocation(terminalStatus);
 	const outcomeCounts = formatOutcomeCounts(
 		stepResults.filter((step) => step.status === "error").length,
 		stepResults.filter((step) => step.status === "aborted").length,
@@ -817,6 +832,7 @@ function finalizeRunError(runState: CommandRunState, error: unknown): FinalizedR
 }
 
 export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore) {
+	pi = store.asyncTasks?.wrap(pi) ?? pi;
 	const execute = async (
 		_toolCallId: string,
 		params: Record<string, any>,
@@ -931,7 +947,7 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 		const hasChain = asyncAction === "chain";
 		const hasSingle = asyncAction === "run" || asyncAction === "continue";
 		const mode: LaunchMode = hasBatch ? "batch" : hasChain ? "chain" : "single";
-		const shouldRunAsync = isInteractiveTuiContext(ctx);
+		const shouldRunAsync = isInteractiveTuiContext(ctx) || store.asyncTasks?.provider.accepting === true;
 		const makeDetails = (
 			modeOverride: LaunchMode = mode,
 			results: SingleResult[] = [],
@@ -1389,6 +1405,7 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 		}
 
 		function registerRunLaunch(config: RunLaunchConfig): CommandRunState {
+			assertAsyncAdmission();
 			const initialDisplayTask = buildSubagentDisplayTaskFallback(config.taskForDisplay);
 			let runState: CommandRunState;
 			if (config.existingRunState) {
@@ -1523,6 +1540,7 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 								resumeSessionId: runState.claudeSessionId,
 								sidecarSessionFile: runState.sessionFile,
 								persistedSessionBaseOffset: runState.persistedSessionBaseOffset,
+								...(store.asyncTasks?.provider.supported ? { asyncRunId: runState.id } : {}),
 								onDiagnostic: createRunDiagnosticSink(pi, runState),
 							},
 						);
@@ -1678,7 +1696,10 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 			void (async () => {
 				try {
 					const finalized = await launchRunInBackground(runState, taskForAgent);
-					if (runState.removed || store.disposed) return;
+					if (runState.removed || store.disposed) {
+						discardRemovedAsyncInvocation(finalized.rawOutput);
+						return;
+					}
 					updateCommandRunsWidget(store);
 
 					if (finalized.result?.exitCode === ESCALATION_EXIT_CODE) {
@@ -1704,12 +1725,16 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 						if (entry) entry.pendingCompletion = makePendingCompletion(completionMessage, true);
 					}
 
-					ctx.ui?.notify?.(
-						`subagent tool run #${runState.id} (${resolvedAgent}) ${terminalLabel}`,
-						terminalLabel === "completed" ? "info" : terminalLabel === "aborted" ? "warning" : "error",
-					);
+					if (terminalLabel !== "aborted" || !isExpectedTrackedCancellation())
+						ctx.ui?.notify?.(
+							`subagent tool run #${runState.id} (${resolvedAgent}) ${terminalLabel}`,
+							terminalLabel === "completed" ? "info" : terminalLabel === "aborted" ? "warning" : "error",
+						);
 				} catch (error: unknown) {
-					if (runState.removed || store.disposed) return;
+					if (runState.removed || store.disposed) {
+						discardRemovedAsyncInvocation(String(error));
+						return;
+					}
 					const finalized = finalizeRunError(runState, error);
 					const terminalLabel = getFinalizedRunLabel(finalized);
 					const errorMessage = buildRunCompletionMessage(finalized);
@@ -1720,10 +1745,11 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 						const entry = store.globalLiveRuns.get(runState.id);
 						if (entry) entry.pendingCompletion = makePendingCompletion(errorMessage, true);
 					}
-					ctx.ui?.notify?.(
-						`subagent tool run #${runState.id} ${terminalLabel}: ${runState.lastLine}`,
-						terminalLabel === "aborted" ? "warning" : "error",
-					);
+					if (terminalLabel !== "aborted" || !isExpectedTrackedCancellation())
+						ctx.ui?.notify?.(
+							`subagent tool run #${runState.id} ${terminalLabel}: ${runState.lastLine}`,
+							terminalLabel === "aborted" ? "warning" : "error",
+						);
 					updateCommandRunsWidget(store);
 				} finally {
 					clearRunAbortState(runState);
@@ -1872,12 +1898,14 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 								content,
 								display: true,
 								details: {
+									...asyncInvocationDetails(),
 									batchId,
 									runIds: batch.runIds,
 									status: batchTerminalStatus === "completed" ? "done" : batchTerminalStatus,
 									runSummaries: orderedRuns.map((run) => buildRunAnalyticsSummary(run)),
 								},
 							};
+							retainAsyncCompletion(message);
 							retireFinishedGroup(store, snapshotBatchGroup(store, batch, batchTerminalStatus));
 							if (isInOriginSession(ctx, batch.originSessionFile)) {
 								pi.sendMessage(message, { deliverAs: "followUp", triggerTurn: true });
@@ -1901,14 +1929,15 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 								updateWidget: false,
 								removalReason: "trim",
 							});
-							ctx.ui?.notify?.(
-								batchTerminalStatus === "error"
-									? `subagent batch ${batchId} finished with errors`
-									: batchTerminalStatus === "aborted"
-										? `subagent batch ${batchId} aborted`
-										: `subagent batch ${batchId} completed`,
-								batchTerminalStatus === "error" ? "error" : batchTerminalStatus === "aborted" ? "warning" : "info",
-							);
+							if (batchTerminalStatus !== "aborted" || !isExpectedTrackedCancellation())
+								ctx.ui?.notify?.(
+									batchTerminalStatus === "error"
+										? `subagent batch ${batchId} finished with errors`
+										: batchTerminalStatus === "aborted"
+											? `subagent batch ${batchId} aborted`
+											: `subagent batch ${batchId} completed`,
+									batchTerminalStatus === "error" ? "error" : batchTerminalStatus === "aborted" ? "warning" : "info",
+								);
 						}
 					} catch (error: unknown) {
 						const finalized = finalizeRunError(runState, error);
@@ -1923,12 +1952,14 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 								content: formatBatchSummary(batchId, orderedRuns, batchTerminalStatus),
 								display: true,
 								details: {
+									...asyncInvocationDetails(),
 									batchId,
 									runIds: batch.runIds,
 									status: batchTerminalStatus,
 									runSummaries: orderedRuns.map((run) => buildRunAnalyticsSummary(run)),
 								},
 							};
+							retainAsyncCompletion(message);
 							retireFinishedGroup(store, snapshotBatchGroup(store, batch, batchTerminalStatus));
 							if (isInOriginSession(ctx, batch.originSessionFile)) {
 								pi.sendMessage(message, { deliverAs: "followUp", triggerTurn: true });
@@ -2247,12 +2278,14 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 							content: formatPipelineSummary(pipelineId, pipeline.stepResults, terminalStatus),
 							display: true,
 							details: {
+								...asyncInvocationDetails(),
 								pipelineId,
 								stepRunIds: pipeline.stepRunIds,
 								status: terminalStatus === "completed" ? "done" : terminalStatus,
 								runSummaries: orderedRuns.map((run) => buildRunAnalyticsSummary(run)),
 							},
 						};
+						retainAsyncCompletion(message);
 						retireFinishedGroup(store, snapshotPipeline(pipeline, terminalStatus));
 						if (isInOriginSession(ctx, pipeline.originSessionFile)) {
 							pi.sendMessage(message, { deliverAs: "followUp", triggerTurn: true });
@@ -2301,5 +2334,20 @@ export function createSubagentToolExecute(pi: ExtensionAPI, store: SubagentStore
 		};
 	};
 
-	return execute;
+	return (...args: Parameters<typeof execute>) => {
+		const parsed = parseSubagentToolCommand(args[1].command, { knownRunIds: [...store.commandRuns.keys()] });
+		if (
+			parsed.type !== "params" ||
+			!["run", "continue", "batch", "chain"].includes(String(parsed.params.asyncAction ?? "run")) ||
+			!store.asyncTasks
+		)
+			return execute(...args);
+		return store.asyncTasks.invoke(
+			String(args[1].command).slice(0, 500),
+			args[0],
+			() => execute(...args),
+			false,
+			args[2],
+		);
+	};
 }