diff --git a/.changeset/cli-stream-input.md b/.changeset/cli-stream-input.md new file mode 100644 index 0000000000..ab61e4fca4 --- /dev/null +++ b/.changeset/cli-stream-input.md @@ -0,0 +1,5 @@ +--- +'@e2b/cli': patch +--- + +`e2b sandbox connect`/`create` send terminal input over a single persistent stream instead of one request per 10ms batch diff --git a/packages/cli/src/terminal.ts b/packages/cli/src/terminal.ts index 242143dbdb..260f23e98d 100644 --- a/packages/cli/src/terminal.ts +++ b/packages/cli/src/terminal.ts @@ -1,7 +1,5 @@ import * as e2b from 'e2b' -const FLUSH_INPUT_INTERVAL_MS = 10 - function getStdoutSize() { return { cols: process.stdout.columns, @@ -36,19 +34,35 @@ export async function spawnConnectedTerminal( ...(opts.envs !== undefined ? { envs: opts.envs } : {}), }) - const inputQueue = new BatchedQueue(async (batch) => { - const combined = Buffer.concat(batch) - await sandbox.pty.sendInput(terminalSession.pid, combined) - }, FLUSH_INPUT_INTERVAL_MS) + const openInput = () => + sandbox.pty.streamInput(terminalSession.pid, { timeoutMs: 0 }) + let input = openInput() + let closed = false + let inputErrorReported = false - const resizeListener = process.stdout.on('resize', () => + const resizeListener = () => sandbox.pty.resize(terminalSession.pid, getStdoutSize()) - ) - const stdinListener = process.stdin.on('data', (data) => { - inputQueue.push(data) - }) - - inputQueue.start() + const sendInput = (data: Buffer, retried = false): void => { + const handle = input + handle.send(data).then( + () => { + inputErrorReported = false + }, + (err) => { + if (closed) return + if (input === handle) input = openInput() + if (!retried) return sendInput(data, true) + if (inputErrorReported) { + return + } + inputErrorReported = true + console.error('Error sending input:', err) + } + ) + } + const stdinListener = (data: Buffer) => sendInput(data) + process.stdout.on('resize', resizeListener) + process.stdin.on('data', stdinListener) // Wait for terminal session to finish try { @@ -66,55 +80,11 @@ export async function spawnConnectedTerminal( throw err } finally { // Cleanup + closed = true process.stdout.write('\n') - resizeListener.destroy() - stdinListener.destroy() - await inputQueue.stop() + process.stdout.off('resize', resizeListener) + process.stdin.off('data', stdinListener) + await input.close().catch(() => {}) process.stdin.setRawMode(false) } } - -class BatchedQueue { - private queue: T[] = [] - private isFlushing = false - private intervalId?: NodeJS.Timeout - - constructor( - private flushHandler: (batch: T[]) => Promise, - private flushIntervalMs: number - ) {} - - push(item: T) { - this.queue.push(item) - } - - start() { - this.intervalId = setInterval(async () => { - if (this.isFlushing) return - - this.isFlushing = true - await this.flush() - this.isFlushing = false - }, this.flushIntervalMs) - } - - async stop() { - if (this.intervalId) { - clearInterval(this.intervalId) - this.intervalId = undefined - } - - await this.flush() - } - - private async flush() { - if (this.queue.length === 0) return - - const batch = this.queue.splice(0, this.queue.length) - try { - await this.flushHandler(batch) - } catch (err) { - console.error('Error sending input:', err) - } - } -} diff --git a/packages/cli/tests/terminal.test.ts b/packages/cli/tests/terminal.test.ts new file mode 100644 index 0000000000..720ea67dcf --- /dev/null +++ b/packages/cli/tests/terminal.test.ts @@ -0,0 +1,186 @@ +import type { Sandbox } from 'e2b' +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest' + +import { spawnConnectedTerminal } from '../src/terminal' + +function createDeferred() { + let resolve!: () => void + const promise = new Promise((done) => { + resolve = done + }) + return { promise, resolve } +} + +function createInputHandle( + sendImplementation: (data: Buffer) => Promise = async () => {} +) { + return { + send: vi.fn<(data: Buffer) => Promise>(async (data: Buffer) => + sendImplementation(data) + ), + close: vi.fn(async () => {}), + } +} + +function sentChunks(handle: ReturnType) { + const calls = handle.send.mock.calls as [Buffer][] + return calls.map(([data]) => data.toString()) +} + +function createSandbox(handles: ReturnType[]) { + const wait = createDeferred() + let nextHandle = 0 + const streamInput = vi.fn( + () => handles[nextHandle++] ?? handles[handles.length - 1] + ) + const sandbox = { + pty: { + create: vi.fn().mockResolvedValue({ + pid: 1, + wait: () => wait.promise, + }), + resize: vi.fn().mockResolvedValue(undefined), + streamInput, + }, + } as unknown as Sandbox + + return { sandbox, streamInput, wait } +} + +const flush = () => new Promise((resolve) => setImmediate(resolve)) +const originalRawModeDescriptor = Object.getOwnPropertyDescriptor( + process.stdin, + 'setRawMode' +) +let setRawModeMock = vi.fn() + +beforeEach(() => { + setRawModeMock = vi.fn() + Object.defineProperty(process.stdin, 'setRawMode', { + configurable: true, + writable: true, + value: setRawModeMock, + }) + vi.spyOn(process.stdout, 'write').mockImplementation((() => true) as never) +}) + +afterEach(() => { + vi.restoreAllMocks() + if (originalRawModeDescriptor) { + Object.defineProperty( + process.stdin, + 'setRawMode', + originalRawModeDescriptor + ) + } else { + Reflect.deleteProperty(process.stdin, 'setRawMode') + } +}) + +describe('spawnConnectedTerminal input stream', () => { + test('sends chunks in order over one persistent stream', async () => { + const handle = createInputHandle() + const { sandbox, streamInput, wait } = createSandbox([handle]) + const terminal = spawnConnectedTerminal(sandbox) + + try { + await flush() + process.stdin.emit('data', Buffer.from('first')) + process.stdin.emit('data', Buffer.from('second')) + await flush() + + expect(streamInput).toHaveBeenCalledTimes(1) + expect(streamInput).toHaveBeenCalledWith(1, { timeoutMs: 0 }) + expect(sentChunks(handle)).toEqual(['first', 'second']) + } finally { + wait.resolve() + await terminal + } + }) + + test('reopens once per failed handle and resends chunks in order', async () => { + const first = createInputHandle(async () => { + throw new Error('first stream failed') + }) + const replacement = createInputHandle() + const { sandbox, streamInput, wait } = createSandbox([first, replacement]) + const terminal = spawnConnectedTerminal(sandbox) + + try { + await flush() + process.stdin.emit('data', Buffer.from('first')) + process.stdin.emit('data', Buffer.from('second')) + await flush() + + expect(streamInput).toHaveBeenCalledTimes(2) + expect(streamInput).toHaveBeenNthCalledWith(1, 1, { timeoutMs: 0 }) + expect(streamInput).toHaveBeenNthCalledWith(2, 1, { timeoutMs: 0 }) + expect(sentChunks(replacement)).toEqual(['first', 'second']) + + process.stdin.emit('data', Buffer.from('later')) + await flush() + + expect(sentChunks(replacement)).toEqual(['first', 'second', 'later']) + } finally { + wait.resolve() + await terminal + } + }) + + test('reports repeated resend failures only once', async () => { + const first = createInputHandle(async () => { + throw new Error('first stream failed') + }) + const replacement = createInputHandle(async () => { + throw new Error('replacement stream failed') + }) + const finalHandle = createInputHandle() + const { sandbox, streamInput, wait } = createSandbox([ + first, + replacement, + finalHandle, + ]) + const consoleError = vi.spyOn(console, 'error').mockImplementation(() => {}) + const terminal = spawnConnectedTerminal(sandbox) + + try { + await flush() + process.stdin.emit('data', Buffer.from('first')) + process.stdin.emit('data', Buffer.from('second')) + await flush() + + expect(streamInput).toHaveBeenCalledTimes(3) + expect(sentChunks(replacement)).toEqual(['first', 'second']) + expect(consoleError).toHaveBeenCalledTimes(1) + expect(consoleError).toHaveBeenCalledWith( + 'Error sending input:', + expect.any(Error) + ) + } finally { + wait.resolve() + await terminal + } + }) + + test('cleans up listeners, closes the input stream, and disables raw mode', async () => { + const handle = createInputHandle() + const { sandbox, wait } = createSandbox([handle]) + const stdinListeners = process.stdin.listenerCount('data') + const resizeListeners = process.stdout.listenerCount('resize') + const terminal = spawnConnectedTerminal(sandbox) + + try { + await flush() + expect(process.stdin.listenerCount('data')).toBe(stdinListeners + 1) + expect(process.stdout.listenerCount('resize')).toBe(resizeListeners + 1) + } finally { + wait.resolve() + await terminal + } + + expect(process.stdin.listenerCount('data')).toBe(stdinListeners) + expect(process.stdout.listenerCount('resize')).toBe(resizeListeners) + expect(handle.close).toHaveBeenCalledOnce() + expect(setRawModeMock).toHaveBeenCalledWith(false) + }) +})