arena/codebase_executor.py

"""Sandboxed command executors for Codebase archives."""

from __future__ import annotations

import os
import shlex
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any, Callable
from uuid import uuid4

from deepagents.backends import LocalShellBackend

from .codebase_archive import (
    CodebaseArchiveStore,
    LocalCodebaseArchiveStore,
    create_codebase_archive_store,
)
from .codebase_handoff import (
    CodebaseHandoff,
    DEFAULT_CODEBASE_HANDOFF_ROOT,
    local_snapshot_workspace,
    pointer_error,
)
from .sandbox_lease import SandboxLeaseManager


DEFAULT_SNAPSHOT_ROOT = DEFAULT_CODEBASE_HANDOFF_ROOT
DEFAULT_VALIDATOR_SANDBOX_ROOT = Path(tempfile.gettempdir()) / "agent-swarm-workbench-validator"


class LocalCodebaseExecutor:
    """Runs validation checks against local Codebase archive materializations."""

    def __init__(self, *, root_dir: Path | None = None, timeout_seconds: int = 20) -> None:
        self.root_dir = root_dir or DEFAULT_SNAPSHOT_ROOT
        self.timeout_seconds = timeout_seconds

    def run(self, archive_pointer: str, command: str) -> tuple[bool, str]:
        workspace = local_snapshot_workspace(self.root_dir, archive_pointer)
        if workspace is None:
            return False, pointer_error(archive_pointer)
        if not workspace.exists():
            return False, f"snapshot workspace not found: {archive_pointer}"

        try:
            result = subprocess.run(
                ["/bin/sh", "-lc", command],
                cwd=workspace,
                capture_output=True,
                text=True,
                timeout=self.timeout_seconds,
                env={"PATH": _executor_path()},
                check=False,
            )
        except subprocess.TimeoutExpired as exc:
            output = (exc.stdout or "") + (exc.stderr or "")
            return False, f"{output}\n[Command timed out after {self.timeout_seconds}s]"

        output = result.stdout
        if result.stderr:
            output = f"{output}[stderr] {result.stderr}"
        status = "succeeded" if result.returncode == 0 else "failed"
        return result.returncode == 0, f"{output}\n[Command {status} with exit code {result.returncode}]"


class CodebaseSandboxExecutor:
    """Restores Codebase archives into a sandbox backend before running checks."""

    def __init__(
        self,
        *,
        backend_factory: Callable[[], Any],
        root_dir: Path | None = None,
        archive_store: CodebaseArchiveStore | None = None,
        timeout_seconds: int = 20,
        sandbox_ttl_seconds: float | None = None,
        sandbox_leases: SandboxLeaseManager | None = None,
    ) -> None:
        self.root_dir = root_dir or DEFAULT_SNAPSHOT_ROOT
        self.archive_store = archive_store or LocalCodebaseArchiveStore(root_dir=self.root_dir)
        self.handoff = CodebaseHandoff(archive_store=self.archive_store, root_dir=self.root_dir)
        self.backend_factory = backend_factory
        self.backend = None
        self.timeout_seconds = timeout_seconds
        self.workspace_dir = "workspace"
        self._restored: set[str] = set()
        self._lease_id = f"validator-{uuid4().hex}"
        self._sandbox_leases = sandbox_leases or SandboxLeaseManager(ttl_seconds=sandbox_ttl_seconds)

    def run(self, archive_pointer: str, command: str) -> tuple[bool, str]:
        self._ensure_backend()
        self._sandbox_leases.touch(self._lease_id)
        try:
            self._restore(archive_pointer)
        except (OSError, ValueError) as exc:
            return False, f"could not restore snapshot: {exc}"

        result = self.backend.execute(f"cd {shlex.quote(self.workspace_dir)} && {command}", timeout=self.timeout_seconds)
        output = getattr(result, "output", "")
        exit_code = getattr(result, "exit_code", None)
        if exit_code is None:
            return True, output
        status = "succeeded" if exit_code == 0 else "failed"
        return exit_code == 0, f"{output}\n[Command {status} with exit code {exit_code}]"

    def close(self) -> None:
        self._sandbox_leases.close(self._lease_id)

    def _open_backend(self) -> None:
        self.backend = self.backend_factory()
        self.workspace_dir = getattr(self.backend, "_validator_workspace", "workspace")
        self._restored = set()
        self._sandbox_leases.start(self._lease_id, self._close_backend)

    def _ensure_backend(self) -> None:
        if self.backend is None:
            self._open_backend()

    def _close_backend(self) -> None:
        if self.backend is None:
            return
        sandbox = getattr(self.backend, "_daytona_sandbox", None)
        cleanup = None
        if sandbox is not None:
            cleanup = getattr(sandbox, "delete", None) or getattr(sandbox, "stop", None)
        if cleanup is None:
            cleanup = getattr(self.backend, "delete", None) or getattr(self.backend, "stop", None)
        if callable(cleanup):
            cleanup()
        self.backend = None
        self._restored = set()

    def _restore(self, archive_pointer: str) -> None:
        if archive_pointer in self._restored:
            return
        self.handoff.restore_to_backend(
            archive_pointer,
            self.backend,
            workspace_dir=self.workspace_dir,
        )
        self._restored.add(archive_pointer)


def create_codebase_executor(
    archive_store: CodebaseArchiveStore | None = None,
    provider: str | None = None,
) -> CodebaseSandboxExecutor:
    provider = (provider or os.getenv("VALIDATOR_SANDBOX_PROVIDER", "local")).lower()
    archive_store = archive_store or create_codebase_archive_store()
    if provider == "local":
        return CodebaseSandboxExecutor(backend_factory=_local_validator_backend, archive_store=archive_store)
    if provider == "daytona":
        return CodebaseSandboxExecutor(backend_factory=_daytona_validator_backend, archive_store=archive_store)
    raise ValueError(f"unsupported validator sandbox provider: {provider}")


def _executor_path() -> str:
    executable_dir = str(Path(sys.executable).parent)
    return os.pathsep.join([executable_dir, "/usr/local/bin", "/usr/bin", "/bin"])


def _local_validator_backend() -> LocalShellBackend:
    root = DEFAULT_VALIDATOR_SANDBOX_ROOT / uuid4().hex
    root.mkdir(parents=True, exist_ok=True)
    backend = LocalShellBackend(root_dir=root, virtual_mode=True, inherit_env=True)
    backend._validator_workspace = "workspace"
    return backend


def _daytona_validator_backend() -> Any:
    from daytona import Daytona
    from langchain_daytona import DaytonaSandbox

    client = Daytona()
    sandbox = client.create()
    backend = DaytonaSandbox(sandbox=sandbox)
    backend._daytona_client = client
    backend._daytona_sandbox = sandbox
    backend._validator_workspace = "/home/daytona/workspace"
    return backend