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/_forgejo_cache.py`` (the ``tools/`` directory is not a package).
"""
from __future__ import annotations

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


# ─── 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:
    """Walk every page of issue comments on ``pr_number`` and return the
    newest timestamp of any comment containing one of the markers in
    ``CLAIM_COMMENT_MARKERS_RECOGNISED`` (or ``None`` if no recognised
    marker exists on any page).

    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.

    Why full pagination: long-lived PRs (e.g. cycle-cap'd reviews) can
    accumulate hundreds of comments. Reading only the first page would
    let an in-flight claim's comment fall off the end and trigger
    spurious early releases by the sweep.
    """
    newest: datetime | None = None
    page = 1
    while True:
        res = get(
            f"/repos/{REPO_OWNER}/{REPO_NAME}/issues/{pr_number}/comments"
            f"?limit=50&page={page}",
            cfg,
        )
        if res["status"] != 200:
            return newest
        chunk = res["body"] or []
        for c in chunk:
            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:
                when = datetime.fromisoformat(created.replace("Z", "+00:00"))
            except ValueError:
                continue
            if newest is None or when > newest:
                newest = when
        if len(chunk) < 50:
            return newest
        page += 1