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

// Adapted from Pi c20cb09772bf4e2590a316cb54514cef76df4293.
// Copyright (c) 2025 Mario Zechner. MIT license: licenses/third-party-notices.txt.
import { randomBytes } from "node:crypto";
import { createWriteStream, type WriteStream } from "node:fs";
import { open, readdir, rm, stat } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, type TruncationResult, truncateTail } from "./truncate.ts";

export interface OutputAccumulatorOptions {
  maxLines?: number;
  maxBytes?: number;
  tempFilePrefix?: string;
}

export interface OutputSnapshot {
  content: string;
  truncation: TruncationResult;
  fullOutputPath?: string;
}

export interface FullOutput {
  content: string;
  /** Whether `content` omits part of the output. */
  truncated: boolean;
}

function defaultTempFilePath(prefix: string): string {
  const id = randomBytes(8).toString("hex");
  return join(tmpdir(), `${prefix}-${id}.log`);
}

/** Keep recent output available to read, and prune expired logs on the next bash invocation. */
export async function cleanupOutputFiles(prefix: string, directory = tmpdir()): Promise<void> {
  const cutoff = Date.now() - 7 * 24 * 60 * 60 * 1000;
  for (const file of await readdir(directory, { withFileTypes: true })) {
    if (!file.isFile() || !file.name.startsWith(`${prefix}-`)) continue;
    if (!/^[a-f0-9]{16}\.log$/.test(file.name.slice(prefix.length + 1))) continue;
    const path = join(directory, file.name);
    try {
      if ((await stat(path)).mtimeMs < cutoff) await rm(path, { force: true });
    } catch {
      // Optional cleanup must not prevent a command from running.
    }
  }
}

function byteLength(text: string): number {
  return Buffer.byteLength(text, "utf-8");
}

/**
 * Incrementally tracks streaming output with bounded memory when producers await `append`.
 *
 * 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 completedLines = 0;
  private totalLines = 0;
  private currentLineBytes = 0;
  private hasOpenLine = false;
  private finished = false;

  private tempFilePath: string | undefined;
  private tempFileStream: WriteStream | undefined;
  private tempFileError: Error | undefined;
  private tempFileDrain: Promise<void> | 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";
  }

  /** Await the returned promise, when present, before appending another chunk. */
  append(data: Buffer): void | Promise<void> {
    if (this.finished) {
      throw new Error("Cannot append to a finished output accumulator");
    }

    if (this.tempFileError) throw this.tempFileError;

    this.totalRawBytes += data.length;
    this.appendDecodedText(this.decoder.decode(data, { stream: true }));

    if (this.tempFileStream || this.shouldUseTempFile()) {
      this.ensureTempFile();
      const stream = this.tempFileStream;
      if (stream && !stream.write(data)) return this.waitForTempFileDrain(stream);
    } 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;
    if (!stream.closed)
      await new Promise<void>((resolve) => {
        stream.once("close", resolve);
        stream.end();
      });
    this.tempFileStream = undefined;
    if (this.tempFileError) {
      if (this.tempFilePath) await rm(this.tempFilePath, { force: true }).catch(() => {});
      throw this.tempFileError;
    }
  }

  /**
   * The complete output, for callers that can take more than the display snapshot. Call after
   * `finish()` and `closeTempFile()`. Output longer than `maxBytes` raw bytes keeps its first and
   * last `maxBytes / 2` bytes around an omission marker.
   */
  async readFullOutput(maxBytes: number): Promise<FullOutput> {
    if (!this.tempFilePath) {
      return { content: new TextDecoder().decode(Buffer.concat(this.rawChunks)), truncated: false };
    }
    const file = await open(this.tempFilePath, "r");
    try {
      const size = (await file.stat()).size;
      if (size <= maxBytes) {
        return { content: new TextDecoder().decode(await file.readFile()), truncated: false };
      }
      const headBytes = Math.floor(maxBytes / 2);
      const tailBytes = maxBytes - headBytes;
      const head = Buffer.alloc(headBytes);
      const tail = Buffer.alloc(tailBytes);
      await file.read(head, 0, headBytes, 0);
      await file.read(tail, 0, tailBytes, size - tailBytes);
      // Cut at character boundaries: streaming decode holds back an incomplete trailing sequence,
      // and the tail skips leading continuation bytes.
      const headText = new TextDecoder().decode(head, { stream: true });
      let tailStart = 0;
      while (tailStart < tail.length && (tail[tailStart] & 0xc0) === 0x80) tailStart++;
      const tailText = new TextDecoder().decode(tail.subarray(tailStart));
      const omitted = size - headBytes - tailBytes;
      return { content: `${headText}\n\n[... ${omitted} bytes omitted ...]\n\n${tailText}`, truncated: true };
    } finally {
      await file.close();
    }
  }

  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;
      this.hasOpenLine = true;
    } else {
      this.completedLines += newlines;
      const tail = text.slice(lastNewline + 1);
      this.currentLineBytes = byteLength(tail);
      this.hasOpenLine = tail.length > 0;
    }
    this.totalLines = this.completedLines + (this.hasOpenLine ? 1 : 0);
  }

  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 waitForTempFileDrain(stream: WriteStream): Promise<void> {
    if (this.tempFileDrain) return this.tempFileDrain;
    const drain = new Promise<void>((resolve, reject) => {
      const finish = () => {
        stream.removeListener("drain", finish);
        stream.removeListener("error", finish);
        stream.removeListener("close", finish);
        this.tempFileDrain = undefined;
        if (this.tempFileError) reject(this.tempFileError);
        else if (stream.destroyed && !stream.writableFinished)
          reject(new Error("Output file closed before draining"));
        else resolve();
      };
      stream.once("drain", finish);
      stream.once("error", finish);
      stream.once("close", finish);
    });
    // Legacy synchronous producers can ignore the promise. closeTempFile still reports disk errors.
    void drain.catch(() => {});
    this.tempFileDrain = drain;
    return drain;
  }

  private ensureTempFile(): void {
    if (this.tempFilePath) {
      return;
    }
    this.tempFilePath = defaultTempFilePath(this.tempFilePrefix);
    this.tempFileStream = createWriteStream(this.tempFilePath, { mode: 0o600 });
    this.tempFileStream.on("error", (error) => {
      this.tempFileError = error;
    });
    for (const chunk of this.rawChunks) {
      this.tempFileStream.write(chunk);
    }
    this.rawChunks = [];
  }
}