loopix-sdk 2.30.0__py3-none-any.whl

loopix/sandbox_async/commands/command_handle.py

@@ -0,0 +1,298 @@
+import asyncio
+import codecs
+import inspect
+from typing import (
+    Optional,
+    Callable,
+    Any,
+    AsyncGenerator,
+    List,
+    Awaitable,
+    Union,
+    Tuple,
+    Coroutine,
+)
+
+from loopix.envd.rpc import ahandle_rpc_exception_with_health
+from loopix.envd.process import process_pb2
+from loopix.exceptions import SandboxException
+from loopix.sandbox.commands.command_handle import (
+    CommandExitException,
+    CommandResult,
+    Stderr,
+    Stdout,
+    PtyOutput,
+)
+from loopix.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)

loopix/sandbox_async/commands/pty.py

@@ -0,0 +1,257 @@
+from typing import Dict, Optional
+
+import loopix_connect
+import httpcore
+import httpx
+
+from packaging.version import Version
+from loopix.envd.process import process_connect, process_pb2
+from loopix.connection_config import (
+    Username,
+    ConnectionConfig,
+    KEEPALIVE_PING_HEADER,
+    KEEPALIVE_PING_INTERVAL_SEC,
+)
+from loopix.exceptions import SandboxException
+from loopix.envd.api import acheck_sandbox_health
+from loopix.envd.rpc import authentication_header, ahandle_rpc_exception_with_health
+from loopix.sandbox.commands.command_handle import PtySize
+from loopix.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=loopix_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, loopix_connect.ConnectException):
+                if e.status == loopix_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
+            ),
+        )