diff --git a/src/main/daemon/daemon-stream-data-batcher.test.ts b/src/main/daemon/daemon-stream-data-batcher.test.ts index 08b4624eb78..da2286d3ed1 100644 --- a/src/main/daemon/daemon-stream-data-batcher.test.ts +++ b/src/main/daemon/daemon-stream-data-batcher.test.ts @@ -1,13 +1,14 @@ import { describe, expect, it, vi } from 'vitest' import type { Socket } from 'net' import { DaemonStreamDataBatcher } from './daemon-stream-data-batcher' +import { createNdjsonParser } from './ndjson' -function createBatcher() { +function createBatcher(options?: ConstructorParameters[1]) { const streamSocket = { destroyed: false, write: vi.fn() } as unknown as Socket & { write: ReturnType } - const batcher = new DaemonStreamDataBatcher(() => ({ streamSocket })) + const batcher = new DaemonStreamDataBatcher(() => ({ streamSocket }), options) return { batcher, streamSocket } } @@ -100,4 +101,32 @@ describe('DaemonStreamDataBatcher', () => { vi.useRealTimers() } }) + + it('writes large stream data as parser-sized NDJSON events', () => { + vi.useFakeTimers() + try { + const maxLineBytes = 256 + const { batcher, streamSocket } = createBatcher({ maxLineBytes }) + const data = 'x'.repeat(maxLineBytes * 3) + const onMessage = vi.fn() + const onError = vi.fn() + const parser = createNdjsonParser(onMessage, onError, { maxLineBytes }) + + batcher.enqueue('client-1', 'session-1', data) + vi.advanceTimersByTime(8) + for (const [line] of streamSocket.write.mock.calls) { + parser.feed(String(line)) + } + + expect(onError).not.toHaveBeenCalled() + expect(onMessage).toHaveBeenCalled() + expect( + onMessage.mock.calls + .map(([message]) => (message as { payload?: { data?: string } }).payload?.data ?? '') + .join('') + ).toBe(data) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/src/main/daemon/daemon-stream-data-batcher.ts b/src/main/daemon/daemon-stream-data-batcher.ts index b9b1f6d92c9..b83f2135097 100644 --- a/src/main/daemon/daemon-stream-data-batcher.ts +++ b/src/main/daemon/daemon-stream-data-batcher.ts @@ -1,5 +1,5 @@ import type { Socket } from 'net' -import { encodeNdjson } from './ndjson' +import { encodeNdjson, NDJSON_MAX_LINE_BYTES } from './ndjson' type StreamDataClient = { streamSocket: Socket | null @@ -20,12 +20,99 @@ type EnqueueOptions = { flushMaxChars?: number } +type DaemonStreamDataBatcherOptions = { + maxLineBytes?: number +} + +function encodeStreamDataEvent(sessionId: string, data: string): string { + return encodeNdjson({ + type: 'event', + event: 'data', + sessionId, + payload: { data } + }) +} + +function streamDataEventLineBytes(sessionId: string, data: string): number { + return Buffer.byteLength(encodeStreamDataEvent(sessionId, data), 'utf8') +} + +function isHighSurrogate(value: number): boolean { + return value >= 0xd800 && value <= 0xdbff +} + +function isLowSurrogate(value: number): boolean { + return value >= 0xdc00 && value <= 0xdfff +} + +function clampToSafeSplitIndex(value: string, start: number, end: number): number { + if (end <= start || end >= value.length) { + return end + } + const prev = value.charCodeAt(end - 1) + const next = value.charCodeAt(end) + return isHighSurrogate(prev) && isLowSurrogate(next) ? end - 1 : end +} + +function nextSafeSplitIndex(value: string, start: number): number { + const next = Math.min(value.length, start + 1) + if ( + next < value.length && + isHighSurrogate(value.charCodeAt(start)) && + isLowSurrogate(value.charCodeAt(next)) + ) { + return next + 1 + } + return next +} + +function splitStreamDataForNdjson(sessionId: string, data: string, maxLineBytes: number): string[] { + if (streamDataEventLineBytes(sessionId, data) <= maxLineBytes) { + return [data] + } + + const chunks: string[] = [] + let start = 0 + while (start < data.length) { + let low = start + 1 + let high = data.length + let best = start + + while (low <= high) { + const rawMid = Math.floor((low + high) / 2) + const mid = clampToSafeSplitIndex(data, start, rawMid) + if (mid <= start) { + low = rawMid + 1 + continue + } + + if (streamDataEventLineBytes(sessionId, data.slice(start, mid)) <= maxLineBytes) { + best = mid + low = rawMid + 1 + } else { + high = rawMid - 1 + } + } + + const end = best > start ? best : nextSafeSplitIndex(data, start) + chunks.push(data.slice(start, end)) + start = end + } + + return chunks +} + export class DaemonStreamDataBatcher { private pendingByClient = new Map() private getClient: (clientId: string) => StreamDataClient | undefined + private maxLineBytes: number - constructor(getClient: (clientId: string) => StreamDataClient | undefined) { + constructor( + getClient: (clientId: string) => StreamDataClient | undefined, + options: DaemonStreamDataBatcherOptions = {} + ) { this.getClient = getClient + this.maxLineBytes = Math.max(1, options.maxLineBytes ?? NDJSON_MAX_LINE_BYTES) } enqueue(clientId: string, sessionId: string, data: string, options: EnqueueOptions = {}): void { @@ -79,14 +166,7 @@ export class DaemonStreamDataBatcher { } for (const entry of batch.queue) { - client.streamSocket.write( - encodeNdjson({ - type: 'event', - event: 'data', - sessionId: entry.sessionId, - payload: { data: entry.data } - }) - ) + this.writeStreamDataEvent(client.streamSocket, entry.sessionId, entry.data) } } @@ -137,14 +217,7 @@ export class DaemonStreamDataBatcher { } for (const entry of flushed) { - client.streamSocket.write( - encodeNdjson({ - type: 'event', - event: 'data', - sessionId: entry.sessionId, - payload: { data: entry.data } - }) - ) + this.writeStreamDataEvent(client.streamSocket, entry.sessionId, entry.data) } } @@ -161,4 +234,13 @@ export class DaemonStreamDataBatcher { this.pendingByClient.delete(id) } } + + private writeStreamDataEvent(streamSocket: Socket, sessionId: string, data: string): void { + // Why: createNdjsonParser rejects oversized lines. Terminal output can + // burst faster than the batch interval, so writer-side chunking prevents + // the daemon from dropping its own stream events at the receiver. + for (const chunk of splitStreamDataForNdjson(sessionId, data, this.maxLineBytes)) { + streamSocket.write(encodeStreamDataEvent(sessionId, chunk)) + } + } }