diff --git a/.changeset/stream-input-sdk.md b/.changeset/stream-input-sdk.md new file mode 100644 index 0000000000..9bb8d3462f --- /dev/null +++ b/.changeset/stream-input-sdk.md @@ -0,0 +1,7 @@ +--- +"e2b": minor +"@e2b/python-sdk": minor +--- + +- JS: `sandbox.commands.streamInput(pid)` / `sandbox.pty.streamInput(pid)` return a `CommandInputHandle` that sends ordered stdin/PTY input over one `StreamInput` RPC. +- Python async: `stream_input(pid)` returns an `AsyncCommandInputHandle`; Python sync: `stream_input(pid, data)` blocks while it sends an iterable of input. diff --git a/packages/js-sdk/src/connectionConfig.ts b/packages/js-sdk/src/connectionConfig.ts index b946cc69b4..54af513b4b 100644 --- a/packages/js-sdk/src/connectionConfig.ts +++ b/packages/js-sdk/src/connectionConfig.ts @@ -8,6 +8,7 @@ import { InvalidArgumentError } from './errors' const supportedDomains = ['e2b.app', 'e2b.dev', 'e2b.pro', 'e2b-staging.dev'] export const REQUEST_TIMEOUT_MS = 60_000 // 60 seconds +export const STREAM_INPUT_TIMEOUT_MS = 60_000 // 60 seconds export const DEFAULT_RETRIES = 3 export const KEEPALIVE_PING_INTERVAL_SEC = 50 // 50 seconds diff --git a/packages/js-sdk/src/envd/streamingTransport.ts b/packages/js-sdk/src/envd/streamingTransport.ts new file mode 100644 index 0000000000..87acd9b80a --- /dev/null +++ b/packages/js-sdk/src/envd/streamingTransport.ts @@ -0,0 +1,103 @@ +import type { Interceptor } from '@connectrpc/connect' +import { createTransport } from '@connectrpc/connect/protocol-connect' +import type { + UniversalClientFn, + UniversalClientRequest, + UniversalClientResponse, +} from '@connectrpc/connect/protocol' + +function asyncIterableToReadableStream( + iterable: AsyncIterable +): ReadableStream { + const iterator = iterable[Symbol.asyncIterator]() + + return new ReadableStream({ + async pull(controller) { + try { + const { done, value } = await iterator.next() + if (done) { + controller.close() + } else { + controller.enqueue(value) + } + } catch (error) { + controller.error(error) + } + }, + async cancel(reason) { + await iterator.return?.(reason) + }, + }) +} + +function readableStreamToAsyncIterable( + stream: ReadableStream | null +): AsyncIterable { + return { + async *[Symbol.asyncIterator]() { + if (!stream) { + return + } + + const reader = stream.getReader() + try { + while (true) { + const { done, value } = await reader.read() + if (done) { + return + } + yield value + } + } finally { + reader.releaseLock() + } + }, + } +} + +function createStreamingHttpClient( + fetch: typeof globalThis.fetch +): UniversalClientFn { + return async ( + request: UniversalClientRequest + ): Promise => { + const response = await fetch(request.url, { + method: request.method, + headers: request.header, + body: request.body + ? asyncIterableToReadableStream(request.body) + : undefined, + signal: request.signal, + duplex: 'half', + } as RequestInit & { duplex: 'half' }) + + return { + status: response.status, + header: response.headers, + body: readableStreamToAsyncIterable(response.body), + trailer: new Headers(), + } + } +} + +export function createClientStreamingTransport({ + baseUrl, + fetch, + interceptors = [], +}: { + baseUrl: string + fetch: typeof globalThis.fetch + interceptors?: Interceptor[] +}) { + return createTransport({ + baseUrl, + httpClient: createStreamingHttpClient(fetch), + useBinaryFormat: false, + interceptors, + acceptCompression: [], + sendCompression: null, + compressMinBytes: 1024, + readMaxBytes: 0xffffffff, + writeMaxBytes: 0xffffffff, + }) +} diff --git a/packages/js-sdk/src/index.ts b/packages/js-sdk/src/index.ts index 5de45f1226..7d2553ae3f 100644 --- a/packages/js-sdk/src/index.ts +++ b/packages/js-sdk/src/index.ts @@ -51,6 +51,8 @@ export type { } from './sandbox/filesystem/watchHandle' export { CommandExitError } from './sandbox/commands/commandHandle' +export { CommandInputHandle } from './sandbox/commands/inputHandle' +export type { CommandStreamInputOpts } from './sandbox/commands/inputHandle' export type { CommandResult, Stdout, diff --git a/packages/js-sdk/src/sandbox/commands/index.ts b/packages/js-sdk/src/sandbox/commands/index.ts index 446147ca8e..b4c6411c48 100644 --- a/packages/js-sdk/src/sandbox/commands/index.ts +++ b/packages/js-sdk/src/sandbox/commands/index.ts @@ -28,7 +28,10 @@ import { import { ENVD_COMMANDS_STDIN, ENVD_ENVD_CLOSE } from '../../envd/versions' import { SandboxError } from '../../errors' import { CommandHandle, CommandResult } from './commandHandle' +import { CommandInputHandle, CommandStreamInputOpts } from './inputHandle' export { Pty } from './pty' +export { CommandInputHandle } +export type { CommandStreamInputOpts } /** * Options for sending a command request. @@ -131,6 +134,7 @@ export interface ProcessInfo { */ export class Commands { protected readonly rpc: Client + private readonly streamingRpc: Client private readonly defaultProcessConnectionTimeout = 60_000 // 60 seconds private readonly envdVersion: string @@ -139,9 +143,11 @@ export class Commands { constructor( transport: Transport, private readonly envdApi: EnvdApiClient, - private readonly connectionConfig: ConnectionConfig + private readonly connectionConfig: ConnectionConfig, + streamingTransport: Transport ) { this.rpc = createClient(ProcessService, transport) + this.streamingRpc = createClient(ProcessService, streamingTransport) this.envdVersion = envdApi.version this.checkHealth = () => checkSandboxHealth(this.envdApi) } @@ -229,6 +235,42 @@ export class Commands { } } + /** + * Open a persistent input stream to a command. + * + * Input sent through this stream is ordered relative to other `send` calls. + * Use this instead of {@link Commands.sendStdin} when ordering matters; each + * `sendStdin` call is a separate request. The stream sends keepalive messages + * while idle. + * + * In browsers, this requires fetch request-body streaming support (Chromium + * over HTTP/2). + * + * @param pid Process ID of the command. + * @param opts Optional stream timeout and abort signal. `timeoutMs` defaults to `STREAM_INPUT_TIMEOUT_MS` (60_000 milliseconds). + * @returns A handle whose `send()` method queues input in order and whose `close()` method ends the stream and waits for envd. + * @throws {NotFoundError} If the command does not exist, surfaced from `send()` or `close()`. + * @throws {TimeoutError} If the stream timeout expires, surfaced from `send()` or `close()`. + * @example + * ```ts + * const input = sandbox.commands.streamInput(pid) + * await input.send('hello\n') + * await input.close() + * ``` + */ + streamInput( + pid: number, + opts: CommandStreamInputOpts = {} + ): CommandInputHandle { + return new CommandInputHandle( + this.streamingRpc, + pid, + 'stdin', + opts, + this.checkHealth + ) + } + /** * Close command stdin. * diff --git a/packages/js-sdk/src/sandbox/commands/inputHandle.ts b/packages/js-sdk/src/sandbox/commands/inputHandle.ts new file mode 100644 index 0000000000..ba46fcb6b3 --- /dev/null +++ b/packages/js-sdk/src/sandbox/commands/inputHandle.ts @@ -0,0 +1,205 @@ +import type { Client } from '@connectrpc/connect' +import { create } from '@bufbuild/protobuf' + +import { + KEEPALIVE_PING_INTERVAL_SEC, + STREAM_INPUT_TIMEOUT_MS, +} from '../../connectionConfig' +import { + Process as ProcessService, + StreamInputRequestSchema, +} from '../../envd/process/process_pb' +import { + handleRpcErrorWithHealthCheck, + SandboxHealthCheck, +} from '../../envd/rpc' +import { SandboxError } from '../../errors' + +const CLOSE = Symbol('close') + +/** + * Options for streaming input to a command or PTY. + */ +export interface CommandStreamInputOpts { + /** + * Maximum duration of the input stream in milliseconds. + * `0` disables the timeout. + * + * @default STREAM_INPUT_TIMEOUT_MS (60_000 milliseconds) + */ + timeoutMs?: number + /** + * Abort signal for the input stream. + */ + signal?: AbortSignal +} + +type InputKind = 'stdin' | 'pty' +type QueuedInput = string | Uint8Array | typeof CLOSE + +/** + * An ordered, long-lived input handle for a command or PTY. + */ +export class CommandInputHandle { + private readonly queue: QueuedInput[] = [] + private readonly completion: Promise + private readonly keepaliveIntervalMs: number + private wake?: () => void + private closing = false + private stopped = false + private failure?: Error + private closePromise?: Promise + + /** + * @hidden + * @internal + */ + constructor( + private readonly rpc: Client, + private readonly pid: number, + private readonly inputKind: InputKind, + opts: CommandStreamInputOpts, + private readonly checkHealth: SandboxHealthCheck, + keepaliveIntervalMs = KEEPALIVE_PING_INTERVAL_SEC * 1000 + ) { + this.keepaliveIntervalMs = keepaliveIntervalMs + this.completion = this.rpc + .streamInput(this.requests(), { + signal: opts.signal, + timeoutMs: opts.timeoutMs ?? STREAM_INPUT_TIMEOUT_MS, + }) + .then(() => { + this.stopped = true + this.notify() + }) + .catch(async (error: unknown) => { + this.failure = await handleRpcErrorWithHealthCheck( + error, + this.checkHealth + ) + this.stopped = true + this.notify() + throw this.failure + }) + this.completion.catch(() => undefined) + } + + /** + * Enqueue input for transmission in call order. + * + * @param data Input data to send. + * @returns A promise that resolves when the input is queued for ordered transmission. + * @throws {NotFoundError} If the command or PTY does not exist. + * @throws {TimeoutError} If the input stream times out. + * @throws {SandboxError} If this handle is closed. + */ + async send(data: string | Uint8Array): Promise { + if (this.failure) { + throw this.failure + } + if (this.closing || this.stopped) { + throw new SandboxError( + 'Cannot send input after the input handle is closed. Open a new one with `sandbox.commands.streamInput(pid)` or `sandbox.pty.streamInput(pid)`.' + ) + } + + this.queue.push(data instanceof Uint8Array ? data.slice() : data) + this.notify() + } + + /** + * Finish the input stream and wait for envd to acknowledge it. + * + * @returns A promise that resolves when envd acknowledges the end of the input stream. + * @throws {NotFoundError} If the command or PTY does not exist. + * @throws {TimeoutError} If the input stream times out. + */ + close(): Promise { + if (this.closePromise) { + return this.closePromise + } + + this.closing = true + this.queue.push(CLOSE) + this.notify() + this.closePromise = this.completion + return this.closePromise + } + + private notify() { + this.wake?.() + } + + private waitForInput(): Promise { + return new Promise((resolve) => { + let settled = false + const timer = setTimeout(() => { + settled = true + this.wake = undefined + resolve(false) + }, this.keepaliveIntervalMs) + + this.wake = () => { + if (settled) { + return + } + settled = true + clearTimeout(timer) + this.wake = undefined + resolve(true) + } + }) + } + + private async *requests() { + yield create(StreamInputRequestSchema, { + event: { + case: 'start', + value: { + process: { + selector: { + case: 'pid', + value: this.pid, + }, + }, + }, + }, + }) + + while (!this.stopped) { + if (this.queue.length > 0) { + const input = this.queue.shift()! + if (input === CLOSE) { + return + } + + const value = + typeof input === 'string' ? new TextEncoder().encode(input) : input + yield create(StreamInputRequestSchema, { + event: { + case: 'data', + value: { + input: { + input: { + case: this.inputKind, + value, + }, + }, + }, + }, + }) + continue + } + + const receivedInput = await this.waitForInput() + if (!receivedInput && !this.stopped && this.queue.length === 0) { + yield create(StreamInputRequestSchema, { + event: { + case: 'keepalive', + value: {}, + }, + }) + } + } + } +} diff --git a/packages/js-sdk/src/sandbox/commands/pty.ts b/packages/js-sdk/src/sandbox/commands/pty.ts index 7fea01131d..e7f930e328 100644 --- a/packages/js-sdk/src/sandbox/commands/pty.ts +++ b/packages/js-sdk/src/sandbox/commands/pty.ts @@ -10,6 +10,7 @@ import { setupRequestController, } from '../../connectionConfig' import { CommandHandle } from './commandHandle' +import { CommandInputHandle, CommandStreamInputOpts } from './inputHandle' import { authenticationHeader, handleRpcErrorWithHealthCheck, @@ -75,6 +76,7 @@ export type PtyConnectOpts = Pick & */ export class Pty { private readonly rpc: Client + private readonly streamingRpc: Client private readonly envdVersion: string private readonly checkHealth: SandboxHealthCheck @@ -83,9 +85,11 @@ export class Pty { constructor( private readonly transport: Transport, private readonly envdApi: EnvdApiClient, - private readonly connectionConfig: ConnectionConfig + private readonly connectionConfig: ConnectionConfig, + streamingTransport: Transport ) { this.rpc = createClient(ProcessService, this.transport) + this.streamingRpc = createClient(ProcessService, streamingTransport) this.envdVersion = envdApi.version this.checkHealth = () => checkSandboxHealth(this.envdApi) } @@ -254,6 +258,42 @@ export class Pty { } } + /** + * Open a persistent input stream to a PTY. + * + * Input sent through this stream is ordered relative to other `send` calls. + * Use this instead of {@link Pty.sendInput} when ordering matters; each + * `sendInput` call is a separate request. The stream sends keepalive messages + * while idle. + * + * In browsers, this requires fetch request-body streaming support (Chromium + * over HTTP/2). + * + * @param pid Process ID of the PTY. + * @param opts Optional stream timeout and abort signal. `timeoutMs` defaults to `STREAM_INPUT_TIMEOUT_MS` (60_000 milliseconds). + * @returns A handle whose `send()` method queues input in order and whose `close()` method ends the stream and waits for envd. + * @throws {NotFoundError} If the PTY does not exist, surfaced from `send()` or `close()`. + * @throws {TimeoutError} If the stream timeout expires, surfaced from `send()` or `close()`. + * @example + * ```ts + * const input = sandbox.pty.streamInput(pid) + * await input.send('echo hello\n') + * await input.close() + * ``` + */ + streamInput( + pid: number, + opts: CommandStreamInputOpts = {} + ): CommandInputHandle { + return new CommandInputHandle( + this.streamingRpc, + pid, + 'pty', + opts, + this.checkHealth + ) + } + /** * Resize PTY. * Call this when the terminal window is resized and the number of columns and rows has changed. diff --git a/packages/js-sdk/src/sandbox/index.ts b/packages/js-sdk/src/sandbox/index.ts index 3af8e5e3d4..12609db666 100644 --- a/packages/js-sdk/src/sandbox/index.ts +++ b/packages/js-sdk/src/sandbox/index.ts @@ -10,6 +10,7 @@ import { EnvdApiClient, handleEnvdApiError } from '../envd/api' import { rejectProxyUnavailableResponse } from '../envd/rpc' import { createEnvdFetch, createEnvdRpcFetch } from '../envd/http2' import { createRpcLogger } from '../logs' +import { createClientStreamingTransport } from '../envd/streamingTransport' import { Commands, Pty } from './commands' import { CommandExitError } from './commands/commandHandle' import { Filesystem } from './filesystem' @@ -164,39 +165,49 @@ export class Sandbox extends SandboxApi { } const envdFetch = createEnvdFetch(this.connectionConfig) const envdRpcFetch = createEnvdRpcFetch(this.connectionConfig) + const interceptors = opts?.logger + ? [createRpcLogger(opts.logger)] + : undefined + + const rpcFetch: typeof fetch = (url, options) => { + // Patch fetch to always use redirect: "follow" + // connect-web doesn't allow to configure redirect option - https://github.com/connectrpc/connect-es/pull/1082 + // connect-web package uses redirect: "error" which is not supported in edge runtimes + // E2B endpoints should be safe to use with redirect: "follow" https://github.com/e2b-dev/E2B/issues/531#issuecomment-2779492867 + + const headers = new Headers({ + 'User-Agent': this.connectionConfig.headers?.['User-Agent'] ?? '', + }) + new Headers(options?.headers).forEach((value, key) => + headers.append(key, value) + ) + new Headers(sandboxHeaders).forEach((value, key) => + headers.append(key, value) + ) - const rpcTransport = createConnectTransport({ - baseUrl: this.envdApiUrl, - useBinaryFormat: false, - interceptors: opts?.logger ? [createRpcLogger(opts.logger)] : undefined, - fetch: (url, options) => { - // Patch fetch to always use redirect: "follow" - // connect-web doesn't allow to configure redirect option - https://github.com/connectrpc/connect-es/pull/1082 - // connect-web package uses redirect: "error" which is not supported in edge runtimes - // E2B endpoints should be safe to use with redirect: "follow" https://github.com/e2b-dev/E2B/issues/531#issuecomment-2779492867 - - const headers = new Headers({ - 'User-Agent': this.connectionConfig.headers?.['User-Agent'] ?? '', - }) - new Headers(options?.headers).forEach((value, key) => - headers.append(key, value) - ) - new Headers(sandboxHeaders).forEach((value, key) => - headers.append(key, value) - ) + if (this.envdAccessToken) { + headers.append('X-Access-Token', this.envdAccessToken) + } - if (this.envdAccessToken) { - headers.append('X-Access-Token', this.envdAccessToken) - } + options = { + ...(options ?? {}), + headers: headers, + redirect: 'follow', + } - options = { - ...(options ?? {}), - headers: headers, - redirect: 'follow', - } + return envdRpcFetch(url, options).then(rejectProxyUnavailableResponse) + } - return envdRpcFetch(url, options).then(rejectProxyUnavailableResponse) - }, + const rpcTransport = createConnectTransport({ + baseUrl: this.envdApiUrl, + useBinaryFormat: false, + interceptors, + fetch: rpcFetch, + }) + const rpcStreamingTransport = createClientStreamingTransport({ + baseUrl: this.envdApiUrl, + fetch: rpcFetch, + interceptors, }) this.envdApi = new EnvdApiClient( @@ -222,9 +233,15 @@ export class Sandbox extends SandboxApi { this.commands = new Commands( rpcTransport, this.envdApi, - this.connectionConfig + this.connectionConfig, + rpcStreamingTransport + ) + this.pty = new Pty( + rpcTransport, + this.envdApi, + this.connectionConfig, + rpcStreamingTransport ) - this.pty = new Pty(rpcTransport, this.envdApi, this.connectionConfig) this.git = new Git(this.commands) } diff --git a/packages/js-sdk/tests/sandbox/commands/streamInput.test.ts b/packages/js-sdk/tests/sandbox/commands/streamInput.test.ts new file mode 100644 index 0000000000..65304bd2bd --- /dev/null +++ b/packages/js-sdk/tests/sandbox/commands/streamInput.test.ts @@ -0,0 +1,46 @@ +import { expect } from 'vitest' + +import { NotFoundError, TimeoutError } from '../../../src' +import { sandboxTest } from '../../setup' + +sandboxTest('stream input to a command in order', async ({ sandbox }) => { + const cmd = await sandbox.commands.run('cat', { + background: true, + stdin: true, + }) + + const input = sandbox.commands.streamInput(cmd.pid) + await input.send('a') + await input.send('b\n') + await input.send(new TextEncoder().encode('c')) + await input.close() + + for (let i = 0; i < 10; i++) { + if (cmd.stdout === 'ab\nc') { + break + } + await new Promise((resolve) => setTimeout(resolve, 200)) + } + + expect(cmd.stdout).toBe('ab\nc') + await cmd.kill() +}) + +sandboxTest('stream input reports a missing process', async ({ sandbox }) => { + const input = sandbox.commands.streamInput(999999) + await input.send('missing').catch(() => undefined) + + await expect(input.close()).rejects.toBeInstanceOf(NotFoundError) +}) + +sandboxTest('stream input respects its timeout', async ({ sandbox }) => { + const cmd = await sandbox.commands.run('cat', { + background: true, + stdin: true, + }) + const input = sandbox.commands.streamInput(cmd.pid, { timeoutMs: 1000 }) + + await new Promise((resolve) => setTimeout(resolve, 2000)) + await expect(input.close()).rejects.toBeInstanceOf(TimeoutError) + await cmd.kill() +}) diff --git a/packages/js-sdk/tests/sandbox/commands/streamInput.unit.test.ts b/packages/js-sdk/tests/sandbox/commands/streamInput.unit.test.ts new file mode 100644 index 0000000000..9eb0f46ca6 --- /dev/null +++ b/packages/js-sdk/tests/sandbox/commands/streamInput.unit.test.ts @@ -0,0 +1,100 @@ +import type { Client } from '@connectrpc/connect' +import { expect, test } from 'vitest' + +import { Process as ProcessService } from '../../../src/envd/process/process_pb' +import { CommandInputHandle } from '../../../src/sandbox/commands/inputHandle' + +test('stream input preserves order and sends keepalives while idle', async () => { + const messages: any[] = [] + let resolveStarted!: () => void + const started = new Promise((resolve) => { + resolveStarted = resolve + }) + const rpc = { + async streamInput(requests: AsyncIterable) { + for await (const request of requests) { + messages.push(request) + if (request.event.case === 'start') { + resolveStarted() + } + } + return {} + }, + } as unknown as Client + + const input = new CommandInputHandle( + rpc, + 42, + 'stdin', + {}, + async () => true, + 50 + ) + + await started + await input.send('a') + await input.send(new TextEncoder().encode('b')) + await new Promise((resolve) => setTimeout(resolve, 125)) + await input.send('c') + await input.close() + + const events = messages.map((message) => message.event.case) + const firstDataIndex = events.indexOf('data') + const lastDataIndex = events.lastIndexOf('data') + const keepaliveIndex = events.indexOf('keepalive') + const data = messages + .filter((message) => message.event.case === 'data') + .map((message) => + new TextDecoder().decode(message.event.value.input.input.value) + ) + + expect(events[0]).toBe('start') + expect( + events.filter((event) => event === 'keepalive').length + ).toBeGreaterThan(0) + expect(keepaliveIndex).toBeGreaterThan(firstDataIndex) + expect(keepaliveIndex).toBeLessThan(lastDataIndex) + expect(data).toEqual(['a', 'b', 'c']) + await expect(input.send('late')).rejects.toThrow( + 'sandbox.commands.streamInput(pid)` or `sandbox.pty.streamInput(pid)' + ) +}) + +test('stream input snapshots mutable byte arrays when enqueued', async () => { + const messages: any[] = [] + let release!: () => void + const gate = new Promise((resolve) => { + release = resolve + }) + const rpc = { + async streamInput(requests: AsyncIterable) { + await gate + for await (const request of requests) { + messages.push(request) + } + return {} + }, + } as unknown as Client + + const input = new CommandInputHandle( + rpc, + 42, + 'stdin', + {}, + async () => true, + 50 + ) + const buffer = new TextEncoder().encode('a') + + await input.send(buffer) + buffer.set(new TextEncoder().encode('b')) + release() + await input.close() + + const data = messages + .filter((message) => message.event.case === 'data') + .map((message) => + new TextDecoder().decode(message.event.value.input.input.value) + ) + expect(data).toEqual(['a']) +}) diff --git a/packages/js-sdk/tests/sandbox/pty/streamInput.test.ts b/packages/js-sdk/tests/sandbox/pty/streamInput.test.ts new file mode 100644 index 0000000000..064b83d519 --- /dev/null +++ b/packages/js-sdk/tests/sandbox/pty/streamInput.test.ts @@ -0,0 +1,33 @@ +import { expect } from 'vitest' + +import { sandboxTest } from '../../setup' + +sandboxTest( + 'stream input to a PTY preserves character order', + async ({ sandbox }) => { + let output = '' + const decoder = new TextDecoder() + const terminal = await sandbox.pty.create({ + cols: 80, + rows: 24, + onData: (data) => { + output += decoder.decode(data, { stream: true }) + }, + }) + + const input = sandbox.pty.streamInput(terminal.pid) + for (const character of 'echo hello-stream\n') { + await input.send(character) + } + await input.close() + + for (let i = 0; i < 10; i++) { + if (/(?:^|[\r\n])hello-stream(?:[\r\n]|$)/.test(output)) { + break + } + await new Promise((resolve) => setTimeout(resolve, 200)) + } + + expect(output).toMatch(/(?:^|[\r\n])hello-stream(?:[\r\n]|$)/) + } +) diff --git a/packages/python-sdk/e2b/__init__.py b/packages/python-sdk/e2b/__init__.py index 53830c87c2..2682a76160 100644 --- a/packages/python-sdk/e2b/__init__.py +++ b/packages/python-sdk/e2b/__init__.py @@ -114,6 +114,7 @@ SnapshotMode, ) from .sandbox_async.commands.command_handle import AsyncCommandHandle +from .sandbox_async.commands.command_input_handle import AsyncCommandInputHandle from .sandbox_async.filesystem.watch_handle import AsyncWatchHandle from .sandbox_async.main import AsyncSandbox from .sandbox_async.paginator import AsyncSandboxPaginator, AsyncSnapshotPaginator @@ -280,6 +281,7 @@ "AsyncSandbox", "AsyncWatchHandle", "AsyncCommandHandle", + "AsyncCommandInputHandle", # Template "Template", "AsyncTemplate", diff --git a/packages/python-sdk/e2b/connection_config.py b/packages/python-sdk/e2b/connection_config.py index efee15ec61..795053b83a 100644 --- a/packages/python-sdk/e2b/connection_config.py +++ b/packages/python-sdk/e2b/connection_config.py @@ -30,6 +30,7 @@ """ REQUEST_TIMEOUT: float = 60.0 # 60 seconds +STREAM_INPUT_TIMEOUT: float = 60.0 # 60 seconds DEFAULT_RETRIES = 3 # Idle bound for every read on the streaming envd file-transfer transport: diff --git a/packages/python-sdk/e2b/envd/interceptors.py b/packages/python-sdk/e2b/envd/interceptors.py index b50300afaa..63883a0778 100644 --- a/packages/python-sdk/e2b/envd/interceptors.py +++ b/packages/python-sdk/e2b/envd/interceptors.py @@ -79,6 +79,24 @@ async def intercept_unary( self._apply(ctx) return await call_next(request, ctx) + def intercept_client_stream_sync( + self, + call_next: Callable[[Iterator[REQ], RequestContext], RES], + request: Iterator[REQ], + ctx: RequestContext, + ) -> RES: + self._apply(ctx) + return call_next(request, ctx) + + async def intercept_client_stream( + self, + call_next: Callable[[AsyncIterator[REQ], RequestContext], Awaitable[RES]], + request: AsyncIterator[REQ], + ctx: RequestContext, + ) -> RES: + self._apply(ctx) + return await call_next(request, ctx) + def intercept_server_stream_sync( self, call_next: Callable[[REQ, RequestContext], Iterator[RES]], @@ -161,6 +179,36 @@ async def intercept_unary( self._log_response(ctx, None) return response + def intercept_client_stream_sync( + self, + call_next: Callable[[Iterator[REQ], RequestContext], RES], + request: Iterator[REQ], + ctx: RequestContext, + ) -> RES: + self._log_request(ctx) + try: + response = call_next(request, ctx) + except Exception as e: + self._log_response(ctx, e) + raise + self._log_response(ctx, None) + return response + + async def intercept_client_stream( + self, + call_next: Callable[[AsyncIterator[REQ], RequestContext], Awaitable[RES]], + request: AsyncIterator[REQ], + ctx: RequestContext, + ) -> RES: + self._log_request(ctx) + try: + response = await call_next(request, ctx) + except Exception as e: + self._log_response(ctx, e) + raise + self._log_response(ctx, None) + return response + def intercept_server_stream_sync( self, call_next: Callable[[REQ, RequestContext], Iterator[RES]], diff --git a/packages/python-sdk/e2b/sandbox/commands/stream_input.py b/packages/python-sdk/e2b/sandbox/commands/stream_input.py new file mode 100644 index 0000000000..c6309587aa --- /dev/null +++ b/packages/python-sdk/e2b/sandbox/commands/stream_input.py @@ -0,0 +1,37 @@ +from typing import Literal, Union + +from protobuf import Oneof + +from e2b.envd.process import process_pb + +InputKind = Literal["stdin", "pty"] + + +def stream_input_start(pid: int) -> process_pb.StreamInputRequest: + return process_pb.StreamInputRequest( + event=Oneof( + "start", + process_pb.StreamInputRequest.StartEvent( + process=process_pb.ProcessSelector(selector=Oneof("pid", pid)) + ), + ) + ) + + +def stream_input_data( + input_kind: InputKind, data: Union[str, bytes] +) -> process_pb.StreamInputRequest: + payload = data.encode("utf-8") if isinstance(data, str) else data + input_value: Oneof[Literal["stdin"], bytes] | Oneof[Literal["pty"], bytes] + if input_kind == "stdin": + input_value = Oneof("stdin", payload) + else: + input_value = Oneof("pty", payload) + return process_pb.StreamInputRequest( + event=Oneof( + "data", + process_pb.StreamInputRequest.DataEvent( + input=process_pb.ProcessInput(input=input_value) + ), + ) + ) diff --git a/packages/python-sdk/e2b/sandbox_async/commands/command.py b/packages/python-sdk/e2b/sandbox_async/commands/command.py index e7a52b55af..67073f4fd4 100644 --- a/packages/python-sdk/e2b/sandbox_async/commands/command.py +++ b/packages/python-sdk/e2b/sandbox_async/commands/command.py @@ -9,6 +9,7 @@ Username, KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, + STREAM_INPUT_TIMEOUT, ) from protobuf import Oneof @@ -25,6 +26,7 @@ from e2b.exceptions import SandboxException from e2b.sandbox.commands.main import ProcessInfo from e2b.sandbox.commands.command_handle import CommandResult +from e2b.sandbox_async.commands.command_input_handle import AsyncCommandInputHandle from e2b.sandbox_async.commands.command_handle import AsyncCommandHandle, Stderr, Stdout from e2b.sandbox_async.utils import OutputHandler @@ -145,6 +147,32 @@ async def send_stdin( except Exception as e: raise await ahandle_rpc_exception_with_health(e, self._check_health) + def stream_input( + self, + pid: int, + *, + timeout: Optional[float] = STREAM_INPUT_TIMEOUT, + ) -> AsyncCommandInputHandle: + """Open an ordered input handle to command stdin. + + :param pid: Process ID of the command. + :param timeout: Timeout in seconds; defaults to ``STREAM_INPUT_TIMEOUT`` + (60 seconds). ``0`` or ``None`` disables the timeout. + :return: Handle for queuing ordered input with ``send()`` and ending the + stream with ``close()``. + :raises NotFoundException: If the command does not exist; surfaced by + the handle's ``send()`` or ``close()``. + :raises TimeoutException: If the stream timeout expires; surfaced by + the handle's ``send()`` or ``close()``. + """ + return AsyncCommandInputHandle( + self._rpc, + pid, + "stdin", + timeout, + self._check_health, + ) + async def close_stdin( self, pid: int, diff --git a/packages/python-sdk/e2b/sandbox_async/commands/command_input_handle.py b/packages/python-sdk/e2b/sandbox_async/commands/command_input_handle.py new file mode 100644 index 0000000000..ea011a25e5 --- /dev/null +++ b/packages/python-sdk/e2b/sandbox_async/commands/command_input_handle.py @@ -0,0 +1,131 @@ +import asyncio +from typing import Any, AsyncGenerator, Awaitable, Callable, Optional, Union + +from protobuf import Oneof + +from e2b.connection_config import KEEPALIVE_PING_INTERVAL_SEC +from e2b.envd.process import process_pb +from e2b.envd.rpc import ahandle_rpc_exception_with_health +from e2b.envd.utils import timeout_to_ms +from e2b.exceptions import SandboxException +from e2b.sandbox.commands.stream_input import ( + InputKind, + stream_input_data, + stream_input_start, +) + + +class _CloseSignal: + pass + + +_CLOSE = _CloseSignal() +_QueueItem = Union[str, bytes, _CloseSignal] + + +class AsyncCommandInputHandle: + """An ordered, long-lived input handle for a command or PTY.""" + + def __init__( + self, + rpc: Any, + pid: int, + input_kind: InputKind, + timeout: Optional[float], + check_health: Callable[[], Awaitable[Optional[bool]]], + _keepalive_interval: float = KEEPALIVE_PING_INTERVAL_SEC, + ) -> None: + self._rpc = rpc + self._pid = pid + self._input_kind = input_kind + self._timeout = timeout + self._check_health = check_health + self._keepalive_interval = _keepalive_interval + self._queue: asyncio.Queue[_QueueItem] = asyncio.Queue() + self._closing = False + self._stopped = False + self._failure: Optional[Exception] = None + self._task = asyncio.create_task(self._run()) + + async def send(self, data: Union[str, bytes]) -> None: + """Queue input for ordered transmission. + + :param data: Input data to send. + :return: ``None`` once the input has been queued. + :raises NotFoundException: If the command or PTY does not exist. + :raises TimeoutException: If the input stream times out. + :raises SandboxException: If this handle is closed. Open a new one with + ``sandbox.commands.stream_input(pid)`` or + ``sandbox.pty.stream_input(pid)``. + """ + if self._failure is not None: + raise self._failure + if self._closing or self._stopped: + raise SandboxException( + "Cannot send input after the input handle is closed. Open a new one " + "with `sandbox.commands.stream_input(pid)` or " + "`sandbox.pty.stream_input(pid)`." + ) + self._queue.put_nowait(data) + + async def close(self) -> None: + """Finish the input stream and wait for envd to acknowledge it. + + :return: ``None`` when envd acknowledges the end of the input stream. + :raises NotFoundException: If the command or PTY does not exist. + :raises TimeoutException: If the input stream times out. + """ + if not self._closing: + self._closing = True + self._queue.put_nowait(_CLOSE) + await asyncio.shield(self._task) + if self._failure is not None: + raise self._failure + + async def __aenter__(self) -> "AsyncCommandInputHandle": + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.close() + + async def _run(self) -> None: + try: + await self._rpc.stream_input( + self._requests(), + timeout_ms=timeout_to_ms(self._timeout), + ) + except Exception as error: + self._failure = await ahandle_rpc_exception_with_health( + error, self._check_health + ) + self._stopped = True + self._discard_pending_input() + else: + self._stopped = True + + def _discard_pending_input(self) -> None: + while not self._queue.empty(): + try: + self._queue.get_nowait() + except asyncio.QueueEmpty: + break + self._queue.put_nowait(_CLOSE) + + async def _requests(self) -> AsyncGenerator[process_pb.StreamInputRequest, None]: + yield stream_input_start(self._pid) + + while not self._stopped: + try: + data = await asyncio.wait_for( + self._queue.get(), timeout=self._keepalive_interval + ) + except asyncio.TimeoutError: + yield process_pb.StreamInputRequest( + event=Oneof("keepalive", process_pb.StreamInputRequest.KeepAlive()) + ) + continue + + if isinstance(data, _CloseSignal): + return + + yield stream_input_data(self._input_kind, data) diff --git a/packages/python-sdk/e2b/sandbox_async/commands/pty.py b/packages/python-sdk/e2b/sandbox_async/commands/pty.py index d1dd3b3b37..d7c1a4b29a 100644 --- a/packages/python-sdk/e2b/sandbox_async/commands/pty.py +++ b/packages/python-sdk/e2b/sandbox_async/commands/pty.py @@ -13,6 +13,7 @@ ConnectionConfig, KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, + STREAM_INPUT_TIMEOUT, ) from e2b.envd.api import acheck_sandbox_health from e2b.envd.rpc import akill_process, ahandle_rpc_exception_with_health @@ -23,6 +24,7 @@ ) from e2b.envd.client_async import as_stream, create_rpc_client, first_event from e2b.sandbox.commands.command_handle import PtySize +from e2b.sandbox_async.commands.command_input_handle import AsyncCommandInputHandle from e2b.sandbox_async.commands.command_handle import ( AsyncCommandHandle, OutputHandler, @@ -108,6 +110,32 @@ async def send_stdin( except Exception as e: raise await ahandle_rpc_exception_with_health(e, self._check_health) + def stream_input( + self, + pid: int, + *, + timeout: Optional[float] = STREAM_INPUT_TIMEOUT, + ) -> AsyncCommandInputHandle: + """Open an ordered input handle to a PTY. + + :param pid: Process ID of the PTY. + :param timeout: Timeout in seconds; defaults to ``STREAM_INPUT_TIMEOUT`` + (60 seconds). ``0`` or ``None`` disables the timeout. + :return: Handle for queuing ordered input with ``send()`` and ending the + stream with ``close()``. + :raises NotFoundException: If the PTY does not exist; surfaced by the + handle's ``send()`` or ``close()``. + :raises TimeoutException: If the stream timeout expires; surfaced by + the handle's ``send()`` or ``close()``. + """ + return AsyncCommandInputHandle( + self._rpc, + pid, + "pty", + timeout, + self._check_health, + ) + async def create( self, size: PtySize, diff --git a/packages/python-sdk/e2b/sandbox_sync/commands/command.py b/packages/python-sdk/e2b/sandbox_sync/commands/command.py index 82e871144a..be10f32dec 100644 --- a/packages/python-sdk/e2b/sandbox_sync/commands/command.py +++ b/packages/python-sdk/e2b/sandbox_sync/commands/command.py @@ -3,6 +3,7 @@ Callable, Dict, Generator, + Iterable, List, Literal, Optional, @@ -19,6 +20,7 @@ Username, KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, + STREAM_INPUT_TIMEOUT, ) from protobuf import Oneof @@ -32,9 +34,10 @@ ) from e2b.envd.client_sync import as_stream, create_rpc_client from e2b.envd.versions import ENVD_COMMANDS_STDIN, ENVD_ENVD_CLOSE -from e2b.exceptions import SandboxException +from e2b.exceptions import InvalidArgumentException, SandboxException from e2b.sandbox.commands.main import ProcessInfo from e2b.sandbox.commands.command_handle import CommandResult +from e2b.sandbox.commands.stream_input import stream_input_data, stream_input_start from e2b.sandbox_sync.commands.command_handle import CommandHandle @@ -154,6 +157,48 @@ def send_stdin( except Exception as e: raise handle_rpc_exception_with_health(e, self._check_health) + def stream_input( + self, + pid: int, + data: Iterable[Union[str, bytes]], + *, + timeout: Optional[float] = STREAM_INPUT_TIMEOUT, + ) -> None: + """Stream input to command stdin. + + Blocks until ``data`` is exhausted; input is sent as the iterable + yields it, and no keepalives are sent while it blocks. + + :param pid: Process ID of the command. + :param data: Iterable of strings or bytes to send in order. + :param timeout: Timeout in seconds; defaults to ``STREAM_INPUT_TIMEOUT`` + (60 seconds). ``0`` or ``None`` disables the timeout. + :return: ``None`` after all input is sent and envd acknowledges the stream. + :raises NotFoundException: If the command does not exist. + :raises TimeoutException: If the stream timeout expires. + :raises SandboxException: If the streaming request fails. + :raises InvalidArgumentException: If ``data`` is a single ``str`` or + ``bytes`` value instead of an iterable of chunks. + """ + if isinstance(data, (str, bytes, bytearray, memoryview)): + raise InvalidArgumentException( + "`data` must be an iterable of str or bytes chunks, not a single " + "chunk; wrap it in a list, e.g. `stream_input(pid, [data])`." + ) + + def requests() -> Generator[process_pb.StreamInputRequest, None, None]: + yield stream_input_start(pid) + for item in data: + yield stream_input_data("stdin", item) + + try: + self._rpc.stream_input( + requests(), + timeout_ms=timeout_to_ms(timeout), + ) + except Exception as error: + raise handle_rpc_exception_with_health(error, self._check_health) + def close_stdin( self, pid: int, diff --git a/packages/python-sdk/e2b/sandbox_sync/commands/pty.py b/packages/python-sdk/e2b/sandbox_sync/commands/pty.py index d0122d349c..d3cc7d1c67 100644 --- a/packages/python-sdk/e2b/sandbox_sync/commands/pty.py +++ b/packages/python-sdk/e2b/sandbox_sync/commands/pty.py @@ -1,4 +1,4 @@ -from typing import Dict, Optional +from typing import Dict, Generator, Iterable, Optional, Union import httpx @@ -13,6 +13,7 @@ ConnectionConfig, KEEPALIVE_PING_HEADER, KEEPALIVE_PING_INTERVAL_SEC, + STREAM_INPUT_TIMEOUT, ) from e2b.envd.api import check_sandbox_health from e2b.envd.rpc import handle_rpc_exception_with_health, kill_process @@ -22,7 +23,9 @@ timeout_to_ms, ) from e2b.envd.client_sync import as_stream, create_rpc_client +from e2b.exceptions import InvalidArgumentException from e2b.sandbox.commands.command_handle import PtySize +from e2b.sandbox.commands.stream_input import stream_input_data, stream_input_start from e2b.sandbox_sync.commands.command_handle import CommandHandle @@ -104,6 +107,48 @@ def send_stdin( except Exception as e: raise handle_rpc_exception_with_health(e, self._check_health) + def stream_input( + self, + pid: int, + data: Iterable[Union[str, bytes]], + *, + timeout: Optional[float] = STREAM_INPUT_TIMEOUT, + ) -> None: + """Stream input to a PTY. + + Blocks until ``data`` is exhausted; input is sent as the iterable + yields it, and no keepalives are sent while it blocks. + + :param pid: Process ID of the PTY. + :param data: Iterable of strings or bytes to send in order. + :param timeout: Timeout in seconds; defaults to ``STREAM_INPUT_TIMEOUT`` + (60 seconds). ``0`` or ``None`` disables the timeout. + :return: ``None`` after all input is sent and envd acknowledges the stream. + :raises NotFoundException: If the PTY does not exist. + :raises TimeoutException: If the stream timeout expires. + :raises SandboxException: If the streaming request fails. + :raises InvalidArgumentException: If ``data`` is a single ``str`` or + ``bytes`` value instead of an iterable of chunks. + """ + if isinstance(data, (str, bytes, bytearray, memoryview)): + raise InvalidArgumentException( + "`data` must be an iterable of str or bytes chunks, not a single " + "chunk; wrap it in a list, e.g. `stream_input(pid, [data])`." + ) + + def requests() -> Generator[process_pb.StreamInputRequest, None, None]: + yield stream_input_start(pid) + for item in data: + yield stream_input_data("pty", item) + + try: + self._rpc.stream_input( + requests(), + timeout_ms=timeout_to_ms(timeout), + ) + except Exception as error: + raise handle_rpc_exception_with_health(error, self._check_health) + def create( self, size: PtySize, diff --git a/packages/python-sdk/tests/async/sandbox_async/commands/test_stream_input.py b/packages/python-sdk/tests/async/sandbox_async/commands/test_stream_input.py new file mode 100644 index 0000000000..e7eec307c0 --- /dev/null +++ b/packages/python-sdk/tests/async/sandbox_async/commands/test_stream_input.py @@ -0,0 +1,91 @@ +import asyncio + +import pytest + +from e2b import AsyncSandbox, NotFoundException, SandboxException, TimeoutException +from e2b.sandbox_async.commands.command_input_handle import AsyncCommandInputHandle + + +async def test_stream_input_to_command(async_sandbox: AsyncSandbox): + cmd = await async_sandbox.commands.run("cat", background=True, stdin=True) + stream = async_sandbox.commands.stream_input(cmd.pid, timeout=0) + await stream.send("a") + await stream.send("b\n") + await stream.send(b"c") + await stream.close() + + await async_sandbox.commands.close_stdin(cmd.pid) + result = await cmd.wait() + assert result.stdout == "ab\nc" + + +async def test_stream_input_reports_missing_command(async_sandbox: AsyncSandbox): + stream = async_sandbox.commands.stream_input(999999) + try: + await stream.send("data") + except NotFoundException: + pass + + with pytest.raises(NotFoundException): + await stream.close() + + +async def test_stream_input_timeout(async_sandbox: AsyncSandbox): + cmd = await async_sandbox.commands.run("cat", background=True, stdin=True) + stream = async_sandbox.commands.stream_input(cmd.pid, timeout=1) + await asyncio.sleep(2) + + with pytest.raises(TimeoutException): + await stream.close() + + +async def test_async_command_input_handle_orders_data_and_sends_keepalives(): + class FakeRpc: + async def stream_input(self, requests, timeout_ms=None): + assert timeout_ms is None + self.requests = [request async for request in requests] + + rpc = FakeRpc() + stream = AsyncCommandInputHandle( + rpc, 42, "stdin", 0, lambda: None, _keepalive_interval=0.01 + ) + await stream.send("first") + await asyncio.sleep(0.04) + await stream.send(b"second") + await stream.close() + + events = [request.event for request in rpc.requests] + assert events[0].field == "start" + assert [event.field for event in events].count("keepalive") >= 1 + data = [event.value.input.input.value for event in events if event.field == "data"] + assert data == [b"first", b"second"] + await stream.close() + with pytest.raises(SandboxException, match="sandbox.commands.stream_input"): + await stream.send("late") + + +async def test_async_command_input_handle_close_survives_cancelled_wait(): + class FakeRpc: + def __init__(self): + self.gate = asyncio.Event() + + async def stream_input(self, requests, timeout_ms=None): + await self.gate.wait() + self.requests = [request async for request in requests] + + rpc = FakeRpc() + handle = AsyncCommandInputHandle(rpc, 42, "stdin", 0, lambda: None) + await handle.send(b"payload") + + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(handle.close(), 0.01) + + rpc.gate.set() + await handle.close() + + data = [ + request.event.value.input.input.value + for request in rpc.requests + if request.event.field == "data" + ] + assert data == [b"payload"] diff --git a/packages/python-sdk/tests/async/sandbox_async/pty/test_stream_input.py b/packages/python-sdk/tests/async/sandbox_async/pty/test_stream_input.py new file mode 100644 index 0000000000..e4b06b7dbc --- /dev/null +++ b/packages/python-sdk/tests/async/sandbox_async/pty/test_stream_input.py @@ -0,0 +1,18 @@ +from e2b import AsyncSandbox +from e2b.sandbox.commands.command_handle import PtySize + + +async def test_stream_input_to_pty(async_sandbox: AsyncSandbox): + output = [] + terminal = await async_sandbox.pty.create( + PtySize(cols=80, rows=24), + on_data=lambda data: output.append(data.decode("utf-8")), + ) + stream = async_sandbox.pty.stream_input(terminal.pid, timeout=0) + + for char in "echo hello-stream\nexit\n": + await stream.send(char) + await stream.close() + await terminal.wait() + + assert any(line.strip() == "hello-stream" for line in "".join(output).splitlines()) diff --git a/packages/python-sdk/tests/sync/sandbox_sync/commands/test_stream_input.py b/packages/python-sdk/tests/sync/sandbox_sync/commands/test_stream_input.py new file mode 100644 index 0000000000..519316ed1b --- /dev/null +++ b/packages/python-sdk/tests/sync/sandbox_sync/commands/test_stream_input.py @@ -0,0 +1,66 @@ +import time +from typing import Any, cast + +import pytest + +from e2b import InvalidArgumentException, NotFoundException, Sandbox, TimeoutException +from e2b.sandbox_sync.commands.command import Commands + + +def test_stream_input_to_command(sandbox: Sandbox): + cmd = sandbox.commands.run("cat", background=True, stdin=True) + + def input_lines(): + yield "first\n" + time.sleep(1) + yield "second\n" + time.sleep(1) + yield b"third\n" + + sandbox.commands.stream_input(cmd.pid, input_lines(), timeout=0) + + sandbox.commands.close_stdin(cmd.pid) + result = cmd.wait() + assert result.stdout.splitlines() == ["first", "second", "third"] + + +def test_stream_input_reports_missing_command(sandbox: Sandbox): + with pytest.raises(NotFoundException): + sandbox.commands.stream_input(999999, ["data"]) + + +def test_stream_input_rejects_single_chunk(sandbox: Sandbox): + with pytest.raises(InvalidArgumentException, match=r"\[data\]"): + sandbox.commands.stream_input(1, cast(Any, b"ls\n")) + + +def test_stream_input_timeout(sandbox: Sandbox): + cmd = sandbox.commands.run("cat", background=True, stdin=True) + + def slow_input(): + yield "first\n" + time.sleep(2) + yield "second\n" + + with pytest.raises(TimeoutException): + sandbox.commands.stream_input(cmd.pid, slow_input(), timeout=1) + + sandbox.commands.close_stdin(cmd.pid) + cmd.wait() + + +def test_stream_input_orders_start_and_data(): + class FakeRpc: + def stream_input(self, requests, timeout_ms=None): + assert timeout_ms is None + self.requests = list(requests) + + rpc = FakeRpc() + commands = Commands.__new__(Commands) + commands._rpc = rpc + commands._check_health = lambda: None + commands.stream_input(42, ["first", "café", b"second"], timeout=0) + events = [request.event for request in rpc.requests] + assert events[0].field == "start" + data = [event.value.input.input.value for event in events if event.field == "data"] + assert data == [b"first", b"caf\xc3\xa9", b"second"] diff --git a/packages/python-sdk/tests/sync/sandbox_sync/pty/test_stream_input.py b/packages/python-sdk/tests/sync/sandbox_sync/pty/test_stream_input.py new file mode 100644 index 0000000000..2e1e106691 --- /dev/null +++ b/packages/python-sdk/tests/sync/sandbox_sync/pty/test_stream_input.py @@ -0,0 +1,24 @@ +from typing import Any, cast + +import pytest + +from e2b import InvalidArgumentException, Sandbox +from e2b.sandbox.commands.command_handle import PtySize + + +def test_stream_input_to_pty(sandbox: Sandbox): + terminal = sandbox.pty.create(PtySize(cols=80, rows=24)) + + def input_chars(): + yield from "echo hello-stream\nexit\n" + + sandbox.pty.stream_input(terminal.pid, input_chars(), timeout=0) + + output = [] + terminal.wait(on_pty=lambda data: output.append(data.decode("utf-8"))) + assert any(line.strip() == "hello-stream" for line in "".join(output).splitlines()) + + +def test_stream_input_rejects_single_chunk(sandbox: Sandbox): + with pytest.raises(InvalidArgumentException, match=r"\[data\]"): + sandbox.pty.stream_input(1, cast(Any, "ls\n"))