Gitea-Tools/worktree_cleanup_audit.py

"""Session-owned worktree cleanup audit and TTL enforcement (#401).

LLM workflows create many session-owned worktrees under ``branches/``
(review, baseline, merge-simulation, issue, and conflict-fix worktrees).
When a run stops early, races a sibling session, hits a validation failure,
or loses shell/cwd state, those worktrees are left behind and later workflow
decisions get harder and riskier.

This module provides:

* ``build_worktree_metadata`` — ownership/purpose metadata for a worktree.
* ``classify_worktree`` — safety-first classification into the audit
  vocabulary (active open PR, active issue work, dirty, clean stale
  removable, detached review leftover, unsafe/unknown).
* ``assess_worktree_removal`` — a per-worktree removal decision with an
  explicit proof and block reasons.
* git-shelling helpers (``list_worktrees``, ``read_worktree_dirty``,
  ``git_worktree_list``, ``remove_worktree``) and ``audit_branches_directory``
  that classify every entry under ``branches/``.

Pure assessment functions take explicit state so they are unit-testable
without a filesystem or network. Only the thin git helpers shell out, and
removal is only ever executed after ``assess_worktree_removal`` proves the
worktree is safe to delete.
"""

from __future__ import annotations

import os
import re
import subprocess
from datetime import datetime, timezone
from typing import Any

PROTECTED_BRANCHES = frozenset({"master", "main", "dev"})
DEFAULT_TTL_HOURS = float(os.environ.get("GITEA_WORKTREE_TTL_HOURS", "24") or 24)

# Workflow types that can create session-owned worktrees.
WORKFLOW_REVIEW = "review"
WORKFLOW_BASELINE = "baseline"
WORKFLOW_MERGE_SIMULATION = "merge_simulation"
WORKFLOW_ISSUE_WORK = "issue_work"
WORKFLOW_CONFLICT_FIX = "conflict_fix"
WORKFLOW_UNKNOWN = "unknown"

# Workflow types whose worktrees are transient and removed automatically at
# successful workflow completion (acceptance criterion 2).
AUTO_REMOVE_ON_SUCCESS = frozenset(
    {WORKFLOW_REVIEW, WORKFLOW_BASELINE, WORKFLOW_MERGE_SIMULATION}
)

# Cleanup-audit classification vocabulary (acceptance criterion 4).
CLASS_ACTIVE_OPEN_PR = "active_open_pr"
CLASS_ACTIVE_ISSUE_WORK = "active_issue_work"
CLASS_DIRTY_LOCAL = "dirty_local_worktree"
CLASS_CLEAN_STALE_REMOVABLE = "clean_stale_removable"
CLASS_DETACHED_REVIEW_LEFTOVER = "detached_review_leftover"
CLASS_UNSAFE_UNKNOWN = "unsafe_unknown"

# Only these two classifications may ever be removed automatically.
REMOVABLE_CLASSES = frozenset(
    {CLASS_CLEAN_STALE_REMOVABLE, CLASS_DETACHED_REVIEW_LEFTOVER}
)

_ISSUE_REF_RE = re.compile(r"issue-(\d+)", re.IGNORECASE)
_ISSUE_BRANCH_PREFIXES = ("feat/", "fix/", "docs/", "chore/")


def infer_workflow_type(path: str | None, branch: str | None = None) -> str:
    """Infer the creating workflow type from a worktree path or branch name."""
    text = f"{path or ''} {branch or ''}".lower()
    if "baseline" in text:
        return WORKFLOW_BASELINE
    if "merge-sim" in text or "merge_sim" in text or "mergesim" in text:
        return WORKFLOW_MERGE_SIMULATION
    if "review" in text or "review-pr" in text:
        return WORKFLOW_REVIEW
    if "conflict" in text:
        return WORKFLOW_CONFLICT_FIX
    if "issue-" in text or (branch or "").startswith(_ISSUE_BRANCH_PREFIXES):
        return WORKFLOW_ISSUE_WORK
    return WORKFLOW_UNKNOWN


def _extract_issue_number(path: str | None, branch: str | None) -> int | None:
    match = _ISSUE_REF_RE.search(f"{path or ''} {branch or ''}")
    return int(match.group(1)) if match else None


def build_worktree_metadata(
    *,
    path: str,
    branch: str | None = None,
    head_sha: str | None = None,
    workflow_type: str | None = None,
    issue_number: int | None = None,
    pr_number: int | None = None,
    creator: str | None = None,
    profile: str | None = None,
    created_at: str | None = None,
    last_used_at: str | None = None,
    cleanup_eligibility: str | None = None,
) -> dict[str, Any]:
    """Return ownership/purpose metadata for a session-owned worktree.

    Covers acceptance criterion 1: path, workflow type, issue/PR number,
    branch/head SHA, creator identity/profile, created and last-used
    timestamps, and cleanup eligibility.
    """
    wt = workflow_type or infer_workflow_type(path, branch)
    issue = issue_number
    if issue is None:
        issue = _extract_issue_number(path, branch)
    return {
        "path": path,
        "workflow_type": wt,
        "issue_number": issue,
        "pr_number": pr_number,
        "branch": branch,
        "head_sha": head_sha,
        "creator": creator,
        "profile": profile,
        "created_at": created_at,
        "last_used_at": last_used_at,
        "auto_remove_on_success": wt in AUTO_REMOVE_ON_SUCCESS,
        "cleanup_eligibility": cleanup_eligibility,
    }


def _parse_timestamp(value: str | None) -> datetime | None:
    if not value:
        return None
    text = str(value).strip()
    if not text:
        return None
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        parsed = datetime.fromisoformat(text)
    except ValueError:
        return None
    if parsed.tzinfo is None:
        return parsed.replace(tzinfo=timezone.utc)
    return parsed


def is_ttl_expired(
    *,
    last_used_at: str | None,
    now: datetime | str | None,
    ttl_hours: float = DEFAULT_TTL_HOURS,
) -> bool:
    """Return True only when the age is known and exceeds ``ttl_hours``.

    Unknown or unparseable timestamps fail safe (not expired) so a worktree
    is never treated as removable merely because its age is unknown.
    """
    last = _parse_timestamp(last_used_at)
    now_dt = now if isinstance(now, datetime) else _parse_timestamp(now)
    if last is None or now_dt is None:
        return False
    if now_dt.tzinfo is None:
        now_dt = now_dt.replace(tzinfo=timezone.utc)
    return (now_dt - last).total_seconds() > ttl_hours * 3600.0


def classify_worktree(
    *,
    workflow_type: str,
    is_dirty: bool,
    has_open_pr: bool = False,
    has_active_lease: bool = False,
    has_active_issue_lock: bool = False,
    is_detached: bool = False,
    branch_gone: bool = False,
    ttl_expired: bool = False,
    is_protected: bool = False,
    metadata_known: bool = True,
) -> str:
    """Classify a worktree, safety-first: any preservation signal wins.

    Dirty, open-PR, leased, active-lock, protected, and unknown states are
    all non-removable and are checked before any removable classification,
    so nothing removable can shadow a preservation signal (criteria 6-8).
    """
    if is_protected:
        # The main checkout / a protected base branch is never removable.
        return CLASS_UNSAFE_UNKNOWN
    if is_dirty:
        return CLASS_DIRTY_LOCAL  # never auto-deleted (criterion 6)
    if has_open_pr:
        return CLASS_ACTIVE_OPEN_PR  # never auto-deleted (criterion 7)
    if has_active_lease:
        return CLASS_ACTIVE_ISSUE_WORK  # never auto-deleted (criterion 8)
    if has_active_issue_lock:
        return CLASS_ACTIVE_ISSUE_WORK
    if not metadata_known or workflow_type == WORKFLOW_UNKNOWN:
        return CLASS_UNSAFE_UNKNOWN  # never auto-deleted without proof

    # Clean, no PR, no lease, no lock, known workflow type.
    if workflow_type in AUTO_REMOVE_ON_SUCCESS:
        if is_detached or branch_gone:
            return CLASS_DETACHED_REVIEW_LEFTOVER
        return CLASS_CLEAN_STALE_REMOVABLE
    # issue_work / conflict_fix: only removable once the TTL has expired.
    if ttl_expired:
        return CLASS_CLEAN_STALE_REMOVABLE
    return CLASS_ACTIVE_ISSUE_WORK


def is_removable(classification: str) -> bool:
    """Return True only for the two auto-removable classifications."""
    return classification in REMOVABLE_CLASSES


def assess_worktree_removal(
    *,
    path: str,
    branch: str | None,
    head_sha: str | None,
    is_dirty: bool,
    has_open_pr: bool,
    has_active_lease: bool,
    classification: str,
) -> dict[str, Any]:
    """Return a per-worktree removal decision with proof (criterion 9).

    A worktree is only safe to remove when it is clean, has no active PR,
    has no active lease, and its classification is auto-removable.
    """
    block_reasons: list[str] = []
    if is_dirty:
        block_reasons.append("worktree has uncommitted changes")
    if has_open_pr:
        block_reasons.append("worktree branch has an open PR")
    if has_active_lease:
        block_reasons.append("worktree has an active lease")
    if not is_removable(classification):
        block_reasons.append(
            f"classification '{classification}' is not auto-removable"
        )
    return {
        "path": path,
        "branch": branch,
        "head_sha": head_sha,
        "classification": classification,
        "clean": not is_dirty,
        "no_active_pr": not has_open_pr,
        "no_active_lease": not has_active_lease,
        "safe_to_remove": not block_reasons,
        "block_reasons": block_reasons,
    }


def plan_success_cleanup(
    *,
    metadata: dict[str, Any],
    is_dirty: bool,
    has_open_pr: bool,
    has_active_lease: bool,
) -> dict[str, Any]:
    """Decide whether a just-completed worktree is removed at success.

    Review/baseline/merge-simulation worktrees are removed automatically at
    successful completion (criterion 2); everything else is preserved and
    reported. Dirty/PR/leased worktrees are always preserved (criteria 6-8).
    """
    workflow_type = metadata.get("workflow_type", WORKFLOW_UNKNOWN)
    if not metadata.get("auto_remove_on_success"):
        return {
            "remove": False,
            "reason": f"workflow type '{workflow_type}' is preserved by policy",
        }
    classification = classify_worktree(
        workflow_type=workflow_type,
        is_dirty=is_dirty,
        has_open_pr=has_open_pr,
        has_active_lease=has_active_lease,
    )
    decision = assess_worktree_removal(
        path=metadata.get("path", ""),
        branch=metadata.get("branch"),
        head_sha=metadata.get("head_sha"),
        is_dirty=is_dirty,
        has_open_pr=has_open_pr,
        has_active_lease=has_active_lease,
        classification=classification,
    )
    return {
        "remove": decision["safe_to_remove"],
        "reason": "clean transient worktree removable at success completion"
        if decision["safe_to_remove"]
        else "; ".join(decision["block_reasons"]),
        "classification": classification,
        "decision": decision,
    }


def cleanup_failure_report(path: str, reason: str) -> dict[str, Any]:
    """Structured leftover-worktree record for the final report (criterion 3)."""
    return {"path": path, "removed": False, "reason": reason}


# --------------------------------------------------------------------------
# git-shelling helpers (only these touch the filesystem)
# --------------------------------------------------------------------------


def parse_worktree_porcelain(text: str) -> list[dict[str, Any]]:
    """Parse ``git worktree list --porcelain`` output into entries."""
    entries: list[dict[str, Any]] = []
    current: dict[str, Any] = {}
    for raw in (text or "").splitlines():
        line = raw.rstrip("\n")
        if not line:
            if current:
                entries.append(current)
                current = {}
            continue
        if line.startswith("worktree "):
            if current:
                entries.append(current)
            current = {
                "path": line[len("worktree ") :].strip(),
                "head": None,
                "branch": None,
                "detached": False,
                "bare": False,
            }
        elif line.startswith("HEAD "):
            current["head"] = line[len("HEAD ") :].strip()
        elif line.startswith("branch "):
            ref = line[len("branch ") :].strip()
            current["branch"] = ref.replace("refs/heads/", "", 1)
        elif line == "detached":
            current["detached"] = True
        elif line == "bare":
            current["bare"] = True
    if current:
        entries.append(current)
    return entries


def list_worktrees(project_root: str) -> list[dict[str, Any]]:
    """Return parsed ``git worktree list`` entries for ``project_root``."""
    result = subprocess.run(
        ["git", "-C", project_root, "worktree", "list", "--porcelain"],
        capture_output=True,
        text=True,
        check=False,
    )
    if result.returncode != 0:
        return []
    return parse_worktree_porcelain(result.stdout)


def git_worktree_list(project_root: str) -> str:
    """Return plain ``git worktree list`` output for final verification (criterion 10)."""
    result = subprocess.run(
        ["git", "-C", project_root, "worktree", "list"],
        capture_output=True,
        text=True,
        check=False,
    )
    return (result.stdout or "").strip()


def read_worktree_dirty(path: str) -> dict[str, Any]:
    """Return dirty state for a worktree path via ``git status --porcelain``."""
    if not path or not os.path.isdir(path):
        return {"exists": False, "dirty": None, "dirty_files": []}
    result = subprocess.run(
        ["git", "-C", path, "status", "--porcelain"],
        capture_output=True,
        text=True,
        check=False,
    )
    files = [ln for ln in (result.stdout or "").splitlines() if ln.strip()]
    return {"exists": True, "dirty": bool(files), "dirty_files": files}


def remove_worktree(project_root: str, path: str) -> dict[str, Any]:
    """Remove a single worktree via ``git worktree remove`` (no ``--force``)."""
    result = subprocess.run(
        ["git", "-C", project_root, "worktree", "remove", path],
        capture_output=True,
        text=True,
        check=False,
    )
    ok = result.returncode == 0
    return {
        "path": path,
        "removed": ok,
        "reason": None if ok else (result.stderr or "git worktree remove failed").strip(),
    }


def _is_under_branches(project_root: str, path: str) -> bool:
    branches_root = os.path.join(os.path.abspath(project_root), "branches")
    return os.path.abspath(path or "").startswith(branches_root + os.sep)


def audit_branches_directory(
    project_root: str,
    *,
    open_pr_branches: set[str] | None = None,
    leased_branches: set[str] | None = None,
    active_issue_branches: set[str] | None = None,
    now: datetime | str | None = None,
    ttl_hours: float = DEFAULT_TTL_HOURS,
) -> dict[str, Any]:
    """Classify every session-owned worktree under ``branches/``.

    Read-only: shells out to git for discovery and dirty state, then applies
    the pure classifier. Returns per-worktree classifications, counts, the
    list of removable candidates, and the ``git worktree list`` proof.
    """
    open_pr_branches = open_pr_branches or set()
    leased_branches = leased_branches or set()
    active_issue_branches = active_issue_branches or set()

    worktrees: list[dict[str, Any]] = []
    for entry in list_worktrees(project_root):
        path = entry.get("path") or ""
        branch = entry.get("branch")
        is_protected = (branch in PROTECTED_BRANCHES) or not _is_under_branches(
            project_root, path
        )
        dirty_state = read_worktree_dirty(path)
        is_dirty = bool(dirty_state.get("dirty"))
        metadata = build_worktree_metadata(
            path=path, branch=branch, head_sha=entry.get("head")
        )
        has_open_pr = bool(branch) and branch in open_pr_branches
        has_active_lease = bool(branch) and branch in leased_branches
        has_active_lock = bool(branch) and branch in active_issue_branches
        ttl_expired = is_ttl_expired(
            last_used_at=metadata.get("last_used_at"), now=now, ttl_hours=ttl_hours
        )
        classification = classify_worktree(
            workflow_type=metadata["workflow_type"],
            is_dirty=is_dirty,
            has_open_pr=has_open_pr,
            has_active_lease=has_active_lease,
            has_active_issue_lock=has_active_lock,
            is_detached=bool(entry.get("detached")),
            branch_gone=branch is None and not entry.get("detached"),
            ttl_expired=ttl_expired,
            is_protected=is_protected,
        )
        metadata["cleanup_eligibility"] = classification
        worktrees.append(
            {
                **metadata,
                "detached": bool(entry.get("detached")),
                "dirty": is_dirty,
                "dirty_files": dirty_state.get("dirty_files", []),
                "has_open_pr": has_open_pr,
                "has_active_lease": has_active_lease,
                "has_active_issue_lock": has_active_lock,
                "is_protected": is_protected,
                "classification": classification,
                "removable": is_removable(classification),
            }
        )

    counts: dict[str, int] = {}
    for wt in worktrees:
        counts[wt["classification"]] = counts.get(wt["classification"], 0) + 1
    removable = [wt for wt in worktrees if wt["removable"]]

    return {
        "project_root": project_root,
        "worktrees": worktrees,
        "counts": counts,
        "removable_candidates": removable,
        "removable_count": len(removable),
        "total": len(worktrees),
        "git_worktree_list": git_worktree_list(project_root),
    }