cleveragents-core/tools/_claim_runtime.py

"""Shared runtime primitives for the auto-agents drivers.

Extracted verbatim from :mod:`tools.merge_drive` per
``docs/development/conflict-drive-plan.md`` § 4 so that
:mod:`tools.merge_drive` and the upcoming :mod:`tools.conflict_drive`
share a single source of truth for:

- HTTP client (``idempotent_get`` / ``state_change`` / ``get`` / ``post``
  / ``patch`` / ``delete`` and ``APIError``);
- Single-instance lock + heartbeat (``SingleInstanceLock`` /
  ``write_heartbeat``);
- Cooperative cancellation (``StopEvent``);
- Forgejo claim labels (``claim_pr`` / ``release_pr_claim`` /
  ``sweep_expired_claims`` / ``sweep_all_expired_claims``) and the
  underlying label add / remove helpers.

The extraction is behaviour-preserving: :mod:`tools.merge_drive`
re-exports the public names so existing callers and unit tests
continue to work. New code should import from ``_claim_runtime``
directly.

This module is loaded via :func:`importlib.util.spec_from_file_location`
to mirror the loader contract used by ``tools/merge_drive.py`` and
``tools/_pipeline_cache.py`` (the ``tools/`` directory is not a package).
"""

from __future__ import annotations

import fcntl
import json
import os
import sys
import time
import urllib.error
import urllib.request
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Protocol


# Spec #6 — ``_find_newest_claim_at`` routes through ``_pr_comments_cache``
# so a heavy PR's claim-marker scan reuses the same delta cache the
# implementer prefetch already maintains. Module-loaded lazily on first
# call: ``_pr_comments_cache`` imports ``_review_fetch``, which loads
# ``_claim_runtime`` (us) at module scope — eager-importing the cache
# here would form a circle. Lazy keeps the import graph acyclic and
# only pays the load cost once per process.
_PR_COMMENTS_CACHE_MOD: Any | None = None


def _get_comments_cache() -> Any:
    """Return the lazily-loaded ``_pr_comments_cache`` module."""
    global _PR_COMMENTS_CACHE_MOD
    if _PR_COMMENTS_CACHE_MOD is None:
        tools_dir = str(Path(__file__).resolve().parent)
        if tools_dir not in sys.path:
            sys.path.insert(0, tools_dir)
        from _loader import (  # type: ignore[import-not-found]
            load_sibling,
        )

        _PR_COMMENTS_CACHE_MOD = load_sibling(
            "_pr_comments_cache", "_pr_comments_cache.py"
        )
    return _PR_COMMENTS_CACHE_MOD


# ─── Forgejo coordinates (read from env at import) ─────────────────────────

REPO_OWNER = os.environ.get("FORGEJO_OWNER", "cleveragents")
REPO_NAME = os.environ.get("FORGEJO_REPO", "cleveragents-core")
ORG_NAME = os.environ.get("FORGEJO_ORG", "cleveragents")
API_BASE = os.environ.get(
    "FORGEJO_API_BASE", "https://git.cleverthis.com/api/v1"
).rstrip("/")


class RuntimeContext(Protocol):
    """Duck-typed context object the runtime helpers expect.

    Both :class:`merge_drive.DriverConfig` and
    :class:`conflict_drive.DriverConfig` satisfy this protocol; the
    runtime layer only needs the four fields below and never depends on
    driver-specific knobs.
    """

    token: str
    request_timeout_s: int
    api_retries: int
    claim_ttl_seconds: int


# ─── HTTP / API ────────────────────────────────────────────────────────────


class APIError(Exception):
    """Raised when an API call returns an unexpected status."""

    def __init__(self, status: int, body: Any, url: str):
        super().__init__(f"HTTP {status} from {url}: {body!r}")
        self.status = status
        self.body = body
        self.url = url


def _do_request(
    method: str,
    path: str,
    cfg: RuntimeContext,
    body: Any | None,
    token: str | None = None,
) -> dict[str, Any]:
    """Single fire-and-forget request. Used by both retry policies."""
    url = f"{API_BASE}{path}" if path.startswith("/") else path
    data = json.dumps(body).encode() if body is not None else None
    headers = {"Authorization": f"token {token or cfg.token}"}
    if data is not None:
        headers["Content-Type"] = "application/json"
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=cfg.request_timeout_s) as resp:
            payload = resp.read()
            ct = resp.headers.get("Content-Type", "")
            return {
                "status": resp.status,
                "body": (
                    json.loads(payload)
                    if payload and "application/json" in ct
                    else payload
                ),
            }
    except urllib.error.HTTPError as e:
        try:
            body_parsed = json.loads(e.read())
        except (ValueError, TypeError):
            body_parsed = ""
        return {"status": e.code, "body": body_parsed}


def idempotent_get(path: str, cfg: RuntimeContext) -> dict[str, Any]:
    """GET with full retry policy: retry on transport errors AND 5xx up to
    ``cfg.api_retries`` times with exponential backoff. 4xx is returned to
    the caller without retry (authoritative).
    """
    last_err: Exception | None = None
    for attempt in range(1, cfg.api_retries + 1):
        try:
            res = _do_request("GET", path, cfg, None)
            if res["status"] >= 500 and attempt < cfg.api_retries:
                time.sleep(2.0 * attempt)
                continue
            return res
        except (urllib.error.URLError, TimeoutError, OSError) as e:
            last_err = e
            if attempt < cfg.api_retries:
                time.sleep(2.0 * attempt)
                continue
    raise RuntimeError(f"network error contacting {path}: {last_err}")


def state_change(
    method: str,
    path: str,
    cfg: RuntimeContext,
    body: Any | None = None,
    *,
    token: str | None = None,
) -> dict[str, Any]:
    """Mutation with the strict state-change retry policy:

    - Network/connection errors: retry **at most once** (one transient hop
      can drop the response after the server accepted; we accept that risk
      because most of our writes are idempotent at the HTTP level given
      ``head_commit_id`` / label-already-present checks).
    - HTTP errors (4xx and 5xx alike): NEVER retried. 4xx is authoritative;
      5xx on a write may have applied state already, so we hand control to
      the caller to decide.
    """
    try:
        return _do_request(method, path, cfg, body, token=token)
    except (urllib.error.URLError, TimeoutError, OSError):
        try:
            return _do_request(method, path, cfg, body, token=token)
        except (urllib.error.URLError, TimeoutError, OSError) as e2:
            raise RuntimeError(
                f"network error after one retry on {method} {path}: {e2}"
            ) from e2


def get(path: str, cfg: RuntimeContext) -> dict[str, Any]:
    return idempotent_get(path, cfg)


def post(path: str, cfg: RuntimeContext, body: Any) -> dict[str, Any]:
    return state_change("POST", path, cfg, body)


def patch(path: str, cfg: RuntimeContext, body: Any) -> dict[str, Any]:
    return state_change("PATCH", path, cfg, body)


def delete(path: str, cfg: RuntimeContext) -> dict[str, Any]:
    return state_change("DELETE", path, cfg)


# ─── Single-instance lock + heartbeat ──────────────────────────────────────


class SingleInstanceLock:
    """``fcntl.flock`` on a stable path so only one driver runs at a time."""

    def __init__(self, path: Path):
        self.path = path
        self._fd: int | None = None

    def acquire(self) -> bool:
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._fd = os.open(str(self.path), os.O_RDWR | os.O_CREAT, 0o644)
        try:
            fcntl.flock(self._fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            os.close(self._fd)
            self._fd = None
            return False
        os.write(self._fd, f"{os.getpid()}\n".encode())
        os.fsync(self._fd)
        return True

    def release(self) -> None:
        if self._fd is None:
            return
        try:
            fcntl.flock(self._fd, fcntl.LOCK_UN)
        finally:
            os.close(self._fd)
            self._fd = None


def write_heartbeat(path: Path) -> None:
    """Atomic write of the current timestamp + pid."""
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_text(
        json.dumps(
            {
                "pid": os.getpid(),
                "timestamp": datetime.now(timezone.utc).isoformat(),
            }
        )
    )
    tmp.replace(path)


# ─── Stop event (SIGTERM/SIGINT propagation into long-running phases) ─────


class StopEvent:
    """Cooperative cancellation token. Set by the signal handler at the
    outer loop boundary; checked inside ``wait_for_ci`` and other long
    waits so SIGTERM doesn't have to block for the full CI timeout.
    """

    def __init__(self) -> None:
        self._stopped = False

    def set(self) -> None:
        self._stopped = True

    def is_set(self) -> bool:
        return self._stopped

    def sleep(self, seconds: float) -> None:
        """Sleep up to ``seconds``, returning early if stop is requested.
        Granularity ~1 s.
        """
        end = time.monotonic() + seconds
        while not self._stopped and time.monotonic() < end:
            time.sleep(min(1.0, max(0.0, end - time.monotonic())))


# ─── Claim labels (auto/claimed-merge add / release / sweep) ──────────────
#
# Per the AGENTS.md "Auto-* label registry": auto/claimed-merge marks the
# PR as currently being processed by the merge driver instance. The label
# MUST be added before any mutation and removed at the end of the cycle
# (success or failure) so operators can see ownership in the Forgejo UI
# and so a crashed driver doesn't leave PRs stuck "claimed forever".
#
# The ``conflict_drive.py`` driver reuses this same label so the merge
# driver and the conflict driver are mutually exclusive on a given PR.

CLAIM_LABEL = "auto/claimed-merge"

# Per-driver claim/release markers (P1-8). Each driver writes its own
# marker so an operator scanning a PR's audit log can immediately tell
# which driver claimed/released without parsing the comment body.
# ``CLAIM_COMMENT_MARKER`` / ``CLAIM_RELEASE_MARKER`` retain the legacy
# merge-driver values so external automation that scrapes them keeps
# working unchanged.
CLAIM_COMMENT_MARKER = "<!-- merge_drive.py: claim -->"
CLAIM_RELEASE_MARKER = "<!-- merge_drive.py: release -->"
CONFLICT_CLAIM_COMMENT_MARKER = "<!-- conflict_drive.py: claim -->"
CONFLICT_RELEASE_MARKER = "<!-- conflict_drive.py: release -->"

# Markers we *recognise* during the sweep. We accept ``claim_pr.ts``'s
# marker so an operator using that helper to manually claim
# ``auto/claimed-merge`` isn't immediately stripped on the next sweep —
# every recognised marker is honoured for TTL purposes regardless of
# which driver wrote it.
CLAIM_COMMENT_MARKERS_RECOGNISED: tuple[str, ...] = (
    CLAIM_COMMENT_MARKER,
    CONFLICT_CLAIM_COMMENT_MARKER,
    "<!-- claim_pr.ts: do-not-edit -->",
)

# Tier 1 mutual-respect contract: the driver sweeps stale claims for every
# auto/claimed-* label, not just its own. This protects against orphaned
# implementer / reviewer claims left behind by a crashed worker so the PR
# does not stay invisible forever to every supervisor (Patch A's
# excludeClaimed filter would otherwise hide it). The driver itself only
# ever WRITES auto/claimed-merge — these other two labels are written by
# claim_pr.ts when invoked from pr-review-worker / implementation-worker.
ALL_CLAIM_LABELS: tuple[str, ...] = (
    CLAIM_LABEL,
    "auto/claimed-implementer",
    "auto/claimed-reviewer",
)


_LABEL_ID_CACHE: dict[str, int] = {}


def lookup_label_id(name: str, cfg: RuntimeContext) -> int | None:
    """Resolve a label name to its numeric id, checking repo then org.
    Cached for the lifetime of the process.
    """
    if name in _LABEL_ID_CACHE:
        return _LABEL_ID_CACHE[name]
    for path in (
        f"/repos/{REPO_OWNER}/{REPO_NAME}/labels",
        f"/orgs/{ORG_NAME}/labels",
    ):
        page = 1
        while True:
            res = get(f"{path}?limit=50&page={page}", cfg)
            if res["status"] != 200:
                break
            chunk = res["body"] or []
            for lbl in chunk:
                if lbl.get("name") == name:
                    label_id = int(lbl.get("id"))
                    _LABEL_ID_CACHE[name] = label_id
                    return label_id
            if len(chunk) < 50:
                break
            page += 1
    return None


def _add_label(pr_number: int, label_name: str, cfg: RuntimeContext) -> bool:
    """Add a label to the PR. Returns True iff the API call succeeded.
    Idempotent at the API level — Forgejo no-ops on already-attached labels.
    """
    label_id = lookup_label_id(label_name, cfg)
    if label_id is None:
        return False
    res = post(
        f"/repos/{REPO_OWNER}/{REPO_NAME}/issues/{pr_number}/labels",
        cfg,
        {"labels": [label_id]},
    )
    return res["status"] in (200, 201, 204)


def _remove_label(pr_number: int, label_name: str, cfg: RuntimeContext) -> bool:
    """Remove a label from the PR. Returns True iff a delete was attempted
    (HTTP 404 from the underlying call is treated as success — already gone).
    """
    label_id = lookup_label_id(label_name, cfg)
    if label_id is None:
        return False
    res = delete(
        f"/repos/{REPO_OWNER}/{REPO_NAME}/issues/{pr_number}/labels/{label_id}",
        cfg,
    )
    return res["status"] in (200, 204, 404)


def claim_pr(
    pr_number: int,
    cfg: RuntimeContext,
    *,
    dry_run: bool = False,
    driver_name: str = "merge_drive.py",
    claim_marker: str = CLAIM_COMMENT_MARKER,
) -> dict[str, Any]:
    """Acquire ``auto/claimed-merge`` and post a TTL claim comment.

    Forgejo's add-labels endpoint is HTTP-idempotent, so this can run on
    an already-claimed PR (e.g. a previous cycle by THIS driver that the
    sweep would have released). The TTL comment is posted unconditionally;
    duplicate claim comments are intentional — they document each cycle's
    take on the PR and the sweep uses the newest one to compute TTL.

    P1-8: ``driver_name`` and ``claim_marker`` parametrise the comment
    body so the conflict driver's claim comments are attributed to
    ``conflict_drive.py`` rather than ``merge_drive.py``. Defaults
    preserve byte-equivalence for the historic merge-driver path.
    """
    if dry_run:
        return {"applied": False, "dry_run": True, "pr": pr_number}
    if not _add_label(pr_number, CLAIM_LABEL, cfg):
        return {"applied": False, "pr": pr_number, "error": "label-not-found"}
    ttl_until = (
        datetime.now(timezone.utc) + timedelta(seconds=cfg.claim_ttl_seconds)
    ).isoformat()
    body = (
        f"{claim_marker}\n\n"
        f"Claimed by `{driver_name}` (pid {os.getpid()}) until `{ttl_until}`.\n\n"
        f"This claim is advisory and will be released when the cycle ends, "
        f"or after the TTL by a sibling driver's expired-claim sweep."
    )
    post(
        f"/repos/{REPO_OWNER}/{REPO_NAME}/issues/{pr_number}/comments",
        cfg,
        {"body": body},
    )
    return {"applied": True, "pr": pr_number, "ttl_until": ttl_until}


def release_pr_claim(
    pr_number: int,
    cfg: RuntimeContext,
    *,
    terminal_state: str,
    detail: str = "",
    op_label_map: dict[str, str] | None = None,
    driver_name: str = "merge_drive.py",
    release_marker: str = CLAIM_RELEASE_MARKER,
) -> dict[str, Any]:
    """Release the ``auto/claimed-merge`` label and apply the operational
    label that matches ``terminal_state`` in ``op_label_map`` (when one is
    provided). Posts a single release comment so the PR's history shows
    what happened.

    ``op_label_map`` is an injected dependency so each driver supplies
    its own state→label table; the merge driver passes
    ``RELEASE_STATE_LABEL_MAP`` (defined in :mod:`merge_drive`), the
    conflict driver passes its own escalation/needs-implementer map.
    Default ``None`` (no operational label) preserves the legacy
    behaviour of pure-release callers.

    P1-8: ``driver_name`` / ``release_marker`` parametrise the comment
    attribution so each driver's release events are immediately
    distinguishable in the PR audit log without parsing the body text.
    """
    _remove_label(pr_number, CLAIM_LABEL, cfg)
    op_label = (op_label_map or {}).get(terminal_state)
    op_applied = False
    if op_label is not None:
        op_applied = _add_label(pr_number, op_label, cfg)
    body = (
        f"{release_marker}\n\n"
        f"Released by `{driver_name}` (pid {os.getpid()}). "
        f"terminal_state=`{terminal_state}`"
        + (f", op_label=`{op_label}`" if op_label else "")
        + (f"\n\nDetail: {detail}" if detail else "")
    )
    post(
        f"/repos/{REPO_OWNER}/{REPO_NAME}/issues/{pr_number}/comments",
        cfg,
        {"body": body},
    )
    return {"released": True, "operational_label": op_label, "op_applied": op_applied}


def sweep_expired_claims(
    cfg: RuntimeContext,
    *,
    session_pr_numbers: set[int],
    label: str = CLAIM_LABEL,
) -> list[int]:
    """Walk every open PR with ``label``; for each one **not** claimed by
    THIS process (per ``session_pr_numbers``) AND whose newest claim
    comment is older than ``cfg.claim_ttl_seconds``, release the label so
    the orphaned PR re-enters the candidate pool. Returns the list of PR
    numbers swept.

    The ``label`` parameter defaults to ``auto/claimed-merge`` so existing
    callers that only sweep the merge driver's own claim are unchanged.
    Pass any value from ``ALL_CLAIM_LABELS`` (e.g. ``auto/claimed-reviewer``)
    to sweep a foreign claim type. ``session_pr_numbers`` only applies to
    the driver's own label; for foreign labels the driver is never the
    owner so callers should pass ``set()``.

    Idempotency: removing a label that's already gone is a no-op (HTTP 404
    treated as success).
    """
    swept: list[int] = []
    cutoff = datetime.now(timezone.utc) - timedelta(seconds=cfg.claim_ttl_seconds)
    page = 1
    while True:
        res = get(
            f"/repos/{REPO_OWNER}/{REPO_NAME}/issues"
            f"?state=open&type=pulls&labels={label}&limit=50&page={page}",
            cfg,
        )
        if res["status"] != 200:
            break
        chunk = res["body"] or []
        for issue in chunk:
            num = int(issue.get("number") or 0)
            if num in session_pr_numbers or num == 0:
                continue
            newest_claim_at = _find_newest_claim_at(num, cfg)
            # No marker found OR last claim is older than TTL → release.
            if newest_claim_at is None or newest_claim_at < cutoff:
                _remove_label(num, label, cfg)
                swept.append(num)
        if len(chunk) < 50:
            break
        page += 1
    return swept


def sweep_all_expired_claims(
    cfg: RuntimeContext, *, session_pr_numbers: set[int]
) -> dict[str, list[int]]:
    """Sweep stale claims across every label in ``ALL_CLAIM_LABELS`` and
    return a per-label list of swept PR numbers.

    For ``auto/claimed-merge`` the driver is the owner and excludes
    PRs in ``session_pr_numbers`` (this cycle's claims). For
    ``auto/claimed-implementer`` and ``auto/claimed-reviewer`` the driver
    never writes the label, so it passes an empty exclusion set — any PR
    whose claim comment is older than the TTL gets released regardless of
    the driver's session state.

    This is the orphan-safety net for the Tier 1 mutual-respect contract:
    if a reviewer or implementer worker crashes between claim and release,
    the merge driver's next cycle releases the label so the PR re-enters
    the relevant supervisor's candidate pool.
    """
    out: dict[str, list[int]] = {}
    for label in ALL_CLAIM_LABELS:
        excl = session_pr_numbers if label == CLAIM_LABEL else set[int]()
        out[label] = sweep_expired_claims(cfg, session_pr_numbers=excl, label=label)
    return out


def _find_newest_claim_at(pr_number: int, cfg: RuntimeContext) -> datetime | None:
    """Return the newest timestamp of any comment on ``pr_number``
    containing one of ``CLAIM_COMMENT_MARKERS_RECOGNISED`` — or ``None``
    if no recognised marker exists in the cached comment timeline.

    Why "any of": ``merge_drive.py`` writes ``<!-- merge_drive.py: claim
    -->`` on every claim, but the operator-facing ``claim_pr.ts`` helper
    can also legitimately set ``auto/claimed-merge`` and writes its own
    marker (``<!-- claim_pr.ts: do-not-edit -->``). The sweep must
    respect either source so a manual claim isn't immediately overridden
    by the next driver cycle.

    Spec #6 — routes through ``_pr_comments_cache.get_pr_comments`` so a
    heavy PR (e.g. #30 with 1440 comments) pays a 29-page walk only on
    the cold-cache cycle. Subsequent cycles delta-fetch only the new
    comments since the last cached cursor, turning the sweep's per-PR
    cost from O(total_comments) into O(new_comments). Cache misses /
    truncated seeds still pay the full walk that cycle but persist the
    result so the next cycle is fast.

    Iterates the cached list newest-first and returns on the first
    marker match — claim comments are typically among the most recent
    posts, so the early exit keeps the scan cheap even on PRs with
    thousands of cached comments.

    Fail-safe on incomplete data: if the cache returns ``completed=False``
    (transient fetch failure or unfinished backfill of a previously
    truncated seed) AND we did not find a marker, we cannot prove the
    claim is stale — releasing the label on partial data would race
    against an in-flight worker whose claim comment is in the un-fetched
    middle. Return ``datetime.now(UTC)`` so the caller's
    ``newest_claim_at < cutoff`` check keeps the label this cycle; the
    next delta will complete the timeline and let us decide on real data.
    """
    cache_mod = _get_comments_cache()
    comments, completed = cache_mod.get_pr_comments(
        cfg,
        pr_number,
        owner=REPO_OWNER,
        repo=REPO_NAME,
    )
    # Comments are oldest-first (Forgejo creation order, preserved by
    # the cache's append-only merge). Scan REVERSED so the first marker
    # hit is the newest one.
    for c in reversed(comments):
        if not isinstance(c, dict):
            continue
        body = c.get("body") or ""
        if not any(m in body for m in CLAIM_COMMENT_MARKERS_RECOGNISED):
            continue
        created = c.get("created_at") or ""
        try:
            return datetime.fromisoformat(created.replace("Z", "+00:00"))
        except ValueError:
            # Malformed timestamp on this match — keep scanning for an
            # older-but-parseable marker rather than treating it as
            # absent and racing the fail-safe.
            continue
    if not completed:
        # No marker found, but the timeline view is incomplete — fail
        # safe by reporting "claimed just now" so the sweep doesn't
        # release a label whose marker may be in the un-fetched range.
        return datetime.now(timezone.utc)
    return None