From 2c35f5bf5ceb8f5a976360d00e35bbc81bae8e5d Mon Sep 17 00:00:00 2001 From: Neil <4138956+nwparker@users.noreply.github.com> Date: Sun, 7 Jun 2026 03:45:24 -0700 Subject: [PATCH] perf: batch relay pty output (#4777) --- src/relay/pty-handler.test.ts | 128 ++++++++++++++++++++++++++++++++ src/relay/pty-handler.ts | 135 +++++++++++++++++++++++++++++++++- 2 files changed, 262 insertions(+), 1 deletion(-) diff --git a/src/relay/pty-handler.test.ts b/src/relay/pty-handler.test.ts index a31d35259e9..f643f0997f7 100644 --- a/src/relay/pty-handler.test.ts +++ b/src/relay/pty-handler.test.ts @@ -212,9 +212,86 @@ describe('PtyHandler', () => { expect(dataCallback).toBeDefined() dataCallback!('hello world') + expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.data', expect.anything()) + vi.advanceTimersByTime(8) expect(dispatcher.notify).toHaveBeenCalledWith('pty.data', { id: 'pty-1', data: 'hello world' }) }) + it('coalesces background PTY output before notifying the client', async () => { + let dataCallback: ((data: string) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn() + }) + + await dispatcher.callRequest('pty.spawn', {}) + dataCallback!('hello ') + dataCallback!('world') + + expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.data', expect.anything()) + vi.advanceTimersByTime(8) + expect(dispatcher.notify).toHaveBeenCalledWith('pty.data', { + id: 'pty-1', + data: 'hello world' + }) + }) + + it('sends recent-input redraw output immediately', async () => { + let dataCallback: ((data: string) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn() + }) + + await dispatcher.callRequest('pty.spawn', {}) + dispatcher.callNotification('pty.data', { id: 'pty-1', data: 'a' }) + dispatcher.notify.mockClear() + + dataCallback!('\x1b[20;2Hredraw') + + expect(dispatcher.notify).toHaveBeenCalledWith('pty.data', { + id: 'pty-1', + data: '\x1b[20;2Hredraw' + }) + vi.advanceTimersByTime(8) + expect(dispatcher.notify).toHaveBeenCalledTimes(1) + }) + + it('drains large relay PTY output in bounded slices', async () => { + let dataCallback: ((data: string) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn() + }) + + await dispatcher.callRequest('pty.spawn', {}) + const firstChunk = 'x'.repeat(16 * 1024) + dataCallback!(`${firstChunk}tail`) + + vi.advanceTimersByTime(8) + expect(dispatcher.notify).toHaveBeenCalledTimes(1) + expect(dispatcher.notify).toHaveBeenNthCalledWith(1, 'pty.data', { + id: 'pty-1', + data: firstChunk + }) + + vi.advanceTimersByTime(1) + expect(dispatcher.notify).toHaveBeenCalledTimes(2) + expect(dispatcher.notify).toHaveBeenNthCalledWith(2, 'pty.data', { + id: 'pty-1', + data: 'tail' + }) + }) + it('returns attach replay instead of notifying when replay notification is suppressed', async () => { let dataCallback: ((data: string) => void) | undefined mockPtySpawn.mockReturnValue({ @@ -235,6 +312,8 @@ describe('PtyHandler', () => { expect(result).toEqual({ replay: 'buffered output' }) expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.replay', expect.anything()) + vi.advanceTimersByTime(8) + expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.data', expect.anything()) }) it('notifies replay on normal attach', async () => { @@ -258,6 +337,8 @@ describe('PtyHandler', () => { id: 'pty-1', data: 'buffered output' }) + vi.advanceTimersByTime(8) + expect(dispatcher.notify).not.toHaveBeenCalledWith('pty.data', expect.anything()) }) it('notifies on PTY exit and removes from map', async () => { @@ -278,6 +359,30 @@ describe('PtyHandler', () => { expect(handler.activePtyCount).toBe(0) }) + it('flushes pending PTY output before notifying exit', async () => { + let dataCallback: ((data: string) => void) | undefined + let exitCallback: ((info: { exitCode: number }) => void) | undefined + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn((cb: (info: { exitCode: number }) => void) => { + exitCallback = cb + }) + }) + + await dispatcher.callRequest('pty.spawn', {}) + dataCallback!('final output') + exitCallback!({ exitCode: 0 }) + + expect(dispatcher.notify).toHaveBeenNthCalledWith(1, 'pty.data', { + id: 'pty-1', + data: 'final output' + }) + expect(dispatcher.notify).toHaveBeenNthCalledWith(2, 'pty.exit', { id: 'pty-1', code: 0 }) + }) + it('writes data to PTY via pty.data notification', async () => { const mockWrite = vi.fn() mockPtySpawn.mockReturnValue({ @@ -320,6 +425,29 @@ describe('PtyHandler', () => { expect(mockKill).toHaveBeenCalledWith('SIGTERM') }) + it('flushes pending PTY output before immediate shutdown cleanup', async () => { + let dataCallback: ((data: string) => void) | undefined + const mockKill = vi.fn() + mockPtySpawn.mockReturnValue({ + ...mockPtyInstance, + kill: mockKill, + onData: vi.fn((cb: (data: string) => void) => { + dataCallback = cb + }), + onExit: vi.fn() + }) + + await dispatcher.callRequest('pty.spawn', {}) + dataCallback!('last words') + await dispatcher.callRequest('pty.shutdown', { id: 'pty-1', immediate: true }) + + expect(dispatcher.notify).toHaveBeenNthCalledWith(1, 'pty.data', { + id: 'pty-1', + data: 'last words' + }) + expect(mockKill).toHaveBeenCalledWith('SIGKILL') + }) + it('notifies pty.exit when graceful shutdown falls back to SIGKILL', async () => { let onExitCb: ((evt: { exitCode: number }) => void) | undefined const mockKill = vi.fn() diff --git a/src/relay/pty-handler.ts b/src/relay/pty-handler.ts index c6c6ae7e6f2..9384967a4db 100644 --- a/src/relay/pty-handler.ts +++ b/src/relay/pty-handler.ts @@ -53,6 +53,10 @@ type ManagedPty = { worktreeId?: string } +type PendingPtyOutput = { + data: string +} + function disposeManagedPty(managed: ManagedPty): void { if (managed.disposed) { return @@ -84,6 +88,14 @@ function disposeManagedPty(managed: ManagedPty): void { } const DEFAULT_GRACE_TIME_MS = DEFAULT_SSH_RELAY_GRACE_PERIOD_SECONDS * 1000 export const REPLAY_BUFFER_MAX = 100 * 1024 +const PTY_OUTPUT_BATCH_INTERVAL_MS = 8 +const PTY_OUTPUT_DRAIN_CONTINUE_MS = 1 +const PTY_OUTPUT_FLUSH_CHUNK_CHARS = 16 * 1024 +const PTY_OUTPUT_FLUSH_MAX_WRITES = 2 +const INTERACTIVE_OUTPUT_WINDOW_MS = 100 +const INTERACTIVE_OUTPUT_MAX_CHARS = 1024 +const INTERACTIVE_REDRAW_MAX_CHARS = PTY_OUTPUT_FLUSH_CHUNK_CHARS +const INTERACTIVE_OUTPUT_BUDGET_CHARS = 32 * 1024 const ALLOWED_SIGNALS = new Set([ 'SIGINT', 'SIGTERM', @@ -130,6 +142,10 @@ export class PtyHandler { private dispatcher: RelayDispatcher private graceTimeMs: number private graceTimer: ReturnType | null = null + private outputFlushTimer: ReturnType | null = null + private pendingOutputByPty = new Map() + private lastInputAtByPty = new Map() + private interactiveOutputCharsByPty = new Map() // Why: external observers need to drop per-pane state when a PTY exits. // Today the relay composes multiple consumers (hook-server cache eviction // and plugin-overlay dir cleanup) into a single callback at the call site @@ -212,7 +228,7 @@ export class PtyHandler { if (managed.buffered.length > REPLAY_BUFFER_MAX) { managed.buffered = managed.buffered.slice(-REPLAY_BUFFER_MAX) } - this.dispatcher.notify('pty.data', { id: managed.id, data }) + this.enqueuePtyOutput(managed.id, data) }) managed.pty.onExit(({ exitCode }: { exitCode: number }) => { if (managed.disposed) { @@ -235,9 +251,11 @@ export class PtyHandler { clearTimeout(managed.killTimer) managed.killTimer = undefined } + this.flushPtyOutput(managed.id) this.dispatcher.notify('pty.exit', { id: managed.id, code: exitCode }) this.notifyExitListener(managed) this.ptys.delete(managed.id) + this.clearPtyFlowState(managed.id) // Why: release the ptmx fd on the natural-exit path. Without this the // node-pty wrapper's _socket stays alive until GC and the master fd // leaks (see docs/fix-pty-fd-leak.md). @@ -288,6 +306,103 @@ export class PtyHandler { }) } + private isLikelyInteractiveRedraw(data: string): boolean { + if (data.length <= INTERACTIVE_OUTPUT_MAX_CHARS) { + return true + } + return data.length <= INTERACTIVE_REDRAW_MAX_CHARS && data.includes('\x1b[') + } + + private shouldSendInteractiveOutputNow(id: string, data: string): boolean { + const lastInputAt = this.lastInputAtByPty.get(id) + const now = performance.now() + if (lastInputAt === undefined || now - lastInputAt > INTERACTIVE_OUTPUT_WINDOW_MS) { + this.interactiveOutputCharsByPty.delete(id) + return false + } + if (!this.isLikelyInteractiveRedraw(data)) { + this.interactiveOutputCharsByPty.set(id, INTERACTIVE_OUTPUT_BUDGET_CHARS) + return false + } + const usedChars = this.interactiveOutputCharsByPty.get(id) ?? 0 + if (usedChars + data.length > INTERACTIVE_OUTPUT_BUDGET_CHARS) { + this.interactiveOutputCharsByPty.set(id, INTERACTIVE_OUTPUT_BUDGET_CHARS) + return false + } + this.interactiveOutputCharsByPty.set(id, usedChars + data.length) + return true + } + + private enqueuePtyOutput(id: string, data: string): void { + const existing = this.pendingOutputByPty.get(id) + const pending = { data: (existing?.data ?? '') + data } + if (this.shouldSendInteractiveOutputNow(id, pending.data)) { + this.pendingOutputByPty.delete(id) + this.clearOutputFlushTimerIfIdle() + // Why: remote agent TUIs redraw around each keystroke. Background relay + // batching should reduce SSH chatter, not add visible input echo delay. + this.dispatcher.notify('pty.data', { id, data: pending.data }) + return + } + this.pendingOutputByPty.set(id, pending) + this.scheduleOutputFlush(PTY_OUTPUT_BATCH_INTERVAL_MS) + } + + private scheduleOutputFlush(delayMs: number): void { + if (this.outputFlushTimer !== null) { + return + } + this.outputFlushTimer = setTimeout(() => this.flushPendingOutput(), delayMs) + } + + private flushPendingOutput(): void { + this.outputFlushTimer = null + let writes = 0 + for (const [id, pending] of Array.from(this.pendingOutputByPty.entries())) { + if (writes >= PTY_OUTPUT_FLUSH_MAX_WRITES) { + break + } + this.pendingOutputByPty.delete(id) + const chunk = pending.data.slice(0, PTY_OUTPUT_FLUSH_CHUNK_CHARS) + const remaining = pending.data.slice(PTY_OUTPUT_FLUSH_CHUNK_CHARS) + if (remaining) { + this.pendingOutputByPty.set(id, { data: remaining }) + } + this.dispatcher.notify('pty.data', { id, data: chunk }) + writes++ + } + if (this.pendingOutputByPty.size > 0 && writes > 0) { + // Why: relay-side output can arrive as a large single PTY chunk. Yield + // between slices so client input and control frames can interleave. + this.scheduleOutputFlush(PTY_OUTPUT_DRAIN_CONTINUE_MS) + } + } + + private flushPtyOutput(id: string): void { + const pending = this.pendingOutputByPty.get(id) + if (!pending) { + return + } + this.pendingOutputByPty.delete(id) + this.dispatcher.notify('pty.data', { id, data: pending.data }) + this.clearOutputFlushTimerIfIdle() + } + + private clearOutputFlushTimerIfIdle(): void { + if (this.pendingOutputByPty.size > 0 || this.outputFlushTimer === null) { + return + } + clearTimeout(this.outputFlushTimer) + this.outputFlushTimer = null + } + + private clearPtyFlowState(id: string): void { + this.pendingOutputByPty.delete(id) + this.lastInputAtByPty.delete(id) + this.interactiveOutputCharsByPty.delete(id) + this.clearOutputFlushTimerIfIdle() + } + private async spawn( params: Record, context?: RequestContext @@ -391,6 +506,11 @@ export class PtyHandler { // restart still replays the full terminal history instead of only output // generated since the previous attach. if (managed.buffered) { + // Why: relay batching may still hold bytes that are already included in + // the full replay buffer. Drop that pending notification before attach + // so reconnect/suppressed replay cannot render the same bytes twice. + this.pendingOutputByPty.delete(id) + this.clearOutputFlushTimerIfIdle() if (params.suppressReplayNotification) { return { replay: managed.buffered } } @@ -407,6 +527,8 @@ export class PtyHandler { } const managed = this.ptys.get(id) if (managed && !managed.disposed) { + this.lastInputAtByPty.set(id, performance.now()) + this.interactiveOutputCharsByPty.set(id, 0) managed.pty.write(data) } } @@ -430,6 +552,7 @@ export class PtyHandler { } if (immediate) { + this.flushPtyOutput(id) managed.pty.kill('SIGKILL') // Why: SIGKILL has already reaped the child; release the ptmx fd on the // same tick. Deferring to onExit leaves a window where the fd is live @@ -446,6 +569,7 @@ export class PtyHandler { // a no-op. this.notifyExitListener(managed) this.ptys.delete(id) + this.clearPtyFlowState(id) } else { managed.pty.kill('SIGTERM') @@ -462,6 +586,7 @@ export class PtyHandler { const still = this.ptys.get(id) if (still && !still.disposed) { still.pty.kill('SIGKILL') + this.flushPtyOutput(id) // Why: emit pty.exit BEFORE disposeManagedPty sets disposed=true. // The natural onExit short-circuits on `managed.disposed`, so // without this notify the renderer never learns the pane is dead @@ -476,6 +601,7 @@ export class PtyHandler { this.notifyExitListener(still) disposeManagedPty(still) this.ptys.delete(id) + this.clearPtyFlowState(id) } }, 5000) } @@ -675,6 +801,13 @@ export class PtyHandler { dispose(): void { this.cancelGraceTimer() + if (this.outputFlushTimer !== null) { + clearTimeout(this.outputFlushTimer) + this.outputFlushTimer = null + } + this.pendingOutputByPty.clear() + this.lastInputAtByPty.clear() + this.interactiveOutputCharsByPty.clear() for (const [, managed] of this.ptys) { if (managed.killTimer) { clearTimeout(managed.killTimer)