"""``muse coord heartbeat`` — extend a reservation's TTL without modifying it. Long-running agents use heartbeats to keep their reservations alive past the original TTL. Each heartbeat call atomically rewrites a single keep-alive file at ``.muse/coordination/heartbeats/.json``, updating ``extended_expires_at`` to ``now + extension_seconds``. The immutable reservation record is never modified. :func:`active_reservations` and :func:`filter_reservations` automatically use ``max(reservation.expires_at, heartbeat.extended_expires_at)`` as the effective expiry, so the heartbeat transparently extends the reservation's liveness window. Usage:: muse coord heartbeat --run-id AGENT-42 muse coord heartbeat --run-id AGENT-42 --extension 7200 muse coord heartbeat --run-id AGENT-42 --json Recommended pattern ------------------- Poll every ``extension_seconds / 2`` to keep a safety margin:: while working: do_work() muse coord heartbeat $RES_ID --run-id $RUN_ID --extension 3600 Exit code 1 (``already_released``) is a hard signal to stop — the reservation was cancelled by another agent or coordinator. Agents should treat it the same as a graceful shutdown request. JSON output schema:: { "status": "ok" | "already_released" | "not_found", "reservation_id": str, "run_id": str, "last_beat_at": str, // ISO 8601 "extended_expires_at": str, // ISO 8601 "ttl_extended_seconds": int, "elapsed_seconds": float } Exit codes:: 0 — heartbeat written successfully 1 — bad arguments (--run-id too long, --extension out of range, bad UUID) or reservation already released (stop sending heartbeats) 4 — reservation not found Flags: ``--run-id ID`` Agent/pipeline identifier for the audit trail (required). Maximum length: 256 characters. ``--extension SECONDS`` Number of seconds from now to set as the new expiry (default: 3600). Range: 1–31 536 000 (1 s to 1 year). """ from __future__ import annotations import argparse import json import sys import time from muse.core.coordination import ( _validate_reservation_id, create_heartbeat, load_all_reservations, load_released_ids, ) from muse.core.errors import ExitCode from muse.core.repo import require_repo from muse.core.validation import clamp_int, sanitize_display # ── Input constraints ───────────────────────────────────────────────────────── #: Maximum byte-length of the ``--run-id`` value. Mirrors the limit in #: ``reserve.py`` and ``release_coord.py`` so audit records stay bounded. _MAX_RUN_ID_LEN: int = 256 #: Maximum value for ``--extension``. Same ceiling as ``reserve.py``'s #: ``--ttl`` so a single heartbeat cannot pin a reservation open indefinitely. _MAX_EXTENSION_SECONDS: int = 31_536_000 # 1 year # ── CLI registration ────────────────────────────────────────────────────────── def register( subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]", ) -> None: """Register the ``heartbeat`` subcommand on *subparsers* (under ``muse coord``). Wires all flags with their defaults, choices, and help text so that ``--help`` output is accurate. Sets ``func`` to :func:`run`. """ parser = subparsers.add_parser( "heartbeat", help="Extend a reservation's TTL by writing a heartbeat.", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "reservation_id", metavar="RESERVATION_ID", help="UUID of the reservation to keep alive.", ) parser.add_argument( "--run-id", required=True, dest="run_id", metavar="ID", help=( "Agent run-id sending the heartbeat (for audit trail). " f"Max {_MAX_RUN_ID_LEN} characters." ), ) parser.add_argument( "--extension", type=int, default=3600, dest="extension_seconds", metavar="SECONDS", help=( f"Seconds from now to set as the new expiry " f"(default: 3600; range: 1–{_MAX_EXTENSION_SECONDS:,})." ), ) parser.add_argument( "--format", "-f", default="text", dest="fmt", choices=("text", "json"), help="Output format: text (default) or json.", ) parser.add_argument( "--json", action="store_const", const="json", dest="fmt", help="Shorthand for --format json.", ) parser.set_defaults(func=run) # ── Command implementation ──────────────────────────────────────────────────── def run(args: argparse.Namespace) -> None: """Write or refresh a heartbeat for an active reservation. Execution order --------------- 1. **Validate inputs** — ``--run-id`` length, ``--extension`` range via :func:`~muse.core.validation.clamp_int`, and UUID syntax of ``reservation_id``. All validation fires before any file I/O; any failure exits :attr:`~muse.core.errors.ExitCode.USER_ERROR` (1). 2. **Released guard** — checks :func:`~muse.core.coordination.load_released_ids` (stem-only scan, no JSON parsing). If the reservation is released, exits 1 with a message telling the agent to stop. 3. **Existence check** — loads the full reservation list and confirms the ID is present. If not, exits :attr:`~muse.core.errors.ExitCode.NOT_FOUND` (4). 4. **Write heartbeat** — calls :func:`~muse.core.coordination.create_heartbeat`, which atomically rewrites ``.muse/coordination/heartbeats/.json``. 5. **Emit output** — text to *stdout* (or compact JSON when ``--format json``). Security -------- * ``--run-id`` is capped at :data:`_MAX_RUN_ID_LEN` to bound audit record size. * ``--extension`` is clamped to ``[1, _MAX_EXTENSION_SECONDS]`` preventing a single heartbeat from locking a reservation open for an unbounded time. * ``reservation_id`` is validated as a UUID before any path is constructed, preventing directory traversal attacks. * All display strings are passed through :func:`~muse.core.validation.sanitize_display`. Args: args: Parsed ``argparse.Namespace`` with attributes ``reservation_id``, ``run_id``, ``extension_seconds``, and ``fmt``. Exit codes: 0 — heartbeat written successfully. 1 — bad arguments, or reservation already released (stop work). 4 — reservation not found. """ t0 = time.monotonic() reservation_id: str = args.reservation_id run_id: str = args.run_id extension_seconds: int = args.extension_seconds fmt: str = args.fmt as_json = fmt == "json" # ── Input validation (before any file I/O) ──────────────────────────────── if len(run_id) > _MAX_RUN_ID_LEN: msg = f"--run-id is too long ({len(run_id)} chars; max {_MAX_RUN_ID_LEN})" if as_json: print(json.dumps({"error": msg, "status": "bad_args"})) else: print(f"❌ {msg}", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) try: extension_seconds = clamp_int( extension_seconds, 1, _MAX_EXTENSION_SECONDS, "--extension" ) except ValueError as exc: msg = str(exc) if as_json: print(json.dumps({"error": msg, "status": "bad_args"})) else: print(f"❌ Invalid --extension: {msg}", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) try: _validate_reservation_id(reservation_id) except ValueError as exc: if as_json: print(json.dumps({"error": str(exc), "status": "bad_id"})) else: print(f"❌ {exc}", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) root = require_repo() # ── Released guard ──────────────────────────────────────────────────────── # Stem-only scan — no JSON parsing, O(released_count). released_ids = load_released_ids(root) if reservation_id in released_ids: elapsed = round(time.monotonic() - t0, 4) if as_json: print(json.dumps({ "status": "already_released", "reservation_id": reservation_id, "elapsed_seconds": elapsed, })) else: rid = sanitize_display(reservation_id) print( f"❌ reservation {rid} is already released — " "stop sending heartbeats", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) # ── Existence check ─────────────────────────────────────────────────────── all_res = load_all_reservations(root) known_ids = {r.reservation_id for r in all_res} if reservation_id not in known_ids: elapsed = round(time.monotonic() - t0, 4) if as_json: print(json.dumps({ "status": "not_found", "reservation_id": reservation_id, "elapsed_seconds": elapsed, })) else: rid = sanitize_display(reservation_id) print(f"❌ reservation {rid} not found", file=sys.stderr) raise SystemExit(ExitCode.NOT_FOUND) # ── Write heartbeat ─────────────────────────────────────────────────────── hb = create_heartbeat(root, reservation_id, run_id, extension_seconds) elapsed = round(time.monotonic() - t0, 4) if as_json: print(json.dumps({ "status": "ok", "reservation_id": hb.reservation_id, "run_id": hb.run_id, "last_beat_at": hb.last_beat_at.isoformat(), "extended_expires_at": hb.extended_expires_at.isoformat(), "ttl_extended_seconds": extension_seconds, "elapsed_seconds": elapsed, })) else: rid = sanitize_display(hb.reservation_id) run_label = sanitize_display(hb.run_id) exp_str = hb.extended_expires_at.isoformat()[:19] print( f"💓 heartbeat {rid} run={run_label}" f" extended until {exp_str}Z (+{extension_seconds}s)" )