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`.

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


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
    return True


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 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:
        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) -> 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.
    """
    stop_timer(job_id, registry_dir)
    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))]
    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", proc.pid, job_id, interval)
    return proc.pid


def _should_stop(job_id: str, registry_dir: str, dispatch_token, attempt_id) -> Optional[str]:
    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) -> 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)
        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("-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)


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