GodopuandClaude Fable 5.1
2b1d590010
Review round 4 (codex R3): two concurrent `start_timer` calls for the same attempt both passed the "no timer" check, both spawned, and the pidfile tracked only one - the untracked timer could still publish a retained `running` after `completed`. A per-job flock (`<job>.heartbeat.lock`, separate from the registry/seq lock so in-flight beats can still take their seq) now makes observe/drain/spawn/publish-pid one ownership operation; the terminal stop takes the same lock; a start that finds the attempt ended (terminal record or raised barrier) does not arm; the timer removes its own pidfile only under the lock so a replacement's pid is never deleted. Regressions: concurrent starts leave exactly one tracked owner that the terminal stop drains; a start after the barrier arms nothing. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VCgu9Fa9iT7TbJUSzdq4Ff
474 lines
20 KiB
Python
474 lines
20 KiB
Python
"""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, so no heartbeat can trail a
|
|
`completed` (review round 2, codex R3): it raises an attempt-scoped barrier
|
|
under the registry lock that allocates seqs (a beat either took its seq before
|
|
it or is refused when it asks for one), stops the timer through the pidfile
|
|
(`<registry>/<job_id>.heartbeat.pid`), and waits for the timer to drain its
|
|
in-flight beats - killing it after `MAM_HEARTBEAT_DRAIN_S` - before it takes
|
|
its own seq.
|
|
|
|
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 contextlib
|
|
import fcntl
|
|
import json
|
|
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
|
|
|
|
|
|
DRAIN_ENV = "MAM_HEARTBEAT_DRAIN_S"
|
|
# Short on purpose: the barrier already keeps any beat from taking a seq after
|
|
# the terminal, so the drain only lets an idle timer exit cleanly. A beat still
|
|
# on the wire is cancelled with the timer (its seq is older than the terminal's)
|
|
# rather than holding back the terminal event - which may be the orchestrator
|
|
# closing the job - for seconds of broker round trips.
|
|
DEFAULT_DRAIN_S = 1.0
|
|
|
|
|
|
def drain_seconds() -> float:
|
|
"""How long a terminal publish waits for the timer to drain (env override)."""
|
|
try:
|
|
return max(0.0, float(os.environ.get(DRAIN_ENV, DEFAULT_DRAIN_S)))
|
|
except ValueError:
|
|
return DEFAULT_DRAIN_S
|
|
|
|
|
|
def lifecycle_lock_path(job_id: str, registry_dir: str) -> str:
|
|
return os.path.join(registry_dir, "{}.heartbeat.lock".format(job_id))
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def lifecycle_lock(job_id: str, registry_dir: str):
|
|
"""The per-job timer ownership lock (review round 4, codex R3).
|
|
|
|
Observing/draining the previous owner, spawning the replacement and
|
|
publishing its pid through the pidfile is one ownership operation, and the
|
|
terminal stop takes part in it, so two concurrent `started --heartbeat`
|
|
calls cannot both arm a timer with only one of them tracked, and a start
|
|
that loses the race against a terminal publish finds the barrier and does
|
|
not arm at all. This is a separate flock from the registry/seq lock on
|
|
purpose: an in-flight beat needs that lock to take its seq, and the drain
|
|
waits for exactly those beats, so holding the seq lock across the drain
|
|
would deadlock."""
|
|
os.makedirs(registry_dir, exist_ok=True)
|
|
with open(lifecycle_lock_path(job_id, registry_dir), "a+") as fh:
|
|
fcntl.flock(fh.fileno(), fcntl.LOCK_EX)
|
|
try:
|
|
yield
|
|
finally:
|
|
fcntl.flock(fh.fileno(), fcntl.LOCK_UN)
|
|
|
|
|
|
def stop_timer(job_id: str, registry_dir: str, drain: Optional[float] = None) -> bool:
|
|
"""Terminate the job's timer if one runs. Returns True when one was signalled.
|
|
|
|
``drain`` (seconds): also wait for the timer to exit - on TERM it stops
|
|
scheduling and joins its in-flight beats - and SIGKILL it if it has not
|
|
gone by then, so none of its beats can still publish afterwards. A terminal
|
|
publish uses this (review round 2, codex R3): "TERM sent" is not "timer
|
|
stopped". Serialised with `start_timer` through the lifecycle lock."""
|
|
with lifecycle_lock(job_id, registry_dir):
|
|
return _stop_timer_locked(job_id, registry_dir, drain)
|
|
|
|
|
|
def _stop_timer_locked(job_id: str, registry_dir: str, drain: Optional[float] = None) -> bool:
|
|
"""`stop_timer` body; the caller holds the lifecycle lock."""
|
|
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
|
|
if drain is not None:
|
|
deadline = time.monotonic() + drain
|
|
while time.monotonic() < deadline and _pid_alive(pid):
|
|
time.sleep(0.1)
|
|
if _pid_alive(pid):
|
|
logger.warning("heartbeat timer %d for %s did not drain within %.0fs; killing it", pid, job_id, drain)
|
|
try:
|
|
os.kill(pid, signal.SIGKILL)
|
|
except OSError:
|
|
pass
|
|
deadline = time.monotonic() + 2.0
|
|
while time.monotonic() < deadline and _pid_alive(pid):
|
|
time.sleep(0.05)
|
|
try:
|
|
os.unlink(path)
|
|
except OSError:
|
|
pass
|
|
return True
|
|
|
|
|
|
def barrier_path(job_id: str, registry_dir: str) -> str:
|
|
return os.path.join(registry_dir, "{}.heartbeat.barrier".format(job_id))
|
|
|
|
|
|
def raise_barrier(job_id: str, registry_dir: str, dispatch_token=None, attempt_id=None) -> None:
|
|
"""Mark this attempt as ending (review round 2, codex R3).
|
|
|
|
Written under the registry lock - the lock `next_seq` takes - so a beat
|
|
either took its seq before this point (and the terminal publisher then
|
|
waits for the timer to drain) or meets the barrier when it asks for a seq
|
|
and is refused without consuming one."""
|
|
with mqtt_common.registry_lock(registry_dir):
|
|
with open(barrier_path(job_id, registry_dir), "w") as fh:
|
|
json.dump({"dispatch_token": dispatch_token, "attempt_id": attempt_id}, fh)
|
|
|
|
|
|
def barrier_blocks(job_id: str, registry_dir: str, dispatch_token=None, attempt_id=None) -> bool:
|
|
"""True when a terminal barrier was raised for this attempt (a later
|
|
attempt, with its own token, is not blocked). Check it under the registry
|
|
lock when the answer must be atomic with a seq."""
|
|
try:
|
|
with open(barrier_path(job_id, registry_dir)) as fh:
|
|
barrier = json.load(fh)
|
|
except (OSError, ValueError):
|
|
return False
|
|
if not isinstance(barrier, dict):
|
|
return True
|
|
if dispatch_token and barrier.get("dispatch_token"):
|
|
return barrier["dispatch_token"] == dispatch_token
|
|
if attempt_id is not None and barrier.get("attempt_id") is not None:
|
|
return int(barrier["attempt_id"]) == int(attempt_id)
|
|
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.
|
|
It is drained as at a terminal event - and killed if it does not exit -
|
|
before its pidfile is handed over (review round 3, codex R3): a beat it had
|
|
already admitted would otherwise still reach the wire after a terminal
|
|
publish that can find only the new timer through the pidfile.
|
|
``watch_pid`` binds the timer to a process (default: the nearest agent-CLI
|
|
ancestor of this process, if any); 0 disables the watch explicitly.
|
|
Returns 0 without arming when the attempt already ended (terminal record
|
|
or raised barrier): a start that lost the race against the terminal publish
|
|
must not leave a timer behind (review round 4, codex R3).
|
|
"""
|
|
if watch_pid is None:
|
|
watch_pid = nearest_agent_cli_ancestor()
|
|
with lifecycle_lock(job_id, registry_dir):
|
|
_stop_timer_locked(job_id, registry_dir, drain=drain_seconds())
|
|
try:
|
|
status = load_job(job_id, registry_dir).get("status")
|
|
except FileNotFoundError:
|
|
status = None
|
|
if status in TERMINAL_STATUSES or barrier_blocks(job_id, registry_dir, dispatch_token, attempt_id):
|
|
logger.info("not arming a heartbeat timer for %s: the attempt already ended", job_id)
|
|
return 0
|
|
return _spawn_timer_locked(job_id, registry_dir, interval, dispatch_token, attempt_id, python, watch_pid)
|
|
|
|
|
|
def _spawn_timer_locked(job_id, registry_dir, interval, dispatch_token, attempt_id, python, watch_pid) -> int:
|
|
"""Spawn the timer and publish its pid; the caller holds the lifecycle lock."""
|
|
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 barrier_blocks(job_id, registry_dir, dispatch_token, attempt_id):
|
|
return "the terminal barrier for this attempt is raised"
|
|
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:
|
|
# _beat=True: admitted under the seq lock, never after this
|
|
# attempt's terminal barrier (review round 2, codex R3)
|
|
rc = publish_event.main(argv, _beat=True)
|
|
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 (terminal barrier, terminal record or 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)
|
|
# Under the lifecycle lock: a replacement may be writing its own pid to
|
|
# this file right now, and that owner's pidfile must not be deleted.
|
|
try:
|
|
with lifecycle_lock(job_id, registry_dir):
|
|
with open(own_pidfile) as fh:
|
|
mine = fh.read().strip() == str(os.getpid())
|
|
if mine:
|
|
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())
|