@oh-my-pi/pi-coding-agent 6.7.67 → 6.8.0

package/src/core/streaming-output.ts

@@ -1,154 +1,210 @@
 import { tmpdir } from "node:os";
 import { join } from "node:path";
+import { sanitizeText } from "@oh-my-pi/pi-utils";
 import { nanoid } from "nanoid";
-import stripAnsi from "strip-ansi";
-import { truncateTail } from "./tools/truncate";
+import { DEFAULT_MAX_BYTES, DEFAULT_MAX_COLUMN } from "./tools/truncate";
 
-/**
- * Sanitize binary output for display/storage.
- * Removes characters that crash string-width or cause display issues:
- * - Control characters (except tab, newline, carriage return)
- * - Lone surrogates
- * - Unicode Format characters (crash string-width due to a bug)
- * - Characters with undefined code points
- */
-export function sanitizeBinaryOutput(str: string): string {
-	// Use Array.from to properly iterate over code points (not code units)
-	// This handles surrogate pairs correctly and catches edge cases where
-	// codePointAt() might return undefined
-	return Array.from(str)
-		.filter((char) => {
-			// Filter out characters that cause string-width to crash
-			// This includes:
-			// - Unicode format characters
-			// - Lone surrogates (already filtered by Array.from)
-			// - Control chars except \t \n \r
-			// - Characters with undefined code points
-
-			const code = char.codePointAt(0);
-
-			// Skip if code point is undefined (edge case with invalid strings)
-			if (code === undefined) return false;
-
-			// Allow tab, newline, carriage return
-			if (code === 0x09 || code === 0x0a || code === 0x0d) return true;
-
-			// Filter out control characters (0x00-0x1F, except 0x09, 0x0a, 0x0x0d)
-			if (code <= 0x1f) return false;
-
-			// Filter out Unicode format characters
-			if (code >= 0xfff9 && code <= 0xfffb) return false;
-
-			return true;
-		})
-		.join("");
+export interface OutputResult {
+	output: string;
+	truncated: boolean;
+	fullOutputPath?: string;
 }
 
-/**
- * Sanitize text output: strip ANSI codes, remove binary garbage, normalize line endings.
- */
-export function sanitizeText(text: string): string {
-	return sanitizeBinaryOutput(stripAnsi(text)).replace(/\r/g, "");
+export interface OutputSinkOptions {
+	allocateFilePath?: () => string;
+	spillThreshold?: number;
+	maxColumn?: number;
+	onLine?: (line: string) => void;
+	onChunk?: (chunk: string) => void;
 }
 
-interface OutputFileSink {
-	write(data: string): number | Promise<number>;
-	end(): void;
+function defaultFilePathAllocator(): string {
+	return join(tmpdir(), `omp-${nanoid()}.log`);
 }
 
-export function createSanitizer(): TransformStream<Uint8Array, string> {
-	const decoder = new TextDecoder();
-	return new TransformStream({
-		transform(chunk, controller) {
-			const text = sanitizeText(decoder.decode(chunk, { stream: true }));
-			if (text) {
-				controller.enqueue(text);
-			}
-		},
-		flush(controller) {
-			const text = sanitizeText(decoder.decode());
-			if (text) {
-				controller.enqueue(text);
+/**
+ * Line-buffered output sink with file spill support.
+ *
+ * Uses a single string buffer with line position tracking.
+ * When memory limit exceeded, spills ~half to file in one batch operation.
+ */
+export class OutputSink {
+	private buffer = "";
+	private lineEnds: number[] = []; // String index after each \n
+
+	private fileSink?: Bun.FileSink;
+	private filePath?: string;
+
+	private readonly allocateFilePath: () => string;
+	private readonly spillThreshold: number;
+	private readonly maxColumn: number;
+	private readonly onLine?: (line: string) => void;
+	private readonly onChunk?: (chunk: string) => void;
+
+	constructor(options?: OutputSinkOptions) {
+		const {
+			allocateFilePath = defaultFilePathAllocator,
+			spillThreshold = DEFAULT_MAX_BYTES,
+			maxColumn = DEFAULT_MAX_COLUMN,
+			onLine,
+			onChunk,
+		} = options ?? {};
+
+		this.allocateFilePath = allocateFilePath;
+		this.spillThreshold = spillThreshold;
+		this.maxColumn = maxColumn;
+		this.onLine = onLine;
+		this.onChunk = onChunk;
+	}
+
+	private pushLine(line: string, term?: string): void {
+		while (line.length > this.maxColumn) {
+			this.pushLine(line.slice(0, this.maxColumn), "--\n");
+			line = line.slice(this.maxColumn);
+		}
+
+		this.buffer += line;
+		if (term) {
+			this.buffer += term;
+		}
+
+		this.lineEnds.push(this.buffer.length);
+		this.onLine?.(line);
+
+		if (this.buffer.length > this.spillThreshold) {
+			this.spillHalf();
+		}
+	}
+
+	private pushChunk(line: string): void {
+		this.onChunk?.(line);
+		this.pushLine(line);
+	}
+
+	private getFileSink(): Bun.FileSink {
+		if (!this.fileSink) {
+			const filePath = this.allocateFilePath();
+			this.filePath = filePath;
+			this.fileSink = Bun.file(filePath).writer();
+		}
+		return this.fileSink;
+	}
+
+	private spillHalf(): void {
+		const target = this.buffer.length >>> 1;
+
+		// Binary search: first line ending >= target
+		let lo = 0;
+		let hi = this.lineEnds.length;
+		while (lo < hi) {
+			const mid = (lo + hi) >>> 1;
+			if (this.lineEnds[mid] < target) {
+				lo = mid + 1;
+			} else {
+				hi = mid;
 			}
-		},
-	});
-}
+		}
+
+		// Clamp: evict at least 1 line, keep at least 1 line
+		const splitIdx = Math.max(1, Math.min(lo, this.lineEnds.length - 1));
+		const splitPos = this.lineEnds[splitIdx - 1];
 
-export async function pumpStream(readable: ReadableStream<Uint8Array>, writer: WritableStreamDefaultWriter<string>) {
-	const reader = readable.pipeThrough(createSanitizer()).getReader();
-	try {
-		while (true) {
-			const { done, value } = await reader.read();
-			if (done) break;
-			await writer.write(value);
+		// Write evicted portion to file
+		this.getFileSink().write(this.buffer.slice(0, splitPos));
+
+		// Truncate buffer, shift line positions
+		this.buffer = this.buffer.slice(splitPos);
+		const remaining = this.lineEnds.length - splitIdx;
+		for (let i = 0; i < remaining; i++) {
+			this.lineEnds[i] = this.lineEnds[i + splitIdx] - splitPos;
 		}
-	} finally {
-		reader.releaseLock();
+		this.lineEnds.length = remaining;
 	}
-}
 
-export interface OutputSinkDump {
-	output: string;
-	truncated: boolean;
-	fullOutputPath?: string;
-}
+	createWritable(): WritableStream<Uint8Array> {
+		const decoder = new TextDecoder("utf-8", { ignoreBOM: true });
+		let buf = "";
+
+		const flushLines = () => {
+			let start = 0;
+			while (true) {
+				const nl = buf.indexOf("\n", start);
+				if (nl === -1) break;
+				this.pushChunk(buf.slice(start, nl + 1));
+				start = nl + 1;
+			}
+			buf = buf.slice(start);
+		};
+
+		const finalize = () => {
+			buf += sanitizeText(decoder.decode());
+			flushLines();
+			buf = buf.trimEnd();
+			if (buf) {
+				this.pushChunk(`${buf}\n`);
+			}
+		};
 
-export class OutputSink {
-	private readonly stream: WritableStream<string>;
-	private readonly chunks: Array<{ text: string; bytes: number }> = [];
-	private chunkBytes = 0;
-	private totalBytes = 0;
-	private fullOutputPath: string | undefined;
-	private fullOutputStream: OutputFileSink | undefined;
-
-	constructor(
-		private readonly spillThreshold: number,
-		private readonly maxBuffer: number,
-		private readonly onChunk?: (text: string) => void,
-	) {
-		this.stream = new WritableStream<string>({
-			write: (text) => {
-				const bytes = Buffer.byteLength(text, "utf-8");
-				this.totalBytes += bytes;
-
-				if (this.totalBytes > this.spillThreshold && !this.fullOutputPath) {
-					this.fullOutputPath = join(tmpdir(), `omp-${nanoid()}.buffer`);
-					const stream = Bun.file(this.fullOutputPath).writer();
-					for (const chunk of this.chunks) {
-						stream.write(chunk.text);
-					}
-					this.fullOutputStream = stream;
-				}
-				this.fullOutputStream?.write(text);
-
-				this.chunks.push({ text, bytes });
-				this.chunkBytes += bytes;
-				while (this.chunkBytes > this.maxBuffer && this.chunks.length > 1) {
-					const removed = this.chunks.shift();
-					if (removed) {
-						this.chunkBytes -= removed.bytes;
-					}
-				}
-
-				this.onChunk?.(text);
+		return new WritableStream<Uint8Array>({
+			write: (chunk) => {
+				buf += sanitizeText(decoder.decode(chunk, { stream: true }));
+				flushLines();
 			},
-			close: () => {
-				this.fullOutputStream?.end();
+			close: finalize,
+			abort: finalize,
+		});
+	}
+
+	createStringWritable(): WritableStream<string> {
+		let buf = "";
+
+		const flushLines = () => {
+			let start = 0;
+			while (true) {
+				const nl = buf.indexOf("\n", start);
+				if (nl === -1) break;
+				this.pushChunk(buf.slice(start, nl + 1));
+				start = nl + 1;
+			}
+			buf = buf.slice(start);
+		};
+
+		const finalize = () => {
+			flushLines();
+			buf = buf.trimEnd();
+			if (buf) {
+				this.pushChunk(`${buf}\n`);
+			}
+		};
+
+		return new WritableStream<string>({
+			write: (chunk) => {
+				buf += sanitizeText(chunk);
+				flushLines();
 			},
+			close: finalize,
+			abort: finalize,
 		});
 	}
 
-	getWriter(): WritableStreamDefaultWriter<string> {
-		return this.stream.getWriter();
+	async close(): Promise<void> {
+		await this.fileSink?.end();
 	}
 
-	dump(annotation?: string): OutputSinkDump {
+	dump(annotation?: string): OutputResult {
+		let output = this.buffer;
 		if (annotation) {
-			const text = `\n\n${annotation}`;
-			this.chunks.push({ text, bytes: Buffer.byteLength(text, "utf-8") });
+			output += `\n${annotation}\n`;
+		}
+		if (!this.filePath) {
+			return { output, truncated: false };
 		}
-		const full = this.chunks.map((chunk) => chunk.text).join("");
-		const { content, truncated } = truncateTail(full);
-		return { output: truncated ? content : full, truncated, fullOutputPath: this.fullOutputPath };
+		this.fileSink!.write(this.buffer);
+		this.fileSink!.flush();
+		return {
+			output,
+			truncated: true,
+			fullOutputPath: this.filePath,
+		};
 	}
 }