cleveragents-core/tools/_pr_clone.py

"""Pre-clone a PR's working copy for the auto-agents dispatchers.

Shared by the reviewer and implementer dispatchers. Reviewer falls
back here when the embedded diff is insufficient; implementer needs
an editable copy under ``/tmp/`` so its worker's ``edit: /tmp/**``
rule covers the source it must change.

Knobs (canonical names; the four shared ones honour their
``REVIEW_DISPATCHER_*`` legacy spellings with a one-shot
deprecation warning):

- ``DISPATCHER_MIRROR_PATH`` / ``DISPATCHER_GIT_FETCH_TIMEOUT_S`` /
  ``DISPATCHER_GIT_CLONE_TIMEOUT_S`` /
  ``DISPATCHER_MIRROR_MAX_STALENESS_S`` -- shared mirror & timeouts.
- ``REVIEW_DISPATCHER_WORKTREE_BASE`` /
  ``IMPLEMENTER_DISPATCHER_WORKTREE_BASE`` -- per-kind worktree base
  (path shape ``pr-{n}-{kind}-{run_tag}`` for tail-the-log diagnostics).
- ``REVIEW_DISPATCHER_DISABLE_PRECLONE`` /
  ``IMPLEMENTER_DISPATCHER_DISABLE_PRECLONE`` -- per-kind kill-switch.
- ``IMPLEMENTER_DISPATCHER_PRECLONE`` -- Phase-3 feature flag,
  default-off until rollout (reviewer has no feature flag).

Read-only with respect to upstream. Credential isolation lives in
:mod:`_pr_clone_creds`.
"""
from __future__ import annotations

import logging
import os
import shutil
import subprocess
import sys
import time
import uuid
from dataclasses import dataclass
from pathlib import Path
from typing import Any, TypedDict

_TOOLS_DIR = str(Path(__file__).resolve().parent)
if _TOOLS_DIR not in sys.path:
    sys.path.insert(0, _TOOLS_DIR)
from _loader import load_sibling as _load_sibling  # noqa: E402  type: ignore[import-not-found]


_pr_clone_creds = _load_sibling("_pr_clone_creds", "_pr_clone_creds.py")
_ensure_askpass_script = _pr_clone_creds._ensure_askpass_script
_git_env = _pr_clone_creds._git_env
_GIT_ENV_PASSTHROUGH = _pr_clone_creds._GIT_ENV_PASSTHROUGH

_logger = logging.getLogger("pr_clone")


# ─── Defaults + legacy-aware env reader ────────────────────────────────────

_DEFAULT_MIRROR_PATH = "/tmp/.cleveragents-mirror.git"
_DEFAULT_FETCH_TIMEOUT_S = 300
_DEFAULT_CLONE_TIMEOUT_S = 900
# Mirror staleness is a CACHE, not a freshness guarantee:
# ``_ensure_mirror`` skips ``git fetch --prune`` when mtime is
# younger than this. 60s amortises refreshes inside a single
# dispatch cycle (default 300s).
_DEFAULT_MIRROR_STALENESS_S = 60

_LEGACY_DEPRECATION_LOGGED: set[str] = set()


def _legacy_or_canonical_env(canonical: str, legacy: str) -> str | None:
    """Read ``canonical``; fall back to ``legacy`` with a one-shot
    deprecation log so operators see the warning once per process,
    not once per call. Returns ``None`` when neither is set."""
    val = os.environ.get(canonical)
    if val:
        return val
    legacy_val = os.environ.get(legacy)
    if legacy_val:
        if legacy not in _LEGACY_DEPRECATION_LOGGED:
            _logger.warning(
                "%s is deprecated and shared between both dispatchers; "
                "please rename to %s",
                legacy,
                canonical,
            )
            _LEGACY_DEPRECATION_LOGGED.add(legacy)
        return legacy_val
    return None


# ─── Per-kind config ───────────────────────────────────────────────────────


class _KindConfig(TypedDict):
    """Shape of each ``_KIND_CONFIG`` entry. Typed so a typo'd key
    is caught by the type checker, not at runtime."""

    worktree_base_env: str
    default_worktree_base: str
    disable_env: str
    enable_env: str | None  # None → no feature flag, always-enabled


# ``enable_env`` is set only for kinds whose pre-clone is gated
# behind a Phase-3 rollout flag (default-off until Phase 4 flips
# it). ``None`` means always-enabled, kill-switch is the only knob.
_KIND_CONFIG: dict[str, _KindConfig] = {
    "review": {
        "worktree_base_env": "REVIEW_DISPATCHER_WORKTREE_BASE",
        "default_worktree_base": "/tmp/cleveragents-review-worktrees",
        "disable_env": "REVIEW_DISPATCHER_DISABLE_PRECLONE",
        "enable_env": None,
    },
    "implementer": {
        "worktree_base_env": "IMPLEMENTER_DISPATCHER_WORKTREE_BASE",
        "default_worktree_base": "/tmp/cleveragents-implementer-worktrees",
        "disable_env": "IMPLEMENTER_DISPATCHER_DISABLE_PRECLONE",
        "enable_env": "IMPLEMENTER_DISPATCHER_PRECLONE",
    },
}


def _kind_cfg(kind: str) -> _KindConfig:
    """Look up ``kind``. A typo'd kind logs ``ERROR`` and falls back
    to ``"review"`` rather than raising so one misconfigured caller
    cannot crash a dispatch cycle. ``ERROR`` (not ``WARNING``) so
    the misroute is loud enough to surface in operator alerting:
    silent fall-back to the reviewer worktree base is hazardous
    enough that the dispatcher contract treats it as a programmer
    bug, not a transient anomaly."""
    if kind not in _KIND_CONFIG:
        _logger.error(
            "unknown kind=%r; expected one of %s; falling back to 'review'",
            kind,
            sorted(_KIND_CONFIG),
        )
        return _KIND_CONFIG["review"]
    return _KIND_CONFIG[kind]


# ─── Paths ─────────────────────────────────────────────────────────────────


_TRUTHY_ENV_VALUES = frozenset({"1", "true", "yes", "on"})


def _mirror_path() -> Path:
    return Path(
        _legacy_or_canonical_env(
            "DISPATCHER_MIRROR_PATH", "REVIEW_DISPATCHER_MIRROR_PATH"
        )
        or _DEFAULT_MIRROR_PATH
    )


def _worktree_base_with_cfg(cfg: _KindConfig) -> Path:
    """Production gate twin: shares one cfg across consumers."""
    return Path(
        os.environ.get(cfg["worktree_base_env"])
        or cfg["default_worktree_base"]
    )


def _worktree_base(kind: str = "review") -> Path:
    """Test/direct-caller wrapper around :func:`_worktree_base_with_cfg`."""
    return _worktree_base_with_cfg(_kind_cfg(kind))


def _env_truthy(name: str) -> bool:
    return os.environ.get(name, "").strip().lower() in _TRUTHY_ENV_VALUES


def _is_preclone_disabled_with_cfg(cfg: _KindConfig) -> bool:
    """Production gate twin: shares one cfg across consumers."""
    return _env_truthy(cfg["disable_env"])


def _is_preclone_feature_enabled_with_cfg(cfg: _KindConfig) -> bool:
    """Production gate twin: shares one cfg across consumers."""
    enable_env = cfg["enable_env"]
    return True if enable_env is None else _env_truthy(enable_env)


def _is_preclone_disabled(kind: str = "review") -> bool:
    """Test/direct-caller wrapper; the production gate path uses
    :func:`_is_preclone_disabled_with_cfg` to avoid a duplicate
    ``_kind_cfg`` lookup."""
    return _is_preclone_disabled_with_cfg(_kind_cfg(kind))


def _is_preclone_feature_enabled(kind: str = "review") -> bool:
    """Test/direct-caller wrapper; the production gate path uses
    :func:`_is_preclone_feature_enabled_with_cfg`."""
    return _is_preclone_feature_enabled_with_cfg(_kind_cfg(kind))


def _fetch_timeout_s() -> int:
    raw = _legacy_or_canonical_env(
        "DISPATCHER_GIT_FETCH_TIMEOUT_S", "REVIEW_DISPATCHER_GIT_FETCH_TIMEOUT_S"
    )
    return int(raw or _DEFAULT_FETCH_TIMEOUT_S)


def _clone_timeout_s() -> int:
    raw = _legacy_or_canonical_env(
        "DISPATCHER_GIT_CLONE_TIMEOUT_S", "REVIEW_DISPATCHER_GIT_CLONE_TIMEOUT_S"
    )
    return int(raw or _DEFAULT_CLONE_TIMEOUT_S)


def _mirror_max_staleness_s() -> int:
    raw = _legacy_or_canonical_env(
        "DISPATCHER_MIRROR_MAX_STALENESS_S", "REVIEW_DISPATCHER_MIRROR_MAX_STALENESS_S"
    )
    return int(raw or _DEFAULT_MIRROR_STALENESS_S)


# ─── Bare mirror lifecycle ─────────────────────────────────────────────────


# ``FORGEJO_URL`` may point at the API root (REST helpers need it);
# git clones from the web root, so strip the API suffix here.
_FORGEJO_API_SUFFIXES = ("/api/v1", "/api")


def _clone_url(cfg: Any) -> str:
    """Build the HTTPS URL git clones from. Credentials are
    delivered via ``GIT_ASKPASS``, never embedded in the URL."""
    base = str(getattr(cfg, "forgejo_url", "")).rstrip("/")
    for suffix in _FORGEJO_API_SUFFIXES:
        if base.endswith(suffix):
            base = base[: -len(suffix)]
            break
    owner = getattr(cfg, "owner", "")
    repo = getattr(cfg, "repo", "")
    return f"{base}/{owner}/{repo}.git"


def _ensure_mirror(cfg: Any) -> Path | None:
    """Create the bare mirror on first use, refresh it otherwise.
    Returns the path on success, ``None`` on failure."""
    mirror = _mirror_path()
    env = _git_env(cfg)
    if not mirror.exists():
        mirror.parent.mkdir(parents=True, exist_ok=True)
        url = _clone_url(cfg)
        _logger.info("creating bare mirror at %s from %s", mirror, url)
        try:
            subprocess.run(
                ["git", "clone", "--mirror", url, str(mirror)],
                env=env,
                check=True,
                timeout=_clone_timeout_s(),
                capture_output=True,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
            _logger.warning("bare-mirror clone failed: %s", exc)
            try:
                if mirror.exists():
                    shutil.rmtree(mirror)
            except OSError:
                pass
            return None
        return mirror
    head = mirror / "HEAD"
    try:
        age = time.time() - head.stat().st_mtime if head.exists() else float("inf")
    except OSError:
        age = float("inf")
    if age >= _mirror_max_staleness_s():
        _logger.info("refreshing bare mirror at %s (age=%.0fs)", mirror, age)
        try:
            subprocess.run(
                ["git", "--git-dir", str(mirror), "fetch", "--prune", "origin"],
                env=env,
                check=True,
                timeout=_fetch_timeout_s(),
                capture_output=True,
            )
        except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
            _logger.warning("mirror fetch failed: %s", exc)
            # Stale-but-usable beats no clone at all.
    return mirror


# ─── Worktree handle ───────────────────────────────────────────────────────


@dataclass
class WorktreeHandle:
    """A pre-cloned working copy for one PR session.

    ``path`` is the worktree the worker can ``read`` (and, for the
    implementer, ``edit``). ``cleanup`` removes it from the
    dispatcher's ``post_session_action`` finally branch. ``kind``
    records which dispatcher prepared the handle so cleanup logs
    can be correlated with the source dispatcher.
    """

    path: Path
    mirror: Path
    head_sha: str
    pr_number: int
    kind: str = "review"

    def cleanup(self, cfg: Any) -> None:
        _logger.info(
            "cleaning up %s worktree at %s for PR #%s",
            self.kind,
            self.path,
            self.pr_number,
        )
        env = _git_env(cfg)
        try:
            subprocess.run(
                [
                    "git",
                    "--git-dir",
                    str(self.mirror),
                    "worktree",
                    "remove",
                    "--force",
                    str(self.path),
                ],
                env=env,
                check=False,
                timeout=60,
                capture_output=True,
            )
        except subprocess.TimeoutExpired:
            _logger.warning(
                "worktree remove timed out for %s PR #%s at %s; falling "
                "back to filesystem rmtree",
                self.kind,
                self.pr_number,
                self.path,
            )
        # Belt-and-suspenders: even if `worktree remove` succeeded
        # the mirror sometimes leaves stale bookkeeping. Prune.
        try:
            subprocess.run(
                ["git", "--git-dir", str(self.mirror), "worktree", "prune"],
                env=env,
                check=False,
                timeout=60,
                capture_output=True,
            )
        except subprocess.TimeoutExpired:
            pass
        try:
            if self.path.exists():
                shutil.rmtree(self.path, ignore_errors=True)
        except OSError:
            pass


# ─── Public API ────────────────────────────────────────────────────────────


def prepare_pr_worktree(
    cfg: Any,
    pr_number: int,
    head_sha: str,
    *,
    kind: str = "review",
) -> WorktreeHandle | None:
    """Clone (or refresh + worktree-add) the PR's head_sha into a
    fresh path under the kind-specific worktree base.

    ``kind`` selects the worktree base, the kill-switch, and the
    optional Phase-3 feature flag. Returns ``None`` when the
    pre-clone is gated off or fails; the caller falls back to
    ``git-isolator-util`` in that case.
    """
    # Resolve cfg ONCE (typo'd kind logs the fall-through ERROR
    # once per call) and normalise ``kind`` to the effective
    # fall-back so path filenames + ``WorktreeHandle.kind`` match
    # the base directory the worktree lives under.
    cfg_entry = _kind_cfg(kind)
    if kind not in _KIND_CONFIG:
        kind = "review"
    if not _is_preclone_feature_enabled_with_cfg(cfg_entry):
        _logger.info(
            "PR #%s %s pre-clone gated off (feature flag %s not set)",
            pr_number,
            kind,
            cfg_entry["enable_env"],
        )
        return None
    if _is_preclone_disabled_with_cfg(cfg_entry):
        _logger.info(
            "PR #%s %s pre-clone disabled via %s; worker will fall back "
            "to git-isolator-util if needed",
            pr_number,
            kind,
            cfg_entry["disable_env"],
        )
        return None
    if not head_sha:
        _logger.warning("PR #%s %s pre-clone skipped: no head_sha", pr_number, kind)
        return None
    mirror = _ensure_mirror(cfg)
    if mirror is None:
        return None
    base = _worktree_base_with_cfg(cfg_entry)
    base.mkdir(parents=True, exist_ok=True)
    # 8-char UUID prevents collisions when two cycles of the same
    # kind hit the same PR.
    run_tag = uuid.uuid4().hex[:8]
    path = base / f"pr-{int(pr_number)}-{kind}-{run_tag}"
    env = _git_env(cfg)
    # Confirm SHA is in mirror; staleness is best-effort. A force-push
    # inside the staleness window can leave ``worktree add`` failing
    # with "fatal: invalid reference" — a targeted fetch closes the race.
    if not _fetch_specific_sha(mirror, head_sha, env):
        return None
    try:
        subprocess.run(
            [
                "git",
                "--git-dir",
                str(mirror),
                "worktree",
                "add",
                "--detach",
                str(path),
                head_sha,
            ],
            env=env,
            check=True,
            timeout=_clone_timeout_s(),
            capture_output=True,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
        _logger.warning(
            "worktree add for PR #%s @ %s failed: %s",
            pr_number,
            head_sha[:12],
            exc,
        )
        return None
    return WorktreeHandle(
        path=path,
        mirror=mirror,
        head_sha=head_sha,
        pr_number=int(pr_number),
        kind=kind,
    )


def _fetch_specific_sha(
    mirror: Path, head_sha: str, env: dict[str, str]
) -> bool:
    """Ensure ``head_sha`` exists in the bare mirror, fetching it if
    necessary. Returns True iff the SHA is present after the call.
    """
    try:
        result = subprocess.run(
            [
                "git",
                "--git-dir",
                str(mirror),
                "cat-file",
                "-e",
                f"{head_sha}^{{commit}}",
            ],
            env=env,
            check=False,
            timeout=10,
            capture_output=True,
        )
    except subprocess.TimeoutExpired:
        result = None  # treat as "missing", attempt the fetch below
    if result is not None and result.returncode == 0:
        return True
    _logger.info(
        "mirror missing %s; targeted fetch", head_sha[:12]
    )
    try:
        subprocess.run(
            [
                "git",
                "--git-dir",
                str(mirror),
                "fetch",
                "origin",
                head_sha,
            ],
            env=env,
            check=True,
            timeout=_fetch_timeout_s(),
            capture_output=True,
        )
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as exc:
        _logger.warning(
            "targeted fetch for %s failed: %s", head_sha[:12], exc
        )
        return False
    return True


__all__ = ("WorktreeHandle", "prepare_pr_worktree")