rpc.py

import base64

from typing import Optional
from packaging.version import Version
from e2b_connect.client import Code, ConnectException

from e2b.exceptions import (
    SandboxException,
    InvalidArgumentException,
    NotFoundException,
    TimeoutException,
    TimeoutType,
    format_sandbox_timeout_exception,
    AuthenticationException,
    RateLimitException,
)
from e2b.connection_config import Username, default_username
from e2b.envd.versions import ENVD_DEFAULT_USER


def handle_rpc_exception(e: Exception):
    if isinstance(e, ConnectException):
        if e.status == Code.invalid_argument:
            return InvalidArgumentException(e.message)
        elif e.status == Code.unauthenticated:
            return AuthenticationException(e.message)
        elif e.status == Code.not_found:
            return NotFoundException(e.message)
        elif e.status == Code.unavailable:
            return format_sandbox_timeout_exception(e.message)
        elif e.status == Code.resource_exhausted:
            return RateLimitException(
                f"{e.message}: Rate limit exceeded, please try again later."
            )
        elif e.status == Code.canceled:
            return TimeoutException(
                f"{e.message}: This error is likely due to exceeding 'request_timeout'. You can pass the request timeout value as an option when making the request.",
                type=TimeoutType.REQUEST,
            )
        elif e.status == Code.deadline_exceeded:
            return TimeoutException(
                f"{e.message}: This error is likely due to exceeding 'timeout' — the total time a long running request (like process or directory watch) can be active. It can be modified by passing 'timeout' when making the request. Use '0' to disable the timeout.",
                type=TimeoutType.EXECUTION,
            )
        else:
            return SandboxException(f"{e.status}: {e.message}")
    else:
        return e


def authentication_header(
    envd_version: Version, user: Optional[Username] = None
) -> dict[str, str]:
    if user is None and envd_version < ENVD_DEFAULT_USER:
        user = default_username

    if not user:
        return {}

    value = f"{user}:"

    encoded = base64.b64encode(value.encode("utf-8")).decode("utf-8")

    return {"Authorization": f"Basic {encoded}"}