arena/codebase_handoff.py

"""Codebase handoff between Swarm workspaces, archives, and Validator sandboxes."""

from __future__ import annotations

import re
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from .codebase_archive import (
    CodebaseArchiveStore,
    LocalCodebaseArchiveStore,
    archive_from_directory,
    archive_from_downloads,
    supports_archive_pointer,
)


DEFAULT_CODEBASE_HANDOFF_ROOT = Path(tempfile.gettempdir()) / "agent-swarm-workbench-codebases"
LOCAL_SNAPSHOT_SCHEME = "local-snapshot://"
CODEBASE_ID_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")


@dataclass
class CodebaseHandoff:
    """Moves a Codebase across Workspace, archive, and Validator sandbox seams."""

    archive_store: CodebaseArchiveStore | None = None
    root_dir: Path | None = None

    def __post_init__(self) -> None:
        root = self.root_dir or DEFAULT_CODEBASE_HANDOFF_ROOT
        if self.archive_store is None:
            self.archive_store = LocalCodebaseArchiveStore(root_dir=root)

    def snapshot_local_workspace(self, source_root: Path, codebase_id: str) -> str:
        _validate_codebase_id(codebase_id)
        return self.archive_store.save(archive_from_directory(codebase_id, source_root / "workspace"))

    def snapshot_backend_workspace(
        self,
        backend: Any,
        codebase_id: str,
        *,
        workspace_dir: str = "/workspace",
    ) -> str:
        _validate_codebase_id(codebase_id)
        workspace_dir = normalize_workspace_dir(workspace_dir)
        glob_result = backend.glob(f"{workspace_dir}/**")
        if getattr(glob_result, "error", None):
            raise RuntimeError("could not list workspace files")

        file_paths = [
            _entry_path(match)
            for match in getattr(glob_result, "matches", [])
            if not _entry_is_dir(match)
        ]
        if not file_paths:
            raise FileNotFoundError("workspace has no files")

        archive_files: list[tuple[str, bytes]] = []
        for item in backend.download_files(file_paths):
            if getattr(item, "error", None):
                raise RuntimeError(f"could not download workspace file: {getattr(item, 'path', '')}")
            source_path = str(getattr(item, "path", ""))
            relative = workspace_relative_path(source_path, workspace_dir=workspace_dir)
            archive_files.append((relative.as_posix(), getattr(item, "content", b"") or b""))

        return self.archive_store.save(archive_from_downloads(codebase_id, archive_files))

    def restore_to_backend(
        self,
        archive_pointer: str,
        backend: Any,
        *,
        workspace_dir: str,
    ) -> int:
        if not supports_archive_pointer(archive_pointer):
            raise ValueError(pointer_error(archive_pointer))

        workspace = (self.root_dir or DEFAULT_CODEBASE_HANDOFF_ROOT) / restore_id(archive_pointer) / "workspace"
        self.archive_store.materialize(archive_pointer, workspace)
        files = [
            (sandbox_workspace_path(workspace_dir, path.relative_to(workspace).as_posix()), path.read_bytes())
            for path in workspace.rglob("*")
            if path.is_file()
        ]
        results = backend.upload_files(files)
        for result in results:
            if getattr(result, "error", None):
                raise RuntimeError(
                    f"could not upload file: {getattr(result, 'path', '')}: {getattr(result, 'error', '')}"
                )
        return len(files)


def local_snapshot_workspace(root_dir: Path, archive_pointer: str) -> Path | None:
    snapshot_id = local_snapshot_id(archive_pointer)
    if snapshot_id is None:
        return None
    return root_dir / snapshot_id / "workspace"


def local_snapshot_id(archive_pointer: str) -> str | None:
    if not archive_pointer.startswith(LOCAL_SNAPSHOT_SCHEME):
        return None
    snapshot_id = archive_pointer.removeprefix(LOCAL_SNAPSHOT_SCHEME)
    if not CODEBASE_ID_PATTERN.fullmatch(snapshot_id):
        return None
    return snapshot_id


def restore_id(archive_pointer: str) -> str:
    if "://" not in archive_pointer:
        raise ValueError("unsupported archive pointer")
    snapshot_id = archive_pointer.split("://", 1)[1]
    _validate_codebase_id(snapshot_id)
    return snapshot_id


def pointer_error(archive_pointer: str) -> str:
    if archive_pointer.startswith(LOCAL_SNAPSHOT_SCHEME):
        return "invalid snapshot id"
    return "unsupported archive pointer"


def normalize_workspace_dir(workspace_dir: str) -> str:
    return workspace_dir.rstrip("/") or "/workspace"


def workspace_relative_path(path: str, *, workspace_dir: str = "/workspace") -> Path:
    prefix = f"{normalize_workspace_dir(workspace_dir)}/"
    if not path.startswith(prefix):
        raise ValueError(f"path outside workspace: {path}")
    relative = path.removeprefix(prefix)
    if not relative or ".." in Path(relative).parts:
        raise ValueError(f"unsafe workspace path: {path}")
    return Path(relative)


def sandbox_workspace_path(workspace_dir: str, relative_path: str) -> str:
    return f"{normalize_workspace_dir(workspace_dir)}/{relative_path}"


def _validate_codebase_id(codebase_id: str) -> None:
    if not CODEBASE_ID_PATTERN.fullmatch(codebase_id):
        raise ValueError("invalid codebase id")


def _entry_path(entry: Any) -> str:
    return str(entry.get("path", "") if isinstance(entry, dict) else getattr(entry, "path", ""))


def _entry_is_dir(entry: Any) -> bool:
    return bool(entry.get("is_dir") if isinstance(entry, dict) else getattr(entry, "is_dir", False))