From 6b18cdbac12edb1d6156b1307cb5133fa0ad1a41 Mon Sep 17 00:00:00 2001 From: Armin Ronacher Date: Mon, 4 May 2026 19:06:07 +0200 Subject: [PATCH] fix(coding-agent): stream bash output incrementally (#4165) fixes #4145 --- packages/coding-agent/src/core/tools/bash.ts | 228 +++++++++--------- .../src/core/tools/output-accumulator.ts | 216 +++++++++++++++++ packages/coding-agent/test/tools.test.ts | 74 +++++- 3 files changed, 399 insertions(+), 119 deletions(-) create mode 100644 packages/coding-agent/src/core/tools/output-accumulator.ts diff --git a/packages/coding-agent/src/core/tools/bash.ts b/packages/coding-agent/src/core/tools/bash.ts index 18762371..c49e8996 100644 --- a/packages/coding-agent/src/core/tools/bash.ts +++ b/packages/coding-agent/src/core/tools/bash.ts @@ -1,7 +1,4 @@ -import { randomBytes } from "node:crypto"; -import { createWriteStream, existsSync } from "node:fs"; -import { tmpdir } from "node:os"; -import { join } from "node:path"; +import { existsSync } from "node:fs"; 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"; - -/** - * 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`); -} +import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, formatSize, type TruncationResult } from "./truncate.js"; 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 | undefined; - let totalBytes = 0; - const chunks: Buffer[] = []; - let chunksBytes = 0; - const maxChunksBytes = DEFAULT_MAX_BYTES * 2; - const ensureTempFile = () => { - if (tempFilePath) return; - tempFilePath = getTempFilePath(); - tempFileStream = createWriteStream(tempFilePath); - for (const chunk of chunks) tempFileStream.write(chunk); - }; + const handleData = (data: Buffer) => { + output.append(data); + scheduleOutputUpdate(); + }; - const handleData = (data: Buffer) => { - totalBytes += data.length; - // Start writing to a temp file once output exceeds the in-memory threshold. - if (totalBytes > DEFAULT_MAX_BYTES) { - ensureTempFile(); - } - // Write to temp file if we have one. - 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, - }, - }); - } - }; + const finishOutput = async () => { + output.finish(); + clearUpdateTimer(); + emitOutputUpdate(); + const snapshot = output.snapshot({ persistIfTruncated: true }); + await output.closeTempFile(); + return snapshot; + }; - ops.exec(spawnContext.command, spawnContext.cwd, { - onData: handleData, - signal, - timeout, - env: spawnContext.env, - }) - .then(({ exitCode }) => { - // Combine the rolling buffer chunks. - const fullBuffer = Buffer.concat(chunks); - const fullOutput = fullBuffer.toString("utf-8"); - // Apply tail truncation for the final display payload. - const truncation = truncateTail(fullOutput); - if (truncation.truncated) { - ensureTempFile(); - } - // Close temp file stream before building the final result. - if (tempFileStream) tempFileStream.end(); - let outputText = truncation.content || "(no output)"; - let details: BashToolDetails | undefined; - if (truncation.truncated) { - // Build truncation details and an actionable notice. - details = { truncation, fullOutputPath: tempFilePath }; - const startLine = truncation.totalLines - truncation.outputLines + 1; - const endLine = truncation.totalLines; - if (truncation.lastLinePartial) { - // Edge case: the last line alone is larger than the byte limit. - const lastLineSize = formatSize(Buffer.byteLength(fullOutput.split("\n").pop() || "", "utf-8")); - outputText += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${tempFilePath}]`; - } else if (truncation.truncatedBy === "lines") { - outputText += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${tempFilePath}]`; - } 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); - } + const formatOutput = (snapshot: Awaited>, emptyText = "(no output)") => { + const truncation = snapshot.truncation; + let text = snapshot.content || emptyText; + let details: BashToolDetails | undefined; + if (truncation.truncated) { + details = { truncation, fullOutputPath: snapshot.fullOutputPath }; + const startLine = truncation.totalLines - truncation.outputLines + 1; + const endLine = truncation.totalLines; + if (truncation.lastLinePartial) { + const lastLineSize = formatSize(output.getLastLineBytes()); + text += `\n\n[Showing last ${formatSize(truncation.outputBytes)} of line ${endLine} (line is ${lastLineSize}). Full output: ${snapshot.fullOutputPath}]`; + } else if (truncation.truncatedBy === "lines") { + text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines}. Full output: ${snapshot.fullOutputPath}]`; + } else { + text += `\n\n[Showing lines ${startLine}-${endLine} of ${truncation.totalLines} (${formatSize(DEFAULT_MAX_BYTES)} limit). Full output: ${snapshot.fullOutputPath}]`; + } + } + return { text, details }; + }; + + const appendStatus = (text: string, status: string) => `${text ? `${text}\n\n` : ""}${status}`; + + try { + let exitCode: number | null; + try { + const result = await ops.exec(spawnContext.command, spawnContext.cwd, { + onData: handleData, + signal, + timeout, + env: spawnContext.env, }); - }); + 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; diff --git a/packages/coding-agent/src/core/tools/output-accumulator.ts b/packages/coding-agent/src/core/tools/output-accumulator.ts new file mode 100644 index 00000000..30f4d663 --- /dev/null +++ b/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 { + if (!this.tempFileStream) { + return; + } + + const stream = this.tempFileStream; + this.tempFileStream = undefined; + + await new Promise((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 = []; + } +} diff --git a/packages/coding-agent/test/tools.test.ts b/packages/coding-agent/test/tools.test.ts index df0c8f62..70e4e54f 100644 --- a/packages/coding-agent/test/tools.test.ts +++ b/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[] = [];