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)

    # Credential-helper invocation that sources the Forgejo token from
    # the ``FORGEJO_TOKEN`` env var at fetch/push time. Keeps the token
    # out of ``.git/config`` (where any agent with filesystem read
    # could ``cat .git/config`` to exfiltrate). The helper script is
    # quoted because git's credential.helper accepts shell strings —
    # see ``git help credentials``. ``test "$1" = get`` ensures we
    # only respond to the get action; store/erase become no-ops.
    _CREDENTIAL_HELPER = (
        '!f() { test "$1" = "get" && '
        'echo "username=x" && '
        'echo "password=${FORGEJO_TOKEN:-}"; }; f'
    )

    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.

        Token-handling: the clone_url MUST NOT contain credentials.
        A local ``credential.helper`` is configured immediately after
        clone so subsequent fetch/push operations source the token
        from ``$FORGEJO_TOKEN`` at runtime — keeps the token out of
        ``.git/config``.

        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)
        # Use -c credential.helper at clone time so the initial fetch
        # can authenticate; then bake the same helper into the local
        # repo config for subsequent git ops.
        _git_run(
            [
                "git",
                "-c",
                f"credential.helper={self._CREDENTIAL_HELPER}",
                "clone",
                "--no-single-branch",
                self.clone_url,
                str(self.worktree_dir),
            ],
            cwd=None,
        )
        _git_run(
            [
                "git",
                "config",
                "--local",
                "credential.helper",
                self._CREDENTIAL_HELPER,
            ],
            cwd=self.worktree_dir,
        )

    def _preserve_unsaved_work(self, reset_target: str) -> str | None:
        """Snapshot any work the imminent ``git reset --hard`` would
        discard — onto a local ``auto-scratch/pr-<N>`` branch — so the
        next attempt can inspect/adopt it.

        Two kinds of work are lost on a hard reset, and both are
        captured here:

        - **Uncommitted edits.** A timed-out implementer attempt leaves
          its edits dirty in the worktree.
        - **Committed-but-unpushed commits.** An implementer can commit
          a finished fix and then fail to *push* it (a push race, a
          worker crash, a canonical-output miss). Those commits sit on
          ``HEAD`` ahead of ``reset_target``; a plain hard reset drops
          them silently. Preserving only dirty work (the
          pre-2026-05-20 behavior) lost exactly the case that
          dead-ended PR #39 in run-25.

        The next attempt can inspect the snapshot
        (``git diff HEAD..auto-scratch/pr-<N>``) and cherry-pick
        anything useful — or ignore it. The branch is local-only and
        force-updated per PR, so it never touches the PR branch and
        does not accumulate.

        Best-effort: never raises — a preservation failure must not
        block the attempt. Returns the branch name when something was
        preserved, else None.
        """
        try:
            status = _git_run(["git", "status", "--porcelain"], cwd=self.worktree_dir)
            dirty = bool(status.strip())
            head = _git_run(
                ["git", "rev-parse", "HEAD"], cwd=self.worktree_dir
            ).strip()
            head_ahead = bool(head) and head != reset_target
            if not dirty and not head_ahead:
                return None  # clean worktree at the reset target — nothing to lose
            branch = f"auto-scratch/pr-{self.identity.entity_number}"
            if dirty:
                # write-tree + commit-tree snapshots the index WITHOUT
                # moving HEAD or any branch; `git add -A` first so the
                # snapshot includes new (untracked) files too. The
                # commit is parented on HEAD, so it captures any
                # committed-ahead commits AND the dirty edits at once.
                _git_run(["git", "add", "-A"], cwd=self.worktree_dir)
                tree = _git_run(["git", "write-tree"], cwd=self.worktree_dir).strip()
                # Explicit identity: a freshly-cloned worktree may have
                # no user.name/user.email, and commit-tree refuses
                # without one.
                commit = _git_run(
                    [
                        "git",
                        "-c",
                        "user.email=auto-scratch@cleveragents.local",
                        "-c",
                        "user.name=cleveragents-auto-scratch",
                        "commit-tree",
                        tree,
                        "-p",
                        "HEAD",
                        "-m",
                        "auto-scratch: unsaved work preserved before "
                        f"reset (PR #{self.identity.entity_number})",
                    ],
                    cwd=self.worktree_dir,
                ).strip()
            else:
                # Clean worktree but HEAD carries unpushed commits —
                # HEAD itself already is the snapshot; no new commit
                # needed.
                commit = head
            _git_run(
                ["git", "branch", "-f", branch, commit],
                cwd=self.worktree_dir,
            )
            logger.info(
                "preserved unsaved work on %s (%s; dirty=%s head_ahead=%s)",
                branch,
                commit[:12],
                dirty,
                head_ahead,
            )
            return branch
        except RuntimeError as exc:
            logger.warning("could not preserve unsaved work: %s", exc)
            return 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
            )

        # 2b. Preserve any work the reset below would discard — both
        # uncommitted residue AND committed-but-unpushed commits from a
        # prior attempt — onto auto-scratch/pr-<N>.
        self._preserve_unsaved_work(expected_head_sha)

        # 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",
]