sandbox_e2b.py

"""E2B-based sandbox for cloud deployment without Docker."""

from __future__ import annotations

from pathlib import Path

from e2b_code_interpreter import Sandbox

from llm_harness.sandbox import TIMEOUT_SECONDS

# Reuse a sandbox across tool calls within a session.
# The caller manages the lifecycle via create/close.
_active_sandbox: Sandbox | None = None


def get_or_create_sandbox(
    workspace: Path | None = None,
    scratch_dir: Path | None = None,
) -> Sandbox:
    """Get the active sandbox, creating one if needed and uploading workspace files."""
    global _active_sandbox
    if _active_sandbox is not None:
        return _active_sandbox

    _active_sandbox = Sandbox.create(timeout=300)

    # Create workspace and scratchpad directories in user-writable home
    _active_sandbox.commands.run("mkdir -p /home/user/workspace /home/user/scratchpad")
    # Symlink to expected paths
    _active_sandbox.commands.run(
        "ln -sf /home/user/workspace /workspace; "
        "ln -sf /home/user/scratchpad /scratchpad",
        user="root",
    )

    # Upload workspace files
    if workspace is not None:
        for file_path in workspace.iterdir():
            if file_path.is_file():
                _active_sandbox.files.write(
                    f"/home/user/workspace/{file_path.name}",
                    file_path.read_bytes(),
                )

    return _active_sandbox


def close_sandbox() -> None:
    global _active_sandbox
    if _active_sandbox is not None:
        _active_sandbox.kill()
        _active_sandbox = None


def _execute(sandbox: Sandbox, code: str, timeout: int) -> dict:
    execution = sandbox.run_code(code, timeout=timeout)

    stdout = "\n".join(
        line if isinstance(line, str) else line.text
        for line in execution.logs.stdout
    )
    stderr = "\n".join(
        line if isinstance(line, str) else line.text
        for line in execution.logs.stderr
    )

    if execution.error:
        stderr += f"\n{execution.error.name}: {execution.error.value}"
        exit_code = 1
    else:
        exit_code = 0

    return {
        "stdout": stdout,
        "stderr": stderr,
        "exit_code": exit_code,
        "timed_out": False,
    }


def run_python(
    code: str,
    *,
    workspace: Path | None = None,
    scratch_dir: Path | None = None,
    timeout: int = TIMEOUT_SECONDS,
) -> dict:
    """Execute Python code in an E2B sandbox. Same interface as sandbox.run_python."""
    sandbox = get_or_create_sandbox(workspace, scratch_dir)

    try:
        return _execute(sandbox, code, timeout)
    except Exception as exc:
        if "sandbox" in str(exc).lower() and "not found" in str(exc).lower():
            # Stale sandbox — recreate and retry
            close_sandbox()
            sandbox = get_or_create_sandbox(workspace, scratch_dir)
            try:
                return _execute(sandbox, code, timeout)
            except TimeoutError:
                return {
                    "stdout": "",
                    "stderr": "Execution timed out.",
                    "exit_code": -1,
                    "timed_out": True,
                }
        if isinstance(exc, TimeoutError):
            return {
                "stdout": "",
                "stderr": "Execution timed out.",
                "exit_code": -1,
                "timed_out": True,
            }
        raise