cleveragents-core/tools/controller/worker/workspace.py

"""Per-PR workspace management.

Per plan v6/v9: each attempt operates in a per-PR umbrella directory
at ``/tmp/cleveragents-controller/pr-{owner}-{repo}-{N}/`` containing
the editable worktree, the ``worker.session`` sidecar, and (later)
attempt log archives.

At attempt start the worker MUST:
1. Ensure workspace + clone (if absent) or reuse (if present)
2. ``git fetch origin --prune``
3. Validate ``input_payload.head_sha == current HEAD on origin/<ref>``
   — if mismatched, raise ``StaleInputError`` so the runner reports
   ``outcome='stale-input'`` (master re-prefetches without pickup
   penalty, per v6).
4. ``git reset --hard <expected_head_sha>`` (cleans residue from any
   previous attempt's mid-rebase / mid-edit state on this workspace).

The workspace is per-PR, not per-attempt — multiple attempts on the
same PR share the same on-disk dir. Crash-recovery: workspace dirs
whose owning attempt is no longer in flight are deleted by the
startup janitor (janitor.py).
"""
from __future__ import annotations

import logging
import os
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path

logger = logging.getLogger(__name__)


# Per plan v6: per-PR umbrella dir.
DEFAULT_WORKSPACE_ROOT = Path(
    os.environ.get(
        "CONTROLLER_WORKSPACE_ROOT",
        "/tmp/cleveragents-controller",
    )
)


class StaleInputError(RuntimeError):
    """Raised when the input_payload's head_sha doesn't match the
    repo's current state. Runner maps this to ``WorkerError(
    outcome='stale-input')`` so the master re-prefetches without
    bumping ``pickup_count``."""

    def __init__(self, expected: str, actual: str, head_ref: str):
        self.expected = expected
        self.actual = actual
        self.head_ref = head_ref
        super().__init__(
            f"stale input: enqueued head_sha={expected!r} for ref "
            f"{head_ref!r}, but current head_sha is {actual!r}"
        )


@dataclass(frozen=True)
class WorkspaceIdentity:
    """The (owner, repo, entity_number) tuple that uniquely names a
    workspace dir. Matches the v6 path convention."""

    owner: str
    repo: str
    entity_number: int

    def dir_name(self) -> str:
        return f"pr-{self.owner}-{self.repo}-{self.entity_number}"


class PerPRWorkspace:
    """One PR's on-disk workspace.

    The workspace dir layout::

        /tmp/cleveragents-controller/pr-{owner}-{repo}-{N}/
            worker.session       (sidecar; written by the worker on spawn)
            worktree/            (editable git checkout)
            attempts/            (per-attempt log archive)

    Tests inject a custom ``workspace_root`` so they don't need to
    write to ``/tmp/``.
    """

    def __init__(
        self,
        identity: WorkspaceIdentity,
        *,
        workspace_root: Path | None = None,
        clone_url: str | None = None,
    ):
        self.identity = identity
        self.workspace_root = workspace_root or DEFAULT_WORKSPACE_ROOT
        self.clone_url = clone_url

    @property
    def workspace_dir(self) -> Path:
        return self.workspace_root / self.identity.dir_name()

    @property
    def worktree_dir(self) -> Path:
        return self.workspace_dir / "worktree"

    @property
    def session_sidecar(self) -> Path:
        return self.workspace_dir / "worker.session"

    @property
    def attempts_dir(self) -> Path:
        return self.workspace_dir / "attempts"

    def ensure_present(self) -> None:
        """Create the workspace skeleton if absent. Idempotent."""
        self.workspace_dir.mkdir(parents=True, exist_ok=True)
        self.attempts_dir.mkdir(exist_ok=True)

    def clone_if_absent(self) -> None:
        """``git clone`` into ``worktree/`` if not already cloned.

        Idempotent — a re-call on an already-cloned worktree is a
        no-op. The clone uses ``--no-single-branch`` so subsequent
        ``git fetch`` can pick up new branches.

        Tests can pre-populate ``worktree/`` (e.g., via ``git init``)
        and skip this call.
        """
        if (self.worktree_dir / ".git").exists():
            return
        if self.clone_url is None:
            raise RuntimeError(
                f"clone_url not set; cannot clone into {self.worktree_dir}"
            )
        self.ensure_present()
        # Remove any partial state before cloning.
        if self.worktree_dir.exists():
            shutil.rmtree(self.worktree_dir)
        _git_run(
            ["git", "clone", "--no-single-branch", self.clone_url, str(self.worktree_dir)],
            cwd=None,
        )

    def fetch_and_validate(
        self, *, expected_head_sha: str, head_ref: str
    ) -> None:
        """Refresh remote state + verify the workspace is at the
        expected head_sha.

        - ``git fetch origin --prune`` to pull latest refs.
        - Read ``origin/<head_ref>``; if it differs from
          ``expected_head_sha`` raise ``StaleInputError``.
        - ``git reset --hard <expected_head_sha>`` + ``git clean -fdx``
          to wipe any worktree residue.

        v6 stale-input fix: the worker reports the mismatch as
        ``outcome='stale-input'`` so the master re-prefetches without
        bumping ``pickup_count``.
        """
        # 1. Fetch.
        _git_run(["git", "fetch", "origin", "--prune"], cwd=self.worktree_dir)

        # 2. Compare expected vs current.
        current = _git_run(
            ["git", "rev-parse", f"origin/{head_ref}"],
            cwd=self.worktree_dir,
        ).strip()
        if current != expected_head_sha:
            raise StaleInputError(
                expected=expected_head_sha, actual=current, head_ref=head_ref
            )

        # 3. Reset worktree cleanly to the expected head.
        _git_run(
            ["git", "reset", "--hard", expected_head_sha],
            cwd=self.worktree_dir,
        )
        _git_run(["git", "clean", "-fdx"], cwd=self.worktree_dir)

    def remove(self) -> None:
        """Recursively delete the workspace dir. Used on terminal
        workflow transitions and by the startup janitor."""
        if self.workspace_dir.exists():
            shutil.rmtree(self.workspace_dir)


def _git_run(cmd: list[str], cwd: Path | None) -> str:
    """Run a git command; return stdout. Raises on non-zero exit."""
    try:
        result = subprocess.run(
            cmd, cwd=str(cwd) if cwd else None,
            capture_output=True, text=True, check=True, timeout=120,
        )
    except subprocess.CalledProcessError as exc:
        raise RuntimeError(
            f"git command failed: {' '.join(cmd)}\n"
            f"  stderr: {exc.stderr.strip()}\n"
            f"  stdout: {exc.stdout.strip()}"
        ) from exc
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError(
            f"git command timed out (120s): {' '.join(cmd)}"
        ) from exc
    return result.stdout


__all__ = [
    "DEFAULT_WORKSPACE_ROOT",
    "PerPRWorkspace",
    "StaleInputError",
    "WorkspaceIdentity",
]