From cef55701885524b531d4e7c6686b6df52b1ed0cc Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Thu, 21 May 2026 21:40:37 -0700 Subject: [PATCH] perf: flush daemon stream data per session (#2595) --- src/main/daemon/daemon-server.test.ts | 52 +++++++++++++++ .../daemon/daemon-stream-data-batcher.test.ts | 29 +++++++++ src/main/daemon/daemon-stream-data-batcher.ts | 63 ++++++++++++++++++- 3 files changed, 142 insertions(+), 2 deletions(-) diff --git a/src/main/daemon/daemon-server.test.ts b/src/main/daemon/daemon-server.test.ts index 0290cd943..df547bad8 100644 --- a/src/main/daemon/daemon-server.test.ts +++ b/src/main/daemon/daemon-server.test.ts @@ -1,3 +1,4 @@ +/* eslint-disable max-lines -- Why: daemon server RPC, auth, stream batching, and shutdown behavior share one socket/client harness; splitting would duplicate setup. */ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' import { connect, type Socket } from 'net' import { tmpdir } from 'os' @@ -15,6 +16,7 @@ function createTestDir(): string { function createMockSubprocess(): SubprocessHandle & { _simulateData: (data: string) => void + _simulateExit: (code: number) => void } { let onDataCb: ((data: string) => void) | null = null let onExitCb: ((code: number) => void) | null = null @@ -34,6 +36,9 @@ function createMockSubprocess(): SubprocessHandle & { dispose: vi.fn(), _simulateData(data: string) { onDataCb?.(data) + }, + _simulateExit(code: number) { + onExitCb?.(code) } } } @@ -280,6 +285,53 @@ describe('DaemonServer', () => { vi.useRealTimers() } }) + + it('flushes pending batched stream output before the exit event', async () => { + vi.useFakeTimers() + try { + let subprocess: ReturnType + server = new DaemonServer({ + socketPath, + tokenPath, + spawnSubprocess: () => { + subprocess = createMockSubprocess() + return subprocess + } + }) + const daemon = server as unknown as DaemonServerPrivate + const controlSocket = { destroy: vi.fn() } as unknown as Socket + const streamSocket = { + destroyed: false, + destroy: vi.fn(), + write: vi.fn() + } as unknown as Socket & { write: ReturnType } + + daemon.clients.set('client-1', { + clientId: 'client-1', + controlSocket, + streamSocket + }) + + await daemon.routeRequest('client-1', { + id: 'req-1', + type: 'createOrAttach', + payload: { sessionId: 'test-session', cols: 80, rows: 24 } + }) + + subprocess!._simulateData('final-output') + subprocess!._simulateExit(42) + + expect(streamSocket.write).toHaveBeenCalledTimes(2) + expect(String(streamSocket.write.mock.calls[0]?.[0])).toContain('"event":"data"') + expect(String(streamSocket.write.mock.calls[0]?.[0])).toContain('"data":"final-output"') + expect(String(streamSocket.write.mock.calls[1]?.[0])).toContain('"event":"exit"') + expect(String(streamSocket.write.mock.calls[1]?.[0])).toContain('"code":42') + vi.advanceTimersByTime(8) + expect(streamSocket.write).toHaveBeenCalledTimes(2) + } finally { + vi.useRealTimers() + } + }) }) describe('authentication', () => { diff --git a/src/main/daemon/daemon-stream-data-batcher.test.ts b/src/main/daemon/daemon-stream-data-batcher.test.ts index 408143d94..08b4624eb 100644 --- a/src/main/daemon/daemon-stream-data-batcher.test.ts +++ b/src/main/daemon/daemon-stream-data-batcher.test.ts @@ -71,4 +71,33 @@ describe('DaemonStreamDataBatcher', () => { vi.useRealTimers() } }) + + it('flushes interactive output for one session while another session has large pending output', () => { + vi.useFakeTimers() + try { + const { batcher, streamSocket } = createBatcher() + const background = 'x'.repeat(2048) + + batcher.enqueue('client-1', 'session-background', background) + batcher.enqueue('client-1', 'session-interactive', 'echo', { + flushImmediately: true, + flushMaxChars: 1024 + }) + + expect(streamSocket.write).toHaveBeenCalledTimes(1) + expect(String(streamSocket.write.mock.calls[0]?.[0])).toContain( + '"sessionId":"session-interactive"' + ) + expect(String(streamSocket.write.mock.calls[0]?.[0])).toContain('"data":"echo"') + + vi.advanceTimersByTime(8) + expect(streamSocket.write).toHaveBeenCalledTimes(2) + expect(String(streamSocket.write.mock.calls[1]?.[0])).toContain( + '"sessionId":"session-background"' + ) + expect(String(streamSocket.write.mock.calls[1]?.[0])).toContain(`"data":"${background}"`) + } finally { + vi.useRealTimers() + } + }) }) diff --git a/src/main/daemon/daemon-stream-data-batcher.ts b/src/main/daemon/daemon-stream-data-batcher.ts index f5631e827..b9b1f6d92 100644 --- a/src/main/daemon/daemon-stream-data-batcher.ts +++ b/src/main/daemon/daemon-stream-data-batcher.ts @@ -50,9 +50,10 @@ export class DaemonStreamDataBatcher { if ( options.flushImmediately === true && - batch.queuedChars <= (options.flushMaxChars ?? Number.POSITIVE_INFINITY) + this.queuedCharsForSession(batch, sessionId) <= + (options.flushMaxChars ?? Number.POSITIVE_INFINITY) ) { - this.flush(clientId) + this.flushSession(clientId, sessionId) return } if (!batch.timer) { @@ -89,6 +90,64 @@ export class DaemonStreamDataBatcher { } } + private queuedCharsForSession(batch: PendingStreamDataBatch, sessionId: string): number { + let chars = 0 + for (const entry of batch.queue) { + if (entry.sessionId === sessionId) { + chars += entry.data.length + } + } + return chars + } + + private flushSession(clientId: string, sessionId: string): void { + const batch = this.pendingByClient.get(clientId) + if (!batch) { + return + } + + const flushed: PendingStreamDataBatch['queue'] = [] + const retained: PendingStreamDataBatch['queue'] = [] + let flushedChars = 0 + for (const entry of batch.queue) { + if (entry.sessionId === sessionId) { + flushed.push(entry) + flushedChars += entry.data.length + } else { + retained.push(entry) + } + } + if (flushed.length === 0) { + return + } + + batch.queue = retained + batch.queuedChars -= flushedChars + if (batch.queue.length === 0) { + if (batch.timer) { + clearTimeout(batch.timer) + batch.timer = null + } + this.pendingByClient.delete(clientId) + } + + const client = this.getClient(clientId) + if (!client?.streamSocket || client.streamSocket.destroyed) { + return + } + + for (const entry of flushed) { + client.streamSocket.write( + encodeNdjson({ + type: 'event', + event: 'data', + sessionId: entry.sessionId, + payload: { data: entry.data } + }) + ) + } + } + clear(clientId?: string): void { const batches = clientId === undefined