fix(coding-agent): stream bash output incrementally (#4165)

fixes #4145

packages/coding-agent/src/core/tools/bash.ts

@@ -1,7 +1,4 @@
-import { randomBytes } from "node:crypto";
+import { existsSync } from "node:fs";
-import { createWriteStream, existsSync } from "node:fs";
-import { tmpdir } from "node:os";
-import { join } from "node:path";
 import type { AgentTool } from "@mariozechner/pi-agent-core";
 import { Container, Text, truncateToWidth } from "@mariozechner/pi-tui";
 import { spawn } from "child_process";
@@ -18,17 +15,10 @@ import {
 	untrackDetachedChildPid,
 } from "../../utils/shell.js";
 import type { ToolDefinition, ToolRenderResultOptions } from "../extensions/types.js";
+import { OutputAccumulator } from "./output-accumulator.js";
 import { getTextOutput, invalidArgText, str } from "./render-utils.js";
 import { wrapToolDefinition } from "./tool-definition-wrapper.js";
-import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult, truncateTail } from "./truncate.js";
+import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult } from "./truncate.js";
-
-/**
- * Generate a unique temp file path for bash output.
- */
-function getTempFilePath(): string {
-	const id = randomBytes(8).toString("hex");
-	return join(tmpdir(), `pi-bash-${id}.log`);
-}
 
 const bashSchema = Type.Object({
 	command: Type.String({ description: "Bash command to execute" }),
@@ -161,6 +151,7 @@ export interface BashToolOptions {
 }
 
 const BASH_PREVIEW_LINES = 5;
+const BASH_UPDATE_THROTTLE_MS = 100;
 
 type BashRenderState = {
 	startedAt: number | undefined;
@@ -292,118 +283,119 @@ export function createBashToolDefinition(
 		) {
 			const resolvedCommand = commandPrefix ? `${commandPrefix}\n${command}` : command;
 			const spawnContext = resolveSpawnContext(resolvedCommand, cwd, spawnHook);
+			const output = new OutputAccumulator({ tempFilePrefix: "pi-bash" });
+			let updateTimer: NodeJS.Timeout | undefined;
+			let updateDirty = false;
+			let lastUpdateAt = 0;
+
+			const emitOutputUpdate = () => {
+				if (!onUpdate || !updateDirty) return;
+				updateDirty = false;
+				lastUpdateAt = Date.now();
+				const snapshot = output.snapshot({ persistIfTruncated: true });
+				onUpdate({
+					content: [{ type: "text", text: snapshot.content || "" }],
+					details: {
+						truncation: snapshot.truncation.truncated ? snapshot.truncation : undefined,
+						fullOutputPath: snapshot.fullOutputPath,
+					},
+				});
+			};
+
+			const clearUpdateTimer = () => {
+				if (updateTimer) {
+					clearTimeout(updateTimer);
+					updateTimer = undefined;
+				}
+			};
+
+			const scheduleOutputUpdate = () => {
+				if (!onUpdate) return;
+				updateDirty = true;
+				const delay = BASH_UPDATE_THROTTLE_MS - (Date.now() - lastUpdateAt);
+				if (delay <= 0) {
+					clearUpdateTimer();
+					emitOutputUpdate();
+					return;
+				}
+				updateTimer ??= setTimeout(() => {
+					updateTimer = undefined;
+					emitOutputUpdate();
+				}, delay);
+			};
+
 			if (onUpdate) {
 				onUpdate({ content: [], details: undefined });
 			}
-			return new Promise((resolve, reject) => {
-				let tempFilePath: string | undefined;
-				let tempFileStream: ReturnType<typeof createWriteStream> | undefined;
-				let totalBytes = 0;
-				const chunks: Buffer[] = [];
-				let chunksBytes = 0;
-				const maxChunksBytes = DEFAULT_MAX_BYTES * 2;
 
-				const ensureTempFile = () => {
+			const handleData = (data: Buffer) => {
-					if (tempFilePath) return;
+				output.append(data);
-					tempFilePath = getTempFilePath();
+				scheduleOutputUpdate();
-					tempFileStream = createWriteStream(tempFilePath);
+			};
-					for (const chunk of chunks) tempFileStream.write(chunk);
-				};
 
-				const handleData = (data: Buffer) => {
+			const finishOutput = async () => {
-					totalBytes += data.length;
+				output.finish();
-					// Start writing to a temp file once output exceeds the in-memory threshold.
+				clearUpdateTimer();
-					if (totalBytes > DEFAULT_MAX_BYTES) {
+				emitOutputUpdate();
-						ensureTempFile();
+				const snapshot = output.snapshot({ persistIfTruncated: true });
-					}
+				await output.closeTempFile();
-					// Write to temp file if we have one.
+				return snapshot;
-					if (tempFileStream) tempFileStream.write(data);
+			};
-					// Keep a rolling buffer of recent output for tail truncation.
-					chunks.push(data);
-					chunksBytes += data.length;
-					// Trim old chunks if the rolling buffer grows too large.
-					while (chunksBytes > maxChunksBytes && chunks.length > 1) {
-						const removed = chunks.shift()!;
-						chunksBytes -= removed.length;
-					}
-					// Stream partial output using the rolling tail buffer.
-					if (onUpdate) {
-						const fullBuffer = Buffer.concat(chunks);
-						const fullText = fullBuffer.toString("utf-8");
-						const truncation = truncateTail(fullText);
-						if (truncation.truncated) {
-							ensureTempFile();
-						}
-						onUpdate({
-							content: [{ type: "text", text: truncation.content || "" }],
-							details: {
-								truncation: truncation.truncated ? truncation : undefined,
-								fullOutputPath: tempFilePath,
-							},
-						});
-					}
-				};
 
-				ops.exec(spawnContext.command, spawnContext.cwd, {
+			const formatOutput = (snapshot: Awaited<ReturnType<typeof finishOutput>>, emptyText = "(no output)") => {
-					onData: handleData,
+				const truncation = snapshot.truncation;
-					signal,
+				let text = snapshot.content || emptyText;
-					timeout,
+				let details: BashToolDetails | undefined;
-					env: spawnContext.env,
+				if (truncation.truncated) {
-				})
+					details = { truncation, fullOutputPath: snapshot.fullOutputPath };
-					.then(({ exitCode }) => {
+					const startLine = truncation.totalLines - truncation.outputLines + 1;
-						// Combine the rolling buffer chunks.
+					const endLine = truncation.totalLines;
-						const fullBuffer = Buffer.concat(chunks);
+					if (truncation.lastLinePartial) {
-						const fullOutput = fullBuffer.toString("utf-8");
+						const lastLineSize = formatSize(output.getLastLineBytes());
-						// Apply tail truncation for the final display payload.
+						text += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${snapshot.fullOutputPath}]`;
-						const truncation = truncateTail(fullOutput);
+					} else if (truncation.truncatedBy === "lines") {
-						if (truncation.truncated) {
+						text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${snapshot.fullOutputPath}]`;
-							ensureTempFile();
+					} else {
-						}
+						text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snapshot.fullOutputPath}]`;
-						// Close temp file stream before building the final result.
+					}
-						if (tempFileStream) tempFileStream.end();
+				}
-						let outputText = truncation.content || "(no output)";
+				return { text, details };
-						let details: BashToolDetails | undefined;
+			};
-						if (truncation.truncated) {
+
-							// Build truncation details and an actionable notice.
+			const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`;
-							details = { truncation, fullOutputPath: tempFilePath };
+
-							const startLine = truncation.totalLines - truncation.outputLines + 1;
+			try {
-							const endLine = truncation.totalLines;
+				let exitCode: number | null;
-							if (truncation.lastLinePartial) {
+				try {
-								// Edge case: the last line alone is larger than the byte limit.
+					const result = await ops.exec(spawnContext.command, spawnContext.cwd, {
-								const lastLineSize = formatSize(Buffer.byteLength(fullOutput.split("\n").pop() || "", "utf-8"));
+						onData: handleData,
-								outputText += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${tempFilePath}]`;
+						signal,
-							} else if (truncation.truncatedBy === "lines") {
+						timeout,
-								outputText += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${tempFilePath}]`;
+						env: spawnContext.env,
-							} else {
-								outputText += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${tempFilePath}]`;
-							}
-						}
-						if (exitCode !== 0 && exitCode !== null) {
-							outputText += `\n\nCommand exited with code ${exitCode}`;
-							reject(new Error(outputText));
-						} else {
-							resolve({ content: [{ type: "text", text: outputText }], details });
-						}
-					})
-					.catch((err: Error) => {
-						// Close temp file stream and include buffered output in the error message.
-						if (tempFileStream) tempFileStream.end();
-						const fullBuffer = Buffer.concat(chunks);
-						let output = fullBuffer.toString("utf-8");
-						if (err.message === "aborted") {
-							if (output) output += "\n\n";
-							output += "Command aborted";
-							reject(new Error(output));
-						} else if (err.message.startsWith("timeout:")) {
-							const timeoutSecs = err.message.split(":")[1];
-							if (output) output += "\n\n";
-							output += `Command timed out after ${timeoutSecs} seconds`;
-							reject(new Error(output));
-						} else {
-							reject(err);
-						}
 					});
-			});
+					exitCode = result.exitCode;
+				} catch (err) {
+					const snapshot = await finishOutput();
+					const { text } = formatOutput(snapshot, "");
+					if (err instanceof Error && err.message === "aborted") {
+						throw new Error(appendStatus(text, "Command aborted"));
+					}
+					if (err instanceof Error && err.message.startsWith("timeout:")) {
+						const timeoutSecs = err.message.split(":")[1];
+						throw new Error(appendStatus(text, `Command timed out after ${timeoutSecs} seconds`));
+					}
+					throw err;
+				}
+
+				const snapshot = await finishOutput();
+				const { text: outputText, details } = formatOutput(snapshot);
+				if (exitCode !== 0 && exitCode !== null) {
+					throw new Error(appendStatus(outputText, `Command exited with code ${exitCode}`));
+				}
+				return { content: [{ type: "text", text: outputText }], details };
+			} finally {
+				clearUpdateTimer();
+			}
 		},
 		renderCall(args, _theme, context) {
 			const state = context.state;

packages/coding-agent/src/core/tools/output-accumulator.ts

@@ -0,0 +1,216 @@
+import { randomBytes } from "node:crypto";
+import { createWriteStream, type WriteStream } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, type TruncationResult, truncateTail } from "./truncate.js";
+
+export interface OutputAccumulatorOptions {
+	maxLines?: number;
+	maxBytes?: number;
+	tempFilePrefix?: string;
+}
+
+export interface OutputSnapshot {
+	content: string;
+	truncation: TruncationResult;
+	fullOutputPath?: string;
+}
+
+function defaultTempFilePath(prefix: string): string {
+	const id = randomBytes(8).toString("hex");
+	return join(tmpdir(), `${prefix}-${id}.log`);
+}
+
+function byteLength(text: string): number {
+	return Buffer.byteLength(text, "utf-8");
+}
+
+/**
+ * Incrementally tracks streaming output with bounded memory.
+ *
+ * Appends decode chunks with a streaming UTF-8 decoder, keeps only a decoded
+ * tail for display snapshots, and opens a temp file when the full output needs
+ * to be preserved.
+ */
+export class OutputAccumulator {
+	private readonly maxLines: number;
+	private readonly maxBytes: number;
+	private readonly maxRollingBytes: number;
+	private readonly tempFilePrefix: string;
+	private readonly decoder = new TextDecoder();
+
+	private rawChunks: Buffer[] = [];
+	private tailText = "";
+	private tailBytes = 0;
+	private tailStartsAtLineBoundary = true;
+	private totalRawBytes = 0;
+	private totalDecodedBytes = 0;
+	private totalLines = 1;
+	private currentLineBytes = 0;
+	private finished = false;
+
+	private tempFilePath: string | undefined;
+	private tempFileStream: WriteStream | undefined;
+
+	constructor(options: OutputAccumulatorOptions = {}) {
+		this.maxLines = options.maxLines ?? DEFAULT_MAX_LINES;
+		this.maxBytes = options.maxBytes ?? DEFAULT_MAX_BYTES;
+		this.maxRollingBytes = Math.max(this.maxBytes * 2, 1);
+		this.tempFilePrefix = options.tempFilePrefix ?? "pi-output";
+	}
+
+	append(data: Buffer): void {
+		if (this.finished) {
+			throw new Error("Cannot append to a finished output accumulator");
+		}
+
+		this.totalRawBytes += data.length;
+		this.appendDecodedText(this.decoder.decode(data, { stream: true }));
+
+		if (this.tempFileStream || this.shouldUseTempFile()) {
+			this.ensureTempFile();
+			this.tempFileStream?.write(data);
+		} else if (data.length > 0) {
+			this.rawChunks.push(data);
+		}
+	}
+
+	finish(): void {
+		if (this.finished) {
+			return;
+		}
+		this.finished = true;
+		this.appendDecodedText(this.decoder.decode());
+		if (this.shouldUseTempFile()) {
+			this.ensureTempFile();
+		}
+	}
+
+	snapshot(options: { persistIfTruncated?: boolean } = {}): OutputSnapshot {
+		const tailTruncation = truncateTail(this.getSnapshotText(), {
+			maxLines: this.maxLines,
+			maxBytes: this.maxBytes,
+		});
+		const truncated = this.totalLines > this.maxLines || this.totalDecodedBytes > this.maxBytes;
+		const truncatedBy = truncated
+			? (tailTruncation.truncatedBy ?? (this.totalDecodedBytes > this.maxBytes ? "bytes" : "lines"))
+			: null;
+		const truncation: TruncationResult = {
+			...tailTruncation,
+			truncated,
+			truncatedBy,
+			totalLines: this.totalLines,
+			totalBytes: this.totalDecodedBytes,
+			maxLines: this.maxLines,
+			maxBytes: this.maxBytes,
+		};
+
+		if (options.persistIfTruncated && truncation.truncated) {
+			this.ensureTempFile();
+		}
+
+		return {
+			content: truncation.content,
+			truncation,
+			fullOutputPath: this.tempFilePath,
+		};
+	}
+
+	async closeTempFile(): Promise<void> {
+		if (!this.tempFileStream) {
+			return;
+		}
+
+		const stream = this.tempFileStream;
+		this.tempFileStream = undefined;
+
+		await new Promise<void>((resolve, reject) => {
+			const onError = (error: Error) => {
+				stream.off("finish", onFinish);
+				reject(error);
+			};
+			const onFinish = () => {
+				stream.off("error", onError);
+				resolve();
+			};
+			stream.once("error", onError);
+			stream.once("finish", onFinish);
+			stream.end();
+		});
+	}
+
+	getLastLineBytes(): number {
+		return this.currentLineBytes;
+	}
+
+	private appendDecodedText(text: string): void {
+		if (text.length === 0) {
+			return;
+		}
+
+		const bytes = byteLength(text);
+		this.totalDecodedBytes += bytes;
+		this.tailText += text;
+		this.tailBytes += bytes;
+		if (this.tailBytes > this.maxRollingBytes * 2) {
+			this.trimTail();
+		}
+
+		let newlines = 0;
+		let lastNewline = -1;
+		for (let i = text.indexOf("\n"); i !== -1; i = text.indexOf("\n", i + 1)) {
+			newlines++;
+			lastNewline = i;
+		}
+		if (newlines === 0) {
+			this.currentLineBytes += bytes;
+		} else {
+			this.totalLines += newlines;
+			this.currentLineBytes = byteLength(text.slice(lastNewline + 1));
+		}
+	}
+
+	private trimTail(): void {
+		const buffer = Buffer.from(this.tailText, "utf-8");
+		if (buffer.length <= this.maxRollingBytes) {
+			this.tailBytes = buffer.length;
+			return;
+		}
+
+		let start = buffer.length - this.maxRollingBytes;
+		while (start < buffer.length && (buffer[start] & 0xc0) === 0x80) {
+			start++;
+		}
+
+		this.tailStartsAtLineBoundary = start === 0 ? this.tailStartsAtLineBoundary : buffer[start - 1] === 0x0a;
+		this.tailText = buffer.subarray(start).toString("utf-8");
+		this.tailBytes = byteLength(this.tailText);
+	}
+
+	private getSnapshotText(): string {
+		if (this.tailStartsAtLineBoundary) {
+			return this.tailText;
+		}
+
+		const firstNewline = this.tailText.indexOf("\n");
+		return firstNewline === -1 ? this.tailText : this.tailText.slice(firstNewline + 1);
+	}
+
+	private shouldUseTempFile(): boolean {
+		return (
+			this.totalRawBytes > this.maxBytes || this.totalDecodedBytes > this.maxBytes || this.totalLines > this.maxLines
+		);
+	}
+
+	private ensureTempFile(): void {
+		if (this.tempFilePath) {
+			return;
+		}
+		this.tempFilePath = defaultTempFilePath(this.tempFilePrefix);
+		this.tempFileStream = createWriteStream(this.tempFilePath);
+		for (const chunk of this.rawChunks) {
+			this.tempFileStream.write(chunk);
+		}
+		this.rawChunks = [];
+	}
+}

packages/coding-agent/test/tools.test.ts

@@ -3,7 +3,7 @@ import { tmpdir } from "os";
 import { join } from "path";
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { executeBashWithOperations } from "../src/core/bash-executor.js";
-import { createBashTool, createLocalBashOperations } from "../src/core/tools/bash.js";
+import { type BashOperations, createBashTool, createLocalBashOperations } from "../src/core/tools/bash.js";
 import { computeEditsDiff } from "../src/core/tools/edit-diff.js";
 import {
 	createEditTool,
@@ -449,6 +449,42 @@ describe("Coding Agent Tools", () => {
 			);
 		});
 
+		it("should include full output path for truncated timeout and abort errors", async () => {
+			for (const testCase of [
+				{ error: "timeout:5", expected: "Command timed out after 5 seconds" },
+				{ error: "aborted", expected: "Command aborted" },
+			]) {
+				const operations: BashOperations = {
+					exec: async (_command, _cwd, { onData }) => {
+						for (let i = 1; i <= 3000; i++) {
+							onData(Buffer.from(`${i}\n`, "utf-8"));
+						}
+						throw new Error(testCase.error);
+					},
+				};
+				const bash = createBashTool(testDir, { operations });
+
+				let error: unknown;
+				try {
+					await bash.execute(`test-call-${testCase.error}`, { command: "chatty-fail" });
+				} catch (err) {
+					error = err;
+				}
+
+				expect(error).toBeInstanceOf(Error);
+				const message = (error as Error).message;
+				expect(message).toContain(testCase.expected);
+				expect(message).toMatch(/\[Showing lines \d+-\d+ of \d+\. Full output: /);
+				expect(message).not.toContain("Full output: undefined");
+				const fullOutputPath = message.match(/Full output: ([^\]\n]+)/)?.[1];
+				expect(fullOutputPath).toBeDefined();
+				expect(existsSync(fullOutputPath!)).toBe(true);
+				const fullOutput = readFileSync(fullOutputPath!, "utf-8");
+				expect(fullOutput).toContain("1\n2\n3");
+				expect(fullOutput).toContain("2998\n2999\n3000");
+			}
+		});
+
 		it("should throw error when cwd does not exist", async () => {
 			const nonexistentCwd = "/this/directory/definitely/does/not/exist/12345";
 
@@ -517,6 +553,42 @@ describe("Coding Agent Tools", () => {
 			expect(getTextOutput(result).trim()).toBe("no-prefix");
 		});
 
+		it("should coalesce streaming updates for chatty output", async () => {
+			const operations: BashOperations = {
+				exec: async (_command, _cwd, { onData }) => {
+					for (let i = 0; i < 5000; i++) {
+						onData(Buffer.from(`line ${i}\n`, "utf-8"));
+					}
+					return { exitCode: 0 };
+				},
+			};
+			const updates: Array<{ content: Array<{ type: string; text?: string }>; details?: unknown }> = [];
+			const bash = createBashTool(testDir, { operations });
+
+			const result = await bash.execute("test-call-chatty-updates", { command: "chatty" }, undefined, (update) =>
+				updates.push(update),
+			);
+
+			expect(updates.length).toBeLessThan(25);
+			expect(getTextOutput(result)).toContain("line 4999");
+		});
+
+		it("should decode UTF-8 characters split across output chunks", async () => {
+			const euro = Buffer.from("€\n", "utf-8");
+			const operations: BashOperations = {
+				exec: async (_command, _cwd, { onData }) => {
+					onData(euro.subarray(0, 1));
+					onData(euro.subarray(1));
+					return { exitCode: 0 };
+				},
+			};
+			const bash = createBashTool(testDir, { operations });
+
+			const result = await bash.execute("test-call-split-utf8", { command: "split-utf8" });
+
+			expect(getTextOutput(result).trim()).toBe("€");
+		});
+
 		it("should expose local bash operations for extension reuse", async () => {
 			const ops = createLocalBashOperations();
 			const chunks: Buffer[] = [];