ai-agent/sidecar/pids.py

"""
Pid-record tracking: which sessions are still running, across sidecar restarts.

Owns the mutable ``_procs`` table (the children we launched ourselves) — other
modules reach it as ``pids._procs``, never by copying the reference.
"""

import json
import os
import pathlib
import subprocess

from config import LOG_DIR

# ---- pid tracking ---------------------------------------------------------
# Each running session is persisted to `logs/<sid>.pid` as a small JSON record
# so the sidecar can re-discover, list and interrupt runs that outlived it. The
# `/proc` start-time pins the record to *this* process, so a recycled PID (a new
# unrelated process that happens to reuse the number) is never mistaken for a
# live session — critical before we send it a signal.


# The runs we launched ourselves, kept only so we can *reap* them. A `claude`
# started with Popen stays our child even though it leads its own session, so
# when it exits it lingers as a zombie until the parent waits on it — and a
# zombie still answers `os.kill(pid, 0)` and still has a `/proc/<pid>/stat`. So
# a finished run kept looking "alive", `/resume` took the idempotency shortcut
# below, launched nothing, and the prompt was silently dropped. We poll (reap)
# them before every liveness check, and treat state `Z` as dead as a belt-and-
# braces guard (a survivor of a previous sidecar instance is reparented to init,
# which reaps it for us).
_procs: dict[str, subprocess.Popen] = {}


def _reap() -> None:
    """Wait on any of our children that have exited, so no zombie lingers."""
    for sid, proc in list(_procs.items()):
        if proc.poll() is not None:  # returncode set ⇒ child reaped
            _procs.pop(sid, None)


# Linux reads /proc directly (no dependency — the host sidecar's venv has only
# fastapi + uvicorn); anywhere else (the macOS worker) falls back to psutil,
# which the worker installer puts in its venv. Both answer the same two
# questions: is the process a zombie, and what is its start-time fingerprint.
_HAVE_PROC = pathlib.Path("/proc/self/stat").exists()


def _psutil_proc(pid: int):
    try:
        import psutil  # noqa: PLC0415 — optional, macOS worker only
        return psutil.Process(pid)
    except Exception:  # ImportError, NoSuchProcess, AccessDenied
        return None


def _proc_state(pid: int) -> str | None:
    """Process state letter ('Z' = zombie, 'S'/'R' = live), from /proc/pid/stat
    on Linux, psutil elsewhere."""
    if not _HAVE_PROC:
        p = _psutil_proc(pid)
        try:
            return None if p is None else (
                "Z" if p.status() == "zombie" else "R")
        except Exception:
            return None
    try:
        stat = pathlib.Path(f"/proc/{pid}/stat").read_text()
        return stat[stat.rindex(")") + 1:].split()[0]
    except (OSError, ValueError, IndexError):
        return None


def _proc_starttime(pid: int) -> int | float | None:
    """Process start-time fingerprint: clock ticks since boot (field 22 of
    /proc/pid/stat) on Linux, psutil's ``create_time()`` (epoch seconds,
    rounded) elsewhere.

    Unique per (pid, boot), so it fingerprints the exact process and lets us
    detect PID reuse. The `comm` field (field 2) can contain spaces and parens,
    so we parse everything after the final ')'."""
    if not _HAVE_PROC:
        p = _psutil_proc(pid)
        try:
            return None if p is None else round(p.create_time(), 2)
        except Exception:
            return None
    try:
        stat = pathlib.Path(f"/proc/{pid}/stat").read_text()
        after = stat[stat.rindex(")") + 1:].split()
        return int(after[19])  # field 22 overall (fields 3.. after comm)
    except (OSError, ValueError, IndexError):
        return None


def _pidfile(sid: str) -> pathlib.Path:
    return LOG_DIR / f"{sid}.pid"


def _read_record(sid: str) -> dict | None:
    """Load a session's pid record, tolerating the legacy plain-integer format."""
    try:
        raw = _pidfile(sid).read_text().strip()
    except OSError:
        return None
    if not raw:
        return None
    try:
        rec = json.loads(raw)
        if isinstance(rec, dict) and rec.get("pid"):
            return rec
    except ValueError:
        pass
    try:  # legacy pidfile: bare pid, no start-time guard
        return {"pid": int(raw)}
    except ValueError:
        return None


def _alive(rec: dict) -> bool:
    """True if the recorded process is still *running* — same process, not a zombie.

    Reaps our finished children first: an exited-but-unwaited child would
    otherwise still pass every check below and be mistaken for a live run."""
    pid = rec.get("pid")
    if not pid:
        return False
    _reap()
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True  # exists but owned by someone else — still 'alive'
    except OSError:
        return False
    if _proc_state(pid) == "Z":
        return False  # exited, just not reaped yet
    want = rec.get("starttime")
    if want is not None:
        now = _proc_starttime(pid)
        if now is not None and now != want:
            return False  # PID was recycled into a different process
    return True


def _live_record(sid: str) -> dict | None:
    """The session's pid record if its run is still going; else None (and the
    stale pidfile is dropped, so finished runs don't accumulate on disk)."""
    rec = _read_record(sid)
    if rec is None:
        return None
    if _alive(rec):
        return rec
    _pidfile(sid).unlink(missing_ok=True)
    return None


def _reconcile_pidfiles() -> None:
    """On (re)start, GC pidfiles whose process is gone and count the survivors.

    The spawned runs are daemons, so after a sidecar restart the disk pidfiles
    are our only record of what's still running — reconcile them into a clean
    view and drop the dead ones so /interrupt never signals a stale PID."""
    if not LOG_DIR.is_dir():
        return
    survived = 0
    for p in sorted(LOG_DIR.glob("*.pid")):
        rec = _read_record(p.stem)
        if rec is None or not _alive(rec):
            p.unlink(missing_ok=True)
        else:
            survived += 1
    print(f"[sidecar] startup reconcile: {survived} session(s) survived restart",
          flush=True)