pi-agent-flow 1.3.0 → 1.4.1

package/package.json

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

package/src/flow.ts

@@ -26,11 +26,15 @@ const isWindows = process.platform === "win32";
 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";
 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";
 export const FLOW_TOOL_OPTIMIZE_ENV = "PI_FLOW_TOOL_OPTIMIZE";
 const PI_OFFLINE_ENV = "PI_OFFLINE";
 
@@ -140,6 +144,7 @@ function buildFlowArgs(
 	maxDepth: number = 0,
 	toolOptimize: boolean = false,
 	structuredOutput: boolean = true,
+	timeoutMs?: number,
 ): string[] {
 	const args: string[] = [
 		"--mode",
@@ -210,11 +215,17 @@ function buildFlowArgs(
 		? `You may delegate to sub-flows (depth ${currentDepth}/${effectiveMaxDepth}).`
 		: `You may NOT delegate to sub-flows (depth limit reached).`;
 
+	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`
+			: "";
+
 	const activation =
 		`\n\n<activation flow="${flow.name}" depth="${currentDepth}" tools="${availableTools}">\n` +
 		`You are a [${flow.name}] agent operating at depth ${currentDepth}.\n` +
 		`Available tools: ${availableTools}.\n` +
 		`${delegationRule}\n` +
+		`${timeBudgetHint}` +
 		`Do not attempt to use any tool outside the available set — it will fail.\n` +
 		`</activation>`;
 
@@ -241,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` +
@@ -406,6 +417,14 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 		forkSessionTmpPath = forkTmp.filePath;
 	}
 
+	// Resolve timeout: explicit option > env var > default (10 min)
+	const envTimeoutRaw = process.env[FLOW_TIMEOUT_ENV];
+	const envTimeout = envTimeoutRaw !== undefined ? (() => {
+		const n = Number(envTimeoutRaw);
+		return Number.isSafeInteger(n) && n >= 0 ? n : null;
+	})() : null;
+	const effectiveTimeout = opts.timeoutMs ?? envTimeout ?? DEFAULT_FLOW_TIMEOUT_MS;
+
 	try {
 		const piArgs = buildFlowArgs(
 			flow,
@@ -416,17 +435,10 @@ export async function runFlow(opts: RunFlowOptions): Promise<SingleResult> {
 			maxDepth,
 			toolOptimize,
 			structuredOutput,
+			effectiveTimeout,
 		);
 		let wasAborted = false;
 
-		// Resolve timeout: explicit option > env var > default (10 min)
-		const envTimeoutRaw = process.env[FLOW_TIMEOUT_ENV];
-		const envTimeout = envTimeoutRaw !== undefined ? (() => {
-			const n = Number(envTimeoutRaw);
-			return Number.isSafeInteger(n) && n >= 0 ? n : null;
-		})() : null;
-		const effectiveTimeout = opts.timeoutMs ?? envTimeout ?? DEFAULT_FLOW_TIMEOUT_MS;
-
 		const exitCode = await new Promise<number>((resolve) => {
 			const nextDepth = Math.max(0, Math.floor(parentDepth)) + 1;
 			const propagatedMaxDepth = Math.max(0, Math.floor(maxDepth));
@@ -446,18 +458,25 @@ 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),
 				},
 			});
 
+			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 = () => {
@@ -468,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)], {
@@ -491,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);
@@ -554,18 +575,73 @@ 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
+				// support for out-of-band messages. Until then, the agent only knows its
+				// budget from the initial prompt; this warning is parent-side only.
+				const warningMs = effectiveTimeout - FLOW_TIME_BUDGET_WARNING_MS;
+				if (warningMs > 0) {
+					const warnTimer = setTimeout(() => {
+						if (didClose || settled) return;
+						const remainingSec = Math.round(FLOW_TIME_BUDGET_WARNING_MS / 1000);
+						const warnMsg = `\n[Flow warning] ${remainingSec}s remaining before hard timeout. The agent should wrap up now.`;
+						result.stderr += warnMsg;
+						// Force an update so the parent UI shows the warning immediately.
+						emitUpdate();
+					}, warningMs);
+					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. */