pi-agent-flow 1.4.1 → 1.4.3

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "pi-agent-flow",
-  "version": "1.4.1",
+  "version": "1.4.3",
   "description": "Flow-state delegation extension for Pi coding agent.",
   "type": "module",
   "main": "src/index.ts",

package/src/ambient.d.ts

@@ -46,6 +46,10 @@ declare module "@mariozechner/pi-coding-agent" {
 		renderCall?: (...args: any[]) => any;
 		renderResult?: (...args: any[]) => any;
 	};
+	/** Test-only exports provided by tests/__mocks__/pi-coding-agent.ts. */
+	export const bashToolExecuteCalls: any[][];
+	export function __setBashToolExecuteImpl(fn: (...args: any[]) => Promise<any>): void;
+	export function __resetBashToolMock(): void;
 }
 
 declare module "@mariozechner/pi-tui" {

package/src/flow.ts

@@ -29,12 +29,14 @@ const DEFAULT_FLOW_TIMEOUT_MS = 10 * 60 * 1000; // 10 minutes
 const FLOW_TIME_BUDGET_WARNING_MS = 2 * 60 * 1000; // warn 2 min before kill
 const FLOW_FINAL_URGE_MS = 30 * 1000; // final urge 30 s before kill
 const REPORTING_GRACE_MS = 10_000; // grace period after timeout for agent to report findings
+const FLOW_TOOL_SUMMARY_GRACE_MS = FLOW_FINAL_URGE_MS; // bash/tool abort lead time so the agent can summarize
 const FLOW_DEPTH_ENV = "PI_FLOW_DEPTH";
 const FLOW_MAX_DEPTH_ENV = "PI_FLOW_MAX_DEPTH";
 const FLOW_STACK_ENV = "PI_FLOW_STACK";
 const FLOW_PREVENT_CYCLES_ENV = "PI_FLOW_PREVENT_CYCLES";
 const FLOW_TIMEOUT_ENV = "PI_FLOW_TIMEOUT_MS";
 const FLOW_DEADLINE_ENV = "PI_FLOW_DEADLINE_MS";
+const FLOW_TOOL_SUMMARY_GRACE_ENV = "PI_FLOW_TOOL_SUMMARY_GRACE_MS";
 export const FLOW_TOOL_OPTIMIZE_ENV = "PI_FLOW_TOOL_OPTIMIZE";
 const PI_OFFLINE_ENV = "PI_OFFLINE";
 
@@ -217,7 +219,7 @@ function buildFlowArgs(
 
 	const timeBudgetHint =
 		timeoutMs && timeoutMs > 0
-			? `Time budget: ${Math.round(timeoutMs / 1000)}s total. You will be hard-killed when time expires — plan and prioritize accordingly.\n`
+			? `Time budget: ${Math.round(timeoutMs / 1000)}s total. Long-running tools may be interrupted near the deadline to preserve final-summary time; if a tool reports [Flow timeout], stop tool use and output structured findings immediately.\n`
 			: "";
 
 	const activation =
@@ -443,6 +445,13 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 			const nextDepth = Math.max(0, Math.floor(parentDepth)) + 1;
 			const propagatedMaxDepth = Math.max(0, Math.floor(maxDepth));
 			const propagatedStack = [...parentFlowStack, normalizedFlowName];
+			const proportionalGraceMs = Math.floor(effectiveTimeout * 0.1);
+			const minimumGraceMs = effectiveTimeout >= 10_000 ? 1_000 : Math.floor(effectiveTimeout / 2);
+			const toolSummaryGraceMs = Math.min(
+				FLOW_TOOL_SUMMARY_GRACE_MS,
+				Math.max(0, effectiveTimeout),
+				Math.max(minimumGraceMs, proportionalGraceMs),
+			);
 			const { command, prefixArgs } = resolveFlowSpawn();
 			const proc = spawn(command, [...prefixArgs, ...piArgs], {
 				cwd: taskCwd ?? cwd,
@@ -458,7 +467,13 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 					[FLOW_PREVENT_CYCLES_ENV]: preventCycles ? "1" : "0",
 					[FLOW_TOOL_OPTIMIZE_ENV]: toolOptimize ? "1" : "0",
 					[PI_OFFLINE_ENV]: "1",
-					[FLOW_DEADLINE_ENV]: String(Date.now() + effectiveTimeout),
+					...(effectiveTimeout > 0 ? {
+						[FLOW_DEADLINE_ENV]: String(Date.now() + effectiveTimeout),
+						[FLOW_TOOL_SUMMARY_GRACE_ENV]: String(toolSummaryGraceMs),
+					} : {
+						[FLOW_DEADLINE_ENV]: "",
+						[FLOW_TOOL_SUMMARY_GRACE_ENV]: "0",
+					}),
 				},
 			});
 
@@ -471,6 +486,7 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 			proc.stdin.on("error", () => {
 				/* ignore broken pipe on fast exits */
 			});
+			proc.stdin.end();
 
 			let abortHandler: (() => void) | undefined;
 			let buffer = "";
@@ -582,7 +598,7 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 				else signal.addEventListener("abort", abortHandler, { once: true });
 			}
 
-			// Execution timeout — two-stage: urge message, then stdin injection, then grace, then hard kill
+			// Execution timeout — two-stage: parent-side warnings, tool-level deadline abort, then grace, then hard kill
 			if (effectiveTimeout > 0) {
 				// Warning timer: notify the parent UI that the child is about to be killed.
 				// NOTE: True mid-flight injection into the child's context requires pi-core
@@ -620,21 +636,7 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 					result.stderr += `\nFlow timed out after ${Math.round(effectiveTimeout / 1000)}s.`;
 					emitUpdate();
 
-					// Stage 1: Send timeout instruction to child via stdin
-					try {
-						if (!stdinEnded) {
-							proc.stdin.write(
-								JSON.stringify({
-									type: "flow_timeout",
-									instruction: "stop and report all findings immediately",
-								}) + "\n",
-							);
-						}
-					} catch {
-						/* ignore broken pipe */
-					}
-
-					// Stage 2: Grace period before hard kill
+					// Grace period before hard kill
 					const graceTimer = setTimeout(() => {
 						if (didClose || settled) return;
 						result.stopReason = "timeout";

package/src/timed-bash.ts

@@ -4,6 +4,13 @@
  * Wraps the built-in bash tool to append execution-time classification
  * to every result. This gives the LLM concrete feedback to self-correct
  * strategy (e.g. switch from bash grep to grep tool, batch commands, etc.).
+ *
+ * Child flows also receive a hard deadline from the parent runner. When a
+ * bash command is still running near that deadline, this wrapper aborts just
+ * the bash tool and returns an explicit instruction to stop using tools and
+ * summarize. That preserves the child agent process long enough to produce
+ * its final structured report instead of being killed while a shell command is
+ * still active.
  */
 
 import { createBashToolDefinition } from "@mariozechner/pi-coding-agent";
@@ -21,6 +28,10 @@ export interface TimingReport {
 	label: string;
 }
 
+const FLOW_DEADLINE_ENV = "PI_FLOW_DEADLINE_MS";
+const FLOW_TOOL_SUMMARY_GRACE_ENV = "PI_FLOW_TOOL_SUMMARY_GRACE_MS";
+const DEFAULT_FLOW_TOOL_SUMMARY_GRACE_MS = 30_000;
+
 /** Classify duration into user-defined tiers with actionable feedback. */
 export function classifyDuration(ms: number): TimingReport {
 	const s = ms / 1000;
@@ -56,6 +67,90 @@ export function formatTimingAppendix(report: TimingReport): string {
 	return `\n\n[Execution time: ${report.label}]`;
 }
 
+function parsePositiveSafeInteger(raw: unknown): number | null {
+	if (typeof raw !== "string" || !raw.trim()) return null;
+	const parsed = Number(raw);
+	return Number.isSafeInteger(parsed) && parsed > 0 ? parsed : null;
+}
+
+function parseNonNegativeSafeInteger(raw: unknown): number | null {
+	if (typeof raw !== "string" || !raw.trim()) return null;
+	const parsed = Number(raw);
+	return Number.isSafeInteger(parsed) && parsed >= 0 ? parsed : null;
+}
+
+function getFlowDeadlineMs(): number | null {
+	return parsePositiveSafeInteger(process.env[FLOW_DEADLINE_ENV]);
+}
+
+function getFlowToolSummaryGraceMs(): number {
+	return parseNonNegativeSafeInteger(process.env[FLOW_TOOL_SUMMARY_GRACE_ENV]) ?? DEFAULT_FLOW_TOOL_SUMMARY_GRACE_MS;
+}
+
+function formatDeadlineAppendix(): string {
+	return "\n\n[Flow timeout] Bash command was interrupted to preserve time for the final flow summary. Stop running tools and return structured findings now.";
+}
+
+function appendTextToToolResult(result: any, text: string): void {
+	const textItem = result?.content?.find?.((c: any) => c.type === "text");
+	if (textItem && typeof textItem.text === "string") {
+		textItem.text += text;
+	} else if (Array.isArray(result?.content)) {
+		result.content.push({ type: "text", text: text.trim() });
+	}
+}
+
+function createDeadlineSignal(parentSignal: AbortSignal | undefined): {
+	signal: AbortSignal | undefined;
+	cleanup: () => void;
+	wasDeadlineAbort: () => boolean;
+} {
+	const deadlineMs = getFlowDeadlineMs();
+	if (!deadlineMs) {
+		return { signal: parentSignal, cleanup: () => undefined, wasDeadlineAbort: () => false };
+	}
+
+	const summaryGraceMs = getFlowToolSummaryGraceMs();
+	const abortAtMs = deadlineMs - summaryGraceMs;
+	const delayMs = abortAtMs - Date.now();
+
+	const controller = new AbortController();
+	let deadlineAbort = false;
+	let timer: NodeJS.Timeout | undefined;
+	let relayParentAbort: (() => void) | undefined;
+
+	const abortForDeadline = () => {
+		if (controller.signal.aborted) return;
+		deadlineAbort = true;
+		controller.abort(new Error("Flow deadline reached while bash command was running."));
+	};
+
+	if (parentSignal?.aborted) {
+		controller.abort(parentSignal.reason);
+	} else if (parentSignal) {
+		relayParentAbort = () => controller.abort(parentSignal.reason);
+		parentSignal.addEventListener("abort", relayParentAbort, { once: true });
+	}
+
+	if (delayMs <= 0) {
+		abortForDeadline();
+	} else {
+		timer = setTimeout(abortForDeadline, delayMs);
+		timer.unref?.();
+	}
+
+	return {
+		signal: controller.signal,
+		cleanup: () => {
+			if (timer) clearTimeout(timer);
+			if (parentSignal && relayParentAbort) {
+				parentSignal.removeEventListener("abort", relayParentAbort);
+			}
+		},
+		wasDeadlineAbort: () => deadlineAbort,
+	};
+}
+
 /**
  * Create a timed bash tool definition that wraps the built-in one.
  * Extensions override built-in tools by name, so this replaces the
@@ -93,11 +188,12 @@ export function createTimedBashToolDefinition(
 			ctx: any,
 		) {
 			const start = Date.now();
+			const deadlineSignal = createDeadlineSignal(signal);
 			try {
 				const result = await original.execute(
 					toolCallId,
 					params,
-					signal,
+					deadlineSignal.signal,
 					onUpdate,
 					ctx,
 				);
@@ -105,13 +201,10 @@ export function createTimedBashToolDefinition(
 				const report = classifyDuration(duration);
 				const appendix = formatTimingAppendix(report);
 
-				const textItem = result.content?.find(
-					(c: any) => c.type === "text",
-				);
-				if (textItem && typeof textItem.text === "string") {
-					textItem.text += appendix;
-				} else if (result.content) {
-					result.content.push({ type: "text", text: appendix.trim() });
+				appendTextToToolResult(result, appendix);
+				if (deadlineSignal.wasDeadlineAbort()) {
+					appendTextToToolResult(result, formatDeadlineAppendix());
+					result.isError = true;
 				}
 				return result;
 			} catch (err: any) {
@@ -119,10 +212,22 @@ export function createTimedBashToolDefinition(
 				const report = classifyDuration(duration);
 				const appendix = formatTimingAppendix(report);
 
+				if (deadlineSignal.wasDeadlineAbort()) {
+					const message = typeof err?.message === "string" && err.message.trim()
+						? `${err.message}${appendix}${formatDeadlineAppendix()}`
+						: `${appendix.trim()}${formatDeadlineAppendix()}`;
+					return {
+						content: [{ type: "text", text: message }],
+						isError: true,
+					};
+				}
+
 				if (err?.message && typeof err.message === "string") {
 					err.message += appendix;
 				}
 				throw err;
+			} finally {
+				deadlineSignal.cleanup();
 			}
 		},
 	};