pi-agent-flow 1.2.0 → 1.2.3

package/src/index.ts

@@ -23,16 +23,22 @@ import { mapFlowConcurrent, runFlow } from "./flow.js";
 import {
 	type SingleResult,
 	type FlowDetails,
+	type CompressedFlowResult,
+	type FileEntry,
+	type CommandEntry,
 	emptyFlowUsage,
 	isFlowError,
 	isFlowSuccess,
+	getFlowOutput,
 } from "./types.js";
+import { extractStructuredOutput } from "./structured-output.js";
 import { createBatchTool, createBatchReadTool } from "./batch.js";
 import {
 	createWebTool,
 	looksLikeUrlPrompt,
 	looksLikeWebSearchPrompt,
 } from "./web-tool.js";
+import { createTimedBashToolDefinition } from "./timed-bash.js";
 
 // ---------------------------------------------------------------------------
 // Limits
@@ -471,6 +477,152 @@ function stripReasoningFromAssistantMessage(message: any): {
 function isJsonEqual(a: any, b: any): boolean {
 	return JSON.stringify(a) === JSON.stringify(b);
 }
+// ---------------------------------------------------------------------------
+// Flow result compression
+// ---------------------------------------------------------------------------
+
+/**
+ * Module-level cache populated after each flow tool execution.
+ * Maps toolCallId → compressed flow result.
+ * Consumed by compressFlowToolResults at fork time.
+ */
+const flowResultCache = new Map<string, CompressedFlowResult[]>();
+
+/**
+ * Build compressed representations of flow results and cache them by toolCallId.
+ */
+function cacheFlowResults(toolCallId: string, results: SingleResult[]): void {
+	for (const result of results) {
+		const so = result.structuredOutput;
+		if (!so) continue;
+		const compressed: CompressedFlowResult = {
+			type: result.type,
+			status: isFlowError(result) ? "failed" : "accomplished",
+		};
+		if (so.files.length > 0) compressed.files = so.files;
+		if (so.commands.length > 0) compressed.commands = so.commands;
+		if (result.errorMessage) compressed.error = result.errorMessage;
+		const existing = flowResultCache.get(toolCallId) ?? [];
+		existing.push(compressed);
+		flowResultCache.set(toolCallId, existing);
+	}
+}
+
+/**
+ * Render a compressed flow result as compact text for child context.
+ */
+function renderCompressedFlowResult(r: CompressedFlowResult): string {
+	const parts: string[] = [`[Flow: ${r.type} ${r.status}]`];
+	if (r.files?.length) {
+		const fileLines = r.files.map((f) => {
+			const role = f.role ? ` (${f.role})` : "";
+			const desc = f.description ? ` — ${f.description}` : "";
+			return `  ${f.path}${role}${desc}`;
+		});
+		parts.push(`Files:\n${fileLines.join("\n")}`);
+	}
+	if (r.commands?.length) {
+		const cmdLines = r.commands.map((c) => {
+			const result = c.result ? ` (${c.result})` : "";
+			const purpose = c.purpose ? ` — ${c.purpose}` : "";
+			return `  ${c.tool ?? "cmd"}: ${c.command}${result}${purpose}`;
+		});
+		parts.push(`Commands:\n${cmdLines.join("\n")}`);
+	}
+	if (r.error) parts.push(`Error: ${r.error}`);
+	return parts.join("\n");
+}
+
+/**
+ * Compress flow tool results in a sanitized session snapshot.
+ *
+ * Scans for tool result messages that correspond to flow invocations
+ * and replaces their content with compact compressed output.
+ */
+export function compressFlowToolResults(snapshot: string, cache: Map<string, CompressedFlowResult[]>): string {
+	if (cache.size === 0) return snapshot;
+
+	const lines = snapshot.trimEnd().split("\n");
+	const result: string[] = [];
+
+	// First pass: map toolCallId → tool name from assistant messages
+	const toolCallIdToName = new Map<string, string>();
+	for (const line of lines) {
+		let entry: any;
+		try { entry = JSON.parse(line); } catch { continue; }
+		if (entry?.type !== "message" || entry.message?.role !== "assistant") continue;
+		const content = entry.message.content;
+		if (!Array.isArray(content)) continue;
+		for (const part of content) {
+			if (part.type === "toolCall" && part.toolCallId && part.name) {
+				toolCallIdToName.set(part.toolCallId, part.name);
+			}
+		}
+	}
+
+	// Second pass: compress flow tool results
+	for (const line of lines) {
+		let entry: any;
+		try { entry = JSON.parse(line); } catch { result.push(line); continue; }
+
+		if (entry?.type !== "message" || entry.message?.role !== "tool") {
+			result.push(line);
+			continue;
+		}
+
+		// Extract toolCallId — either from message-level or content-level toolResult
+		let toolCallId: string | undefined;
+		if (typeof entry.message.toolCallId === "string") {
+			toolCallId = entry.message.toolCallId;
+		} else if (Array.isArray(entry.message.content)) {
+			for (const part of entry.message.content) {
+				if (part.type === "toolResult" && part.toolCallId) {
+					toolCallId = part.toolCallId;
+					break;
+				}
+			}
+		}
+
+		if (!toolCallId) { result.push(line); continue; }
+
+		const toolName = toolCallIdToName.get(toolCallId);
+		if (toolName !== "flow") { result.push(line); continue; }
+
+		const compressed = cache.get(toolCallId);
+		if (!compressed || compressed.length === 0) { result.push(line); continue; }
+
+		const rendered = compressed.map(renderCompressedFlowResult).join("\n\n");
+
+		// Replace content in the tool result message
+		if (typeof entry.message.toolCallId === "string") {
+			// Format 1: toolCallId at message level, content is text array
+			entry = {
+				...entry,
+				message: {
+					...entry.message,
+					content: [{ type: "text", text: rendered }],
+				},
+			};
+		} else {
+			// Format 2: toolCallId inside content array
+			entry = {
+				...entry,
+				message: {
+					...entry.message,
+					content: entry.message.content.map((part: any) =>
+						part.type === "toolResult" && part.toolCallId === toolCallId
+							? { ...part, content: rendered }
+							: part,
+					),
+				},
+			};
+		}
+
+		result.push(JSON.stringify(entry));
+	}
+
+	return `${result.join("\n")}\n`;
+}
 
 /**
  * Sanitize a fork session snapshot JSONL to remove only non-inheritable
@@ -533,7 +685,8 @@ function sanitizeForkSnapshot(snapshot: string | null): string | null {
 		sanitizedLines.push(changed ? JSON.stringify(entry) : line);
 	}
 
-	return `${sanitizedLines.join("\n")}\n`;
+	const sanitized = `${sanitizedLines.join("\n")}\n`;
+	return compressFlowToolResults(sanitized, flowResultCache);
 }
 
 function computeActiveTools(optimize: boolean): string[] {
@@ -582,6 +735,8 @@ export default function (pi: ExtensionAPI) {
 
 	// toolOptimize: CLI flag > env var > settings.json > default (true)
 	let toolOptimize = true;
+	// structuredOutput: settings.json > default (true)
+	let structuredOutput = true;
 	const envToolOptimize = process.env[FLOW_TOOL_OPTIMIZE_ENV];
 	if (envToolOptimize !== undefined) {
 		const parsed = parseBoolean(envToolOptimize);
@@ -613,6 +768,9 @@ export default function (pi: ExtensionAPI) {
 			if (typeof flowSettings.toolOptimize === "boolean") {
 				toolOptimize = flowSettings.toolOptimize;
 			}
+			if (typeof flowSettings.structuredOutput === "boolean") {
+				structuredOutput = flowSettings.structuredOutput;
+			}
 		}
 
 		// Only restrict tools for the main orchestrator (depth 0).
@@ -627,6 +785,12 @@ export default function (pi: ExtensionAPI) {
 			pi.registerTool(createBatchReadTool());
 			pi.registerTool(createBatchTool());
 		}
+
+		// Override built-in bash with timed wrapper so the LLM sees execution-time classification.
+		const timedBash = createTimedBashToolDefinition(ctx.cwd);
+		if (timedBash) {
+			pi.registerTool(timedBash);
+		}
 	});
 
 	// Re-apply active tools every turn to survive registry refreshes.
@@ -776,8 +940,9 @@ flow [type] accomplished
 			].join("\n"),
 			parameters: FlowParams,
 
-			async execute(_toolCallId, params, signal, onUpdate, ctx) {
+			async execute(toolCallId, params, signal, onUpdate, ctx) {
 				const discovery = discoverFlows(ctx.cwd, "all");
+				flowResultCache.clear();
 				const { flows } = discovery;
 				const makeDetails = makeFlowDetailsFactory(discovery.projectFlowsDir);
 
@@ -940,6 +1105,7 @@ flow [type] accomplished
 							maxDepth: effectiveMaxDepth,
 							preventCycles,
 							toolOptimize,
+							structuredOutput,
 							model: candidateModel,
 							signal,
 							onUpdate: (partial) => {
@@ -970,6 +1136,7 @@ flow [type] accomplished
 					return result;
 				});
 
+				cacheFlowResults(toolCallId, results);
 				// Build tool result with FULL flow output — no truncation
 				const successCount = results.filter((r) => isFlowSuccess(r)).length;
 				const flowReports = results.map((r) => {

package/src/render.ts

@@ -241,6 +241,42 @@ function renderFlowExpanded(
 	// Flow report (structured output)
 	container.addChild(new Spacer(1));
 	container.addChild(new Text(theme.fg("muted", sectionHeader("report")), 0, 0));
+
+	// Structured output summary (compact badge when available)
+	if (r.structuredOutput) {
+		const so = r.structuredOutput;
+		const statusColor = so.status === "complete" ? "success" : so.status === "partial" ? "warning" : "error";
+		container.addChild(new Text(
+			`${theme.fg(statusColor, `[${so.status}]`)} ${theme.fg("dim", so.summary)}`,
+			0, 0,
+		));
+		if (so.files.length > 0) {
+			container.addChild(new Text(theme.fg("dim", `Files: ${so.files.map((f) => f.path).join(", ")}`), 0, 0));
+		}
+		if (so.commands?.length > 0) {
+			const cmdLabels = so.commands.map((c) => {
+				const short = c.command.length > 30 ? c.command.slice(0, 30) + "..." : c.command;
+				return `${c.tool ?? "cmd"}: ${short}`;
+			});
+			container.addChild(new Text(theme.fg("dim", `Commands: ${cmdLabels.join(", ")}`), 0, 0));
+		}
+		if (so.notDone.length > 0) {
+			const notDoneText = so.notDone.map((item) => {
+				const details = [
+					item.reason ? `reason: ${item.reason}` : undefined,
+					item.blocker ? `blocker: ${item.blocker}` : undefined,
+					item.nextStep ? `next: ${item.nextStep}` : undefined,
+				].filter(Boolean).join("; ");
+				return details ? `${item.item} (${details})` : item.item;
+			}).join("; ");
+			container.addChild(new Text(theme.fg("dim", `Not Done: ${notDoneText}`), 0, 0));
+		}
+		if (so.nextSteps.length > 0) {
+			container.addChild(new Text(theme.fg("dim", `Next: ${so.nextSteps.join("; ")}`), 0, 0));
+		}
+		container.addChild(new Spacer(1));
+	}
+
 	if (flowOutput) {
 		container.addChild(new Markdown(flowOutput.trim(), 0, 0, mdTheme));
 	} else {
@@ -294,6 +330,9 @@ function renderFlowCollapsed(
 	if (r.exitCode === -1 && streamingText) {
 		const logContent = tailText(streamingText, contentBudget(10));
 		container.addChild(new TruncatedText(`${theme.fg("dim", "└─ msg:")} ${theme.fg("dim", logContent)}`, 0, 0));
+	} else if (r.structuredOutput?.summary) {
+		const logContent = truncateChars(r.structuredOutput.summary, contentBudget(10));
+		container.addChild(new TruncatedText(`${theme.fg("dim", "└─ msg:")} ${theme.fg("dim", logContent)}`, 0, 0));
 	} else if (flowOutput) {
 		const logContent = tailText(flowOutput, contentBudget(10));
 		container.addChild(new TruncatedText(`${theme.fg("dim", "└─ msg:")} ${theme.fg("dim", logContent)}`, 0, 0));

package/src/structured-output.ts

@@ -0,0 +1,97 @@
+/**
+ * Structured JSON output extraction from flow responses.
+ *
+ * Parses a JSON code block from the end of the flow's final assistant text
+ * and validates it against the FlowStructuredOutput schema.
+ */
+
+import type { Action, CommandEntry, FileEntry, FlowStructuredOutput, NotDoneItem } from "./types.js";
+
+type FlowStatus = FlowStructuredOutput["status"];
+
+type StructuredOutputRecord = {
+	version: string;
+	status: FlowStatus;
+	summary: string;
+	files?: FileEntry[];
+	actions?: Action[];
+	commands?: CommandEntry[];
+	notDone?: NotDoneItem[];
+	nextSteps?: string[];
+	reasoning?: string[];
+	notes?: string[];
+	extensions?: Record<string, unknown>;
+};
+
+const VALID_STATUSES: FlowStatus[] = ["complete", "partial", "blocked", "failed"];
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+function isOptionalArray(record: Record<string, unknown>, key: string): boolean {
+	return record[key] === undefined || Array.isArray(record[key]);
+}
+
+/** Minimum required fields to consider parsed JSON a valid structured output. */
+function isValidStructuredOutput(obj: unknown): obj is StructuredOutputRecord {
+	if (!isRecord(obj)) return false;
+	return (
+		typeof obj.version === "string" &&
+		typeof obj.status === "string" &&
+		VALID_STATUSES.includes(obj.status as FlowStatus) &&
+		typeof obj.summary === "string" &&
+		isOptionalArray(obj, "files") &&
+		isOptionalArray(obj, "actions") &&
+		isOptionalArray(obj, "commands") &&
+		isOptionalArray(obj, "notDone") &&
+		isOptionalArray(obj, "nextSteps") &&
+		isOptionalArray(obj, "reasoning") &&
+		isOptionalArray(obj, "notes")
+	);
+}
+
+/**
+ * Extract a structured JSON output block from the end of an assistant's text.
+ *
+ * Looks for a final ```json ... ``` code block, parses it, and validates
+ * against the FlowStructuredOutput schema. Returns undefined when the block
+ * is missing, malformed, or fails validation.
+ */
+export function extractStructuredOutput(text: string): FlowStructuredOutput | undefined {
+	if (!text) return undefined;
+
+	// Find the last ```json ... ``` block in the text.
+	// Scan backward from the end to handle multiple blocks.
+	const allMatches = [...text.matchAll(/```json\s*([\s\S]*?)\s*```/g)];
+	const match = allMatches.length > 0 ? allMatches[allMatches.length - 1] : null;
+	if (!match || !match[1]) return undefined;
+
+	const jsonStr = match[1].trim();
+	if (!jsonStr) return undefined;
+
+	let parsed: unknown;
+	try {
+		parsed = JSON.parse(jsonStr);
+	} catch {
+		return undefined;
+	}
+
+	if (!isValidStructuredOutput(parsed)) return undefined;
+
+	// Sanitize: trim string fields and normalize omitted arrays to [] for
+	// backward compatibility with earlier structured-output prompts.
+	return {
+		version: parsed.version.trim(),
+		status: parsed.status,
+		summary: parsed.summary.trim(),
+		files: parsed.files ?? [],
+		actions: parsed.actions ?? [],
+		commands: parsed.commands ?? [],
+		notDone: parsed.notDone ?? [],
+		nextSteps: parsed.nextSteps ?? [],
+		reasoning: parsed.reasoning ?? [],
+		notes: parsed.notes ?? [],
+		...(parsed.extensions !== undefined ? { extensions: parsed.extensions } : {}),
+	};
+}

package/src/timed-bash.ts

@@ -0,0 +1,129 @@
+/**
+ * Timed Bash Tool Wrapper
+ *
+ * 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.).
+ */
+
+import { createBashToolDefinition } from "@mariozechner/pi-coding-agent";
+
+export type TimingTier =
+	| "normal"
+	| "avg"
+	| "long"
+	| "extreme_long"
+	| "very_long";
+
+export interface TimingReport {
+	tier: TimingTier;
+	seconds: number;
+	label: string;
+}
+
+/** Classify duration into user-defined tiers with actionable feedback. */
+export function classifyDuration(ms: number): TimingReport {
+	const s = ms / 1000;
+	if (s < 10) {
+		return { tier: "normal", seconds: s, label: `${s.toFixed(1)}s (normal)` };
+	}
+	if (s < 30) {
+		return { tier: "avg", seconds: s, label: `${s.toFixed(1)}s (avg) — user feedback: consider improving the current commands or find a better solution` };
+	}
+	if (s < 60) {
+		return {
+			tier: "long",
+			seconds: s,
+			label: `${s.toFixed(1)}s (long) — user feedback: consider improving the whole scripts`,
+		};
+	}
+	if (s < 300) {
+		return {
+			tier: "extreme_long",
+			seconds: s,
+			label: `${s.toFixed(1)}s (extreme long) — user feedback: should consider to improve the whole scripts`,
+		};
+	}
+	return {
+		tier: "very_long",
+		seconds: s,
+		label: `${(s / 60).toFixed(1)}min (very long) — user feedback: consider to improve, only run when everything tested with other means`,
+	};
+}
+
+/** Format the timing appendix that gets appended to bash output. */
+export function formatTimingAppendix(report: TimingReport): string {
+	return `\n\n[Execution time: ${report.label}]`;
+}
+
+/**
+ * Create a timed bash tool definition that wraps the built-in one.
+ * Extensions override built-in tools by name, so this replaces the
+ * default `bash` tool transparently.
+ *
+ * Returns `null` if the underlying `createBashToolDefinition` is not
+ * available (e.g. test environment or incompatible CLI version).
+ */
+export function createTimedBashToolDefinition(
+	cwd: string,
+	options?: {
+		shellPath?: string;
+		commandPrefix?: string;
+		operations?: any;
+		spawnHook?: any;
+	},
+): any {
+	let original: any;
+	try {
+		original = createBashToolDefinition(cwd, options);
+	} catch {
+		return null;
+	}
+	if (!original || typeof original.execute !== "function") {
+		return null;
+	}
+
+	return {
+		...original,
+		async execute(
+			toolCallId: string,
+			params: { command: string; timeout?: number },
+			signal: AbortSignal,
+			onUpdate: any,
+			ctx: any,
+		) {
+			const start = Date.now();
+			try {
+				const result = await original.execute(
+					toolCallId,
+					params,
+					signal,
+					onUpdate,
+					ctx,
+				);
+				const duration = Date.now() - start;
+				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() });
+				}
+				return result;
+			} catch (err: any) {
+				const duration = Date.now() - start;
+				const report = classifyDuration(duration);
+				const appendix = formatTimingAppendix(report);
+
+				if (err?.message && typeof err.message === "string") {
+					err.message += appendix;
+				}
+				throw err;
+			}
+		},
+	};
+}

package/src/types.ts

@@ -18,6 +18,102 @@ export interface UsageStats {
 	smoothedTps?: number;
 }
 
+/** Structured file entry in a flow's output. */
+export interface FileEntry {
+	/** Path to the file, relative or absolute. */
+	path: string;
+	/** Semantic role of this file in the flow's work. */
+	role?: "reference" | "read" | "modified" | "created" | "deleted" | "test";
+	/** Why this file matters (1 sentence). */
+	description?: string;
+	/** Short excerpt or snippet (not full content). */
+	snippet?: string;
+	/** Specific line ranges of interest. */
+	ranges?: Array<{
+		start: number;
+		end: number;
+		/** Free-form label like "bug", "fix", "ref", "added". */
+		label?: string;
+	}>;
+}
+
+/** Structured command/tool invocation entry in a flow's output. */
+export interface CommandEntry {
+	/** The command string or tool call representation. */
+	command: string;
+	/** Tool used: bash, grep, find, ls, batch, read, write, edit, flow, web. */
+	tool?: string;
+	/** Working directory or target scope (file path, URL, etc.). */
+	target?: string;
+	/** Whether it succeeded, failed, or was aborted. */
+	result?: "success" | "failure" | "partial" | "aborted";
+	/** High-level output or error excerpt (not full stdout). */
+	output?: string;
+	/** Why this command was run (1 sentence). */
+	purpose?: string;
+}
+
+/** Compressed representation of a flow result for child context inheritance. */
+export interface CompressedFlowResult {
+	/** Flow type (scout, build, debug, etc.). */
+	type: string;
+	/** Execution outcome. */
+	status: "accomplished" | "failed" | "aborted";
+	/** Files touched, read, or referenced. */
+	files?: FileEntry[];
+	/** Commands or tool calls executed. */
+	commands?: CommandEntry[];
+	/** Error message for failed/aborted flows. */
+	error?: string;
+}
+
+/** Action performed or attempted by a flow. */
+export interface Action {
+	type: string;
+	description: string;
+	target?: string;
+	result?: "success" | "failure" | "partial" | "skipped";
+	evidence?: string;
+}
+
+/** Incomplete, skipped, blocked, or deferred work reported by a flow. */
+export interface NotDoneItem {
+	/** The unfinished item. */
+	item: string;
+	/** Why the item was not completed. */
+	reason?: string;
+	/** Concrete blocker preventing completion, when applicable. */
+	blocker?: string;
+	/** Suggested follow-up for this item. */
+	nextStep?: string;
+}
+
+/** Structured JSON output from a flow run. */
+export interface FlowStructuredOutput {
+	/** Schema version for forward compatibility. */
+	version: string;
+	/** Overall completion status. */
+	status: "complete" | "partial" | "blocked" | "failed";
+	/** 1–3 sentence summary of what was accomplished. */
+	summary: string;
+	/** Files touched, read, or referenced. */
+	files: FileEntry[];
+	/** Actions performed or attempted. */
+	actions: Action[];
+	/** Commands or tool calls executed during the flow. */
+	commands: CommandEntry[];
+	/** Incomplete, skipped, blocked, or deferred work. */
+	notDone: NotDoneItem[];
+	/** Recommended next steps or follow-up flows. */
+	nextSteps: string[];
+	/** Reasoning chains, hypotheses, inferences made during the flow. */
+	reasoning: string[];
+	/** Observations, warnings, caveats, side notes. */
+	notes: string[];
+	/** Escape hatch for flow-specific data (audit findings, debug root cause, etc.). */
+	extensions?: Record<string, unknown>;
+}
+
 /** Result of a single flow invocation. */
 export interface SingleResult {
 	type: string;
@@ -34,6 +130,8 @@ export interface SingleResult {
 	sawAgentEnd?: boolean;
 	/** Live in-progress text for status rendering; not part of the final flow report. */
 	streamingText?: string;
+	/** Structured JSON output parsed from the flow's final response. */
+	structuredOutput?: FlowStructuredOutput;
 }
 
 /** Metadata attached to every tool result for rendering. */