diff --git a/src/shared/agent-status-osc-split-frame-scan-budget.test.ts b/src/shared/agent-status-osc-split-frame-scan-budget.test.ts new file mode 100644 index 00000000000..fda2e9b8f3e --- /dev/null +++ b/src/shared/agent-status-osc-split-frame-scan-budget.test.ts @@ -0,0 +1,75 @@ +import { describe, expect, it, vi } from 'vitest' +import { createAgentStatusOscProcessor } from './agent-status-osc' + +/** Total characters swept by terminator/prefix searches across every chunk of a feed. */ +function feedWithScanBudget(chunks: string[]) { + let searchedChars = 0 + const indexOf = String.prototype.indexOf + const spy = vi.spyOn(String.prototype, 'indexOf').mockImplementation(function ( + this: string, + search, + from = 0 + ) { + const found = indexOf.call(this, search, from) + searchedChars += (found === -1 ? this.length : found + String(search).length) - Number(from) + return found + }) + try { + const process = createAgentStatusOscProcessor() + const results = chunks.map((chunk) => process(chunk)) + return { results, searchedChars } + } finally { + spy.mockRestore() + } +} + +describe('OSC 9999 split-frame scan budget', () => { + it('keeps per-chunk work flat as the split frame accumulates', () => { + // One unterminated marker whose payload arrives one character at a time. + const feedOf = (chunkCount: number): string[] => [ + '\x1b]9999;{"state":"working","prompt":"', + ...Array(chunkCount).fill('x') + ] + + const small = feedWithScanBudget(feedOf(2000)) + const large = feedWithScanBudget(feedOf(4000)) + + expect(small.results.every((result) => result.payloads.length === 0)).toBe(true) + // Re-scanning the accumulation would quadruple the budget when the feed doubles. + expect(large.searchedChars).toBeLessThan(small.searchedChars * 3) + }) + + it.each(['\x07', '\x1b\\'])( + 'matches whole-string parsing when split at every offset with terminator %j', + (terminator) => { + const stream = `head\x1b]9999;{"state":"working","prompt":"p"}${terminator}tail` + const whole = createAgentStatusOscProcessor()(stream) + + for (let split = 1; split < stream.length; split += 1) { + const process = createAgentStatusOscProcessor() + const first = process(stream.slice(0, split)) + const second = process(stream.slice(split)) + expect({ + cleanData: first.cleanData + second.cleanData, + payloads: [...first.payloads, ...second.payloads] + }).toEqual({ cleanData: whole.cleanData, payloads: whole.payloads }) + } + } + ) + + it('finds a string terminator straddling the resume boundary', () => { + const process = createAgentStatusOscProcessor() + // The ESC lands as the last character of the carried frame; the backslash arrives next. + expect(process('\x1b]9999;{"state":"working"}\x1b').payloads).toEqual([]) + expect(process('\\rest').payloads).toMatchObject([{ state: 'working' }]) + }) + + it('still parses a payload that completes many chunks later', () => { + const process = createAgentStatusOscProcessor() + process('\x1b]9999;{"state":"wor') + for (const chunk of ['k', 'i', 'n', 'g']) { + expect(process(chunk).payloads).toEqual([]) + } + expect(process('"}\x07done').payloads).toMatchObject([{ state: 'working' }]) + }) +}) diff --git a/src/shared/agent-status-osc.ts b/src/shared/agent-status-osc.ts index adb0c0eee47..3b13d2de5bc 100644 --- a/src/shared/agent-status-osc.ts +++ b/src/shared/agent-status-osc.ts @@ -57,6 +57,9 @@ function findAgentStatusTerminator( export function createAgentStatusOscProcessor(): (data: string) => ProcessedAgentStatusChunk { const MAX_PENDING = 64 * 1024 let pending = '' + // How much of `pending` already failed a terminator search, so a frame split across + // many chunks re-scans only the new bytes instead of the whole accumulation. + let pendingSearched = 0 return (data: string): ProcessedAgentStatusChunk => { // Ordinary terminal output is by far the common case. Keep it on the @@ -76,7 +79,9 @@ export function createAgentStatusOscProcessor(): (data: string) => ProcessedAgen } const combined = pending + data + const resumeFrom = pendingSearched pending = '' + pendingSearched = 0 const payloads: ParsedAgentStatusPayload[] = [] let lastPayloadCleanOffset: number | null = null @@ -100,12 +105,16 @@ export function createAgentStatusOscProcessor(): (data: string) => ProcessedAgen cleanData += combined.slice(cursor, start) const payloadStart = start + OSC_AGENT_STATUS_PREFIX.length - const terminator = findAgentStatusTerminator(combined, payloadStart, nextTerminator) + // Minus one so a `\x1b\\` straddling the previous chunk boundary is still found. + const searchFrom = + start === 0 && resumeFrom > 0 ? Math.max(payloadStart, resumeFrom - 1) : payloadStart + const terminator = findAgentStatusTerminator(combined, searchFrom, nextTerminator) if (terminator === null) { const candidate = combined.slice(start) // Own the frame so it stops pinning the consumed chunk it was sliced from. pending = candidate.length > MAX_PENDING ? '' : ownRetainedString(candidate) + pendingSearched = pending.length break }