@ryan_nookpi/pi-extension-subagent 0.5.11 → 0.5.13

package/README.md

@@ -41,7 +41,7 @@ Equivalent AI tool call:
 { "command": "subagent run worker --isolated -- implement the requested change and run tests" }
 ```
 
-Tool launches are asynchronous. Wait for the automatic completion or failure follow-up instead of polling immediately.
+Tool launches are asynchronous. Wait for the automatic completion or failure follow-up instead of polling immediately. Pi drops queued messages when you press Escape; if a completion was still queued when the run stops, it is added to the session without starting a new turn, and the agent sees it with your next message.
 
 ## Agent definitions
 

package/delivery-guard.ts

@@ -0,0 +1,103 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+type SendMessage = ExtensionAPI["sendMessage"];
+type Message = Parameters<SendMessage>[0];
+
+const GUARD_ID_KEY = "deliveryGuardId";
+
+function guardId(message: unknown): string | undefined {
+	const details = (message as { details?: unknown } | undefined)?.details;
+	if (!details || typeof details !== "object") return undefined;
+	const id = (details as Record<string, unknown>)[GUARD_ID_KEY];
+	return typeof id === "string" ? id : undefined;
+}
+
+/**
+ * Pi keeps steer and follow-up messages in an in-memory queue until the agent loop picks them up, and
+ * Escape clears that queue without restoring custom messages. A completion queued while the agent was
+ * busy would vanish. This tracks every turn-triggering message queued during a run, and when the run
+ * settles without the message reaching the model, appends it to the session without starting a turn:
+ * the user sees it now, and the model sees it on the next prompt.
+ *
+ * Extensions cannot tell Escape from other aborts: Pi's pending count covers only typed messages, not
+ * extension messages. An abort while the model streams keeps the original queued, and Pi delivers it on
+ * the next run after the restored copy. Both carry the same guard id, so the model input keeps only the first.
+ */
+export function guardQueuedDeliveries(pi: ExtensionAPI): ExtensionAPI {
+	let context: ExtensionContext | undefined;
+	let nextId = 0;
+	const queued = new Map<string, Message>();
+
+	const busy = () => {
+		try {
+			return context ? !context.isIdle() : false;
+		} catch {
+			return false;
+		}
+	};
+
+	const restored = new Set<string>();
+
+	pi.on("agent_start", (_event, ctx) => {
+		context = ctx;
+	});
+	pi.on("message_end", (event) => {
+		const id = guardId(event.message);
+		if (id) queued.delete(id);
+	});
+	pi.on("agent_settled", (_event, ctx) => {
+		context = ctx;
+		if (queued.size === 0) return;
+		const lost = [...queued.entries()];
+		queued.clear();
+		for (const [id, message] of lost) {
+			restored.add(id);
+			pi.sendMessage(message, { triggerTurn: false });
+		}
+	});
+	pi.on("context", (event) => {
+		if (restored.size === 0) return;
+		const seen = new Set<string>();
+		let dropped = false;
+		const messages = event.messages.filter((message) => {
+			const id = guardId(message);
+			if (!id || !restored.has(id)) return true;
+			if (!seen.has(id)) {
+				seen.add(id);
+				return true;
+			}
+			dropped = true;
+			return false;
+		});
+		return dropped ? { messages } : undefined;
+	});
+	pi.on("session_shutdown", () => {
+		queued.clear();
+		restored.clear();
+		context = undefined;
+	});
+
+	return new Proxy(pi, {
+		get(target, property, receiver) {
+			if (property !== "sendMessage") return Reflect.get(target, property, receiver);
+			const send: SendMessage = (message, options) => {
+				// An idle send starts its own run with the message as the prompt, so nothing can drop it.
+				if (options?.triggerTurn !== true || options.deliverAs === "nextTurn" || !busy()) {
+					target.sendMessage(message, options);
+					return;
+				}
+				const id = String(++nextId);
+				const details = message.details && typeof message.details === "object" ? message.details : {};
+				const tagged = { ...message, details: { ...details, [GUARD_ID_KEY]: id } } as Message;
+				queued.set(id, tagged);
+				try {
+					target.sendMessage(tagged, options);
+				} catch (error) {
+					queued.delete(id);
+					throw error;
+				}
+			};
+			return send;
+		},
+	});
+}

package/index.ts

@@ -30,6 +30,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 { guardQueuedDeliveries } from "./delivery-guard.js";
 import { SUBAGENT_COMMANDS, SUBAGENT_SHORTCUTS, type SubagentCommandName } from "./registration-manifest.js";
 import { handleRunningReminderContext, ReminderAnchors } from "./running-reminder.js";
 
@@ -41,7 +42,9 @@ interface SubagentCore {
 	lifecycle: typeof import("./lifecycle.js");
 }
 
-export default function (pi: ExtensionAPI) {
+export default function (host: ExtensionAPI) {
+	// Completions are queued while the agent is busy; the guard restores any the queue drops on Escape.
+	const pi = guardQueuedDeliveries(host);
 	const asyncTasks = new SubagentAsyncTasks(pi);
 	let core: SubagentCore | null = null;
 	let corePromise: Promise<SubagentCore> | null = null;

package/lifecycle.ts

@@ -31,6 +31,13 @@ function reconcileRunWithPersistedSession(run: CommandRunState): void {
 		run.lastActivityAt = snapshot.latestActivityAt;
 	}
 	if (!snapshot.isTerminal) return;
+	// A live Pi child can still retry a persisted error. Only its runner's
+	// completion marker proves that failure recovery has finished.
+	if (
+		!snapshot.completionMarker &&
+		(snapshot.terminalStopReason === "error" || snapshot.terminalStopReason === "aborted")
+	)
+		return;
 
 	const exitCode =
 		snapshot.completionMarker?.exitCode ??

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@ryan_nookpi/pi-extension-subagent",
-  "version": "0.5.11",
+  "version": "0.5.13",
   "description": "Asynchronous subagent delegation for pi with run, batch, chain, and continuation workflows.",
   "license": "MIT",
   "repository": {
@@ -32,6 +32,7 @@
     "config.ts",
     "constants.ts",
     "context-limits.ts",
+    "delivery-guard.ts",
     "diagnostics.ts",
     "display-task.ts",
     "escalation.ts",

package/runner.ts

@@ -762,6 +762,9 @@ async function runPiAgent(
 			let lastProcessSignal: NodeJS.Signals | null = null;
 			let lastEventAt = Date.now();
 			let sawAgentEnd = false;
+			// Covers both backoff and the resumed attempt, until Pi reports recovery
+			// or a final agent_end. Raw history can retain errors throughout this time.
+			let retryPending = false;
 			let settleReason = "unknown";
 			let unparsedStdoutCount = 0;
 			const unparsedStdoutTail: string[] = [];
@@ -813,10 +816,24 @@ async function runPiAgent(
 				}, 5000);
 			};
 
+			const clearCompletionFallbacks = () => {
+				if (agentEndFallbackTimer) clearTimeout(agentEndFallbackTimer);
+				if (terminalMessageFallbackTimer) clearTimeout(terminalMessageFallbackTimer);
+				agentEndFallbackTimer = undefined;
+				terminalMessageFallbackTimer = undefined;
+			};
+
 			const syncFromPersistedSession = (allowResolve: boolean): boolean => {
 				if (!sessionFile) return false;
 				const snapshot = readPersistedSessionSnapshot(sessionFile, { startOffset: persistedSessionBaseOffset });
-				if (snapshot.terminalStopReason && !currentResult.stopReason) {
+				const unconfirmedFailure =
+					!snapshot.completionMarker &&
+					(snapshot.terminalStopReason === "error" || snapshot.terminalStopReason === "aborted");
+				if (
+					snapshot.terminalStopReason &&
+					!currentResult.stopReason &&
+					(procExited || settled || (!retryPending && !unconfirmedFailure))
+				) {
 					currentResult.stopReason = snapshot.terminalStopReason;
 				}
 				if (snapshot.messages.length > currentResult.messages.length) {
@@ -824,7 +841,18 @@ async function runPiAgent(
 					currentResult.liveText = undefined;
 					emitUpdate();
 				}
-				if (!allowResolve || !snapshot.isTerminal || settled || procExited || wasAborted) return false;
+				// Persistence can beat stdout. An error response alone is not proof
+				// that Pi has finished: agent_end/auto_retry_start decide its recovery.
+				if (
+					!allowResolve ||
+					!snapshot.isTerminal ||
+					settled ||
+					procExited ||
+					wasAborted ||
+					retryPending ||
+					unconfirmedFailure
+				)
+					return false;
 
 				const forcedCode =
 					snapshot.completionMarker?.exitCode ??
@@ -855,15 +883,44 @@ async function runPiAgent(
 					return;
 				}
 				lastEventAt = Date.now();
+				// Once the local context guard settles the run, late shutdown events
+				// must not replace its failure with the child's last toolUse response.
+				if (contextGuardTripped) return;
+
+				if (event.type === "auto_retry_start") {
+					retryPending = true;
+					sawAgentEnd = false;
+					clearCompletionFallbacks();
+					return;
+				}
+
+				if (event.type === "auto_retry_end") {
+					retryPending = false;
+					if (event.success === false) {
+						currentResult.stopReason = "error";
+						if (event.finalError) currentResult.errorMessage = event.finalError;
+						sawAgentEnd = true;
+						writeCompletionMarkerOnce(1);
+						scheduleAgentEndForceResolve();
+					} else if (currentResult.stopReason === "error" || currentResult.stopReason === "aborted") {
+						// Successful recovery may arrive without its message_end on stdout.
+						// Wait for the successful agent_end instead of reusing the old error.
+						currentResult.stopReason = undefined;
+						currentResult.errorMessage = undefined;
+					} else if (currentResult.stopReason && currentResult.stopReason !== "toolUse") {
+						writeCompletionMarkerOnce(0);
+						scheduleTerminalMessageForceResolve();
+					}
+					return;
+				}
 
 				if (event.type === "agent_start" || event.type === "turn_start") {
 					sawAgentEnd = false;
+					currentResult.stopReason = undefined;
+					currentResult.errorMessage = undefined;
 					currentResult.liveThinking = undefined;
 					currentResult.thoughtText = undefined;
-					if (terminalMessageFallbackTimer) {
-						clearTimeout(terminalMessageFallbackTimer);
-						terminalMessageFallbackTimer = undefined;
-					}
+					clearCompletionFallbacks();
 					return;
 				}
 
@@ -873,21 +930,32 @@ async function runPiAgent(
 					// its objects are freshly deserialized, so reference-based deduplication
 					// would append every prior message again.
 					const eventMessages = (event.messages ?? []) as Message[];
-					const terminalMessage = [...eventMessages].reverse().find((msg) => Boolean((msg as any).stopReason));
+					const terminalMessage = [...eventMessages]
+						.reverse()
+						.find((msg) => msg.role === "assistant" && Boolean(msg.stopReason));
 					if (terminalMessage) {
 						const fingerprint = JSON.stringify(terminalMessage);
 						const alreadyCaptured = currentResult.messages.some((msg) => JSON.stringify(msg) === fingerprint);
 						if (!alreadyCaptured) currentResult.messages.push(terminalMessage);
-						if (!currentResult.stopReason) {
-							currentResult.stopReason = (terminalMessage as any).stopReason;
-							if ((terminalMessage as any).errorMessage) {
-								currentResult.errorMessage = (terminalMessage as any).errorMessage;
-							}
-						}
+						currentResult.stopReason = (terminalMessage as any).stopReason;
+						currentResult.errorMessage = (terminalMessage as any).errorMessage;
 					} else if (currentResult.messages.length === 0 && eventMessages.length > 0) {
 						currentResult.messages.push(eventMessages[eventMessages.length - 1]);
 					}
-					if (currentResult.stopReason && currentResult.stopReason !== "toolUse") {
+					if (event.willRetry === true) {
+						retryPending = true;
+						sawAgentEnd = false;
+						clearCompletionFallbacks();
+						return;
+					}
+					retryPending = false;
+					if (
+						currentResult.stopReason &&
+						currentResult.stopReason !== "toolUse" &&
+						// Older Pi versions omit willRetry; defer their failure marker
+						// until fallback so auto_retry_start can still cancel completion.
+						(currentResult.stopReason !== "error" || event.willRetry === false)
+					) {
 						writeCompletionMarkerOnce(
 							currentResult.stopReason === "error" || currentResult.stopReason === "aborted" ? 1 : 0,
 						);
@@ -977,7 +1045,7 @@ async function runPiAgent(
 						}
 						if (!currentResult.model && msg.model) currentResult.model = msg.model;
 						if (msg.stopReason) currentResult.stopReason = msg.stopReason;
-						if (msg.errorMessage) currentResult.errorMessage = msg.errorMessage;
+						currentResult.errorMessage = msg.errorMessage;
 						currentResult.liveThinking = undefined;
 
 						// Extract thoughtText from thinking block first line only
@@ -989,8 +1057,15 @@ async function runPiAgent(
 							}
 						}
 						const terminalStopReason = (msg as any).stopReason;
-						if (terminalStopReason && terminalStopReason !== "toolUse") {
-							writeCompletionMarkerOnce(terminalStopReason === "error" || terminalStopReason === "aborted" ? 1 : 0);
+						// Error message_end precedes Pi's retry decision. Do not publish
+						// completion or arm a timeout while it may still recover.
+						if (
+							terminalStopReason &&
+							terminalStopReason !== "toolUse" &&
+							terminalStopReason !== "error" &&
+							!retryPending
+						) {
+							writeCompletionMarkerOnce(terminalStopReason === "aborted" ? 1 : 0);
 							scheduleTerminalMessageForceResolve();
 						}
 					}
@@ -1062,11 +1137,18 @@ async function runPiAgent(
 			};
 
 			function scheduleTerminalMessageForceResolve() {
-				if (!currentResult.stopReason || currentResult.stopReason === "toolUse" || settled || procExited) return;
+				if (
+					!currentResult.stopReason ||
+					currentResult.stopReason === "toolUse" ||
+					retryPending ||
+					settled ||
+					procExited
+				)
+					return;
 				if (terminalMessageFallbackTimer) clearTimeout(terminalMessageFallbackTimer);
 
 				terminalMessageFallbackTimer = setTimeout(() => {
-					if (settled || procExited || wasAborted) return;
+					if (settled || procExited || wasAborted || retryPending) return;
 
 					const forcedCode = currentResult.stopReason === "error" || currentResult.stopReason === "aborted" ? 1 : 0;
 
@@ -1082,12 +1164,12 @@ async function runPiAgent(
 			// (e.g. lingering extension timers/transports). In that case, force
 			// resolve after a short quiet period so runs do not remain "running" forever.
 			function scheduleAgentEndForceResolve() {
-				if (!sawAgentEnd || settled || procExited) return;
+				if (!sawAgentEnd || retryPending || settled || procExited) return;
 				if (agentEndFallbackTimer) clearTimeout(agentEndFallbackTimer);
 
 				const marker = lastEventAt;
 				agentEndFallbackTimer = setTimeout(() => {
-					if (settled || procExited || wasAborted) return;
+					if (settled || procExited || wasAborted || retryPending) return;
 					if (lastEventAt !== marker) return;
 
 					const forcedCode = currentResult.stopReason === "error" || currentResult.stopReason === "aborted" ? 1 : 0;