pi-agent-flow 1.4.0 → 1.4.1

package/package.json

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

package/src/flow.ts

@@ -27,6 +27,8 @@ const SIGKILL_TIMEOUT_MS = 5000;
 const AGENT_END_GRACE_MS = 2000;
 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_DEPTH_ENV = "PI_FLOW_DEPTH";
 const FLOW_MAX_DEPTH_ENV = "PI_FLOW_MAX_DEPTH";
 const FLOW_STACK_ENV = "PI_FLOW_STACK";
@@ -250,7 +252,7 @@ function buildFlowArgs(
 			`    { "type": "read", "description": "what was done", "target": "file.ts", "result": "success", "evidence": "output or proof" }\n` +
 			`  ],\n` +
 			`  "commands": [\n` +
-			`    { "command": "curl -s -X POST https://api.example.com/v1/data -H 'Authorization: Bearer token'", "tool": "bash" }\n` +
+			`    { "command": "curl -s -X POST https://api.example.com/v1/data -H 'Authorization: Bearer token'", "tool": "bash", "executionTime": "1.2s (normal)" }\n` +
 			`  ],\n` +
 			`  "notDone": [\n` +
 			`    { "item": "unfinished work", "reason": "why it was not completed", "blocker": "blocking issue if any", "nextStep": "specific follow-up" }\n` +
@@ -460,15 +462,21 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 				},
 			});
 
+			let stdinEnded = false;
+			const endStdin = () => {
+				if (stdinEnded) return;
+				stdinEnded = true;
+				try { proc.stdin.end(); } catch { /* ignore */ }
+			};
 			proc.stdin.on("error", () => {
 				/* ignore broken pipe on fast exits */
 			});
-			proc.stdin.end();
 
+			let abortHandler: (() => void) | undefined;
 			let buffer = "";
 			let didClose = false;
 			let settled = false;
-			let abortHandler: (() => void) | undefined;
+			let timeoutFired = false;
 			let semanticCompletionTimer: NodeJS.Timeout | undefined;
 
 			const clearSemanticCompletionTimer = () => {
@@ -479,6 +487,7 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 			};
 
 			const terminateChild = () => {
+				endStdin();
 				if (isWindows) {
 					if (proc.pid !== undefined) {
 						const killer = spawn("taskkill", ["/T", "/F", "/PID", String(proc.pid)], {
@@ -502,6 +511,7 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 			const finish = (code: number) => {
 				if (settled) return;
 				settled = true;
+				endStdin();
 				clearSemanticCompletionTimer();
 				if (signal && abortHandler) {
 					signal.removeEventListener("abort", abortHandler);
@@ -565,13 +575,14 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 				abortHandler = () => {
 					if (didClose || settled) return;
 					wasAborted = true;
+					endStdin();
 					terminateChild();
 				};
 				if (signal.aborted) abortHandler();
 				else signal.addEventListener("abort", abortHandler, { once: true });
 			}
 
-			// Execution timeout — kill child if it runs too long
+			// Execution timeout — two-stage: urge message, then stdin injection, 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
@@ -590,10 +601,47 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 					warnTimer.unref();
 				}
 
+				// Final urge timer: stronger warning 30 s before hard timeout
+				const urgeMs = effectiveTimeout - FLOW_FINAL_URGE_MS;
+				if (urgeMs > 0) {
+					const urgeTimer = setTimeout(() => {
+						if (didClose || settled) return;
+						const remainingSec = Math.round(FLOW_FINAL_URGE_MS / 1000);
+						const urgeMsg = `\n[Flow warning] ${remainingSec}s remaining before hard timeout. Stop all work and output your structured findings.`;
+						result.stderr += urgeMsg;
+						emitUpdate();
+					}, urgeMs);
+					urgeTimer.unref();
+				}
+
 				const timeoutTimer = setTimeout(() => {
 					if (didClose || settled) return;
+					timeoutFired = true;
 					result.stderr += `\nFlow timed out after ${Math.round(effectiveTimeout / 1000)}s.`;
-					terminateChild();
+					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
+					const graceTimer = setTimeout(() => {
+						if (didClose || settled) return;
+						result.stopReason = "timeout";
+						result.errorMessage = `Flow timed out after ${Math.round(effectiveTimeout / 1000)}s.`;
+						terminateChild();
+					}, REPORTING_GRACE_MS);
+					graceTimer.unref();
 				}, effectiveTimeout);
 				timeoutTimer.unref();
 			}

package/src/structured-output.ts

@@ -110,8 +110,30 @@ export function enrichStructuredOutputCommands(
 	structuredOutput: FlowStructuredOutput,
 	messages: Message[],
 ): FlowStructuredOutput {
+	// Build a map of toolCallId -> execution time string from tool result messages.
+	const timingMap = new Map<string, string>();
+	for (const msg of messages) {
+		if (msg.role !== "tool" || !Array.isArray(msg.content)) continue;
+		const id =
+			(msg as unknown as { toolCallId?: string }).toolCallId ||
+			(msg as unknown as { tool_call_id?: string }).tool_call_id ||
+			"";
+		if (!id) continue;
+		const text = msg.content
+			.filter((c: { type: string; text?: string }) => c.type === "text" && typeof c.text === "string")
+			.map((c: { text: string }) => c.text)
+			.join("");
+		const match = text.match(/\[Execution time: ([^\]]+)\]/);
+		if (match) timingMap.set(id, match[1]);
+	}
+
 	// Collect actual verbatim bash commands (including those inside batch ops)
-	const actualBashCommands: string[] = [];
+	// alongside their toolCallId so we can look up execution time.
+	interface BashEntry {
+		command: string;
+		toolCallId?: string;
+	}
+	const actualBashEntries: BashEntry[] = [];
 	for (const msg of messages) {
 		if (msg.role !== "assistant" || !Array.isArray(msg.content)) continue;
 		for (const part of msg.content) {
@@ -121,10 +143,14 @@ export function enrichStructuredOutputCommands(
 				(part as unknown as { toolName?: string }).toolName ||
 				"";
 			const args = part.arguments || part.input || {};
+			const toolCallId =
+				(part as unknown as { toolCallId?: string }).toolCallId ||
+				(part as unknown as { tool_call_id?: string }).tool_call_id ||
+				"";
 
 			if (name === "bash") {
 				const cmd = (args.command as string) || "";
-				if (cmd) actualBashCommands.push(cmd);
+				if (cmd) actualBashEntries.push({ command: cmd, toolCallId });
 			} else if (name === "batch") {
 				const ops = Array.isArray(args.o)
 					? args.o
@@ -137,20 +163,30 @@ export function enrichStructuredOutputCommands(
 					if (!op) continue;
 					const opType = (op.o ?? op.op) as string;
 					if (opType === "bash" && op.command) {
-						actualBashCommands.push(op.command as string);
+						// Batch-nested bash ops don't have individual toolCallIds,
+						// so execution time can't be correlated.
+						actualBashEntries.push({ command: op.command as string });
 					}
 				}
 			}
 		}
 	}
 
-	if (actualBashCommands.length === 0) return structuredOutput;
+	if (actualBashEntries.length === 0) return structuredOutput;
 
 	// Replace paraphrased bash commands with actual verbatim ones, in order.
 	let bashIdx = 0;
 	const enrichedCommands = structuredOutput.commands.map((cmd) => {
-		if (cmd.tool === "bash" && bashIdx < actualBashCommands.length) {
-			return { command: actualBashCommands[bashIdx++], tool: "bash" };
+		if (cmd.tool === "bash" && bashIdx < actualBashEntries.length) {
+			const entry = actualBashEntries[bashIdx++];
+			const executionTime = entry.toolCallId
+				? timingMap.get(entry.toolCallId)
+				: undefined;
+			return {
+				command: entry.command,
+				tool: "bash",
+				...(executionTime !== undefined ? { executionTime } : {}),
+			};
 		}
 		return { command: cmd.command, tool: cmd.tool };
 	});

package/src/types.ts

@@ -43,6 +43,8 @@ export interface CommandEntry {
 	command: string;
 	/** Tool used: bash, grep, find, ls, batch, read, write, edit, flow, web. */
 	tool?: string;
+	/** Execution time classification from the timed bash wrapper (e.g. "3.5s (normal)"). */
+	executionTime?: string;
 }
 
 /** Compressed representation of a flow result for child context inheritance. */
@@ -176,7 +178,7 @@ export function isFlowComplete(r: Pick<SingleResult, "messages" | "sawAgentEnd">
 export function isFlowSuccess(r: SingleResult): boolean {
 	if (r.exitCode === -1) return false;
 	if (isFlowComplete(r)) return true;
-	return r.exitCode === 0 && r.stopReason !== "error" && r.stopReason !== "aborted";
+	return r.exitCode === 0 && r.stopReason !== "error" && r.stopReason !== "aborted" && r.stopReason !== "timeout";
 }
 
 /** Whether a result represents an error. */