fix(coding-agent): stream bash output incrementally (#4165)
fixes #4145
This commit is contained in:
Armin Ronacher
@@ -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<typeof createWriteStream> | 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<ReturnType<typeof finishOutput>>, 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;
|
||||
|
||||
216
packages/coding-agent/src/core/tools/output-accumulator.ts
Normal file
216
packages/coding-agent/src/core/tools/output-accumulator.ts
Normal file
@@ -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 = [];
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user