From c3372aeadcccbf22cf7e7e63f98838ddef2879bf Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:30:40 -0700 Subject: [PATCH] fix(ai-vault): bound streamed remote JSONL records (#20700) --- .../ai-vault/remote-session-content-lines.ts | 4 ++- .../remote-session-large-transcripts.test.ts | 29 +++++++++++++++ .../remote-session-stream-lifecycle.test.ts | 20 +++++++++++ .../transcript-stream-lines.test.ts | 35 ++++++++++++++++++- .../native-chat/transcript-stream-lines.ts | 32 ++++++++++++++--- 5 files changed, 114 insertions(+), 6 deletions(-) diff --git a/src/main/ai-vault/remote-session-content-lines.ts b/src/main/ai-vault/remote-session-content-lines.ts index d62e8988018..16786704847 100644 --- a/src/main/ai-vault/remote-session-content-lines.ts +++ b/src/main/ai-vault/remote-session-content-lines.ts @@ -4,6 +4,8 @@ import { throwIfAiVaultScanCancelled } from './ai-vault-scan-cancellation' export type RemoteSessionContent = string | AsyncIterable +const MAX_REMOTE_SESSION_RECORD_BYTES = 10 * 1024 * 1024 + const REMOTE_CONTENT_YIELD_LINE_COUNT = 200 const REMOTE_CONTENT_YIELD_CHAR_COUNT = 256 * 1024 @@ -80,7 +82,7 @@ export async function* streamedSessionContentLines( ): AsyncGenerator { let count = 0 let chars = 0 - for await (const record of splitTranscriptStreamLines(bytes)) { + for await (const record of splitTranscriptStreamLines(bytes, MAX_REMOTE_SESSION_RECORD_BYTES)) { throwIfAiVaultScanCancelled(signal) const line = record.line.endsWith('\r') && (record.terminated || signal) diff --git a/src/main/ai-vault/remote-session-large-transcripts.test.ts b/src/main/ai-vault/remote-session-large-transcripts.test.ts index 62d8b8c5ed1..df05700c2cf 100644 --- a/src/main/ai-vault/remote-session-large-transcripts.test.ts +++ b/src/main/ai-vault/remote-session-large-transcripts.test.ts @@ -17,6 +17,35 @@ const jsonl = (rows: unknown[]) => `${rows.map((row) => JSON.stringify(row)).joi const filler = jsonl([{ type: 'irrelevant_event', payload: 'x'.repeat(1024) }]).repeat(11000) describe('large remote history through real relay filesystem', () => { + it('reports an oversized record without losing healthy sessions or publishing a partial session', async () => { + const home = await mkdtemp(join(tmpdir(), 'orca-history-record-limit-')) + try { + const directory = join(home, '.codex', 'sessions') + await mkdir(directory, { recursive: true }) + const metadata = (id: string) => + jsonl([{ type: 'session_meta', payload: { id, cwd: '/repo' } }]) + const badPath = join(directory, 'bad.jsonl') + await writeFile(badPath, metadata('bad') + 'x'.repeat(11 * 1024 * 1024)) + await writeFile(join(directory, 'good.jsonl'), metadata('good')) + const result = await scanRemoteAiVaultSessions({ + provider: createRelayAiVaultFilesystemProvider(), + executionHostId: 'ssh:record-limit', + remoteHome: home, + hostPlatform: platform, + unlimited: true + }) + expect(result.sessions.map((session) => session.sessionId)).toEqual(['good']) + expect(result.issues).toEqual([ + expect.objectContaining({ + path: badPath, + message: 'Session transcript record exceeds 10485760 byte limit' + }) + ]) + } finally { + await rm(home, { recursive: true, force: true }) + } + }) + it('lists a large Codex rollout with middle messages and usage intact', async () => { const home = await mkdtemp(join(tmpdir(), 'orca-history-17744-')) try { diff --git a/src/main/ai-vault/remote-session-stream-lifecycle.test.ts b/src/main/ai-vault/remote-session-stream-lifecycle.test.ts index 46668e3574e..2b8e95ea2a4 100644 --- a/src/main/ai-vault/remote-session-stream-lifecycle.test.ts +++ b/src/main/ai-vault/remote-session-stream-lifecycle.test.ts @@ -4,6 +4,26 @@ import { readStreamedSessionDocument } from './session-document-stream' import { limitRemoteScanFilesystemConcurrency } from './remote-session-scan-concurrency' describe('stream lifetime and retained document work', () => { + it('aborts a newline-free record at the byte ceiling and closes the source', async () => { + let closed = false + let reads = 0 + async function* bytes() { + const chunk = Buffer.alloc(1024 * 1024, 'x') + try { + for (; reads < 100;) { + reads++ + yield chunk + } + } finally { + closed = true + } + } + const lines = streamedSessionContentLines(bytes()) + await expect(lines.next()).rejects.toThrow('record exceeds 10485760 byte limit') + expect(reads).toBe(11) + expect(closed).toBe(true) + }) + it('releases the source when a line consumer finishes early', async () => { let closed = false async function* bytes() { diff --git a/src/main/native-chat/transcript-stream-lines.test.ts b/src/main/native-chat/transcript-stream-lines.test.ts index 48feedfbd98..25ac54d904a 100644 --- a/src/main/native-chat/transcript-stream-lines.test.ts +++ b/src/main/native-chat/transcript-stream-lines.test.ts @@ -1,6 +1,6 @@ import { Readable } from 'node:stream' import { describe, expect, it, vi } from 'vitest' -import { decodeTranscriptStream } from './transcript-stream-lines' +import { decodeTranscriptStream, splitTranscriptStreamLines } from './transcript-stream-lines' const decode = (line: string, id: string) => ({ id, @@ -171,3 +171,36 @@ describe('decodeTranscriptStream', () => { expect(result.consumedBytes).toBe(Buffer.byteLength(complete, 'utf8')) }) }) + +describe('bounded transcript records', () => { + async function collect(chunks: (Buffer | string)[], limit: number) { + const records: string[] = [] + for await (const record of splitTranscriptStreamLines(Readable.from(chunks), limit)) { + records.push(record.line) + } + return records + } + + it.each(['', '\n', '\nnext\n'])('rejects an oversized record ending in %j', async (ending) => { + await expect(collect(['1234', `5${ending}`], 4)).rejects.toThrow('record exceeds 4 byte limit') + await expect(collect([`12345${ending}`], 4)).rejects.toThrow('record exceeds 4 byte limit') + }) + + it('resets the byte budget per record and accepts the exact limit', async () => { + expect(await collect(['1234\n123', '4\n1234'], 4)).toEqual(['1234', '1234', '1234']) + }) + + it('counts UTF-8 bytes across split codepoints', async () => { + const bytes = Buffer.from('😀é') + const chunks = [bytes.subarray(0, 2), bytes.subarray(2, 5), bytes.subarray(5)] + expect((await collect(chunks, 6))[0]).toBe('😀é') + await expect(collect(chunks, 5)).rejects.toThrow('record exceeds 5 byte limit') + expect((await collect(['\ud83d', '\ude00'], 4))[0]).toBe('😀') + }) + + it('checks the decoder tail before emitting it', async () => { + await expect(collect([Buffer.from([0x61, 0xf0, 0x9f])], 3)).rejects.toThrow( + 'record exceeds 3 byte limit' + ) + }) +}) diff --git a/src/main/native-chat/transcript-stream-lines.ts b/src/main/native-chat/transcript-stream-lines.ts index 4a58b4e180f..ce22822b76c 100644 --- a/src/main/native-chat/transcript-stream-lines.ts +++ b/src/main/native-chat/transcript-stream-lines.ts @@ -42,12 +42,13 @@ export async function decodeTranscriptStream( type TranscriptLine = { line: string; byteLength: number; terminated: boolean } export async function* splitTranscriptStreamLines( - stream: AsyncIterable + stream: AsyncIterable, + maxRecordBytes = Infinity ): AsyncGenerator { let records: TranscriptLine[] = [] const framer = createTranscriptLineFramer((line, byteLength, terminated) => { records.push({ line, byteLength, terminated }) - }) + }, maxRecordBytes) for await (const chunk of stream) { framer.write(chunk) for (const record of records) { @@ -63,10 +64,12 @@ export async function* splitTranscriptStreamLines( /** Frame chunks synchronously so native decoding avoids a promise per record. */ function createTranscriptLineFramer( - emit: (line: string, byteLength: number, terminated: boolean) => void + emit: (line: string, byteLength: number, terminated: boolean) => void, + maxRecordBytes = Infinity ): { write(chunk: Buffer | string): void; end(): void } { const decoder = new StringDecoder('utf8') let pending: string[] = [] + let pendingBytes = 0 return { write, end } function write(chunk: Buffer | string): void { @@ -75,23 +78,44 @@ function createTranscriptLineFramer( let newlineIndex = text.indexOf('\n') while (newlineIndex !== -1) { let segment = text.slice(lineStart, newlineIndex + 1) + checkRecordBytes(segment.slice(0, -1)) if (pending.length > 0) { pending.push(segment) segment = pending.join('') pending = [] } + pendingBytes = 0 emit(segment.slice(0, -1), Buffer.byteLength(segment, 'utf8'), true) lineStart = newlineIndex + 1 newlineIndex = text.indexOf('\n', lineStart) } if (lineStart < text.length) { - pending.push(text.slice(lineStart)) + const segment = text.slice(lineStart) + checkRecordBytes(segment) + pending.push(segment) + } + } + + function checkRecordBytes(segment: string): void { + if (maxRecordBytes === Infinity) { + return + } + pendingBytes += Buffer.byteLength(segment, 'utf8') + const previous = pending.at(-1) + // Separately encoded surrogate halves become one four-byte codepoint when joined. + if (previous && /[\uD800-\uDBFF]$/.test(previous) && /^[\uDC00-\uDFFF]/.test(segment)) { + pendingBytes -= 2 + } + if (pendingBytes > maxRecordBytes) { + pending = [] + throw new Error(`Session transcript record exceeds ${maxRecordBytes} byte limit`) } } function end(): void { const tail = decoder.end() if (tail) { + checkRecordBytes(tail) pending.push(tail) } const line = pending.join('')