arena/codebase_archive.py

"""Codebase archive storage for Run and Validator workspaces.

Archives are content-addressed by the Run id, but hidden behind a small store
interface. Local development writes files to disk; hosted deployments can
persist compact archives in Redis so validation survives process restarts.
"""

from __future__ import annotations

import base64
import io
import json
import os
import shutil
import tempfile
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Protocol

from pydantic import BaseModel, Field, field_validator, model_validator

from .redis_client import RedisClient, RedisStoreError, UpstashRedisRestClient


LOCAL_ARCHIVE_SCHEME = "local-snapshot://"
REDIS_ARCHIVE_SCHEME = "redis-archive://"
DEFAULT_ARCHIVE_ROOT = Path(tempfile.gettempdir()) / "agent-swarm-workbench-codebases"
ARCHIVE_FILE_MAX_BYTES = 10_000_000
ARCHIVE_TOTAL_MAX_BYTES = 50_000_000
ARCHIVE_IGNORED_DIR_NAMES = frozenset({"__pycache__"})
ARCHIVE_IGNORED_SUFFIXES = (".pyc", ".pyo")


class CodebaseArchiveError(ValueError):
    pass


class ArchiveFile(BaseModel):
    path: str = Field(min_length=1)
    content_base64: str
    size_bytes: int = Field(ge=0, le=ARCHIVE_FILE_MAX_BYTES)

    @field_validator("path")
    @classmethod
    def validate_path(cls, value: str) -> str:
        return _safe_relative_path(value)


class CodebaseArchiveSnapshot(BaseModel):
    archive_id: str = Field(min_length=1)
    files: list[ArchiveFile] = Field(default_factory=list)

    @model_validator(mode="after")
    def validate_total_size(self) -> "CodebaseArchiveSnapshot":
        total = sum(file.size_bytes for file in self.files)
        if total > ARCHIVE_TOTAL_MAX_BYTES:
            raise ValueError("codebase archive exceeds total size cap")
        return self


class CodebaseArchiveStore(Protocol):
    def save(self, archive: CodebaseArchiveSnapshot) -> str:
        ...

    def load(self, archive_pointer: str) -> CodebaseArchiveSnapshot:
        ...

    def materialize(self, archive_pointer: str, destination: Path) -> None:
        ...


class LocalCodebaseArchiveStore:
    def __init__(self, *, root_dir: Path | None = None) -> None:
        self.root_dir = root_dir or DEFAULT_ARCHIVE_ROOT

    def save(self, archive: CodebaseArchiveSnapshot) -> str:
        destination = self.root_dir / archive.archive_id / "workspace"
        if destination.exists():
            shutil.rmtree(destination)
        destination.mkdir(parents=True, exist_ok=True)
        _write_archive_files(archive, destination)
        return f"{LOCAL_ARCHIVE_SCHEME}{archive.archive_id}"

    def load(self, archive_pointer: str) -> CodebaseArchiveSnapshot:
        archive_id = _pointer_id(archive_pointer, LOCAL_ARCHIVE_SCHEME)
        workspace = self.root_dir / archive_id / "workspace"
        if not workspace.exists():
            raise FileNotFoundError(f"snapshot workspace not found: {archive_id}")
        return archive_from_directory(archive_id, workspace)

    def materialize(self, archive_pointer: str, destination: Path) -> None:
        archive = self.load(archive_pointer)
        if destination.exists():
            shutil.rmtree(destination)
        destination.mkdir(parents=True, exist_ok=True)
        _write_archive_files(archive, destination)


@dataclass
class RedisCodebaseArchiveStore:
    client: RedisClient
    key_prefix: str = "arena"
    key_namespace: str = "codebase-archive"

    def save(self, archive: CodebaseArchiveSnapshot) -> str:
        key = self._key(archive.archive_id)
        self.client.execute("SET", key, archive.model_dump_json())
        return f"{REDIS_ARCHIVE_SCHEME}{archive.archive_id}"

    def load(self, archive_pointer: str) -> CodebaseArchiveSnapshot:
        archive_id = _pointer_id(archive_pointer, REDIS_ARCHIVE_SCHEME)
        payload = self.client.execute("GET", self._key(archive_id))
        if payload is None:
            raise FileNotFoundError(f"codebase archive not found: {archive_id}")
        try:
            return CodebaseArchiveSnapshot.model_validate_json(str(payload))
        except ValueError as exc:
            raise RedisStoreError("codebase archive is invalid") from exc

    def materialize(self, archive_pointer: str, destination: Path) -> None:
        archive = self.load(archive_pointer)
        if destination.exists():
            shutil.rmtree(destination)
        destination.mkdir(parents=True, exist_ok=True)
        _write_archive_files(archive, destination)

    def _key(self, archive_id: str) -> str:
        return f"{self.key_prefix.rstrip(':')}:{self.key_namespace}:{archive_id}"


def create_codebase_archive_store() -> CodebaseArchiveStore:
    provider = (
        os.getenv("CODEBASE_ARCHIVE_PROVIDER")
        or os.getenv("MATCH_STORE_PROVIDER", "local")
    )
    provider = provider.lower()
    if provider in {"local", "memory", "inmemory", "in-memory"}:
        return LocalCodebaseArchiveStore()
    if provider in {"redis", "upstash"}:
        url = os.getenv("UPSTASH_REDIS_REST_URL", "")
        token = os.getenv("UPSTASH_REDIS_REST_TOKEN", "")
        if not url or not token:
            raise ValueError(
                "UPSTASH_REDIS_REST_URL and UPSTASH_REDIS_REST_TOKEN are required "
                "when CODEBASE_ARCHIVE_PROVIDER=redis"
            )
        return RedisCodebaseArchiveStore(client=UpstashRedisRestClient(url=url, token=token))
    raise ValueError(f"Unsupported codebase archive provider: {provider}")


def archive_from_directory(archive_id: str, workspace: Path) -> CodebaseArchiveSnapshot:
    if not workspace.exists():
        raise FileNotFoundError(f"workspace not found: {workspace}")
    files = [
        _archive_file(relative_path, path.read_bytes())
        for path in sorted(workspace.rglob("*"))
        for relative_path in [path.relative_to(workspace).as_posix()]
        if path.is_file()
        if not _is_ignored_archive_path(relative_path)
    ]
    return CodebaseArchiveSnapshot(archive_id=archive_id, files=files)


def archive_from_downloads(
    archive_id: str,
    downloads: list[tuple[str, bytes]],
) -> CodebaseArchiveSnapshot:
    return CodebaseArchiveSnapshot(
        archive_id=archive_id,
        files=[
            _archive_file(path, content)
            for path, content in downloads
            if not _is_ignored_archive_path(path)
        ],
    )


def archive_to_zip_bytes(archive: CodebaseArchiveSnapshot) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as zipped:
        for file in sorted(archive.files, key=lambda item: item.path):
            zipped.writestr(
                _safe_relative_path(file.path),
                base64.b64decode(file.content_base64.encode("ascii")),
            )
    return buffer.getvalue()


def supports_archive_pointer(archive_pointer: str) -> bool:
    return archive_pointer.startswith(LOCAL_ARCHIVE_SCHEME) or archive_pointer.startswith(REDIS_ARCHIVE_SCHEME)


def archive_to_json(archive: CodebaseArchiveSnapshot) -> str:
    return json.dumps(archive.model_dump(mode="json"), sort_keys=True)


def _archive_file(path: str, content: bytes) -> ArchiveFile:
    safe_path = _safe_relative_path(path)
    return ArchiveFile(
        path=safe_path,
        content_base64=base64.b64encode(content).decode("ascii"),
        size_bytes=len(content),
    )


def _is_ignored_archive_path(path: str) -> bool:
    relative = Path(path)
    return (
        any(part in ARCHIVE_IGNORED_DIR_NAMES for part in relative.parts)
        or relative.suffix in ARCHIVE_IGNORED_SUFFIXES
    )


def _write_archive_files(archive: CodebaseArchiveSnapshot, destination: Path) -> None:
    for file in archive.files:
        relative = Path(_safe_relative_path(file.path))
        target = destination / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        target.write_bytes(base64.b64decode(file.content_base64.encode("ascii")))


def _safe_relative_path(path: str) -> str:
    relative = Path(path)
    if relative.is_absolute() or not path or ".." in relative.parts:
        raise CodebaseArchiveError(f"unsafe archive path: {path}")
    return relative.as_posix()


def _pointer_id(archive_pointer: str, scheme: str) -> str:
    if not archive_pointer.startswith(scheme):
        raise CodebaseArchiveError("unsupported archive pointer")
    archive_id = archive_pointer.removeprefix(scheme)
    if not archive_id or "/" in archive_id or ".." in archive_id:
        raise CodebaseArchiveError("invalid archive pointer")
    return archive_id