chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,396 @@
|
||||
from typing import Dict, List, Literal, Optional, Union, overload
|
||||
|
||||
import e2b_connect
|
||||
import httpcore
|
||||
import httpx
|
||||
from packaging.version import Version
|
||||
from e2b.connection_config import (
|
||||
ConnectionConfig,
|
||||
Username,
|
||||
KEEPALIVE_PING_HEADER,
|
||||
KEEPALIVE_PING_INTERVAL_SEC,
|
||||
)
|
||||
from e2b.envd.process import process_connect, process_pb2
|
||||
from e2b.envd.api import acheck_sandbox_health
|
||||
from e2b.envd.rpc import authentication_header, ahandle_rpc_exception_with_health
|
||||
from e2b.envd.versions import ENVD_COMMANDS_STDIN, ENVD_ENVD_CLOSE
|
||||
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_handle import AsyncCommandHandle, Stderr, Stdout
|
||||
from e2b.sandbox_async.utils import OutputHandler
|
||||
|
||||
|
||||
class Commands:
|
||||
"""
|
||||
Module for executing commands in the sandbox.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
envd_api_url: str,
|
||||
connection_config: ConnectionConfig,
|
||||
pool: httpcore.AsyncConnectionPool,
|
||||
envd_version: Version,
|
||||
envd_api: httpx.AsyncClient,
|
||||
) -> None:
|
||||
self._connection_config = connection_config
|
||||
self._envd_version = envd_version
|
||||
self._check_health = lambda: acheck_sandbox_health(envd_api)
|
||||
self._rpc = process_connect.ProcessClient(
|
||||
envd_api_url,
|
||||
# TODO: Fix and enable compression again — the headers compression is not solved for streaming.
|
||||
# compressor=e2b_connect.GzipCompressor,
|
||||
async_pool=pool,
|
||||
json=True,
|
||||
headers=connection_config.sandbox_headers,
|
||||
logger=connection_config.logger,
|
||||
)
|
||||
|
||||
async def list(
|
||||
self,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> List[ProcessInfo]:
|
||||
"""
|
||||
Lists all running commands and PTY sessions.
|
||||
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: List of running commands and PTY sessions
|
||||
"""
|
||||
try:
|
||||
res = await self._rpc.alist(
|
||||
process_pb2.ListRequest(),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
return [
|
||||
ProcessInfo(
|
||||
pid=p.pid,
|
||||
tag=p.tag if p.HasField("tag") else None,
|
||||
cmd=p.config.cmd,
|
||||
args=list(p.config.args),
|
||||
envs=dict(p.config.envs),
|
||||
cwd=p.config.cwd if p.config.HasField("cwd") else None,
|
||||
)
|
||||
for p in res.processes
|
||||
]
|
||||
except Exception as e:
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def kill(
|
||||
self,
|
||||
pid: int,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Kill a running command specified by its process ID.
|
||||
It uses `SIGKILL` signal to kill the command.
|
||||
|
||||
:param pid: Process ID of the command. You can get the list of processes using `sandbox.commands.list()`
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: `True` if the command was killed, `False` if the command was not found
|
||||
"""
|
||||
try:
|
||||
await self._rpc.asend_signal(
|
||||
process_pb2.SendSignalRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
signal=process_pb2.Signal.SIGNAL_SIGKILL,
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
if isinstance(e, e2b_connect.ConnectException):
|
||||
if e.status == e2b_connect.Code.not_found:
|
||||
return False
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def send_stdin(
|
||||
self,
|
||||
pid: int,
|
||||
data: Union[str, bytes],
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Send data to command stdin.
|
||||
|
||||
:param pid Process ID of the command. You can get the list of processes using `sandbox.commands.list()`.
|
||||
:param data: Data to send to the command
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
try:
|
||||
await self._rpc.asend_input(
|
||||
process_pb2.SendInputRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
input=process_pb2.ProcessInput(
|
||||
stdin=data.encode() if isinstance(data, str) else data,
|
||||
),
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def close_stdin(
|
||||
self,
|
||||
pid: int,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Close the command stdin.
|
||||
|
||||
This signals EOF to the command. The command must have been started with `stdin=True`.
|
||||
|
||||
:param pid Process ID of the command. You can get the list of processes using `sandbox.commands.list()`.
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
if self._envd_version < ENVD_ENVD_CLOSE:
|
||||
raise SandboxException(
|
||||
f"Sandbox envd version {self._envd_version} doesn't support closing stdin. "
|
||||
f"Please rebuild your template to pick up the latest sandbox version."
|
||||
)
|
||||
|
||||
try:
|
||||
await self._rpc.aclose_stdin(
|
||||
process_pb2.CloseStdinRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
@overload
|
||||
async def run(
|
||||
self,
|
||||
cmd: str,
|
||||
background: Union[Literal[False], None] = None,
|
||||
envs: Optional[Dict[str, str]] = None,
|
||||
user: Optional[Username] = None,
|
||||
cwd: Optional[str] = None,
|
||||
on_stdout: Optional[OutputHandler[Stdout]] = None,
|
||||
on_stderr: Optional[OutputHandler[Stderr]] = None,
|
||||
stdin: Optional[bool] = None,
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> CommandResult:
|
||||
"""
|
||||
Start a new command and wait until it finishes executing.
|
||||
|
||||
:param cmd: Command to execute
|
||||
:param background: **`False` if the command should be executed in the foreground**, `True` if the command should be executed in the background
|
||||
:param envs: Environment variables used for the command
|
||||
:param user: User to run the command as
|
||||
:param cwd: Working directory to run the command
|
||||
:param on_stdout: Callback for command stdout output
|
||||
:param on_stderr: Callback for command stderr output
|
||||
:param stdin: If `True`, the command will have a stdin stream that you can send data to using `sandbox.commands.send_stdin()`
|
||||
:param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: `CommandResult` result of the command execution
|
||||
"""
|
||||
...
|
||||
|
||||
@overload
|
||||
async def run(
|
||||
self,
|
||||
cmd: str,
|
||||
background: Literal[True],
|
||||
envs: Optional[Dict[str, str]] = None,
|
||||
user: Optional[Username] = None,
|
||||
cwd: Optional[str] = None,
|
||||
on_stdout: Optional[OutputHandler[Stdout]] = None,
|
||||
on_stderr: Optional[OutputHandler[Stderr]] = None,
|
||||
stdin: Optional[bool] = None,
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> AsyncCommandHandle:
|
||||
"""
|
||||
Start a new command and return a handle to interact with it.
|
||||
|
||||
:param cmd: Command to execute
|
||||
:param background: `False` if the command should be executed in the foreground, **`True` if the command should be executed in the background**
|
||||
:param envs: Environment variables used for the command
|
||||
:param user: User to run the command as
|
||||
:param cwd: Working directory to run the command
|
||||
:param on_stdout: Callback for command stdout output
|
||||
:param on_stderr: Callback for command stderr output
|
||||
:param stdin: If `True`, the command will have a stdin stream that you can send data to using `sandbox.commands.send_stdin()`
|
||||
:param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: `AsyncCommandHandle` handle to interact with the running command
|
||||
"""
|
||||
...
|
||||
|
||||
async def run(
|
||||
self,
|
||||
cmd: str,
|
||||
background: Union[bool, None] = None,
|
||||
envs: Optional[Dict[str, str]] = None,
|
||||
user: Optional[Username] = None,
|
||||
cwd: Optional[str] = None,
|
||||
on_stdout: Optional[OutputHandler[Stdout]] = None,
|
||||
on_stderr: Optional[OutputHandler[Stderr]] = None,
|
||||
stdin: Optional[bool] = None,
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
):
|
||||
# Check version for stdin support
|
||||
if stdin is False and self._envd_version < ENVD_COMMANDS_STDIN:
|
||||
raise SandboxException(
|
||||
f"Sandbox envd version {self._envd_version} can't specify stdin, it's always turned on. "
|
||||
f"Please rebuild your template if you need this feature."
|
||||
)
|
||||
|
||||
# Default to `False`
|
||||
stdin = stdin or False
|
||||
|
||||
proc = await self._start(
|
||||
cmd,
|
||||
envs,
|
||||
user,
|
||||
cwd,
|
||||
stdin,
|
||||
timeout,
|
||||
request_timeout,
|
||||
on_stdout=on_stdout,
|
||||
on_stderr=on_stderr,
|
||||
)
|
||||
|
||||
return proc if background else await proc.wait()
|
||||
|
||||
async def _start(
|
||||
self,
|
||||
cmd: str,
|
||||
envs: Optional[Dict[str, str]],
|
||||
user: Optional[Username],
|
||||
cwd: Optional[str],
|
||||
stdin: bool,
|
||||
timeout: Optional[float],
|
||||
request_timeout: Optional[float],
|
||||
on_stdout: Optional[OutputHandler[Stdout]],
|
||||
on_stderr: Optional[OutputHandler[Stderr]],
|
||||
) -> AsyncCommandHandle:
|
||||
events = self._rpc.astart(
|
||||
process_pb2.StartRequest(
|
||||
process=process_pb2.ProcessConfig(
|
||||
cmd="/bin/bash",
|
||||
envs=envs,
|
||||
args=["-l", "-c", cmd],
|
||||
cwd=cwd,
|
||||
),
|
||||
stdin=stdin,
|
||||
),
|
||||
headers={
|
||||
**authentication_header(self._envd_version, user),
|
||||
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
|
||||
},
|
||||
timeout=timeout,
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
start_event = await events.__anext__()
|
||||
|
||||
if not start_event.HasField("event"):
|
||||
raise SandboxException(
|
||||
f"Failed to start process: expected start event, got {start_event}"
|
||||
)
|
||||
|
||||
pid = start_event.event.start.pid
|
||||
return AsyncCommandHandle(
|
||||
pid=pid,
|
||||
handle_kill=lambda: self.kill(pid),
|
||||
events=events,
|
||||
on_stdout=on_stdout,
|
||||
on_stderr=on_stderr,
|
||||
handle_send_stdin=lambda data, request_timeout=None: self.send_stdin(
|
||||
pid, data, request_timeout
|
||||
),
|
||||
handle_close_stdin=lambda request_timeout=None: self.close_stdin(
|
||||
pid, request_timeout
|
||||
),
|
||||
check_health=self._check_health,
|
||||
)
|
||||
except Exception as e:
|
||||
try:
|
||||
await events.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def connect(
|
||||
self,
|
||||
pid: int,
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
on_stdout: Optional[OutputHandler[Stdout]] = None,
|
||||
on_stderr: Optional[OutputHandler[Stderr]] = None,
|
||||
) -> AsyncCommandHandle:
|
||||
"""
|
||||
Connects to a running command.
|
||||
You can use `AsyncCommandHandle.wait()` to wait for the command to finish and get execution results.
|
||||
|
||||
:param pid: Process ID of the command to connect to. You can get the list of processes using `sandbox.commands.list()`
|
||||
:param request_timeout: Request timeout in **seconds**
|
||||
:param timeout: Timeout for the command connection in **seconds**. Using `0` will not limit the command connection time
|
||||
:param on_stdout: Callback for command stdout output
|
||||
:param on_stderr: Callback for command stderr output
|
||||
|
||||
:return: `AsyncCommandHandle` handle to interact with the running command
|
||||
"""
|
||||
events = self._rpc.aconnect(
|
||||
process_pb2.ConnectRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
),
|
||||
headers={
|
||||
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
|
||||
},
|
||||
timeout=timeout,
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
start_event = await events.__anext__()
|
||||
|
||||
if not start_event.HasField("event"):
|
||||
raise SandboxException(
|
||||
f"Failed to connect to process: expected start event, got {start_event}"
|
||||
)
|
||||
|
||||
pid = start_event.event.start.pid
|
||||
return AsyncCommandHandle(
|
||||
pid=pid,
|
||||
handle_kill=lambda: self.kill(pid),
|
||||
events=events,
|
||||
on_stdout=on_stdout,
|
||||
on_stderr=on_stderr,
|
||||
handle_send_stdin=lambda data, request_timeout=None: self.send_stdin(
|
||||
pid, data, request_timeout
|
||||
),
|
||||
handle_close_stdin=lambda request_timeout=None: self.close_stdin(
|
||||
pid, request_timeout
|
||||
),
|
||||
check_health=self._check_health,
|
||||
)
|
||||
except Exception as e:
|
||||
try:
|
||||
await events.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
@@ -0,0 +1,298 @@
|
||||
import asyncio
|
||||
import codecs
|
||||
import inspect
|
||||
from typing import (
|
||||
Optional,
|
||||
Callable,
|
||||
Any,
|
||||
AsyncGenerator,
|
||||
List,
|
||||
Awaitable,
|
||||
Union,
|
||||
Tuple,
|
||||
Coroutine,
|
||||
)
|
||||
|
||||
from e2b.envd.rpc import ahandle_rpc_exception_with_health
|
||||
from e2b.envd.process import process_pb2
|
||||
from e2b.exceptions import SandboxException
|
||||
from e2b.sandbox.commands.command_handle import (
|
||||
CommandExitException,
|
||||
CommandResult,
|
||||
Stderr,
|
||||
Stdout,
|
||||
PtyOutput,
|
||||
)
|
||||
from e2b.sandbox_async.utils import OutputHandler
|
||||
|
||||
|
||||
class AsyncCommandHandle:
|
||||
"""
|
||||
Command execution handle.
|
||||
|
||||
It provides methods for waiting for the command to finish, retrieving stdout/stderr, and killing the command.
|
||||
"""
|
||||
|
||||
@property
|
||||
def pid(self):
|
||||
"""
|
||||
Command process ID.
|
||||
"""
|
||||
return self._pid
|
||||
|
||||
@property
|
||||
def stdout(self):
|
||||
"""
|
||||
Command stdout output.
|
||||
"""
|
||||
return "".join(self._stdout_chunks)
|
||||
|
||||
@property
|
||||
def stderr(self):
|
||||
"""
|
||||
Command stderr output.
|
||||
"""
|
||||
return "".join(self._stderr_chunks)
|
||||
|
||||
@property
|
||||
def error(self):
|
||||
"""
|
||||
Command execution error message.
|
||||
"""
|
||||
if self._result is None:
|
||||
return None
|
||||
return self._result.error
|
||||
|
||||
@property
|
||||
def exit_code(self):
|
||||
"""
|
||||
Command execution exit code.
|
||||
|
||||
`0` if the command finished successfully.
|
||||
|
||||
It is `None` if the command is still running.
|
||||
"""
|
||||
if self._result is None:
|
||||
return None
|
||||
return self._result.exit_code
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pid: int,
|
||||
handle_kill: Callable[[], Coroutine[Any, Any, bool]],
|
||||
events: AsyncGenerator[
|
||||
Union[process_pb2.StartResponse, process_pb2.ConnectResponse], Any
|
||||
],
|
||||
on_stdout: Optional[OutputHandler[Stdout]] = None,
|
||||
on_stderr: Optional[OutputHandler[Stderr]] = None,
|
||||
on_pty: Optional[OutputHandler[PtyOutput]] = None,
|
||||
handle_send_stdin: Optional[
|
||||
Callable[[Union[str, bytes], Optional[float]], Coroutine[Any, Any, None]]
|
||||
] = None,
|
||||
handle_close_stdin: Optional[
|
||||
Callable[[Optional[float]], Coroutine[Any, Any, None]]
|
||||
] = None,
|
||||
check_health: Optional[Callable[[], Awaitable[Optional[bool]]]] = None,
|
||||
):
|
||||
self._pid = pid
|
||||
self._handle_kill = handle_kill
|
||||
self._handle_send_stdin = handle_send_stdin
|
||||
self._handle_close_stdin = handle_close_stdin
|
||||
self._check_health = check_health
|
||||
self._events = events
|
||||
|
||||
self._stdout_chunks: List[str] = []
|
||||
self._stderr_chunks: List[str] = []
|
||||
|
||||
self._stdout_decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
|
||||
self._stderr_decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
|
||||
|
||||
self._on_stdout = on_stdout
|
||||
self._on_stderr = on_stderr
|
||||
self._on_pty = on_pty
|
||||
|
||||
self._result: Optional[CommandResult] = None
|
||||
self._iteration_exception: Optional[Exception] = None
|
||||
|
||||
self._wait = asyncio.create_task(self._handle_events())
|
||||
|
||||
def _flush_decoders(
|
||||
self,
|
||||
) -> List[Union[Tuple[Stdout, None, None], Tuple[None, Stderr, None]]]:
|
||||
"""
|
||||
Flush any bytes still buffered in the stream decoders.
|
||||
|
||||
Incomplete trailing UTF-8 sequences are emitted as replacement
|
||||
characters, matching the per-chunk decoding behavior.
|
||||
"""
|
||||
events: List[Union[Tuple[Stdout, None, None], Tuple[None, Stderr, None]]] = []
|
||||
out = self._stdout_decoder.decode(b"", final=True)
|
||||
if out:
|
||||
self._stdout_chunks.append(out)
|
||||
events.append((out, None, None))
|
||||
err = self._stderr_decoder.decode(b"", final=True)
|
||||
if err:
|
||||
self._stderr_chunks.append(err)
|
||||
events.append((None, err, None))
|
||||
return events
|
||||
|
||||
async def _iterate_events(
|
||||
self,
|
||||
) -> AsyncGenerator[
|
||||
Union[
|
||||
Tuple[Stdout, None, None],
|
||||
Tuple[None, Stderr, None],
|
||||
Tuple[None, None, PtyOutput],
|
||||
],
|
||||
None,
|
||||
]:
|
||||
try:
|
||||
async for event in self._events:
|
||||
if event.event.HasField("data"):
|
||||
if event.event.data.stdout:
|
||||
out = self._stdout_decoder.decode(event.event.data.stdout)
|
||||
if out:
|
||||
self._stdout_chunks.append(out)
|
||||
yield out, None, None
|
||||
if event.event.data.stderr:
|
||||
out = self._stderr_decoder.decode(event.event.data.stderr)
|
||||
if out:
|
||||
self._stderr_chunks.append(out)
|
||||
yield None, out, None
|
||||
if event.event.data.pty:
|
||||
yield None, None, event.event.data.pty
|
||||
if event.event.HasField("end"):
|
||||
# Flush trailing decoder bytes into the accumulators and
|
||||
# record the result before yielding the flushed chunks, so a
|
||||
# consumer that stops iterating on the first flushed chunk
|
||||
# still observes the exit code.
|
||||
flushed = list(self._flush_decoders())
|
||||
self._result = CommandResult(
|
||||
stdout="".join(self._stdout_chunks),
|
||||
stderr="".join(self._stderr_chunks),
|
||||
exit_code=event.event.end.exit_code,
|
||||
error=event.event.end.error,
|
||||
)
|
||||
for f in flushed:
|
||||
yield f
|
||||
except Exception:
|
||||
# The stream raised before an end event (e.g. disconnect or RPC
|
||||
# failure). Flush any bytes still buffered in the decoders so
|
||||
# incomplete trailing sequences surface as replacement characters
|
||||
# instead of being silently dropped, then re-raise so the error is
|
||||
# still surfaced by the consumer.
|
||||
for flushed in self._flush_decoders():
|
||||
yield flushed
|
||||
raise
|
||||
|
||||
# If the stream closed without an end event (e.g. disconnect or a
|
||||
# dropped connection), flush any bytes still buffered in the decoders
|
||||
# so incomplete trailing sequences surface as replacement characters
|
||||
# instead of being silently dropped.
|
||||
if self._result is None:
|
||||
for flushed in self._flush_decoders():
|
||||
yield flushed
|
||||
|
||||
async def disconnect(self) -> None:
|
||||
"""
|
||||
Disconnects from the command.
|
||||
|
||||
The command is not killed, but SDK stops receiving events from the command.
|
||||
You can reconnect to the command using `sandbox.commands.connect` method.
|
||||
"""
|
||||
self._wait.cancel()
|
||||
await asyncio.wait([self._wait])
|
||||
try:
|
||||
await self._events.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
async def _handle_events(self):
|
||||
try:
|
||||
async for stdout, stderr, pty in self._iterate_events():
|
||||
if stdout is not None and self._on_stdout:
|
||||
cb = self._on_stdout(stdout)
|
||||
if inspect.isawaitable(cb):
|
||||
await cb
|
||||
elif stderr is not None and self._on_stderr:
|
||||
cb = self._on_stderr(stderr)
|
||||
if inspect.isawaitable(cb):
|
||||
await cb
|
||||
elif pty is not None and self._on_pty:
|
||||
cb = self._on_pty(pty)
|
||||
if inspect.isawaitable(cb):
|
||||
await cb
|
||||
except StopAsyncIteration:
|
||||
pass
|
||||
except Exception as e:
|
||||
self._iteration_exception = await ahandle_rpc_exception_with_health(
|
||||
e, self._check_health
|
||||
)
|
||||
|
||||
async def wait(self) -> CommandResult:
|
||||
"""
|
||||
Wait for the command to finish and return the result.
|
||||
If the command exits with a non-zero exit code, it throws a `CommandExitException`.
|
||||
|
||||
:return: `CommandResult` result of command execution
|
||||
"""
|
||||
await self._wait
|
||||
if self._iteration_exception:
|
||||
raise self._iteration_exception
|
||||
|
||||
if self._result is None:
|
||||
raise Exception("Command ended without an end event")
|
||||
|
||||
if self._result.exit_code != 0:
|
||||
raise CommandExitException(
|
||||
stdout="".join(self._stdout_chunks),
|
||||
stderr="".join(self._stderr_chunks),
|
||||
exit_code=self._result.exit_code,
|
||||
error=self._result.error,
|
||||
)
|
||||
|
||||
return self._result
|
||||
|
||||
async def kill(self) -> bool:
|
||||
"""
|
||||
Kills the command.
|
||||
|
||||
It uses `SIGKILL` signal to kill the command
|
||||
|
||||
:return: `True` if the command was killed successfully, `False` if the command was not found
|
||||
"""
|
||||
result = await self._handle_kill()
|
||||
return result
|
||||
|
||||
async def send_stdin(
|
||||
self,
|
||||
data: Union[str, bytes],
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Send data to the command stdin.
|
||||
|
||||
The command must have been started with `stdin=True`.
|
||||
|
||||
:param data: Data to send to the command
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
if self._handle_send_stdin is None:
|
||||
raise SandboxException(
|
||||
"Sending stdin is not supported for this command handle."
|
||||
)
|
||||
await self._handle_send_stdin(data, request_timeout)
|
||||
|
||||
async def close_stdin(self, request_timeout: Optional[float] = None) -> None:
|
||||
"""
|
||||
Close the command stdin.
|
||||
|
||||
This signals EOF to the command. The command must have been started with `stdin=True`.
|
||||
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
if self._handle_close_stdin is None:
|
||||
raise SandboxException(
|
||||
"Closing stdin is not supported for this command handle."
|
||||
)
|
||||
await self._handle_close_stdin(request_timeout)
|
||||
@@ -0,0 +1,257 @@
|
||||
from typing import Dict, Optional
|
||||
|
||||
import e2b_connect
|
||||
import httpcore
|
||||
import httpx
|
||||
|
||||
from packaging.version import Version
|
||||
from e2b.envd.process import process_connect, process_pb2
|
||||
from e2b.connection_config import (
|
||||
Username,
|
||||
ConnectionConfig,
|
||||
KEEPALIVE_PING_HEADER,
|
||||
KEEPALIVE_PING_INTERVAL_SEC,
|
||||
)
|
||||
from e2b.exceptions import SandboxException
|
||||
from e2b.envd.api import acheck_sandbox_health
|
||||
from e2b.envd.rpc import authentication_header, ahandle_rpc_exception_with_health
|
||||
from e2b.sandbox.commands.command_handle import PtySize
|
||||
from e2b.sandbox_async.commands.command_handle import (
|
||||
AsyncCommandHandle,
|
||||
OutputHandler,
|
||||
PtyOutput,
|
||||
)
|
||||
|
||||
|
||||
class Pty:
|
||||
"""
|
||||
Module for interacting with PTYs (pseudo-terminals) in the sandbox.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
envd_api_url: str,
|
||||
connection_config: ConnectionConfig,
|
||||
pool: httpcore.AsyncConnectionPool,
|
||||
envd_version: Version,
|
||||
envd_api: httpx.AsyncClient,
|
||||
) -> None:
|
||||
self._connection_config = connection_config
|
||||
self._envd_version = envd_version
|
||||
self._check_health = lambda: acheck_sandbox_health(envd_api)
|
||||
self._rpc = process_connect.ProcessClient(
|
||||
envd_api_url,
|
||||
# TODO: Fix and enable compression again — the headers compression is not solved for streaming.
|
||||
# compressor=e2b_connect.GzipCompressor,
|
||||
async_pool=pool,
|
||||
json=True,
|
||||
headers=connection_config.sandbox_headers,
|
||||
logger=connection_config.logger,
|
||||
)
|
||||
|
||||
async def kill(
|
||||
self,
|
||||
pid: int,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> bool:
|
||||
"""
|
||||
Kill PTY.
|
||||
|
||||
:param pid: Process ID of the PTY
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: `true` if the PTY was killed, `false` if the PTY was not found
|
||||
"""
|
||||
try:
|
||||
await self._rpc.asend_signal(
|
||||
process_pb2.SendSignalRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
signal=process_pb2.Signal.SIGNAL_SIGKILL,
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
return True
|
||||
except Exception as e:
|
||||
if isinstance(e, e2b_connect.ConnectException):
|
||||
if e.status == e2b_connect.Code.not_found:
|
||||
return False
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def send_stdin(
|
||||
self,
|
||||
pid: int,
|
||||
data: bytes,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Send input to a PTY.
|
||||
|
||||
:param pid: Process ID of the PTY
|
||||
:param data: Input data to send
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
try:
|
||||
await self._rpc.asend_input(
|
||||
process_pb2.SendInputRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
input=process_pb2.ProcessInput(
|
||||
pty=data,
|
||||
),
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def create(
|
||||
self,
|
||||
size: PtySize,
|
||||
on_data: OutputHandler[PtyOutput],
|
||||
user: Optional[Username] = None,
|
||||
cwd: Optional[str] = None,
|
||||
envs: Optional[Dict[str, str]] = None,
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> AsyncCommandHandle:
|
||||
"""
|
||||
Start a new PTY (pseudo-terminal).
|
||||
|
||||
:param size: Size of the PTY
|
||||
:param on_data: Callback to handle PTY data
|
||||
:param user: User to use for the PTY
|
||||
:param cwd: Working directory for the PTY
|
||||
:param envs: Environment variables for the PTY
|
||||
:param timeout: Timeout for the PTY in **seconds**
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: Handle to interact with the PTY
|
||||
"""
|
||||
envs = dict(envs) if envs else {}
|
||||
envs.setdefault("TERM", "xterm-256color")
|
||||
envs.setdefault("LANG", "C.UTF-8")
|
||||
envs.setdefault("LC_ALL", "C.UTF-8")
|
||||
events = self._rpc.astart(
|
||||
process_pb2.StartRequest(
|
||||
process=process_pb2.ProcessConfig(
|
||||
cmd="/bin/bash",
|
||||
envs=envs,
|
||||
args=["-i", "-l"],
|
||||
cwd=cwd,
|
||||
),
|
||||
pty=process_pb2.PTY(
|
||||
size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols)
|
||||
),
|
||||
),
|
||||
headers={
|
||||
**authentication_header(self._envd_version, user),
|
||||
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
|
||||
},
|
||||
timeout=timeout,
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
start_event = await events.__anext__()
|
||||
|
||||
if not start_event.HasField("event"):
|
||||
raise SandboxException(
|
||||
f"Failed to start process: expected start event, got {start_event}"
|
||||
)
|
||||
|
||||
return AsyncCommandHandle(
|
||||
pid=start_event.event.start.pid,
|
||||
handle_kill=lambda: self.kill(start_event.event.start.pid),
|
||||
events=events,
|
||||
on_pty=on_data,
|
||||
check_health=self._check_health,
|
||||
)
|
||||
except Exception as e:
|
||||
try:
|
||||
await events.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def connect(
|
||||
self,
|
||||
pid: int,
|
||||
on_data: OutputHandler[PtyOutput],
|
||||
timeout: Optional[float] = 60,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> AsyncCommandHandle:
|
||||
"""
|
||||
Connect to a running PTY.
|
||||
|
||||
:param pid: Process ID of the PTY to connect to. You can get the list of running PTYs using `sandbox.pty.list()`.
|
||||
:param on_data: Callback to handle PTY data
|
||||
:param timeout: Timeout for the PTY connection in **seconds**. Using `0` will not limit the connection time
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
|
||||
:return: Handle to interact with the PTY
|
||||
"""
|
||||
events = self._rpc.aconnect(
|
||||
process_pb2.ConnectRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
),
|
||||
timeout=timeout,
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
headers={
|
||||
KEEPALIVE_PING_HEADER: str(KEEPALIVE_PING_INTERVAL_SEC),
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
start_event = await events.__anext__()
|
||||
|
||||
if not start_event.HasField("event"):
|
||||
raise SandboxException(
|
||||
f"Failed to connect to process: expected start event, got {start_event}"
|
||||
)
|
||||
|
||||
return AsyncCommandHandle(
|
||||
pid=start_event.event.start.pid,
|
||||
handle_kill=lambda: self.kill(start_event.event.start.pid),
|
||||
events=events,
|
||||
on_pty=on_data,
|
||||
check_health=self._check_health,
|
||||
)
|
||||
except Exception as e:
|
||||
try:
|
||||
await events.aclose()
|
||||
except Exception:
|
||||
pass
|
||||
raise await ahandle_rpc_exception_with_health(e, self._check_health)
|
||||
|
||||
async def resize(
|
||||
self,
|
||||
pid: int,
|
||||
size: PtySize,
|
||||
request_timeout: Optional[float] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Resize PTY.
|
||||
Call this when the terminal window is resized and the number of columns and rows has changed.
|
||||
|
||||
:param pid: Process ID of the PTY
|
||||
:param size: New size of the PTY
|
||||
:param request_timeout: Timeout for the request in **seconds**
|
||||
"""
|
||||
await self._rpc.aupdate(
|
||||
process_pb2.UpdateRequest(
|
||||
process=process_pb2.ProcessSelector(pid=pid),
|
||||
pty=process_pb2.PTY(
|
||||
size=process_pb2.PTY.Size(rows=size.rows, cols=size.cols),
|
||||
),
|
||||
),
|
||||
request_timeout=self._connection_config.get_request_timeout(
|
||||
request_timeout
|
||||
),
|
||||
)
|
||||
Reference in New Issue
Block a user