@ferris1225/pi-subagents 4.2.13 → 4.3.1

package/src/recovery.ts

@@ -5,6 +5,7 @@ import { existsSync } from "node:fs";
 import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises";
 import { dirname, join } from "node:path";
 import { stripVTControlCharacters } from "node:util";
+import { getSubagentsRoot } from "./spawn.ts";
 import { removeWorktreeGroup, worktreeGroupDir, type WorktreeFinalization } from "./worktree.ts";
 
 export const RECOVERY_MANIFEST_FILE_NAME = "pi-subagents-recovery.json";
@@ -27,7 +28,7 @@ interface RecoveryManifest {
 }
 
 export function getRecoveryManifestPath(configPath: string): string {
-	return join(dirname(configPath), RECOVERY_MANIFEST_FILE_NAME);
+	return join(getSubagentsRoot(configPath), RECOVERY_MANIFEST_FILE_NAME);
 }
 
 function normalizeRecord(value: unknown): RecoveryRecord | undefined {
@@ -46,21 +47,45 @@ function normalizeRecord(value: unknown): RecoveryRecord | undefined {
 	};
 }
 
-export async function readRecoveryRecords(configPath: string): Promise<RecoveryRecord[]> {
+interface RecoveryManifestRead {
+	valid: boolean;
+	records: RecoveryRecord[];
+}
+
+async function readManifest(path: string): Promise<RecoveryManifestRead> {
 	try {
-		const parsed = JSON.parse(await readFile(getRecoveryManifestPath(configPath), "utf8")) as {
-			records?: unknown;
+		const parsed = JSON.parse(await readFile(path, "utf8")) as { records?: unknown };
+		if (!Array.isArray(parsed.records)) return { valid: false, records: [] };
+		return {
+			valid: true,
+			records: parsed.records.flatMap((record) => {
+				const normalized = normalizeRecord(record);
+				return normalized ? [normalized] : [];
+			}),
 		};
-		if (!Array.isArray(parsed.records)) return [];
-		return parsed.records.flatMap((record) => {
-			const normalized = normalizeRecord(record);
-			return normalized ? [normalized] : [];
-		});
 	} catch {
-		return [];
+		return { valid: false, records: [] };
 	}
 }
 
+export async function readRecoveryRecords(configPath: string): Promise<RecoveryRecord[]> {
+	return (await readManifest(getRecoveryManifestPath(configPath))).records;
+}
+
+/** Move the previous agent-root manifest into the internal-state root without
+ * dropping retained artifact pointers. Invalid legacy files stay untouched. */
+export async function relocateRecoveryManifest(configPath: string): Promise<void> {
+	const legacyPath = join(dirname(configPath), RECOVERY_MANIFEST_FILE_NAME);
+	const currentPath = getRecoveryManifestPath(configPath);
+	if (legacyPath === currentPath || !existsSync(legacyPath)) return;
+	await withFileMutationQueue(legacyPath, async () => {
+		const legacy = await readManifest(legacyPath);
+		if (!legacy.valid) return;
+		await persistRecoveryRecords(configPath, legacy.records);
+		await rm(legacyPath, { force: true });
+	});
+}
+
 function recoveryKey(record: RecoveryRecord): string {
 	return `${record.runId}\0${record.worktreePath ?? ""}\0${record.patchPath ?? ""}\0${record.error ?? ""}`;
 }

package/src/rpc-run.ts

@@ -397,6 +397,40 @@ interface RpcResponse {
 	data?: unknown;
 }
 
+interface RpcSessionUsage {
+	input: number;
+	output: number;
+	cacheRead: number;
+	cacheWrite: number;
+	cost: number;
+	contextTokens: number;
+}
+
+function sessionUsage(data: unknown): RpcSessionUsage | undefined {
+	if (!data || typeof data !== "object") return undefined;
+	const value = data as { tokens?: Record<string, unknown>; cost?: unknown; contextUsage?: { tokens?: unknown } };
+	if (!value.tokens || typeof value.tokens !== "object") return undefined;
+	const number = (candidate: unknown): number =>
+		typeof candidate === "number" && Number.isFinite(candidate) ? candidate : 0;
+	return {
+		input: number(value.tokens.input),
+		output: number(value.tokens.output),
+		cacheRead: number(value.tokens.cacheRead),
+		cacheWrite: number(value.tokens.cacheWrite),
+		cost: number(value.cost),
+		contextTokens: number(value.contextUsage?.tokens),
+	};
+}
+
+function applySessionUsageDelta(result: RpcSingleResult, before: RpcSessionUsage, after: RpcSessionUsage): void {
+	result.usage.input = Math.max(0, after.input - before.input);
+	result.usage.output = Math.max(0, after.output - before.output);
+	result.usage.cacheRead = Math.max(0, after.cacheRead - before.cacheRead);
+	result.usage.cacheWrite = Math.max(0, after.cacheWrite - before.cacheWrite);
+	result.usage.cost = Math.max(0, after.cost - before.cost);
+	result.usage.contextTokens = after.contextTokens;
+}
+
 class RpcCommandRejectedError extends Error {
 	constructor(message: string) {
 		super(message);
@@ -523,6 +557,8 @@ export async function runRpcAgentAttempt(options: RunRpcAttemptOptions): Promise
 	let abortSettlement: Deferred<void> | undefined;
 	let droppedQueuedCount = 0;
 	let initialPromptResolved = false;
+	let usageBaseline: RpcSessionUsage | undefined;
+	let usageSettlementStarted = false;
 	const initialPrompt = deferred<{ accepted: boolean; error?: Error }>();
 	const pendingRequests = new Map<string, PendingRequest>();
 	const outcome = deferred<void>();
@@ -635,6 +671,23 @@ export async function runRpcAgentAttempt(options: RunRpcAttemptOptions): Promise
 		});
 	};
 
+	const settleRunWithUsage = async (): Promise<void> => {
+		if (usageSettlementStarted) return;
+		usageSettlementStarted = true;
+		try {
+			if (usageBaseline) {
+				const response = await send({ type: "get_session_stats" });
+				const finalUsage = sessionUsage(response.data);
+				if (finalUsage) applySessionUsageDelta(result, usageBaseline, finalUsage);
+				if (finalUsage) emit({ kind: "usage", usage: { ...result.usage }, model: result.model });
+			}
+		} catch {
+			/* message events remain the generation-safe accounting fallback */
+		} finally {
+			settleRun();
+		}
+	};
+
 	const waitForAbortSettlement = (): Deferred<void> => {
 		if (abortSettlement) throw new Error("Another RPC abort transition is already in progress.");
 		abortSettlement = deferred<void>();
@@ -846,7 +899,7 @@ export async function runRpcAgentAttempt(options: RunRpcAttemptOptions): Promise
 				stable.resolve();
 				return;
 			}
-			settleRun();
+			void settleRunWithUsage();
 		}
 	};
 
@@ -951,6 +1004,11 @@ export async function runRpcAgentAttempt(options: RunRpcAttemptOptions): Promise
 			};
 			try {
 				await send({ type: "get_state" }, readyTimeoutMs);
+				try {
+					usageBaseline = sessionUsage((await send({ type: "get_session_stats" })).data);
+				} catch {
+					/* older or degraded RPC children fall back to message usage */
+				}
 			} catch (error) {
 				const handshakeError = error instanceof Error ? error : new Error(String(error));
 				if (!control?.isStopRequested()) {

package/src/runtime.ts

@@ -12,7 +12,6 @@ import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-a
 import { rmSync } from "node:fs";
 import { resolveSubagentConcurrency, BackgroundTaskQueue } from "./background.ts";
 import {
-	completionGroupTriggersTurn,
 	createCompletionBatcher,
 	formatActiveRunsFooter,
 	formatCompletionMessage,
@@ -54,6 +53,10 @@ export interface SubagentThread {
 	requestedThinkingLevel?: ThinkingLevel;
 	isolation: IsolationMode;
 	worktree?: WorktreeIsolation;
+	/** Durable restoration failure that permanently blocks continuation. */
+	resumeUnavailableReason?: string;
+	/** Original durable evidence retained when its worktree handle is unavailable. */
+	restorationRecord?: ThreadRecord;
 	state: ThreadState;
 	control: RpcRunControl;
 	queueController?: AbortController;
@@ -104,9 +107,17 @@ export interface SubagentRuntime {
 	 * one-time session-start notice. */
 	restoredRunIds: number[];
 	restoredNotified: boolean;
-	/** Deliver a batch of completion messages to the main window, waking it only
-	 * when the batch needs a turn. */
+	/** Deliver a batch of completion messages as a waking follow-up. */
 	sendCompletionGroup: (items: CompletionMessageItem[]) => void;
+	/** Claim the sole delivery route before a generation can settle. */
+	claimRunDelivery: (runId: number, route: "background" | "await") => void;
+	/** Publish a terminal completion through its claimed route. Immediate failures
+	 * flush older successful batches first. */
+	publishRunCompletion: (runId: number, item: CompletionMessageItem, immediate: boolean) => void;
+	/** Mark awaited results as returned in the tool response. */
+	completeAwaitDelivery: (runIds: readonly number[]) => void;
+	/** Transfer aborted awaited calls back to completion delivery. */
+	fallbackAwaitDelivery: (runIds: readonly number[]) => void;
 	completionBatcher: CompletionBatcher<CompletionMessageItem>;
 	/** Abort controllers per active run, so subagent_stop can cancel a run in-turn. */
 	runControllers: Map<number, AbortController>;
@@ -138,6 +149,11 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 	let compactionInFlight = false;
 	let heldCompletions: CompletionMessageItem[] = [];
 
+	const runDeliveries = new Map<number, {
+		route: "background" | "await";
+		completion?: CompletionMessageItem;
+		immediate: boolean;
+	}>();
 	const runtime: SubagentRuntime = {
 		configPath,
 		backgroundQueue,
@@ -148,6 +164,9 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 		restoredNotified: false,
 		sendCompletionGroup: (items) => {
 			if (!runtime.sessionActive || items.length === 0) return;
+			// Direct (immediate-failure/stop) delivery must follow successes already
+			// held by the debounce batcher. A batcher's own emit sees an empty batch.
+			runtime.completionBatcher?.flush();
 			if (compactionInFlight) {
 				heldCompletions.push(...items);
 				return;
@@ -170,18 +189,41 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 				content: formatCompletionMessage(items) + formatActiveRunsFooter(active),
 				display: true,
 			};
-			if (completionGroupTriggersTurn(items)) {
-				// steer: the result is injected after the current tool call even mid-turn,
-				// or starts a new turn when idle. followUp would sit in the queue until the
-				// whole turn ends — a main agent waiting for the result (sleep/poll) would
-				// never see it delivered, which is exactly the "returned but never woken"
-				// failure mode.
-				pi.sendMessage(message, { deliverAs: "steer", triggerTurn: true });
+			// Follow-ups never interrupt an active parent lane. triggerTurn wakes an
+			// idle parent immediately, while a streaming parent receives the result
+			// only after its current tool/assistant lane settles.
+			pi.sendMessage(message, { deliverAs: "followUp", triggerTurn: true });
+		},
+		claimRunDelivery: (runId, route) => {
+			runDeliveries.set(runId, { route, immediate: false });
+		},
+		publishRunCompletion: (runId, item, immediate) => {
+			const delivery = runDeliveries.get(runId) ?? { route: "background" as const, immediate: false };
+			delivery.completion = item;
+			delivery.immediate = immediate;
+			runDeliveries.set(runId, delivery);
+			if (delivery.route === "await") return;
+			runDeliveries.delete(runId);
+			if (immediate) {
+				runtime.sendCompletionGroup([item]);
 			} else {
-				// No-wake delivery: nextTurn rides along with the next user turn and can
-				// never start a continuation by itself. followUp would auto-continue
-				// whenever pi is already streaming, defeating the opt-out.
-				pi.sendMessage(message, { deliverAs: "nextTurn" });
+				runtime.completionBatcher.push(item);
+			}
+		},
+		completeAwaitDelivery: (runIds) => {
+			for (const runId of runIds) {
+				const delivery = runDeliveries.get(runId);
+				if (delivery?.route === "await") runDeliveries.delete(runId);
+			}
+		},
+		fallbackAwaitDelivery: (runIds) => {
+			for (const runId of runIds) {
+				const delivery = runDeliveries.get(runId);
+				if (!delivery || delivery.route !== "await") continue;
+				delivery.route = "background";
+				if (delivery.completion) {
+					runtime.publishRunCompletion(runId, delivery.completion, delivery.immediate);
+				}
 			}
 		},
 		completionBatcher: undefined as unknown as CompletionBatcher<CompletionMessageItem>,
@@ -279,6 +321,18 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 			const records: ThreadRecord[] = [];
 			for (const thread of runtime.threads.values()) {
 				if (thread.retired) continue;
+				if (thread.resumeUnavailableReason) {
+					// Restoration failures retain their durable evidence and session until
+					// an explicit destructive stop retires them.
+					if (thread.restorationRecord) {
+						records.push({
+							...thread.restorationRecord,
+							updatedAt: Date.now(),
+							elapsedMs: thread.elapsedMs,
+						});
+					}
+					continue;
+				}
 				const previous = previousStates.get(thread.id) ?? thread.state;
 				let state: "parked" | "completed" | "failed";
 				if (previous === "completed" || previous === "failed") {
@@ -313,6 +367,7 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 			runtime.settledRuns.clear();
 			runtime.settledListeners.clear();
 			runtime.runControllers.clear();
+			runDeliveries.clear();
 			// sessionDirs entries still referenced by records stay owned by the
 			// manifest; the next process re-registers them at restore.
 			runtime.sessionDirs.clear();
@@ -329,13 +384,15 @@ export function createRuntime(pi: ExtensionAPI, configPath: string): SubagentRun
 		compactionInFlight = true;
 	});
 	const releaseHeldCompletions = (): void => {
+		// Drain the debounce while the compaction gate is still closed so newer
+		// pending successes append after completions already held by that gate.
+		runtime.completionBatcher.flush();
 		compactionInFlight = false;
 		if (heldCompletions.length === 0) return;
 		const items = heldCompletions;
 		heldCompletions = [];
-		// Re-enter the normal path now that the gate is open: the active-runs
-		// footer and the steer/nextTurn choice must reflect delivery time, not
-		// the moment the items were held.
+		// Re-enter the normal path now that the gate is open so the active-runs
+		// footer reflects delivery time, not the moment the items were held.
 		runtime.sendCompletionGroup(items);
 	};
 	pi.on("session_compact", releaseHeldCompletions);