scalebox-sdk 0.1.0__py3-none-any.whl

scalebox/sandbox_async/sandbox_api.py

@@ -0,0 +1,365 @@
+import datetime
+import urllib.parse
+from typing import Dict, List, Optional
+
+from packaging.version import Version
+
+from ..api import AsyncApiClient, SandboxCreateResponse, handle_api_exception
+from ..api.client.api.sandboxes import (
+    delete_sandboxes_sandbox_id,
+    get_sandboxes,
+    get_sandboxes_sandbox_id,
+    get_sandboxes_sandbox_id_metrics,
+    post_sandboxes,
+    post_sandboxes_sandbox_id_timeout,
+)
+from ..api.client.models import Error, NewSandbox, PostSandboxesSandboxIDTimeoutBody
+from ..connection_config import ConnectionConfig, ProxyTypes
+from ..exceptions import SandboxException, TemplateException
+from ..sandbox.sandbox_api import (
+    ListedSandbox,
+    SandboxApiBase,
+    SandboxInfo,
+    SandboxMetrics,
+    SandboxQuery,
+)
+
+
+class SandboxApi(SandboxApiBase):
+    @classmethod
+    async def list(
+        cls,
+        api_key: Optional[str] = None,
+        query: Optional[SandboxQuery] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+    ) -> List[ListedSandbox]:
+        """
+        List all running sandboxes.
+
+        :param api_key: API key to use for authentication, defaults to `E2B_API_KEY` environment variable
+        :param query: Filter the list of sandboxes, e.g. by metadata `SandboxQuery(metadata={"key": "value"})`, if there are multiple filters they are combined with AND.
+        :param domain: Domain to use for the request, only relevant for self-hosted environments
+        :param debug: Enable debug mode, all requested are then sent to localhost
+        :param request_timeout: Timeout for the request in **seconds**
+        :param headers: Additional headers to send with the request
+        :param proxy: Proxy to use for the request
+
+        :return: List of running sandboxes
+        """
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            request_timeout=request_timeout,
+            headers=headers,
+            proxy=proxy,
+        )
+
+        # Convert filters to the format expected by the API
+        metadata = None
+        if query:
+            if query.metadata:
+                quoted_metadata = {
+                    urllib.parse.quote(k): urllib.parse.quote(v)
+                    for k, v in query.metadata.items()
+                }
+                metadata = urllib.parse.urlencode(quoted_metadata)
+
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
+            res = await get_sandboxes.asyncio_detailed(
+                client=api_client,
+                metadata=metadata,
+            )
+
+        if res.status_code >= 300:
+            raise handle_api_exception(res)
+
+        if res.parsed is None:
+            return []
+
+        return [
+            ListedSandbox(
+                sandbox_id=sandbox.sandbox_id,
+                template_id=sandbox.template_id,
+                name=sandbox.alias if isinstance(sandbox.alias, str) else None,
+                metadata=(
+                    sandbox.metadata if isinstance(sandbox.metadata, dict) else {}
+                ),
+                state=sandbox.state,
+                cpu_count=sandbox.cpu_count,
+                memory_mb=sandbox.memory_mb,
+                started_at=sandbox.started_at,
+                end_at=sandbox.end_at,
+            )
+            for sandbox in res.parsed
+        ]
+
+    @classmethod
+    async def _cls_get_info(
+        cls,
+        sandbox_id: str,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+    ) -> SandboxInfo:
+        """
+        Get the sandbox info.
+        :param sandbox_id: Sandbox ID
+        :param api_key: API key to use for authentication, defaults to `E2B_API_KEY` environment variable
+        :param domain: Domain to use for the request, defaults to `E2B_DOMAIN` environment variable
+        :param debug: Debug mode, defaults to `E2B_DEBUG` environment variable
+        :param request_timeout: Timeout for the request in **seconds**
+        :param headers: Additional headers to send with the request
+        :param proxy: Proxy to use for the request
+
+        :return: Sandbox info
+        """
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            request_timeout=request_timeout,
+            headers=headers,
+            proxy=proxy,
+        )
+
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
+            res = await get_sandboxes_sandbox_id.asyncio_detailed(
+                sandbox_id,
+                client=api_client,
+            )
+
+            if res.status_code >= 300:
+                raise handle_api_exception(res)
+
+            if res.parsed is None:
+                raise Exception("Body of the request is None")
+
+            return SandboxInfo(
+                sandbox_id=res.parsed.sandbox_id,
+                sandbox_domain=res.parsed.domain,
+                template_id=res.parsed.template_id,
+                name=res.parsed.alias if isinstance(res.parsed.alias, str) else None,
+                metadata=(
+                    res.parsed.metadata if isinstance(res.parsed.metadata, dict) else {}
+                ),
+                started_at=res.parsed.started_at,
+                end_at=res.parsed.end_at,
+                envd_version=res.parsed.envd_version,
+                _envd_access_token=res.parsed.envd_access_token,
+            )
+
+    @classmethod
+    async def _cls_kill(
+        cls,
+        sandbox_id: str,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+    ) -> bool:
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            request_timeout=request_timeout,
+            headers=headers,
+            proxy=proxy,
+        )
+
+        if config.debug:
+            # Skip killing the sandbox in debug mode
+            return True
+
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
+            res = await delete_sandboxes_sandbox_id.asyncio_detailed(
+                sandbox_id,
+                client=api_client,
+            )
+
+            if res.status_code == 404:
+                return False
+
+            if res.status_code >= 300:
+                raise handle_api_exception(res)
+
+            return True
+
+    @classmethod
+    async def _cls_set_timeout(
+        cls,
+        sandbox_id: str,
+        timeout: int,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+    ) -> None:
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            request_timeout=request_timeout,
+            headers=headers,
+            proxy=proxy,
+        )
+
+        if config.debug:
+            # Skip setting the timeout in debug mode
+            return
+
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
+            res = await post_sandboxes_sandbox_id_timeout.asyncio_detailed(
+                sandbox_id,
+                client=api_client,
+                body=PostSandboxesSandboxIDTimeoutBody(timeout=timeout),
+            )
+
+            if res.status_code >= 300:
+                raise handle_api_exception(res)
+
+    @classmethod
+    async def _create_sandbox(
+        cls,
+        template: str,
+        timeout: int,
+        metadata: Optional[Dict[str, str]] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        secure: Optional[bool] = None,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+        allow_internet_access: Optional[bool] = True,
+    ) -> SandboxCreateResponse:
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            request_timeout=request_timeout,
+            headers=headers,
+            proxy=proxy,
+        )
+
+        async with AsyncApiClient(
+            config,
+            limits=SandboxApiBase._limits,
+        ) as api_client:
+            res = await post_sandboxes.asyncio_detailed(
+                body=NewSandbox(
+                    template_id=template,
+                    metadata=metadata or {},
+                    timeout=timeout,
+                    env_vars=env_vars or {},
+                    secure=secure or False,
+                    allow_internet_access=allow_internet_access,
+                    is_async=False,
+                ),
+                client=api_client,
+            )
+
+            if res.status_code >= 300:
+                raise handle_api_exception(res)
+
+            if res.parsed is None:
+                raise Exception("Body of the request is None")
+
+            # if Version(res.parsed.envd_version) < Version("0.1.0"):
+            #     await SandboxApi._cls_kill(res.parsed.sandbox_id)
+            #     raise TemplateException(
+            #         "You need to update the template to use the new SDK. "
+            #         "You can do this by running `e2b template build` in the directory with the template."
+            #     )
+
+            return SandboxCreateResponse(
+                sandbox_id=res.parsed.sandbox_id,
+                sandbox_domain=res.parsed.domain,
+                envd_version=res.parsed.envd_version,
+                envd_access_token=res.parsed.envd_access_token,
+            )
+
+    @classmethod
+    async def _cls_get_metrics(
+        cls,
+        sandbox_id: str,
+        start: Optional[datetime.datetime] = None,
+        end: Optional[datetime.datetime] = None,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        debug: Optional[bool] = None,
+        request_timeout: Optional[float] = None,
+        headers: Optional[Dict[str, str]] = None,
+        proxy: Optional[ProxyTypes] = None,
+    ) -> List[SandboxMetrics]:
+        config = ConnectionConfig(
+            api_key=api_key,
+            domain=domain,
+            debug=debug,
+            headers=headers,
+            request_timeout=request_timeout,
+            proxy=proxy,
+        )
+
+        if config.debug:
+            # Skip getting the metrics in debug mode
+            return []
+
+        async with AsyncApiClient(
+            config,
+            limits=cls._limits,
+        ) as api_client:
+            res = await get_sandboxes_sandbox_id_metrics.asyncio_detailed(
+                sandbox_id,
+                start=int(start.timestamp() * 1000) if start else None,
+                end=int(end.timestamp() * 1000) if end else None,
+                client=api_client,
+            )
+
+            if res.status_code >= 300:
+                raise handle_api_exception(res)
+
+            if res.parsed is None:
+                return []
+
+            # Check if res.parse is Error
+            if isinstance(res.parsed, Error):
+                raise SandboxException(f"{res.parsed.message}: Request failed")
+
+            # Convert to typed SandboxMetrics objects
+            return [
+                SandboxMetrics(
+                    cpu_count=metric.cpu_count,
+                    cpu_used_pct=metric.cpu_used_pct,
+                    disk_total=metric.disk_total,
+                    disk_used=metric.disk_used,
+                    mem_total=metric.mem_total,
+                    mem_used=metric.mem_used,
+                    timestamp=metric.timestamp,
+                )
+                for metric in res.parsed
+            ]

scalebox/sandbox_async/utils.py

@@ -0,0 +1,7 @@
+from typing import Awaitable, Callable, TypeVar, Union
+
+T = TypeVar("T")
+OutputHandler = Union[
+    Callable[[T], None],
+    Callable[[T], Awaitable[None]],
+]

scalebox/sandbox_sync/__init__.py

@@ -0,0 +1,2 @@
+from .commands.command_handle import CommandExitException, CommandHandle, CommandResult
+from .main import Sandbox

scalebox/sandbox_sync/commands/command.py

@@ -0,0 +1,300 @@
+from typing import Callable, Dict, List, Literal, Optional, Union, overload
+
+import httpcore
+import urllib3
+
+from ... import csx_connect
+from ...connection_config import (
+    KEEPALIVE_PING_HEADER,
+    KEEPALIVE_PING_INTERVAL_SEC,
+    ConnectionConfig,
+    Username,
+)
+from ...exceptions import SandboxException
+from ...generated import api_pb2, api_pb2_connect
+from ...generated.rpc import authentication_header, handle_rpc_exception
+from ...sandbox.commands.command_handle import CommandResult
+from ...sandbox.commands.main import ProcessInfo
+from ...sandbox_sync.commands.command_handle import CommandHandle
+
+
+class Commands:
+    """
+    Module for executing commands in the sandbox.
+    """
+
+    def __init__(
+        self,
+        envd_api_url: str,
+        connection_config: ConnectionConfig,
+        pool: urllib3.PoolManager,
+    ) -> None:
+        self._connection_config = connection_config
+        self._rpc = api_pb2_connect.ProcessClient(
+            envd_api_url,
+            http_client=pool,
+        )
+        self._headers = connection_config.headers
+        self._pool = pool
+
+    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 = self._rpc.list(
+                api_pb2.ListRequest(),
+                self._headers,
+                timeout_seconds=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+            return [
+                ProcessInfo(
+                    pid=p.pid,
+                    tag=p.tag,
+                    cmd=p.config.cmd,
+                    args=list(p.config.args),
+                    envs=dict(p.config.envs),
+                    cwd=p.config.cwd,
+                )
+                for p in res.processes
+            ]
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def kill(
+        self,
+        pid: int,
+        request_timeout: Optional[float] = None,
+    ) -> bool:
+        """
+        Kills 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:
+            self._rpc.send_signal(
+                api_pb2.SendSignalRequest(
+                    process=api_pb2.ProcessSelector(pid=pid),
+                    signal=api_pb2.Signal.SIGNAL_SIGKILL,
+                ),
+                self._headers,
+                timeout_seconds=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+            return True
+        except Exception as e:
+            if "not found" in str(e):
+                return False
+            raise handle_rpc_exception(e)
+
+    def send_stdin(
+        self,
+        pid: int,
+        data: str,
+        request_timeout: Optional[float] = 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:
+            self._rpc.send_input(
+                api_pb2.SendInputRequest(
+                    process=api_pb2.ProcessSelector(pid=pid),
+                    input=api_pb2.ProcessInput(
+                        stdin=data.encode(),
+                    ),
+                ),
+                self._headers,
+                timeout_seconds=self._connection_config.get_request_timeout(
+                    request_timeout
+                ),
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    @overload
+    def run(
+        self,
+        cmd: str,
+        background: Union[Literal[False], None] = None,
+        envs: Optional[Dict[str, str]] = None,
+        user: Username = "user",
+        cwd: Optional[str] = None,
+        on_stdout: Optional[Callable[[str], None]] = None,
+        on_stderr: Optional[Callable[[str], None]] = 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 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
+    def run(
+        self,
+        cmd: str,
+        background: Literal[True],
+        envs: Optional[Dict[str, str]] = None,
+        user: Username = "user",
+        cwd: Optional[str] = None,
+        on_stdout: None = None,
+        on_stderr: None = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ) -> CommandHandle:
+        """
+        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 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: `CommandHandle` handle to interact with the running command
+        """
+        ...
+
+    def run(
+        self,
+        cmd: str,
+        background: Union[bool, None] = None,
+        envs: Optional[Dict[str, str]] = None,
+        user: Username = "user",
+        cwd: Optional[str] = None,
+        on_stdout: Optional[Callable[[str], None]] = None,
+        on_stderr: Optional[Callable[[str], None]] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ):
+        if background:
+            cmd += " &"
+        proc = self._start(
+            cmd,
+            envs,
+            user,
+            cwd,
+            timeout,
+            request_timeout,
+        )
+
+        return (
+            proc
+            if background
+            else proc.wait(
+                on_stdout=on_stdout,
+                on_stderr=on_stderr,
+            )
+        )
+
+    def _start(
+        self,
+        cmd: str,
+        envs: Optional[Dict[str, str]] = None,
+        user: Username = "user",
+        cwd: Optional[str] = None,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ):
+        events = self._rpc.start(
+            api_pb2.StartRequest(
+                process=api_pb2.ProcessConfig(
+                    cmd="/bin/bash",
+                    envs=envs,
+                    args=["-l","-c", cmd],
+                    cwd=cwd,
+                ),
+            ),
+            self._headers,
+            timeout_seconds=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )
+
+        try:
+            start_event = events.__next__()
+
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to start process: expected start event, got {start_event}"
+                )
+
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)
+
+    def connect(
+        self,
+        pid: int,
+        timeout: Optional[float] = 60,
+        request_timeout: Optional[float] = None,
+    ):
+        """
+        Connects to a running command.
+        You can use `CommandHandle.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 timeout: Timeout for the connection in **seconds**. Using `0` will not limit the connection time
+        :param request_timeout: Timeout for the request in **seconds**
+
+        :return: `CommandHandle` handle to interact with the running command
+        """
+        events = self._rpc.connect(
+            api_pb2.ConnectRequest(
+                process=api_pb2.ProcessSelector(pid=pid),
+            ),
+            self._headers,
+            timeout_seconds=self._connection_config.get_request_timeout(
+                request_timeout
+            ),
+        )
+        try:
+            start_event = events.__next__()
+            if not start_event.HasField("event"):
+                raise SandboxException(
+                    f"Failed to connect to process: expected start event, got {start_event}"
+                )
+            return CommandHandle(
+                pid=start_event.event.start.pid,
+                handle_kill=lambda: self.kill(start_event.event.start.pid),
+                events=events,
+            )
+        except Exception as e:
+            raise handle_rpc_exception(e)