Merge PR #8040: fix(tui,utils): remove O(N²) hot paths when streaming long tool-call args (@AmecoCoding)

packages/ai/src/providers/devin.ts

@@ -125,9 +125,9 @@ export const streamDevin: StreamFunction<"devin-agent"> = (
 		const toolBlocks = new Map<string, ToolCall>();
 		const toolPartialJson = new Map<string, string>();
 		// Last-parsed argument-buffer length per tool-call id — bounds the
-		// mid-stream parse work to O(N) via `parseStreamingJsonThrottled`; the
-		// authoritative final parse still runs unconditionally in the toolcall_end
-		// loop below.
+		// mid-stream parse work to O(N log N) via `parseStreamingJsonThrottled`;
+		// the authoritative final parse still runs unconditionally in the
+		// toolcall_end loop below.
 		const toolLastParseLen = new Map<string, number>();
 		let activeToolCallId: string | undefined;
 		let latestStopReason = StopReason.UNSPECIFIED;

packages/coding-agent/src/tools/write.ts

@@ -1401,6 +1401,68 @@ function normalizeDisplayText(text: unknown): string {
  */
 const WRITE_GUTTER_MIN_WIDTH = 3;
 
+/**
+ * Per-component streaming line index for {@link formatStreamingContent}.
+ * Keyed on the ToolExecutionComponent's persistent render-state object (the
+ * `options` argument renderers receive on every rebuild), so the entry lives
+ * exactly as long as the component and never leaks across tool calls.
+ *
+ * Why: streamed write content is append-only, but the formatter used to
+ * normalize + `split("\n")` the ENTIRE accumulated payload on every reveal
+ * tick — O(n) per tick, O(n²) per stream, which was a measurable main-thread
+ * stall on long writes (and multiplied across concurrent subagent writes).
+ * Tracking the newline count incrementally and extracting only the tail
+ * window makes each tick O(delta + preview lines).
+ */
+interface WriteStreamingLineIndex {
+	/** Full content string as of the last scan. */
+	content: string;
+	/** `1 + count("\n")` over {@link content}. */
+	lineCount: number;
+}
+
+const writeStreamingLineIndex = new WeakMap<object, WriteStreamingLineIndex>();
+
+/** Total logical line count of `content`, resuming from the cached prefix scan when append-only. */
+function streamingTotalLines(streamKey: object | undefined, content: string): number {
+	if (streamKey === undefined) {
+		let lines = 1;
+		for (let i = 0; i < content.length; i++) if (content.charCodeAt(i) === 10) lines++;
+		return lines;
+	}
+	let entry = writeStreamingLineIndex.get(streamKey);
+	if (entry !== undefined && content.startsWith(entry.content)) {
+		let lines = entry.lineCount;
+		for (let i = entry.content.length; i < content.length; i++) if (content.charCodeAt(i) === 10) lines++;
+		entry.content = content;
+		entry.lineCount = lines;
+		return lines;
+	}
+	let lines = 1;
+	for (let i = 0; i < content.length; i++) if (content.charCodeAt(i) === 10) lines++;
+	entry = { content, lineCount: lines };
+	writeStreamingLineIndex.set(streamKey, entry);
+	return lines;
+}
+
+/**
+ * Raw offset just after the (totalLines - previewLines)-th newline — i.e. the
+ * start of the last `previewLines` logical lines — scanning back from the end.
+ * Returns 0 when the whole content fits in the window. Equivalent to
+ * `content.split("\n").slice(-previewLines).join("\n")` without materializing
+ * the full line array.
+ */
+function tailWindowStart(content: string, previewLines: number): number {
+	let newlinesSeen = 0;
+	for (let i = content.length - 1; i >= 0; i--) {
+		if (content.charCodeAt(i) === 10) {
+			newlinesSeen++;
+			if (newlinesSeen === previewLines) return i + 1;
+		}
+	}
+	return 0;
+}
+
 function formatStreamingContent(
 	content: string,
 	expanded: boolean,
@@ -1408,19 +1470,31 @@ function formatStreamingContent(
 	uiTheme: Theme,
 	spinnerFrame?: number,
 	cache?: RenderedStringCache,
+	streamKey?: object,
 ): string {
 	if (!content) return "";
 	const bodyText = cachedRenderedString(cache, uiTheme, expanded, language ?? "", content, () => {
-		const lines = normalizeDisplayText(content).split("\n");
-		const totalLines = lines.length;
 		// Collapsed: follow the streaming edge with a bounded tail window so the box
 		// stays short enough not to strand its scrolled-off head above the viewport
 		// while the block is volatile. `Ctrl+O` (expanded) lifts the cap for a
 		// deliberate full view — matching the eval streaming preview.
-		const startIndex = expanded ? 0 : Math.max(0, totalLines - WRITE_STREAMING_PREVIEW_LINES);
-		const visibleLines = lines.slice(startIndex);
+		let totalLines: number;
+		let startIndex: number;
+		let visibleText: string;
+		if (expanded) {
+			visibleText = normalizeDisplayText(content);
+			totalLines = 1;
+			for (let i = 0; i < visibleText.length; i++) if (visibleText.charCodeAt(i) === 10) totalLines++;
+			startIndex = 0;
+		} else {
+			totalLines = streamingTotalLines(streamKey, content);
+			startIndex = Math.max(0, totalLines - WRITE_STREAMING_PREVIEW_LINES);
+			const tail =
+				startIndex === 0 ? content : content.slice(tailWindowStart(content, WRITE_STREAMING_PREVIEW_LINES));
+			visibleText = tail.replace(/\r/g, "");
+		}
 		const hidden = startIndex;
-		const highlighted = highlightCode(visibleLines.join("\n"), language);
+		const highlighted = highlightCode(visibleText, language);
 		const lineNumberWidth = Math.max(WRITE_GUTTER_MIN_WIDTH, String(totalLines).length);
 
 		let text = "\n\n";
@@ -1518,7 +1592,12 @@ export const writeToolRenderer = {
 			},
 			uiTheme,
 		);
-		const content = normalizeDisplayText(args.content);
+		// Raw content, not normalizeDisplayText(args.content): the collapsed
+		// streaming path normalizes only its tail window, so a full-payload
+		// normalize on every reveal tick would re-introduce the O(n²) streaming
+		// cost formatStreamingContent avoids. Non-string content still falls
+		// back to the normalizing stringify.
+		const content = typeof args.content === "string" ? args.content : normalizeDisplayText(args.content);
 		const streamingCache = createRenderedStringCache();
 		return framedBlock(uiTheme, width => {
 			const body = content
@@ -1529,6 +1608,10 @@ export const writeToolRenderer = {
 						uiTheme,
 						options?.spinnerFrame,
 						streamingCache,
+						// `options` is the ToolExecutionComponent's persistent
+						// render-state object — a stable identity across reveal ticks
+						// that keys the incremental line index.
+						options,
 					)
 				: "";
 			const bodyLines = body ? body.split("\n") : [];

packages/coding-agent/test/write-streaming-incremental.test.ts

@@ -0,0 +1,138 @@
+import { describe, expect, it } from "bun:test";
+import * as themeModule from "@oh-my-pi/pi-coding-agent/modes/theme/theme";
+import { writeToolRenderer } from "@oh-my-pi/pi-coding-agent/tools/write";
+
+const stripAnsi = (s: string): string => s.replace(/\[[0-9;]*m/g, "");
+const hasLine = (lines: readonly string[], n: number): boolean =>
+	new RegExp(`\\bline ${n}\\b`).test(stripAnsi(lines.join("\n")));
+
+/**
+ * Reference algorithm: the pre-incremental formatter normalized the whole
+ * payload, split every line, and sliced the tail window. The incremental
+ * collapsed path must produce byte-identical rows for the same content.
+ */
+function referenceWindow(content: string): { total: number; start: number; visible: string[] } {
+	const lines = content.replace(/\r/g, "").split("\n");
+	const total = lines.length;
+	const start = Math.max(0, total - 12);
+	return { total, start, visible: lines.slice(start) };
+}
+
+describe("write streaming preview incremental line tracking", () => {
+	let initialized = false;
+
+	async function getUiTheme() {
+		if (!initialized) {
+			await themeModule.initTheme();
+			initialized = true;
+		}
+		const uiTheme = (await themeModule.getThemeByName("dark")) ?? (await themeModule.getThemeByName("light"));
+		if (!uiTheme) throw new Error("expected an initialized theme");
+		return uiTheme;
+	}
+
+	function renderCollapsed(content: string, options: { expanded: boolean; isPartial: boolean; spinnerFrame: number }) {
+		return getUiTheme().then(uiTheme => {
+			const component = writeToolRenderer.renderCall({ path: "/tmp/inc.ts", content }, options, uiTheme);
+			if (!component) throw new Error("expected a rendered component for a non-xdev write path");
+			return component.render(120);
+		});
+	}
+
+	it("tracks an append-only stream through one shared render-state object", async () => {
+		// The reveal loop rebuilds via renderCall once per tick with the SAME
+		// persistent options object; simulate growth 5 → 12 → 13 → 25 → 40 lines.
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		const allLines = Array.from({ length: 40 }, (_, i) => `line ${i + 1}`);
+
+		for (const count of [5, 12, 13, 25, 40]) {
+			const content = allLines.slice(0, count).join("\n");
+			const rendered = await renderCollapsed(content, options);
+			const { total, start } = referenceWindow(content);
+			expect(total).toBe(count);
+			// Window shows exactly lines start+1..total with correct numbering.
+			expect(hasLine(rendered, total)).toBe(true);
+			if (start > 0) {
+				expect(hasLine(rendered, start)).toBe(false);
+				expect(hasLine(rendered, start + 1)).toBe(true);
+				expect(stripAnsi(rendered.join("\n"))).toContain(`${start} earlier line`);
+			} else {
+				expect(hasLine(rendered, 1)).toBe(true);
+				expect(stripAnsi(rendered.join("\n"))).not.toContain("earlier line");
+			}
+		}
+	});
+
+	it("matches the split-based reference window across a size battery", async () => {
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		for (const count of [1, 2, 3, 11, 12, 13, 40, 41]) {
+			// Fresh options per size: each tool call gets its own render state.
+			const content = Array.from({ length: count }, (_, i) => `line ${i + 1}`).join("\n");
+			const rendered = stripAnsi((await renderCollapsed(content, options)).join("\n"));
+			const { total, start, visible } = referenceWindow(content);
+			expect(total).toBe(count);
+			for (let i = 0; i < visible.length; i++) {
+				const lineNum = start + i + 1;
+				expect(rendered).toContain(`${lineNum}`);
+				expect(rendered).toContain(visible[i]!);
+			}
+			if (start > 0) expect(rendered).toContain(`… (${start} earlier line${start === 1 ? "" : "s"})`);
+		}
+	});
+
+	it("normalizes CRLF only in the rendered tail, with correct line numbers", async () => {
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		const content = Array.from({ length: 20 }, (_, i) => `line ${i + 1}`).join("\r\n");
+		const rendered = await renderCollapsed(content, options);
+		const text = stripAnsi(rendered.join("\n"));
+		expect(text).not.toContain("\r");
+		// 20 lines → window is lines 9..20.
+		expect(text).toContain("… (8 earlier lines)");
+		expect(hasLine(rendered, 8)).toBe(false);
+		expect(hasLine(rendered, 9)).toBe(true);
+		expect(hasLine(rendered, 20)).toBe(true);
+	});
+
+	it("counts a trailing newline as a final empty row, matching the reference", async () => {
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		const content = `${Array.from({ length: 13 }, (_, i) => `line ${i + 1}`).join("\n")}\n`;
+		const rendered = await renderCollapsed(content, options);
+		const { total, start } = referenceWindow(content);
+		expect(total).toBe(14);
+		expect(start).toBe(2);
+		const text = stripAnsi(rendered.join("\n"));
+		expect(text).toContain("… (2 earlier lines)");
+		expect(hasLine(rendered, 13)).toBe(true);
+		expect(hasLine(rendered, 2)).toBe(false);
+	});
+
+	it("resets cleanly when streamed content is not append-only", async () => {
+		// A restarted stream reuses the component's render state with a buffer
+		// that no longer extends the previous one; the index must not corrupt
+		// the count or the window.
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		const first = Array.from({ length: 20 }, (_, i) => `alpha ${i + 1}`).join("\n");
+		await renderCollapsed(first, options);
+
+		const restarted = "beta 1\nbeta 2";
+		const rendered = await renderCollapsed(restarted, options);
+		const text = stripAnsi(rendered.join("\n"));
+		expect(text).not.toContain("earlier line");
+		expect(text).toContain("beta 1");
+		expect(text).toContain("beta 2");
+		expect(text).not.toContain("alpha");
+	});
+
+	it("resumes append tracking across a CR boundary without miscounting", async () => {
+		const options = { expanded: false, isPartial: true, spinnerFrame: 0 };
+		const part1 = "line 1\r\nline 2\r";
+		const part2 = "line 1\r\nline 2\r\nline 3\r\nline 4";
+		await renderCollapsed(part1, options);
+		const rendered = await renderCollapsed(part2, options);
+		const { total } = referenceWindow(part2);
+		expect(total).toBe(4);
+		expect(hasLine(rendered, 4)).toBe(true);
+		expect(hasLine(rendered, 1)).toBe(true);
+		expect(stripAnsi(rendered.join("\n"))).not.toContain("earlier line");
+	});
+});

packages/utils/src/json-parse.ts

@@ -579,8 +579,8 @@ export function parseStreamingJson<T = Record<string, unknown>>(partialJson: str
 
 /**
  * Default minimum byte growth before `parseStreamingJsonThrottled` will
- * re-parse a streaming tool-call argument buffer. Bounds the mid-stream
- * partial-parse cost from quadratic to linear in N.
+ * re-parse a streaming tool-call argument buffer. Acts as the floor of the
+ * geometric gate — see {@link parseStreamingJsonThrottled}.
  */
 export const STREAMING_JSON_PARSE_MIN_GROWTH = 256;
 
@@ -589,14 +589,23 @@ export const STREAMING_JSON_PARSE_MIN_GROWTH = 256;
  *
  * Tool calls arrive as a long sequence of small deltas — calling
  * `parseStreamingJson(buffer)` on every delta re-parses the entire buffer
- * each time, giving O(N²) work in the total buffer length. Throttling skips
- * the re-parse until at least `minGrowthBytes` of new content has arrived
- * since the last successful parse, bounding mid-stream cost to O(N).
+ * each time, giving O(N²) work in the total buffer length. A fixed re-parse
+ * floor alone does NOT fix this: with `minGrowthBytes` constant, a buffer of
+ * length N is parsed N/minGrowthBytes times at an average cost of N/2, which
+ * is still O(N²) (the constant just shrinks). Long `write` payloads — where
+ * the buffer is the whole file — made this the dominant main-thread stall
+ * during streaming.
+ *
+ * Instead the gate scales geometrically: once the buffer is large, a re-parse
+ * requires growth proportional to the current length (`len / 32`, floored at
+ * `minGrowthBytes`). Parse points then form a geometric progression, so a
+ * buffer of length N is parsed O(log N) times for O(N log N) total work,
+ * while small buffers keep the snappy fixed-cadence updates.
  *
  * Each provider tracks the last parsed length on its tool-call block, so the
  * final `toolcall_end` parse (which providers already perform unconditionally)
  * is the authoritative full parse — the throttle only delays mid-stream UI
- * updates by at most `minGrowthBytes` of accumulated partial content.
+ * updates, by at most ~3% of the accumulated content for large buffers.
  *
  * @returns the parsed object plus the new `parsedLen` to persist; or `null`
  *          when the buffer has not grown enough to warrant a re-parse.
@@ -607,7 +616,9 @@ export function parseStreamingJsonThrottled<T = Record<string, unknown>>(
 	minGrowthBytes: number = STREAMING_JSON_PARSE_MIN_GROWTH,
 ): { value: T; parsedLen: number } | null {
 	const len = partialJson?.length ?? 0;
-	if (len === 0 || (lastParsedLen > 0 && len - lastParsedLen < minGrowthBytes)) return null;
+	if (len === 0) return null;
+	const growth = Math.max(minGrowthBytes, len >> 5);
+	if (lastParsedLen > 0 && len - lastParsedLen < growth) return null;
 	return { value: parseStreamingJson<T>(partialJson), parsedLen: len };
 }
 

packages/utils/test/parse-streaming-json-throttled.test.ts

@@ -68,4 +68,63 @@ describe("parseStreamingJsonThrottled (F5)", () => {
 		expect(parseStreamingJsonThrottled(undefined, 0, 256)).toBeNull();
 		expect(parseStreamingJsonThrottled("", 0, 256)).toBeNull();
 	});
+
+	it("geometric gate: large buffers re-parse O(log N) times, not O(N/minGrowth)", () => {
+		// 512KB of args delivered as 1KB deltas — the long-`write`-payload case.
+		// A fixed 256-byte gate would re-parse ~2048 times; the geometric gate
+		// (len/32 above the floor) must land in the low hundreds at most.
+		const payload = `{"q":"${"x".repeat(512 * 1024)}"}`;
+		let lastParsedLen = 0;
+		let parseCalls = 0;
+
+		for (let i = 1; i <= payload.length; i += 1024) {
+			const slice = payload.slice(0, i);
+			const throttled = parseStreamingJsonThrottled<Record<string, unknown>>(slice, lastParsedLen);
+			if (throttled) {
+				parseCalls++;
+				lastParsedLen = throttled.parsedLen;
+			}
+		}
+
+		expect(parseCalls).toBeGreaterThan(0);
+		expect(parseCalls).toBeLessThan(200);
+		// Fixed-cadence equivalent for scale: payload/256 ≈ 2049 parses.
+		expect(parseCalls).toBeLessThan(payload.length / STREAMING_JSON_PARSE_MIN_GROWTH / 8);
+	});
+
+	it("geometric gate keeps mid-stream snapshots fresh within ~1/32 of the buffer", () => {
+		// After any settled point, the unparsed tail is bounded by len/32, so UI
+		// built on parsed args lags the raw stream by ~3%, never by kilobytes.
+		const payload = `{"q":"${"x".repeat(256 * 1024)}"}`;
+		let lastParsedLen = 0;
+		let maxLagRatio = 0;
+
+		for (let i = 1; i <= payload.length; i += 1024) {
+			const throttled = parseStreamingJsonThrottled<Record<string, unknown>>(payload.slice(0, i), lastParsedLen);
+			if (throttled) lastParsedLen = throttled.parsedLen;
+			if (lastParsedLen > 0) maxLagRatio = Math.max(maxLagRatio, (i - lastParsedLen) / i);
+		}
+
+		expect(maxLagRatio).toBeLessThan(1 / 16);
+	});
+
+	it("geometric gate preserves fixed-cadence behavior for small buffers", () => {
+		// Below len/32 == minGrowthBytes the floor dominates, so a 5KB stream
+		// re-parses at the same ~256-byte cadence as before the change.
+		const payload = `{"q":"${"x".repeat(5000)}"}`;
+		let lastParsedLen = 0;
+		let parseCalls = 0;
+
+		for (let i = 1; i <= payload.length; i++) {
+			const throttled = parseStreamingJsonThrottled<Record<string, unknown>>(payload.slice(0, i), lastParsedLen);
+			if (throttled) {
+				parseCalls++;
+				lastParsedLen = throttled.parsedLen;
+			}
+		}
+
+		// 5108/256 ≈ 20 — identical bound to the pre-geometric contract above.
+		expect(parseCalls).toBeLessThanOrEqual(25);
+		expect(parseCalls).toBeGreaterThan(15);
+	});
 });