perf: batch relay pty output (#4777)

This commit is contained in:
Neil
2026-06-07 03:45:24 -07:00
committed by GitHub
parent 722c9c6697
commit 2c35f5bf5c
2 changed files with 262 additions and 1 deletions
+128
View File
@@ -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()
+134 -1
View File
@@ -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<typeof setTimeout> | null = null
private outputFlushTimer: ReturnType<typeof setTimeout> | null = null
private pendingOutputByPty = new Map<string, PendingPtyOutput>()
private lastInputAtByPty = new Map<string, number>()
private interactiveOutputCharsByPty = new Map<string, number>()
// 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<string, unknown>,
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)