diff --git a/src/main/ai-vault/session-scanner-parse-cache.test.ts b/src/main/ai-vault/session-scanner-parse-cache.test.ts index 6ba5e3266d6..8434f26bac1 100644 --- a/src/main/ai-vault/session-scanner-parse-cache.test.ts +++ b/src/main/ai-vault/session-scanner-parse-cache.test.ts @@ -128,6 +128,45 @@ describe('parseAgentSessionFileCached', () => { expect(incremental?.totalTokens).toBe(420) }) + it('parses an oversized record without quadratic carry copying', async () => { + const root = await makeTempDir() + const path = join(root, 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee.jsonl') + // One tool result far larger than a stream chunk. Re-joining the held-over + // partial line per chunk copies O(record^2); the piece list joins once. + const recordBytes = 4 * 1024 * 1024 + await writeFile( + path, + `${[ + userRecord(0, 'question'), + assistantRecord(1, 'x'.repeat(recordBytes)), + assistantRecord(2, 'tail answer') + ].join('\n')}\n` + ) + + const originalConcat = Buffer.concat + let concatenatedBytes = 0 + Buffer.concat = ((list: readonly Uint8Array[], totalLength?: number) => { + const joined = originalConcat(list as Uint8Array[], totalLength) + concatenatedBytes += joined.length + return joined + }) as typeof Buffer.concat + try { + const stats = createSessionParseStats() + const parsed = await parseAgentSessionFileCached( + await claudeCandidate(path), + process.platform, + stats + ) + expect(parsed).not.toBeNull() + } finally { + Buffer.concat = originalConcat + } + + // Linear joins the record about once; the quadratic form copied many times + // that, growing with the square of the record size. + expect(concatenatedBytes).toBeLessThan(recordBytes * 4) + }) + it('shows a trailing unterminated line without double-counting it later', async () => { const root = await makeTempDir() const path = join(root, 'aaaaaaaa-bbbb-4ccc-8ddd-eeeeeeeeeeee.jsonl') diff --git a/src/main/ai-vault/session-scanner-parse-cache.ts b/src/main/ai-vault/session-scanner-parse-cache.ts index b7aab770432..b3e4448a410 100644 --- a/src/main/ai-vault/session-scanner-parse-cache.ts +++ b/src/main/ai-vault/session-scanner-parse-cache.ts @@ -311,12 +311,28 @@ async function consumeCompleteJsonlLines(args: { }): Promise { let consumedThrough = args.start let bytesRead = 0 - let remainder: Buffer | null = null + // Why a piece list: re-joining the partial line with every chunk made one + // oversized record (a big tool result) cost O(record^2). Joining once, when a + // newline finally arrives, keeps it linear. + let remainderParts: Buffer[] = [] + let remainderLength = 0 const stream = createReadStream(args.path, { start: args.start }) for await (const chunk of stream as AsyncIterable) { bytesRead += chunk.length - const data = remainder ? Buffer.concat([remainder, chunk]) : chunk + // Why check the chunk alone: the pieces held over are all mid-line, so none + // of them contains a newline. + if (!chunk.includes(NEWLINE_BYTE)) { + remainderParts.push(chunk) + remainderLength += chunk.length + continue + } + const data = + remainderLength > 0 + ? Buffer.concat([...remainderParts, chunk], remainderLength + chunk.length) + : chunk + remainderParts = [] + remainderLength = 0 let lineStart = 0 let newlineIndex = data.indexOf(NEWLINE_BYTE, lineStart) while (newlineIndex !== -1) { @@ -329,13 +345,19 @@ async function consumeCompleteJsonlLines(args: { newlineIndex = data.indexOf(NEWLINE_BYTE, lineStart) } consumedThrough += lineStart - // Copy the tail so retaining it doesn't pin the whole chunk buffer. - remainder = lineStart < data.length ? Buffer.from(data.subarray(lineStart)) : null + if (lineStart < data.length) { + // Copy the tail so retaining it doesn't pin the whole chunk buffer. + remainderParts = [Buffer.from(data.subarray(lineStart))] + remainderLength = data.length - lineStart + } } + const trailingPartialLine = + remainderLength > 0 ? Buffer.concat(remainderParts, remainderLength).toString('utf-8') : null + return { consumedThrough, - trailingPartialLine: remainder && remainder.length > 0 ? remainder.toString('utf-8') : null, + trailingPartialLine, bytesRead } }