mirror of
https://github.com/stablyai/orca.git
synced 2026-09-22 00:02:31 +00:00
perf: batch relay pty output (#4777)
This commit is contained in:
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user