perf: flush daemon stream data per session (#2595)

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<typeof createMockSubprocess>
+        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<typeof vi.fn> }
+
+        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', () => {

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()
+    }
+  })
 })

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