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_proc_starttime, 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, *, expected_starttime: int | None = None) -> bool:
    """Cross-platform best-effort ``kill -0`` equivalent.

    Phase 1k+ refinement: when ``expected_starttime`` is provided
    (Linux only), the process's current starttime is compared against
    it. Mismatch = PID was reused by an unrelated process; treat as
    "not the process we wanted" → return False so the janitor doesn't
    SIGKILL the impostor. None (cross-platform or starttime unknown)
    falls back to the legacy PID-only check.
    """
    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
    if expected_starttime is not None:
        actual = read_proc_starttime(pid)
        if actual is not None and actual != expected_starttime:
            # PID was reused by an unrelated process.
            return False
    return True


def _kill_with_grace(
    pid: int, *, grace_s: float = 5.0, expected_starttime: int | None = None
) -> bool:
    """SIGTERM + wait grace + SIGKILL — but only if the PID still
    matches ``expected_starttime`` (Linux). Returns True iff a signal
    was actually delivered to the process; False if there was nothing
    to kill (already dead, PID reused for another process, etc.).

    Phase 1k+ refinement: ``expected_starttime`` defends against PID
    reuse. If the process at this PID has a different starttime, we
    treat it as "already gone" (the original died; OS reused the PID
    for something we have no business signaling) and return False
    without signaling — the metric for "sessions killed" stays
    accurate.

    Zombie-reaping (post-SIGKILL kill -0 still reporting alive) is
    the original parent's responsibility; we don't loop on it.
    """
    if not _pid_alive(pid, expected_starttime=expected_starttime):
        return False
    try:
        os.kill(pid, signal.SIGTERM)
    except ProcessLookupError:
        # Race: process exited between our liveness check and the
        # SIGTERM. No signal was actually delivered.
        return False
    deadline = time.monotonic() + grace_s
    while time.monotonic() < deadline:
        if not _pid_alive(pid, expected_starttime=expected_starttime):
            return True  # SIGTERM did the job.
        time.sleep(0.1)
    try:
        os.kill(pid, signal.SIGKILL)
    except ProcessLookupError:
        # Process died between SIGTERM and SIGKILL; the SIGTERM was
        # delivered (we got past the first os.kill) so it counts.
        return True
    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,
                        expected_starttime=sidecar.subprocess_starttime,
                    ):
                        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,
                    expected_starttime=sidecar.subprocess_starttime,
                ):
                    report.sessions_killed += 1
        else:
            # No callback → SIGKILL fallback.
            if _kill_with_grace(
                sidecar.subprocess_pid,
                expected_starttime=sidecar.subprocess_starttime,
            ):
                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",
]