diff --git a/packages/coding-agent/src/modes/session-worker/private-framing.ts b/packages/coding-agent/src/modes/session-worker/private-framing.ts index 4caf143f6..b4fe53b5b 100644 --- a/packages/coding-agent/src/modes/session-worker/private-framing.ts +++ b/packages/coding-agent/src/modes/session-worker/private-framing.ts @@ -51,7 +51,12 @@ export function encodePrivateFrame( } export class PrivateFrameDecoder { - private buffered: Buffer = Buffer.alloc(0); + // Incoming chunks are kept as segments and only joined when a whole frame has + // arrived. Concatenating into one accumulator on every chunk is O(n^2) in the + // frame size, which a multi-megabyte response split across socket reads feels + // as seconds: the session tree of a long transcript, for example. + private segments: Buffer[] = []; + private pendingBytes = 0; constructor( private readonly validateHeader: PrivateFrameHeaderValidator, @@ -59,20 +64,21 @@ export class PrivateFrameDecoder { ) {} get bufferedBytes(): number { - return this.buffered.length; + return this.pendingBytes; } push(chunk: Uint8Array): PrivateFrame[] { if (chunk.length > 0) { const buffer = Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk); - this.buffered = this.buffered.length === 0 ? buffer : Buffer.concat([this.buffered, buffer]); + this.segments.push(buffer); + this.pendingBytes += buffer.length; } const frames: PrivateFrame[] = []; - let offset = 0; - while (this.buffered.length - offset >= FRAME_PREFIX_BYTES) { - const headerLength = this.buffered.readUInt32BE(offset); - const payloadLength = this.buffered.readUInt32BE(offset + 4); + while (this.pendingBytes >= FRAME_PREFIX_BYTES) { + const prefix = this.peek(FRAME_PREFIX_BYTES); + const headerLength = prefix.readUInt32BE(0); + const payloadLength = prefix.readUInt32BE(4); assertFrameLength("header length", headerLength, this.limits.maxHeaderBytes); assertFrameLength("payload length", payloadLength, this.limits.maxPayloadBytes); if (headerLength === 0) { @@ -80,15 +86,18 @@ export class PrivateFrameDecoder { } const frameLength = FRAME_PREFIX_BYTES + headerLength + payloadLength; - if (this.buffered.length - offset < frameLength) { + if (this.pendingBytes < frameLength) { break; } - const headerStart = offset + FRAME_PREFIX_BYTES; + // Join the frame's segments exactly once (a whole frame that arrived in + // one chunk is taken as-is). + const frame = this.take(frameLength); + const headerStart = FRAME_PREFIX_BYTES; const payloadStart = headerStart + headerLength; let decoded: unknown; try { - decoded = JSON.parse(this.buffered.toString("utf8", headerStart, payloadStart)); + decoded = JSON.parse(frame.toString("utf8", headerStart, payloadStart)); } catch (error) { throw new Error( `Invalid private frame header JSON: ${error instanceof Error ? error.message : String(error)}`, @@ -100,20 +109,55 @@ export class PrivateFrameDecoder { frames.push({ header: decoded, - payload: Buffer.from(this.buffered.subarray(payloadStart, payloadStart + payloadLength)), + payload: Buffer.from(frame.subarray(payloadStart, payloadStart + payloadLength)), }); - offset += frameLength; } + return frames; + } - if (offset > 0) { - this.buffered = Buffer.from(this.buffered.subarray(offset)); + /** Copy of the next `count` buffered bytes, without consuming them. */ + private peek(count: number): Buffer { + const first = this.segments[0]; + if (first !== undefined && first.length >= count) { + return first.subarray(0, count); } - return frames; + const head: Buffer[] = []; + let remaining = count; + for (const segment of this.segments) { + const take = Math.min(segment.length, remaining); + head.push(take === segment.length ? segment : segment.subarray(0, take)); + remaining -= take; + if (remaining === 0) break; + } + return Buffer.concat(head, count); + } + + /** Remove and return the next `count` buffered bytes as one contiguous buffer. */ + private take(count: number): Buffer { + const head: Buffer[] = []; + let remaining = count; + let consumed = 0; + while (remaining > 0) { + const segment = this.segments[consumed]!; + if (segment.length <= remaining) { + head.push(segment); + remaining -= segment.length; + consumed++; + } else { + head.push(segment.subarray(0, remaining)); + // Keep the unconsumed tail as a copy so the caller's chunk is released. + this.segments[consumed] = Buffer.from(segment.subarray(remaining)); + remaining = 0; + } + } + this.segments = this.segments.slice(consumed); + this.pendingBytes -= count; + return head.length === 1 ? head[0]! : Buffer.concat(head, count); } finish(): void { - if (this.buffered.length !== 0) { - throw new Error(`Private frame channel ended with ${this.buffered.length} incomplete bytes`); + if (this.pendingBytes !== 0) { + throw new Error(`Private frame channel ended with ${this.pendingBytes} incomplete bytes`); } } } diff --git a/packages/coding-agent/test/session-worker-private-framing.test.ts b/packages/coding-agent/test/session-worker-private-framing.test.ts index ebb6116f3..40cce94ea 100644 --- a/packages/coding-agent/test/session-worker-private-framing.test.ts +++ b/packages/coding-agent/test/session-worker-private-framing.test.ts @@ -60,6 +60,45 @@ describe("private worker framing", () => { ); }); + it("decodes a large frame split across many small chunks without quadratic buffering", () => { + // A response can be many megabytes (a long session's tree, a transcript + // chunk). Rebuilding the accumulator per chunk made that O(n^2); this stays + // linear, so the bound is generous by orders of magnitude. + const payload = Buffer.alloc(8 * 1024 * 1024, 7); + const frame = encodePrivateFrame({ type: "response", requestId: "big" }, payload); + const decoder = new PrivateFrameDecoder(isTestHeader); + const frames = []; + + const started = Date.now(); + const chunkBytes = 1024; + for (let offset = 0; offset < frame.length; offset += chunkBytes) { + frames.push(...decoder.push(frame.subarray(offset, offset + chunkBytes))); + } + decoder.finish(); + const elapsedMs = Date.now() - started; + + expect(frames).toHaveLength(1); + expect(frames[0]!.header).toEqual({ type: "response", requestId: "big" }); + expect(frames[0]!.payload.length).toBe(payload.length); + expect(frames[0]!.payload[payload.length - 1]).toBe(7); + expect(elapsedMs).toBeLessThan(2000); + }); + + it("keeps unconsumed bytes from a partially consumed chunk", () => { + const first = encodePrivateFrame({ type: "event", requestId: "one" }, Buffer.from([1, 2, 3])); + const second = encodePrivateFrame({ type: "event", requestId: "two" }, Buffer.from([4, 5])); + const decoder = new PrivateFrameDecoder(isTestHeader); + + expect(decoder.push(Buffer.concat([first, second.subarray(0, 3)]))).toEqual([ + { header: { type: "event", requestId: "one" }, payload: Buffer.from([1, 2, 3]) }, + ]); + expect(decoder.bufferedBytes).toBe(3); + expect(decoder.push(second.subarray(3))).toEqual([ + { header: { type: "event", requestId: "two" }, payload: Buffer.from([4, 5]) }, + ]); + expect(decoder.bufferedBytes).toBe(0); + }); + it("reports an incomplete trailing frame", () => { const decoder = new PrivateFrameDecoder(isTestHeader); decoder.push(encodePrivateFrame({ type: "event" }, Buffer.from("body")).subarray(0, 9));