heartbeat_coord.py
file-level
1
files
1
commits
0
hotspots
0
🧊 dead
0
💥 blast risk
| 1 | """``muse coord heartbeat`` — extend a reservation's TTL without modifying it. |
| 2 | |
| 3 | Long-running agents use heartbeats to keep their reservations alive past the |
| 4 | original TTL. Each heartbeat call atomically rewrites a single keep-alive file |
| 5 | at ``.muse/coordination/heartbeats/<reservation-id>.json``, updating |
| 6 | ``extended_expires_at`` to ``now + extension_seconds``. |
| 7 | |
| 8 | The immutable reservation record is never modified. :func:`active_reservations` |
| 9 | and :func:`filter_reservations` automatically use |
| 10 | ``max(reservation.expires_at, heartbeat.extended_expires_at)`` as the effective |
| 11 | expiry, so the heartbeat transparently extends the reservation's liveness window. |
| 12 | |
| 13 | Usage:: |
| 14 | |
| 15 | muse coord heartbeat <reservation-id> --run-id AGENT-42 |
| 16 | muse coord heartbeat <reservation-id> --run-id AGENT-42 --extension 7200 |
| 17 | muse coord heartbeat <reservation-id> --run-id AGENT-42 --json |
| 18 | |
| 19 | Recommended pattern |
| 20 | ------------------- |
| 21 | Poll every ``extension_seconds / 2`` to keep a safety margin:: |
| 22 | |
| 23 | while working: |
| 24 | do_work() |
| 25 | muse coord heartbeat $RES_ID --run-id $RUN_ID --extension 3600 |
| 26 | |
| 27 | Exit code 1 (``already_released``) is a hard signal to stop — the reservation |
| 28 | was cancelled by another agent or coordinator. Agents should treat it the same |
| 29 | as a graceful shutdown request. |
| 30 | |
| 31 | JSON output schema:: |
| 32 | |
| 33 | { |
| 34 | "status": "ok" | "already_released" | "not_found", |
| 35 | "reservation_id": str, |
| 36 | "run_id": str, |
| 37 | "last_beat_at": str, // ISO 8601 |
| 38 | "extended_expires_at": str, // ISO 8601 |
| 39 | "ttl_extended_seconds": int, |
| 40 | "elapsed_seconds": float |
| 41 | } |
| 42 | |
| 43 | Exit codes:: |
| 44 | |
| 45 | 0 — heartbeat written successfully |
| 46 | 1 — bad arguments (--run-id too long, --extension out of range, bad UUID) |
| 47 | or reservation already released (stop sending heartbeats) |
| 48 | 4 — reservation not found |
| 49 | |
| 50 | Flags: |
| 51 | |
| 52 | ``--run-id ID`` |
| 53 | Agent/pipeline identifier for the audit trail (required). |
| 54 | Maximum length: 256 characters. |
| 55 | |
| 56 | ``--extension SECONDS`` |
| 57 | Number of seconds from now to set as the new expiry (default: 3600). |
| 58 | Range: 1–31 536 000 (1 s to 1 year). |
| 59 | """ |
| 60 | |
| 61 | from __future__ import annotations |
| 62 | |
| 63 | import argparse |
| 64 | import json |
| 65 | import sys |
| 66 | import time |
| 67 | |
| 68 | from muse.core.coordination import ( |
| 69 | _validate_reservation_id, |
| 70 | create_heartbeat, |
| 71 | load_all_reservations, |
| 72 | load_released_ids, |
| 73 | ) |
| 74 | from muse.core.errors import ExitCode |
| 75 | from muse.core.repo import require_repo |
| 76 | from muse.core.validation import clamp_int, sanitize_display |
| 77 | |
| 78 | # ── Input constraints ───────────────────────────────────────────────────────── |
| 79 | |
| 80 | #: Maximum byte-length of the ``--run-id`` value. Mirrors the limit in |
| 81 | #: ``reserve.py`` and ``release_coord.py`` so audit records stay bounded. |
| 82 | _MAX_RUN_ID_LEN: int = 256 |
| 83 | |
| 84 | #: Maximum value for ``--extension``. Same ceiling as ``reserve.py``'s |
| 85 | #: ``--ttl`` so a single heartbeat cannot pin a reservation open indefinitely. |
| 86 | _MAX_EXTENSION_SECONDS: int = 31_536_000 # 1 year |
| 87 | |
| 88 | |
| 89 | # ── CLI registration ────────────────────────────────────────────────────────── |
| 90 | |
| 91 | |
| 92 | def register( |
| 93 | subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]", |
| 94 | ) -> None: |
| 95 | """Register the ``heartbeat`` subcommand on *subparsers* (under ``muse coord``). |
| 96 | |
| 97 | Wires all flags with their defaults, choices, and help text so that |
| 98 | ``--help`` output is accurate. Sets ``func`` to :func:`run`. |
| 99 | """ |
| 100 | parser = subparsers.add_parser( |
| 101 | "heartbeat", |
| 102 | help="Extend a reservation's TTL by writing a heartbeat.", |
| 103 | description=__doc__, |
| 104 | formatter_class=argparse.RawDescriptionHelpFormatter, |
| 105 | ) |
| 106 | parser.add_argument( |
| 107 | "reservation_id", |
| 108 | metavar="RESERVATION_ID", |
| 109 | help="UUID of the reservation to keep alive.", |
| 110 | ) |
| 111 | parser.add_argument( |
| 112 | "--run-id", |
| 113 | required=True, |
| 114 | dest="run_id", |
| 115 | metavar="ID", |
| 116 | help=( |
| 117 | "Agent run-id sending the heartbeat (for audit trail). " |
| 118 | f"Max {_MAX_RUN_ID_LEN} characters." |
| 119 | ), |
| 120 | ) |
| 121 | parser.add_argument( |
| 122 | "--extension", |
| 123 | type=int, |
| 124 | default=3600, |
| 125 | dest="extension_seconds", |
| 126 | metavar="SECONDS", |
| 127 | help=( |
| 128 | f"Seconds from now to set as the new expiry " |
| 129 | f"(default: 3600; range: 1–{_MAX_EXTENSION_SECONDS:,})." |
| 130 | ), |
| 131 | ) |
| 132 | parser.add_argument( |
| 133 | "--format", "-f", |
| 134 | default="text", |
| 135 | dest="fmt", |
| 136 | choices=("text", "json"), |
| 137 | help="Output format: text (default) or json.", |
| 138 | ) |
| 139 | parser.add_argument( |
| 140 | "--json", |
| 141 | action="store_const", |
| 142 | const="json", |
| 143 | dest="fmt", |
| 144 | help="Shorthand for --format json.", |
| 145 | ) |
| 146 | parser.set_defaults(func=run) |
| 147 | |
| 148 | |
| 149 | # ── Command implementation ──────────────────────────────────────────────────── |
| 150 | |
| 151 | |
| 152 | def run(args: argparse.Namespace) -> None: |
| 153 | """Write or refresh a heartbeat for an active reservation. |
| 154 | |
| 155 | Execution order |
| 156 | --------------- |
| 157 | 1. **Validate inputs** — ``--run-id`` length, ``--extension`` range via |
| 158 | :func:`~muse.core.validation.clamp_int`, and UUID syntax of |
| 159 | ``reservation_id``. All validation fires before any file I/O; any |
| 160 | failure exits :attr:`~muse.core.errors.ExitCode.USER_ERROR` (1). |
| 161 | 2. **Released guard** — checks :func:`~muse.core.coordination.load_released_ids` |
| 162 | (stem-only scan, no JSON parsing). If the reservation is released, exits |
| 163 | 1 with a message telling the agent to stop. |
| 164 | 3. **Existence check** — loads the full reservation list and confirms the |
| 165 | ID is present. If not, exits |
| 166 | :attr:`~muse.core.errors.ExitCode.NOT_FOUND` (4). |
| 167 | 4. **Write heartbeat** — calls |
| 168 | :func:`~muse.core.coordination.create_heartbeat`, which atomically |
| 169 | rewrites ``.muse/coordination/heartbeats/<id>.json``. |
| 170 | 5. **Emit output** — text to *stdout* (or compact JSON when ``--format json``). |
| 171 | |
| 172 | Security |
| 173 | -------- |
| 174 | * ``--run-id`` is capped at :data:`_MAX_RUN_ID_LEN` to bound audit record size. |
| 175 | * ``--extension`` is clamped to ``[1, _MAX_EXTENSION_SECONDS]`` preventing |
| 176 | a single heartbeat from locking a reservation open for an unbounded time. |
| 177 | * ``reservation_id`` is validated as a UUID before any path is constructed, |
| 178 | preventing directory traversal attacks. |
| 179 | * All display strings are passed through |
| 180 | :func:`~muse.core.validation.sanitize_display`. |
| 181 | |
| 182 | Args: |
| 183 | args: Parsed ``argparse.Namespace`` with attributes ``reservation_id``, |
| 184 | ``run_id``, ``extension_seconds``, and ``fmt``. |
| 185 | |
| 186 | Exit codes: |
| 187 | 0 — heartbeat written successfully. |
| 188 | 1 — bad arguments, or reservation already released (stop work). |
| 189 | 4 — reservation not found. |
| 190 | """ |
| 191 | t0 = time.monotonic() |
| 192 | reservation_id: str = args.reservation_id |
| 193 | run_id: str = args.run_id |
| 194 | extension_seconds: int = args.extension_seconds |
| 195 | fmt: str = args.fmt |
| 196 | as_json = fmt == "json" |
| 197 | |
| 198 | # ── Input validation (before any file I/O) ──────────────────────────────── |
| 199 | |
| 200 | if len(run_id) > _MAX_RUN_ID_LEN: |
| 201 | msg = f"--run-id is too long ({len(run_id)} chars; max {_MAX_RUN_ID_LEN})" |
| 202 | if as_json: |
| 203 | print(json.dumps({"error": msg, "status": "bad_args"})) |
| 204 | else: |
| 205 | print(f"❌ {msg}", file=sys.stderr) |
| 206 | raise SystemExit(ExitCode.USER_ERROR) |
| 207 | |
| 208 | try: |
| 209 | extension_seconds = clamp_int( |
| 210 | extension_seconds, 1, _MAX_EXTENSION_SECONDS, "--extension" |
| 211 | ) |
| 212 | except ValueError as exc: |
| 213 | msg = str(exc) |
| 214 | if as_json: |
| 215 | print(json.dumps({"error": msg, "status": "bad_args"})) |
| 216 | else: |
| 217 | print(f"❌ Invalid --extension: {msg}", file=sys.stderr) |
| 218 | raise SystemExit(ExitCode.USER_ERROR) |
| 219 | |
| 220 | try: |
| 221 | _validate_reservation_id(reservation_id) |
| 222 | except ValueError as exc: |
| 223 | if as_json: |
| 224 | print(json.dumps({"error": str(exc), "status": "bad_id"})) |
| 225 | else: |
| 226 | print(f"❌ {exc}", file=sys.stderr) |
| 227 | raise SystemExit(ExitCode.USER_ERROR) |
| 228 | |
| 229 | root = require_repo() |
| 230 | |
| 231 | # ── Released guard ──────────────────────────────────────────────────────── |
| 232 | # Stem-only scan — no JSON parsing, O(released_count). |
| 233 | released_ids = load_released_ids(root) |
| 234 | if reservation_id in released_ids: |
| 235 | elapsed = round(time.monotonic() - t0, 4) |
| 236 | if as_json: |
| 237 | print(json.dumps({ |
| 238 | "status": "already_released", |
| 239 | "reservation_id": reservation_id, |
| 240 | "elapsed_seconds": elapsed, |
| 241 | })) |
| 242 | else: |
| 243 | rid = sanitize_display(reservation_id) |
| 244 | print( |
| 245 | f"❌ reservation {rid} is already released — " |
| 246 | "stop sending heartbeats", |
| 247 | file=sys.stderr, |
| 248 | ) |
| 249 | raise SystemExit(ExitCode.USER_ERROR) |
| 250 | |
| 251 | # ── Existence check ─────────────────────────────────────────────────────── |
| 252 | all_res = load_all_reservations(root) |
| 253 | known_ids = {r.reservation_id for r in all_res} |
| 254 | if reservation_id not in known_ids: |
| 255 | elapsed = round(time.monotonic() - t0, 4) |
| 256 | if as_json: |
| 257 | print(json.dumps({ |
| 258 | "status": "not_found", |
| 259 | "reservation_id": reservation_id, |
| 260 | "elapsed_seconds": elapsed, |
| 261 | })) |
| 262 | else: |
| 263 | rid = sanitize_display(reservation_id) |
| 264 | print(f"❌ reservation {rid} not found", file=sys.stderr) |
| 265 | raise SystemExit(ExitCode.NOT_FOUND) |
| 266 | |
| 267 | # ── Write heartbeat ─────────────────────────────────────────────────────── |
| 268 | hb = create_heartbeat(root, reservation_id, run_id, extension_seconds) |
| 269 | |
| 270 | elapsed = round(time.monotonic() - t0, 4) |
| 271 | if as_json: |
| 272 | print(json.dumps({ |
| 273 | "status": "ok", |
| 274 | "reservation_id": hb.reservation_id, |
| 275 | "run_id": hb.run_id, |
| 276 | "last_beat_at": hb.last_beat_at.isoformat(), |
| 277 | "extended_expires_at": hb.extended_expires_at.isoformat(), |
| 278 | "ttl_extended_seconds": extension_seconds, |
| 279 | "elapsed_seconds": elapsed, |
| 280 | })) |
| 281 | else: |
| 282 | rid = sanitize_display(hb.reservation_id) |
| 283 | run_label = sanitize_display(hb.run_id) |
| 284 | exp_str = hb.extended_expires_at.isoformat()[:19] |
| 285 | print( |
| 286 | f"💓 heartbeat {rid} run={run_label}" |
| 287 | f" extended until {exp_str}Z (+{extension_seconds}s)" |
| 288 | ) |