letta-server/letta/services/tool_sandbox/e2b_sandbox.py

import asyncio
import time
from typing import TYPE_CHECKING, Any, Dict, Optional

from e2b.sandbox.commands.command_handle import CommandExitException
from e2b_code_interpreter import AsyncSandbox

from letta.log import get_logger
from letta.otel.tracing import log_event, trace_method
from letta.schemas.agent import AgentState
from letta.schemas.enums import SandboxType
from letta.schemas.sandbox_config import SandboxConfig
from letta.schemas.tool import Tool
from letta.schemas.tool_execution_result import ToolExecutionResult
from letta.services.helpers.tool_parser_helper import parse_stdout_best_effort
from letta.services.tool_sandbox.base import AsyncToolSandboxBase
from letta.services.tool_sandbox.typescript_generator import parse_typescript_result
from letta.types import JsonDict
from letta.utils import get_friendly_error_msg

logger = get_logger(__name__)

if TYPE_CHECKING:
    from e2b_code_interpreter import Execution


class AsyncToolSandboxE2B(AsyncToolSandboxBase):
    METADATA_CONFIG_STATE_KEY = "config_state"

    def __init__(
        self,
        tool_name: str,
        args: JsonDict,
        user,
        tool_id: str,
        agent_id: Optional[str] = None,
        project_id: Optional[str] = None,
        force_recreate: bool = True,
        tool_object: Optional[Tool] = None,
        sandbox_config: Optional[SandboxConfig] = None,
        sandbox_env_vars: Optional[Dict[str, Any]] = None,
    ):
        super().__init__(
            tool_name,
            args,
            user,
            tool_id=tool_id,
            agent_id=agent_id,
            project_id=project_id,
            tool_object=tool_object,
            sandbox_config=sandbox_config,
            sandbox_env_vars=sandbox_env_vars,
        )
        self.force_recreate = force_recreate

    @trace_method
    async def run(
        self,
        agent_state: Optional[AgentState] = None,
        additional_env_vars: Optional[Dict] = None,
    ) -> ToolExecutionResult:
        await self._init_async()
        if self.provided_sandbox_config:
            sbx_config = self.provided_sandbox_config
        else:
            sbx_config = await self.sandbox_config_manager.get_or_create_default_sandbox_config_async(
                sandbox_type=SandboxType.E2B, actor=self.user
            )
        # TODO: So this defaults to force recreating always
        # TODO: Eventually, provision one sandbox PER agent, and that agent re-uses that one specifically
        e2b_sandbox = await self.create_e2b_sandbox_with_metadata_hash(sandbox_config=sbx_config)

        logger.info(f"E2B Sandbox configurations: {sbx_config}")
        logger.info(f"E2B Sandbox ID: {e2b_sandbox.sandbox_id}")

        # TODO: This only makes sense if we re-use sandboxes
        # # Since this sandbox was used, we extend its lifecycle by the timeout
        # await sbx.set_timeout(sbx_config.get_e2b_config().timeout)

        # Get environment variables for the sandbox
        envs = await self._gather_env_vars(agent_state, additional_env_vars, sbx_config.id, is_local=False)
        code = await self.generate_execution_script(agent_state=agent_state)

        # Determine language based on tool source_type
        is_typescript = self.is_typescript_tool()
        language = "ts" if is_typescript else None

        try:
            logger.info(f"E2B execution started for ID {e2b_sandbox.sandbox_id}: {self.tool_name} (language={language or 'python'})")
            log_event(
                "e2b_execution_started",
                {
                    "tool": self.tool_name,
                    "sandbox_id": e2b_sandbox.sandbox_id,
                    "code": code,
                    "env_vars": envs,
                    "language": language or "python",
                },
            )
            start_time = time.perf_counter()
            try:
                execution = await e2b_sandbox.run_code(code, envs=envs, language=language)
            except asyncio.CancelledError:
                execution_time = time.perf_counter() - start_time
                logger.info(f"E2B execution cancelled for ID {e2b_sandbox.sandbox_id}: {self.tool_name} (took {execution_time:.2f}s)")
                log_event(
                    "e2b_execution_cancelled",
                    {"tool": self.tool_name, "sandbox_id": e2b_sandbox.sandbox_id, "execution_time_seconds": execution_time},
                )
                raise Exception("Execution cancelled. Transient failure, please retry.")

            execution_time = time.perf_counter() - start_time
            logger.info(f"E2B execution completed in {execution_time:.2f}s for sandbox {e2b_sandbox.sandbox_id}, tool: {self.tool_name}")

            if execution.results:
                # Use appropriate result parser based on language
                if is_typescript:
                    func_return, agent_state = parse_typescript_result(execution.results[0].text)
                else:
                    func_return, agent_state = parse_stdout_best_effort(execution.results[0].text)
                logger.info(f"E2B execution succeeded for ID {e2b_sandbox.sandbox_id}: {self.tool_name} (took {execution_time:.2f}s)")
                log_event(
                    "e2b_execution_succeeded",
                    {
                        "tool": self.tool_name,
                        "sandbox_id": e2b_sandbox.sandbox_id,
                        "func_return": func_return,
                        "execution_time_seconds": execution_time,
                    },
                )
            elif execution.error:
                logger.warning(f"Tool {self.tool_name} raised a {execution.error.name}: {execution.error.value}")
                logger.warning(f"Traceback from e2b sandbox: \n{execution.error.traceback}")
                func_return = get_friendly_error_msg(
                    function_name=self.tool_name, exception_name=execution.error.name, exception_message=execution.error.value
                )
                execution.logs.stderr.append(execution.error.traceback)
                logger.warning(f"E2B execution failed for ID {e2b_sandbox.sandbox_id}: {self.tool_name} (took {execution_time:.2f}s)")
                log_event(
                    "e2b_execution_failed",
                    {
                        "tool": self.tool_name,
                        "sandbox_id": e2b_sandbox.sandbox_id,
                        "error_type": execution.error.name,
                        "error_message": execution.error.value,
                        "func_return": func_return,
                        "execution_time_seconds": execution_time,
                    },
                )
            else:
                logger.warning(f"E2B execution empty for ID {e2b_sandbox.sandbox_id}: {self.tool_name} (took {execution_time:.2f}s)")
                log_event(
                    "e2b_execution_empty",
                    {
                        "tool": self.tool_name,
                        "sandbox_id": e2b_sandbox.sandbox_id,
                        "status": "no_results_no_error",
                        "execution_time_seconds": execution_time,
                    },
                )
                raise ValueError(f"Tool {self.tool_name} returned execution with None")

            return ToolExecutionResult(
                func_return=func_return,
                agent_state=agent_state,
                stdout=execution.logs.stdout,
                stderr=execution.logs.stderr,
                status="error" if execution.error else "success",
                sandbox_config_fingerprint=sbx_config.fingerprint(),
            )
        finally:
            logger.info(f"E2B sandbox {e2b_sandbox.sandbox_id} killed")
            await e2b_sandbox.kill()

    @staticmethod
    def parse_exception_from_e2b_execution(e2b_execution: "Execution") -> Exception:
        builtins_dict = __builtins__ if isinstance(__builtins__, dict) else vars(__builtins__)
        # Dynamically fetch the exception class from builtins, defaulting to Exception if not found
        exception_class = builtins_dict.get(e2b_execution.error.name, Exception)
        return exception_class(e2b_execution.error.value)

    @trace_method
    async def create_e2b_sandbox_with_metadata_hash(self, sandbox_config: SandboxConfig) -> "AsyncSandbox":
        state_hash = sandbox_config.fingerprint()
        e2b_config = sandbox_config.get_e2b_config()

        log_event(
            "e2b_sandbox_create_started",
            {
                "sandbox_fingerprint": state_hash,
                "e2b_config": e2b_config.model_dump(),
            },
        )

        if e2b_config.template:
            sbx = await AsyncSandbox.create(sandbox_config.get_e2b_config().template, metadata={self.METADATA_CONFIG_STATE_KEY: state_hash})
        else:
            sbx = await AsyncSandbox.create(
                metadata={self.METADATA_CONFIG_STATE_KEY: state_hash}, **e2b_config.model_dump(exclude={"pip_requirements"})
            )

        log_event(
            "e2b_sandbox_create_finished",
            {
                "sandbox_id": sbx.sandbox_id,
                "sandbox_fingerprint": state_hash,
            },
        )

        if e2b_config.pip_requirements:
            for package in e2b_config.pip_requirements:
                log_event(
                    "e2b_pip_install_started",
                    {
                        "sandbox_id": sbx.sandbox_id,
                        "package": package,
                    },
                )
                try:
                    await sbx.commands.run(f"pip install {package}")
                    log_event(
                        "e2b_pip_install_finished",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package,
                        },
                    )
                except CommandExitException as e:
                    error_msg = f"Failed to install sandbox pip requirement '{package}' in E2B sandbox. This may be due to package version incompatibility with the E2B environment. Error: {e}"
                    logger.error(error_msg)
                    log_event(
                        "e2b_pip_install_failed",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package,
                            "error": str(e),
                        },
                    )
                    raise RuntimeError(error_msg) from e

        # Install tool-specific pip requirements
        if self.tool and self.tool.pip_requirements:
            for pip_requirement in self.tool.pip_requirements:
                package_str = str(pip_requirement)
                log_event(
                    "tool_pip_install_started",
                    {
                        "sandbox_id": sbx.sandbox_id,
                        "package": package_str,
                        "tool_name": self.tool.name,
                    },
                )
                try:
                    await sbx.commands.run(f"pip install {package_str}")
                    log_event(
                        "tool_pip_install_finished",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package_str,
                            "tool_name": self.tool.name,
                        },
                    )
                except CommandExitException as e:
                    error_msg = f"Failed to install tool pip requirement '{package_str}' for tool '{self.tool.name}' in E2B sandbox. This may be due to package version incompatibility with the E2B environment. Consider updating the package version or removing the version constraint. Error: {e}"
                    logger.error(error_msg)
                    log_event(
                        "tool_pip_install_failed",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package_str,
                            "tool_name": self.tool.name,
                            "error": str(e),
                        },
                    )
                    # Kill sandbox to prevent resource leak
                    await sbx.kill()
                    raise RuntimeError(error_msg) from e

        # Install tool-specific npm requirements (for TypeScript tools)
        if self.tool and self.tool.npm_requirements:
            for npm_requirement in self.tool.npm_requirements:
                package_str = str(npm_requirement)
                log_event(
                    "tool_npm_install_started",
                    {
                        "sandbox_id": sbx.sandbox_id,
                        "package": package_str,
                        "tool_name": self.tool.name,
                    },
                )
                try:
                    await sbx.commands.run(f"npm install {package_str}")
                    log_event(
                        "tool_npm_install_finished",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package_str,
                            "tool_name": self.tool.name,
                        },
                    )
                except CommandExitException as e:
                    error_msg = f"Failed to install tool npm requirement '{package_str}' for tool '{self.tool.name}' in E2B sandbox. This may be due to package version incompatibility. Error: {e}"
                    logger.error(error_msg)
                    log_event(
                        "tool_npm_install_failed",
                        {
                            "sandbox_id": sbx.sandbox_id,
                            "package": package_str,
                            "tool_name": self.tool.name,
                            "error": str(e),
                        },
                    )
                    # Kill sandbox to prevent resource leak
                    await sbx.kill()
                    raise RuntimeError(error_msg) from e

        # Auto-install @letta-ai/letta-client for TypeScript tools (similar to letta_client for Python)
        if self.is_typescript_tool():
            letta_client_package = "@letta-ai/letta-client"
            log_event(
                "letta_client_npm_install_started",
                {
                    "sandbox_id": sbx.sandbox_id,
                    "package": letta_client_package,
                    "tool_name": self.tool.name if self.tool else None,
                },
            )
            try:
                await sbx.commands.run(f"npm install {letta_client_package}")
                log_event(
                    "letta_client_npm_install_finished",
                    {
                        "sandbox_id": sbx.sandbox_id,
                        "package": letta_client_package,
                        "tool_name": self.tool.name if self.tool else None,
                    },
                )
            except CommandExitException as e:
                # Log warning but don't fail - the client is optional for tools that don't need it
                logger.warning(f"Failed to install {letta_client_package} in E2B sandbox: {e}")
                log_event(
                    "letta_client_npm_install_failed",
                    {
                        "sandbox_id": sbx.sandbox_id,
                        "package": letta_client_package,
                        "tool_name": self.tool.name if self.tool else None,
                        "error": str(e),
                    },
                )

        return sbx

    def use_top_level_await(self) -> bool:
        """
        E2B sandboxes run in a Jupyter-like environment with an active event loop,
        so they support top-level await.
        """
        return True

    @staticmethod
    async def list_running_e2b_sandboxes():
        # List running sandboxes and access metadata.
        return await AsyncSandbox.list()