cleveragents-core/tools/controller/worker/janitor.py

"""Worker startup janitor: orphan-workspace + orphan-session cleanup.

Plan v6: at worker startup, before entering the main loop, the worker
scans the workspace root for orphan sidecars:

For each ``{workspace_root}/pr-*/worker.session``:
1. Read the sidecar; if missing/malformed → just delete the workspace.
2. If the recorded subprocess_pid is alive AND no DB row is held by
   this instance_id → orphan. Try OpenCode session-cancel API;
   SIGKILL fallback. Delete sidecar.
3. Delete the workspace dir.

The DB query for "is there a row held by instance_id X" is the
authoritative liveness check; the local pid is the fallback signal.

For this skeleton, the OpenCode-cancel-API call is parameterised as
``cancel_callback``. Production wires it to a real HTTP POST against
``{opencode_url}/sessions/{session_id}/cancel``; tests inject a fake
that records calls.
"""
from __future__ import annotations

import logging
import os
import shutil
import signal
import time
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path

from sqlalchemy import text
from sqlalchemy.engine import Engine

from .session_sidecar import WorkerSession, read_sidecar
from .workspace import DEFAULT_WORKSPACE_ROOT

logger = logging.getLogger(__name__)


# Callback the production worker injects to ask OpenCode to cancel a
# session. Returns True iff the cancel succeeded (or the session
# already didn't exist). False → fall back to SIGKILL.
CancelCallback = Callable[[WorkerSession], bool]


@dataclass
class JanitorReport:
    """Summary of one sweep — used by tests + by the worker's
    structured logging."""

    workspaces_scanned: int = 0
    sidecars_orphan: int = 0
    sessions_cancelled_via_api: int = 0
    sessions_killed: int = 0
    workspaces_removed: int = 0
    errors: list[str] = None  # type: ignore[assignment]

    def __post_init__(self):
        if self.errors is None:
            self.errors = []


def _pid_alive(pid: int) -> bool:
    """Cross-platform best-effort ``kill -0`` equivalent."""
    if pid <= 0:
        return False
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        # Permission denied means the pid exists but we can't signal
        # it. From orphan-detection's standpoint: it's still alive.
        return True
    return True


def _kill_with_grace(pid: int, *, grace_s: float = 5.0) -> bool:
    """SIGTERM + wait grace + SIGKILL.

    Returns True iff we issued at least one kill signal that the
    kernel accepted (ESRCH = nothing to kill; that counts as "already
    dead" → True). After SIGKILL is delivered we DON'T re-check
    ``_pid_alive`` because zombie processes (parent hasn't reaped)
    still report alive via ``kill -0``; reaping zombies isn't the
    janitor's responsibility.
    """
    if not _pid_alive(pid):
        return True
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        return True
    deadline = time.monotonic() + grace_s
    while time.monotonic() < deadline:
        if not _pid_alive(pid):
            return True
        time.sleep(0.1)
    try:
        os.kill(pid, signal.SIGKILL)
    except ProcessLookupError:
        return True
    # SIGKILL was delivered. Zombie-reaping is the parent's job; we
    # consider our work done.
    return True


def sweep_orphans(
    engine: Engine,
    *,
    workspace_root: Path | None = None,
    cancel_callback: CancelCallback | None = None,
) -> JanitorReport:
    """Scan ``workspace_root`` for orphan sidecars + clean them up.

    A sidecar is orphaned iff:
    - No DB row has ``status='in_progress'`` AND
      ``locked_by_instance == sidecar.instance_id``.

    For each orphan: try OpenCode cancel API → SIGKILL → delete
    sidecar → delete workspace dir.
    """
    root = workspace_root or DEFAULT_WORKSPACE_ROOT
    report = JanitorReport()
    if not root.is_dir():
        return report

    # Pre-query the DB once for the set of live instance_ids on
    # in_progress rows. Cheap; lets us skip per-sidecar DB hits.
    live_instances = _query_live_instances(engine)

    for entry in sorted(root.iterdir()):
        if not entry.is_dir() or not entry.name.startswith("pr-"):
            continue
        report.workspaces_scanned += 1
        sidecar_path = entry / "worker.session"

        sidecar = read_sidecar(sidecar_path)
        if sidecar is None:
            # No sidecar (or malformed) → workspace is orphan with no
            # session info; just delete.
            _safe_rmtree(entry, report)
            continue

        if sidecar.instance_id in live_instances:
            # Active worker owns this. Leave it alone.
            continue

        # Orphan: clean it up.
        report.sidecars_orphan += 1
        if cancel_callback is not None:
            try:
                if cancel_callback(sidecar):
                    report.sessions_cancelled_via_api += 1
                else:
                    if _kill_with_grace(sidecar.subprocess_pid):
                        report.sessions_killed += 1
            except Exception as exc:  # noqa: BLE001 — best-effort cleanup
                report.errors.append(
                    f"cancel_callback raised for {sidecar.instance_id}: {exc}"
                )
                if _kill_with_grace(sidecar.subprocess_pid):
                    report.sessions_killed += 1
        else:
            # No callback → SIGKILL fallback.
            if _kill_with_grace(sidecar.subprocess_pid):
                report.sessions_killed += 1

        _safe_rmtree(entry, report)

    return report


def _query_live_instances(engine: Engine) -> set[str]:
    """Return the set of instance_id values currently holding
    in_progress rows."""
    with engine.connect() as conn:
        rows = conn.execute(
            text(
                "SELECT DISTINCT locked_by_instance FROM workflow_attempts "
                "WHERE status = 'in_progress' AND locked_by_instance IS NOT NULL"
            )
        ).all()
    return {row[0] for row in rows}


def _safe_rmtree(path: Path, report: JanitorReport) -> None:
    try:
        shutil.rmtree(path)
        report.workspaces_removed += 1
    except OSError as exc:
        report.errors.append(f"rmtree {path}: {exc}")


__all__ = [
    "CancelCallback",
    "JanitorReport",
    "sweep_orphans",
]