multi-agent-mux/.agents/skills/lib_py/delegation/heartbeat.py

"""heartbeat.py — the agent-side automatic heartbeat timer (Issue #18).

The liveness contract obliges an agent to publish `progress` at least every N
seconds while it works. Relying on the brief to say so meant an agent that
forgot, or one whose turn was one long tool call, was closed as unresponsive
while it was fine. This timer publishes the heartbeat for it.

`publish_event.py --event started --heartbeat N` starts one detached timer per
job; it runs in its own session so the agent's shell may return at once. The
timer stops on its own when

  * the job record reaches a terminal status (completed / error / cancelled),
  * the record's dispatch_token or attempt_id changes (the job was
    re-dispatched: this attempt's heartbeats would be refused anyway),
  * the job record disappears, or
  * its lifetime exceeds the job's wall-clock timeout plus one period.

A terminal publish also stops it directly through the pidfile
(`<registry>/<job_id>.heartbeat.pid`) so no heartbeat can trail a `completed`.

The timer is bound to the agent's life (review B2): it watches one process -
`--watch-pid`, or by default the nearest ancestor of the publishing process
that is an agent CLI (claude, codex, agy, ...; the agent runs the start
command from its shell tool, so that ancestor is the agent itself) - and stops
within one period once that process is gone. A crashed or killed agent
therefore stops heartbeating and is closed after K x N like any silent one;
an agent that is alive but stuck keeps heartbeating, which is the contract's
deliberate trade-off (liveness of the process, not of the work).

Each beat is an ordinary `publish_event` call in this process: it takes the
next seq, is signed with the job's auth_token, carries the dispatch token and
attempt id the timer was started with, and is recorded in events.log before
the wire publish - exactly what the orchestrator's disk and wire liveness
checks read.
"""
from __future__ import annotations

import argparse
import logging
import os
import signal
import subprocess
import sys
import threading
import time
from typing import List, Optional

from lib_py.delegation import mqtt_common
from lib_py.delegation.mqtt_common import default_registry_dir, load_job

logger = logging.getLogger("delegate_job.heartbeat")

TERMINAL_STATUSES = ("completed", "error", "cancelled")
HEARTBEAT_DETAIL = "heartbeat"
MAX_IN_FLIGHT = 8


# The agent CLIs, as lib_py.agents.base names them. Repeated here rather than
# imported: lib_py.delegation may not import a sibling lib_py package (a
# frozen loop binds delegation to a different tree, see __init__).
AGENT_CLI_NAMES = ("claude", "agy", "hermes", "grok", "opencode", "codex")


def _is_agent_cli(command: str) -> bool:
    """Some token before the first option has an agent basename
    (`claude ...`, `node .../claude ...`, `/bin/sh .../agy ...`)."""
    for tok in command.split():
        if tok.startswith("-"):
            break
        if os.path.basename(tok) in AGENT_CLI_NAMES:
            return True
    return False


def _process_table():
    """{pid: (ppid, command)} from ps; empty when ps is unavailable."""
    try:
        out = subprocess.run(["ps", "-axo", "pid=,ppid=,command="], capture_output=True, text=True).stdout
    except OSError:
        return {}
    table = {}
    for line in out.splitlines():
        parts = line.split(None, 2)
        if len(parts) < 2:
            continue
        try:
            pid, ppid = int(parts[0]), int(parts[1])
        except ValueError:
            continue
        table[pid] = (ppid, parts[2] if len(parts) > 2 else "")
    return table


def nearest_agent_cli_ancestor(start_pid: Optional[int] = None) -> Optional[int]:
    """PID of the closest ancestor (start_pid included) that is an agent CLI.

    The agent executes the injected start command through its shell tool, so
    the chain is agent CLI -> shell -> python publish_event; the agent is the
    first agent-CLI process met walking up. None when there is no such
    ancestor (a test, a plain shell, a remote runner)."""
    table = _process_table()
    pid = start_pid if start_pid is not None else os.getpid()
    seen = set()
    while pid and pid > 1 and pid not in seen:
        seen.add(pid)
        entry = table.get(pid)
        if entry is None:
            return None
        ppid, command = entry
        if _is_agent_cli(command):
            return pid
        pid = ppid
    return None


def pidfile_path(job_id: str, registry_dir: str) -> str:
    return os.path.join(registry_dir, "{}.heartbeat.pid".format(job_id))


def _pid_alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    try:
        state = subprocess.run(["ps", "-o", "stat=", "-p", str(pid)],
                               capture_output=True, text=True).stdout.strip()
    except OSError:
        return True
    return bool(state) and not state.startswith("Z")


def running_timer(job_id: str, registry_dir: str) -> Optional[int]:
    """PID of a live timer for this job, or None."""
    try:
        with open(pidfile_path(job_id, registry_dir)) as fh:
            pid = int(fh.read().strip() or 0)
    except (OSError, ValueError):
        return None
    if pid > 0 and _pid_alive(pid):
        return pid
    return None


def _is_timer_process(pid: int, job_id: str) -> bool:
    """The pid really is this job's timer (a stale pidfile after a crash may
    name a recycled pid that now belongs to something else)."""
    entry = _process_table().get(pid)
    if entry is None:
        return False
    command = entry[1]
    return "lib_py.delegation.heartbeat" in command and job_id in command


def stop_timer(job_id: str, registry_dir: str) -> bool:
    """Terminate the job's timer if one runs. Returns True when one was signalled."""
    pid = running_timer(job_id, registry_dir)
    path = pidfile_path(job_id, registry_dir)
    if pid is None or not _is_timer_process(pid, job_id):
        try:
            os.unlink(path)
        except OSError:
            pass
        return False
    try:
        os.kill(pid, signal.SIGTERM)
    except OSError:
        pass
    try:
        os.unlink(path)
    except OSError:
        pass
    return True


def start_timer(job_id: str, registry_dir: str, interval: int,
                dispatch_token: Optional[str] = None, attempt_id: Optional[int] = None,
                python: Optional[str] = None, watch_pid: Optional[int] = None) -> int:
    """Spawn the detached timer process and return its pid.

    A timer that already runs for the job is stopped first: a second `started`
    for the same job (a retry of the same command) must not double the rate.
    ``watch_pid`` binds the timer to a process (default: the nearest agent-CLI
    ancestor of this process, if any); 0 disables the watch explicitly.
    """
    stop_timer(job_id, registry_dir)
    if watch_pid is None:
        watch_pid = nearest_agent_cli_ancestor()
    cmd: List[str] = [python or sys.executable, "-m", "lib_py.delegation.heartbeat",
                      "--job", job_id, "--registry-dir", registry_dir,
                      "--interval", str(int(interval))]
    if dispatch_token:
        cmd += ["--expect-dispatch-token", dispatch_token]
    if attempt_id is not None:
        cmd += ["--attempt-id", str(int(attempt_id))]
    if watch_pid:
        cmd += ["--watch-pid", str(int(watch_pid))]
    env = dict(os.environ)
    # The timer imports lib_py from the same tree as this process.
    pkg_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
    env["PYTHONPATH"] = pkg_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "")
    log_path = os.path.join(registry_dir, "{}.heartbeat.log".format(job_id))
    with open(log_path, "ab") as log:
        proc = subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=log, stderr=subprocess.STDOUT,
                                start_new_session=True, env=env, cwd=os.getcwd())
    with open(pidfile_path(job_id, registry_dir), "w") as fh:
        fh.write(str(proc.pid))
    logger.info("heartbeat timer pid %d for job %s every %ss (watching pid %s)",
                proc.pid, job_id, interval, watch_pid or "none")
    return proc.pid


def _should_stop(job_id: str, registry_dir: str, dispatch_token, attempt_id,
                 watch_pid: Optional[int] = None) -> Optional[str]:
    if watch_pid and not _pid_alive(watch_pid):
        return "watched agent process {} is gone".format(watch_pid)
    try:
        rec = load_job(job_id, registry_dir)
    except FileNotFoundError:
        return "job record gone"
    if rec.get("status") in TERMINAL_STATUSES:
        return "job is {}".format(rec.get("status"))
    if dispatch_token and rec.get("dispatch_token") != dispatch_token:
        return "dispatch token changed (re-dispatched)"
    if attempt_id is not None and rec.get("attempt_id", rec.get("iteration", 1)) != attempt_id:
        return "attempt changed (re-dispatched)"
    return None


def run(job_id: str, registry_dir: str, interval: int, dispatch_token=None,
        attempt_id=None, max_lifetime: Optional[float] = None,
        watch_pid: Optional[int] = None) -> int:
    """The timer loop. Runs until a stop condition holds."""
    from lib_py.delegation import publish_event  # late import: the CLI imports us

    interval = max(1, int(interval))
    if max_lifetime is None:
        try:
            max_lifetime = float(load_job(job_id, registry_dir).get("timeout_sec", 3600)) + interval
        except FileNotFoundError:
            return 1
    stop = {"flag": False}

    def _on_term(_signum, _frame):
        stop["flag"] = True

    signal.signal(signal.SIGTERM, _on_term)
    signal.signal(signal.SIGINT, _on_term)

    started = time.monotonic()
    own_pidfile = pidfile_path(job_id, registry_dir)
    # Fixed cadence: the next beat is due `interval` after the previous one
    # was *due*, not after its publish returned, so a slow broker (seconds per
    # ACK) does not stretch the period the contract promises. Each beat is
    # recorded on disk before its wire publish, so the on-disk cadence holds
    # even when the wire lags. Sleep first: `started` was just published.
    next_due = time.monotonic() + interval
    in_flight: List[threading.Thread] = []
    while not stop["flag"]:
        while time.monotonic() < next_due and not stop["flag"]:
            time.sleep(min(0.5, max(0.0, next_due - time.monotonic())))
        if stop["flag"]:
            break
        next_due += interval
        if next_due < time.monotonic():
            next_due = time.monotonic() + interval  # fell far behind: resync
        reason = _should_stop(job_id, registry_dir, dispatch_token, attempt_id, watch_pid)
        if reason:
            logger.info("heartbeat timer for %s stopping: %s", job_id, reason)
            break
        if time.monotonic() - started > max_lifetime:
            logger.info("heartbeat timer for %s stopping: lifetime %.0fs exceeded", job_id, max_lifetime)
            break
        argv = ["--registry-dir", registry_dir, "--job", job_id, "--event", "progress",
                "--detail", HEARTBEAT_DETAIL, "--attempts", "2"]
        if dispatch_token:
            argv += ["--expect-dispatch-token", dispatch_token]
        if attempt_id is not None:
            argv += ["--attempt-id", str(attempt_id)]
        # Each beat publishes in its own thread: the local record (seq, HMAC,
        # events.log) is taken at once, and the wire publish - seconds per ACK
        # on a slow broker - must not delay the next beat. In-flight beats are
        # capped so a dead broker cannot pile up threads without bound.
        in_flight[:] = [t for t in in_flight if t.is_alive()]
        if len(in_flight) >= MAX_IN_FLIGHT:
            logger.warning("heartbeat for %s skipped: %d publishes still in flight", job_id, len(in_flight))
            continue

        def _beat(argv=argv):
            try:
                rc = publish_event.main(argv)
            except Exception as exc:  # pragma: no cover - defensive
                logger.warning("heartbeat publish raised: %s", exc)
                rc = 2
            if rc == 4:
                logger.info("heartbeat publish refused (stale attempt); timer for %s will stop", job_id)
                stop["flag"] = True

        t = threading.Thread(target=_beat, daemon=True)
        t.start()
        in_flight.append(t)
    for t in in_flight:
        t.join(timeout=15)
    try:
        with open(own_pidfile) as fh:
            if fh.read().strip() == str(os.getpid()):
                os.unlink(own_pidfile)
    except (OSError, ValueError):
        pass
    return 0


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="automatic heartbeat timer for a delegated job")
    parser.add_argument("--job", required=True)
    parser.add_argument("--registry-dir", default=default_registry_dir())
    parser.add_argument("--interval", type=int, required=True, help="heartbeat period N in seconds")
    parser.add_argument("--expect-dispatch-token", default=None)
    parser.add_argument("--attempt-id", type=int, default=None)
    parser.add_argument("--max-lifetime", type=float, default=None)
    parser.add_argument("--watch-pid", type=int, default=None,
                        help="stop once this process (the agent) is gone")
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args(argv)
    mqtt_common.ensure_env_loaded()
    mqtt_common.setup_logging(logging.DEBUG if args.verbose else logging.INFO)
    return run(args.job, args.registry_dir, args.interval, args.expect_dispatch_token,
               args.attempt_id, args.max_lifetime, args.watch_pid)


if __name__ == "__main__":
    sys.exit(main())