#!/usr/bin/env python3 """job_subscriber.py — the single entry point for observing Job events. Subscribes to one job's ``/events`` (or, with ``--wait-any``, the events of every running/pending job in the registry), prints one line to stdout per accepted event, and exits on a terminal event or a timeout. Design points (all flagged in the PLAN review): - terminal state machine: ``completed``/``error`` is acted on exactly once per job, so QoS-1 duplicates or an ``error``-after-``completed`` reorder are safe. - dual timeouts: a wall-clock ``--timeout`` (total budget, started at subscribe time so a cold start can't hang forever) AND an idle ``--idle-timeout`` (no new event for N seconds). - defensive parsing: undecodable payloads, ``schema_version`` mismatches, and ``job_id`` values we did not subscribe for are logged and dropped. stdout = event lines only. Diagnostics go to stderr via logging. Exit codes: 0 all watched jobs reached ``completed`` 1 any watched job reached ``error`` 2 no event for the idle window (K x N) - the wrapper labels it deadline_timeout 3 broker unreachable 4 the wall-clock budget (--timeout) ran out - the wrapper labels it wall_clock_timeout """ from __future__ import annotations import argparse import hashlib import json import logging import queue import sys import time from typing import Any, Dict, List, Optional, Set, Tuple from lib_py.delegation import mqtt_common from lib_py.delegation import registry from lib_py.delegation.mqtt_common import ( SCHEMA_VERSION, broker_config_from_job, default_registry_dir, load_job, make_client, ) logger = logging.getLogger("delegate_job.job_subscriber") TERMINAL_EVENTS = ("completed", "error") TERMINAL_STATUSES = ("completed", "error", "cancelled") STATE_EVENT = "state" # Exit codes. The wrapper mirrors them as SUB_RC_* (a test keeps the two equal). RC_COMPLETED = 0 RC_ERROR = 1 # Issue #33: which clock ended the wait. The wrapper maps both to a timeout # (caller rc 2) and labels the close from them: rc 2 -> deadline_timeout # (no event for the idle window, K x N), rc 4 -> wall_clock_timeout (the job # budget, --timeout, ran out while events may still have been arriving). RC_NO_EVENT = 2 RC_BROKER_UNAVAILABLE = 3 RC_WALL_CLOCK = 4 def _format_line(topic: str, payload: Dict[str, Any]) -> str: source_tag = f" [{payload['source']}]" if payload.get("source") else "" return ( f"{payload.get('timestamp','-')} " f"job={payload.get('job_id','?')} " f"seq={payload.get('seq','?')} " f"{payload.get('event','?') + source_tag:<20} " f"{payload.get('detail','')}" ) def _check_disk_fallback( pending: Set[str], registry_dir: str, terminal: Dict[str, str], expected_dispatch_tokens: Optional[Dict[str, Optional[str]]] = None, expected_attempts: Optional[Dict[str, Optional[int]]] = None, ) -> None: """Check local on-disk registry and audit logs for terminal status (B-15).""" for jid in list(pending): disk_status = None rec_token = None rec_attempt = None rec_seq = None try: job_rec = load_job(jid, registry_dir) disk_status = job_rec.get("status") rec_token = job_rec.get("dispatch_token") rec_attempt = job_rec.get("attempt_id", job_rec.get("iteration")) rec_seq = job_rec.get("transition_seq") except Exception: job_rec = {} if not disk_status or disk_status not in ("completed", "error", "cancelled"): # read_logged_status returns the status.json dict; the status is its # "status" field (Issue #57, comment 5: comparing the dict matched nothing). try: logged = mqtt_common.read_logged_status(jid, mqtt_common.get_logs_dir(registry_dir)) except Exception: logged = None if logged: # The audit mirror is written after the registry commits, so a # record that is not terminal may sit next to the previous # attempt's terminal status. A versioned record (the engine's # transition_seq) takes the audit status only when the audit is # at least as new; an audit status without a sequence cannot # prove which attempt it describes, so a versioned record # ignores it. An unversioned record (never transitioned, or # MAM_ENGINE_CAS=0) and a missing record keep the old rule. aud_seq = logged.get("transition_seq") try: stale = rec_seq is not None and (aud_seq is None or int(aud_seq) < int(rec_seq)) except (TypeError, ValueError): stale = True # a damaged status.json proves nothing either if stale: logger.warning("disk fallback ignored audit status %r for %s: transition_seq %r is older " "than the record's %r", logged.get("status"), jid, aud_seq, rec_seq) else: disk_status = logged.get("status") if disk_status in ("completed", "error", "cancelled"): if expected_dispatch_tokens: exp_token = expected_dispatch_tokens.get(jid) if exp_token is not None and rec_token != exp_token: logger.warning("disk fallback ignored terminal status for %s due to dispatch_token mismatch (%r != %r)", jid, rec_token, exp_token) continue if expected_attempts: exp_att = expected_attempts.get(jid) if exp_att is not None and rec_attempt is not None and rec_attempt < exp_att: logger.warning("disk fallback ignored terminal status for %s due to attempt_id older than expected (%r < %r)", jid, rec_attempt, exp_att) continue synth_event = "completed" if disk_status == "completed" else "error" synth_payload = { "job_id": jid, "event": synth_event, "seq": -1, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "detail": f"terminal status '{disk_status}' resolved via disk-fallback", "source": "disk-fallback", } synth_topic = f"mam/{jid}/events" print(_format_line(synth_topic, synth_payload), flush=True) terminal[jid] = synth_event pending.discard(jid) class _Watcher: """Holds the shared queue + the set of job_ids we accept events for.""" def __init__(self, expected_job_ids: Set[str], expected_tokens: Dict[str, Optional[str]], expected_seqs: Dict[str, int], expected_dispatch_tokens: Optional[Dict[str, Optional[str]]] = None, expected_attempts: Optional[Dict[str, Optional[int]]] = None, logs_dir: Optional[str] = None): self.events: "queue.Queue[Tuple[str, Dict[str, Any]]]" = queue.Queue() self.logs_dir = logs_dir # Issue #57: the audit log next to the registry being watched self.expected = set(expected_job_ids) self.tokens = expected_tokens # job_id -> expected auth_token (or None) self.last_seq = dict(expected_seqs) self.dispatch_tokens = dict(expected_dispatch_tokens or {}) self.attempts = dict(expected_attempts or {}) # Intent identities already accepted, per job. A replayed notification is # re-published with a *fresh* wire sequence, so the monotonic sequence # guard below cannot recognise it: to that guard a retry looks like a new # event. The intent identity is what stays stable across a retry, so # deduplication has to happen on that. # # This is per-process state. Across a restart the disk fallback resolves # terminal status from the registry, and a duplicate terminal is refused # by the transition boundary as an idempotent ack, so a replay after # restart cannot double-apply either. self.seen_intents = {} # Issue #18: jobs for which at least one wire event was accepted. A seq # gap is only meaningful after that - the initial last_seq comes from # the registry and may be behind the wire (or ahead of a fresh remote). self.accepted_any: Set[str] = set() self.state_seq: Dict[str, int] = {} def _note_seq_gap(self, jid: str, expected: int, got: int, via: str, topic: str) -> None: """Issue #18: a seq the subscriber never saw is a detected loss.""" logger.warning("seq gap for job %s: expected %d, got %d (via %s)", jid, expected, got, via) mqtt_common.append_event(jid, { "event": "seq_gap", "expected_seq": expected, "seq": got, "via": via, "topic": topic, }, logs_dir=self.logs_dir) def _on_state(self, msg, payload: Dict[str, Any], jid: str) -> None: """A retained `/state` message (Issue #18): learn the job's status and last seq at (re)connect without waiting for the next heartbeat. Signed and token-checked like an event. Delivered to the main loop as a state item so the idle window resets and a terminal status ends the wait.""" seq = payload.get("seq") if not isinstance(seq, int): logger.warning("drop state for job %s: missing or invalid seq", jid) return # A state message never advances the event watermark: MQTT orders # messages per connection, not across topics, so a state that overtakes # its own event must not make the monotonic guard drop that event. last = max(self.last_seq.get(jid, 0), self.state_seq.get(jid, 0)) if jid in self.accepted_any and seq > last + 1: self._note_seq_gap(jid, last + 1, seq, "state", msg.topic) # Review round 2 (codex R4): a state renews liveness only when it carries # news - a seq beyond every event and state already accepted. The same # retained snapshot re-delivered on every reconnect, or an older one # replayed, still recovers a terminal status in the main loop, but it is # no evidence that the agent is alive now. renews = seq > last if seq > self.state_seq.get(jid, 0): self.state_seq[jid] = seq mqtt_common.append_event(jid, { "event": "received", "source_event": STATE_EVENT, "seq": seq, "topic": msg.topic, "timestamp": payload.get("timestamp"), "detail": "status={} last_event={}".format(payload.get("status"), payload.get("last_event")), }, logs_dir=self.logs_dir) self.events.put((msg.topic, dict(payload, _renews_liveness=renews))) def on_message(self, _client, _userdata, msg) -> None: # --- defensive parsing ------------------------------------------- try: payload = json.loads(msg.payload.decode("utf-8")) except (UnicodeDecodeError, json.JSONDecodeError) as exc: logger.warning("drop unparseable payload on %s: %s", msg.topic, exc) return if not isinstance(payload, dict): logger.warning("drop non-object payload on %s", msg.topic) return if payload.get("schema_version") != SCHEMA_VERSION: logger.warning("drop event with schema_version=%r (expected %d)", payload.get("schema_version"), SCHEMA_VERSION) return jid = payload.get("job_id") if jid not in self.expected: logger.warning("drop event for unexpected job_id=%r on %s", jid, msg.topic) return # --- production auth check: data.auth_token must match if expected --- expected_token = self.tokens.get(jid) if not mqtt_common.verify_hmac(payload, expected_token): logger.warning("drop event for job %s: HMAC verify failed", jid) return # --- CAS dispatch token & attempt check: drop stale attempt events --- expected_disp_tok = self.dispatch_tokens.get(jid) if expected_disp_tok is not None: data = payload.get("data") ev_tok = data.get("dispatch_token") if isinstance(data, dict) else None if ev_tok != expected_disp_tok: logger.warning("drop event for job %s: dispatch_token %r does not match expected %r", jid, ev_tok, expected_disp_tok) return expected_att = self.attempts.get(jid) if expected_att is not None: data = payload.get("data") ev_att = data.get("attempt_id") if isinstance(data, dict) else None if ev_att is not None and isinstance(ev_att, int) and ev_att < expected_att: logger.warning("drop event for job %s: attempt_id %r is older than expected %r", jid, ev_att, expected_att) return if payload.get("event") == STATE_EVENT: self._on_state(msg, payload, jid) return # --- replay attack defense: check monotonic sequence --- seq = payload.get("seq") if seq is None or not isinstance(seq, int): logger.warning("drop event for job %s: missing or invalid seq", jid) return data = payload.get("data") intent = data.get("notification_intent") if isinstance(data, dict) else None identity = None if intent is not None: # Present but not an object: it cannot yield a stable identity, so the # replay guard below could not protect it. Refuse rather than admit a # notification we would be unable to recognise on its retry. if not isinstance(intent, dict): logger.warning("drop event for job %s: notification_intent is %s, not an object", jid, type(intent).__name__) return identity = (intent.get("seq"), intent.get("event"), intent.get("timestamp")) try: already_seen = identity in self.seen_intents.setdefault(jid, set()) except TypeError: # A malformed sender can put an unhashable value in a field. That # must not escape defensive parsing as an exception: a callback # raising here would take down the network loop for every job. logger.warning("drop event for job %s: notification intent %r has no stable identity", jid, intent) return if already_seen: logger.info("drop event for job %s: notification intent %s already accepted", jid, identity) return if seq <= self.last_seq.get(jid, 0): logger.warning("drop event for job %s: seq %d is not monotonically increasing (last %d)", jid, seq, self.last_seq.get(jid, 0)) return if jid in self.accepted_any and seq > self.last_seq.get(jid, 0) + 1: self._note_seq_gap(jid, self.last_seq.get(jid, 0) + 1, seq, "events", msg.topic) self.last_seq[jid] = seq self.accepted_any.add(jid) # Recorded only once the event is actually accepted. A late packet that the # sequence guard refuses must not mark its intent as seen, or the genuine # retry that follows would be mistaken for the replay of an event this # subscriber never delivered. if identity is not None: self.seen_intents[jid].add(identity) # Persistent audit log from the *subscriber's* vantage point: every event # that survives defensive parsing is recorded here, including ones a # different host published. This is the external-observer record that # backstops the publisher's own "published" line if it never wrote one. # Which dispatch this record vouches for. `dispatch_verified` says the # token check above ran (the record had a dispatch_token and the # event's data.dispatch_token matched it) after HMAC verification; on # its own it names no dispatch, and this file is appended to across a # job's attempts, so a reader waiting for a later dispatch would take # an earlier attempt's record as its acknowledgement (review round 4). # `dispatch_fingerprint` is SHA-256 of the token that was verified - # the same identity the local path compares in events.log, minus the # token itself, which the audit redactor masks by key name and which a # 128-bit uuid4 nonce is not recoverable from. `attempt_id` is kept for # readers; it is not the identity, since a token can be re-minted # within an attempt. A reader that demands a token must match the # fingerprint; a record without one proves nothing about identity. mqtt_common.append_event(jid, { "event": "received", "source_event": payload.get("event"), "seq": payload.get("seq"), "topic": msg.topic, "timestamp": payload.get("timestamp"), "detail": payload.get("detail", ""), "dispatch_verified": expected_disp_tok is not None, "dispatch_fingerprint": (hashlib.sha256(expected_disp_tok.encode("utf-8")).hexdigest() if expected_disp_tok else None), "attempt_id": (data.get("attempt_id") if isinstance(data, dict) and data.get("attempt_id") is not None else self.attempts.get(jid)), }, logs_dir=self.logs_dir) self.events.put((msg.topic, payload)) def _sync_remote_status(jid: str, status: str, registry_dir: str, dispatch_token: Optional[str]) -> None: """Issue #18: mirror a wire event into the orchestrator's registry for a job whose agent shares no disk with it. Goes through the same transition boundary as the publisher's own status sync; a refused (late/duplicate) transition is logged, never raised.""" try: res = mqtt_common.transition_job_status( jid, registry_dir, _source="job_subscriber", status=status, expected_dispatch_token=dispatch_token) if not res.committed and not res.idempotent: logger.warning("remote status sync refused for %s -> %s: %s", jid, status, res.reason) except Exception as exc: # pragma: no cover - best effort logger.warning("remote status sync failed for %s: %s", jid, exc) def _collect_jobs(args) -> List[Dict[str, Any]]: """Resolve the list of job records this invocation should watch.""" if args.wait_any: jobs = [r for r in registry.list_jobs(args.registry_dir) if r.get("status") in ("pending", "running")] if not jobs: logger.error("no pending/running jobs to wait for") return jobs job = load_job(args.job, args.registry_dir) # raises FileNotFoundError return [job] def _run_subscriber(args) -> int: try: jobs = _collect_jobs(args) except FileNotFoundError as exc: logger.error("%s", exc) return RC_NO_EVENT # nothing to watch: reported like silence, as before #33 if not jobs: return RC_NO_EVENT expected_ids: Set[str] = {j["job_id"] for j in jobs} tokens = {j["job_id"]: j.get("auth_token") for j in jobs} seqs = {} dispatch_tokens = {} attempts = {} for j in jobs: jid = j["job_id"] last_seq = int(j.get("last_seq", 0)) seqs[jid] = max(0, last_seq - 1) dispatch_tokens[jid] = j.get("dispatch_token") attempts[jid] = j.get("attempt_id", j.get("iteration", 1)) if args.expect_dispatch_token is not None and args.job: dispatch_tokens[args.job] = args.expect_dispatch_token if args.attempt_id is not None and args.job: attempts[args.job] = args.attempt_id logs_dir = mqtt_common.get_logs_dir(args.registry_dir) watcher = _Watcher(expected_ids, tokens, seqs, dispatch_tokens, attempts, logs_dir=logs_dir) # Resolve timeouts from CLI, falling back to the (first) job's settings. base_job = jobs[0] wall_timeout = args.timeout if args.timeout is not None else float(base_job.get("timeout_sec", 3600)) idle_timeout = args.idle_timeout if args.idle_timeout is not None else float(base_job.get("idle_timeout_sec", 120)) # Issue #18: until the first accepted event the "last event" is the # injection itself, which happens after this subscriber started and may # take up to the ack timeout T to produce `started` (a remote runner has # to receive the dispatch first). The idle budget before the first event # is therefore T + K x N; from the first event on it is K x N. ack_timeout = args.ack_timeout if args.ack_timeout is not None else float(base_job.get("ack_timeout_sec", 0) or 0) first_event_seen = False # All watched jobs share a broker in practice; connect using the first # job's broker and subscribe to each job's events topic. config = broker_config_from_job(base_job) client = make_client("subscriber", config) client.on_message = watcher.on_message # Issue #18: a remote job (no shared disk with the agent) has nobody else # to move its registry record - the subscriber does it from the wire. remote_sync = bool(args.remote) or any(j.get("transport") == "remote" for j in jobs) subscribed_topics = [] state_topics = [] for job in jobs: prefix = job.get("topic_prefix") or mqtt_common.topic_prefix_for(job["job_id"]) subscribed_topics.append(f"{prefix}/events") state_topics.append(f"{prefix}/state") connects = {"count": 0} def on_connect(_c, _u, _flags, reason_code, _props): if mqtt_common.reason_code_value(reason_code) != 0: logger.error("broker connection failed: rc=%s", reason_code) return connects["count"] += 1 if connects["count"] > 1: # Issue #18: paho reconnected on its own; re-subscribing here is # what makes the retained state message arrive again at once. logger.warning("broker reconnected (connect #%d); re-subscribing", connects["count"]) for jid in expected_ids: mqtt_common.append_event(jid, {"event": "reconnected", "connect_count": connects["count"]}, logs_dir=logs_dir) for topic in subscribed_topics + state_topics: _c.subscribe(topic, qos=1) logger.info("subscribed to %s", topic) def on_disconnect(_c, _u, _flags, reason_code, _props): rc = mqtt_common.reason_code_value(reason_code) if rc != 0: logger.warning("broker disconnected (rc=%s); will retry reconnect", reason_code) start = time.monotonic() wall_deadline = start + wall_timeout last_event = start last_disk_check = 0.0 retention_checked = not args.expect_retention connected = False subscribed = False def on_subscribe(_c, _u, mid, granted_qos, _props=None): nonlocal last_event, subscribed subscribed = True last_event = time.monotonic() for topic in subscribed_topics: print(f"SUBSCRIBED {topic}", flush=True) client.on_connect = on_connect client.on_disconnect = on_disconnect client.on_subscribe = on_subscribe client.reconnect_delay_set(min_delay=1, max_delay=16) terminal: Dict[str, str] = {} # job_id -> "completed"/"error" pending: Set[str] = set(expected_ids) # Check if all jobs are already in terminal state on disk before connecting _check_disk_fallback(pending, args.registry_dir, terminal, dispatch_tokens, attempts) if not pending: if any(state == "error" for state in terminal.values()): return RC_ERROR return RC_COMPLETED try: mqtt_common.with_retry( lambda: client.connect(config.host, config.port, config.keepalive), attempts=3, base_delay=0.2, max_delay=1.0 )() client.loop_start() connected = True last_event = time.monotonic() except Exception as exc: logger.warning("broker connection failed (%s); checking disk fallback", exc) _check_disk_fallback(pending, args.registry_dir, terminal, dispatch_tokens, attempts) if not pending: if any(state == "error" for state in terminal.values()): return RC_ERROR return RC_COMPLETED logger.error("broker connection failed: %s", exc) return RC_BROKER_UNAVAILABLE try: while pending: now = time.monotonic() if now >= wall_deadline: _check_disk_fallback(pending, args.registry_dir, terminal, dispatch_tokens, attempts) if not pending: break logger.error("wall-clock timeout (%.0fs); still pending: %s", wall_timeout, ", ".join(sorted(pending))) return RC_WALL_CLOCK if not subscribed: last_event = now idle_left = (idle_timeout if first_event_seen else idle_timeout + ack_timeout) - (now - last_event) if idle_left <= 0: _check_disk_fallback(pending, args.registry_dir, terminal, dispatch_tokens, attempts) if not pending: break logger.error("idle timeout (%.0fs, no events); still pending: %s", idle_timeout, ", ".join(sorted(pending))) return RC_NO_EVENT wait = min(wall_deadline - now, idle_left, 1.0) try: topic, payload = watcher.events.get(timeout=wait) except queue.Empty: if not retention_checked and (now - start) > 3.0: logger.warning("--expect-retention set but no retained " "terminal event observed yet") retention_checked = True # Step 2 (B-15): Local disk fallback throttled to every 3.0s if (now - last_disk_check) >= 3.0: last_disk_check = now _check_disk_fallback(pending, args.registry_dir, terminal, dispatch_tokens, attempts) continue retention_checked = True if payload.pop("_renews_liveness", True): last_event = time.monotonic() first_event_seen = True jid = payload["job_id"] event = payload.get("event") if event == STATE_EVENT: # Issue #18: retained job state. It renewed the idle window # above only if it advanced the seq (review round 2, codex R4: # a re-delivered or older snapshot is not liveness); a terminal # status ends the wait exactly like the terminal event. status = payload.get("status") if status in TERMINAL_STATUSES and jid not in terminal: synth = "completed" if status == "completed" else "error" payload = dict(payload, event=synth, source="state", detail="terminal status '{}' recovered from retained state".format(status)) print(_format_line(topic, payload), flush=True) terminal[jid] = synth pending.discard(jid) if remote_sync: _sync_remote_status(jid, status, args.registry_dir, dispatch_tokens.get(jid)) continue print(_format_line(topic, payload), flush=True) if remote_sync and event in ("started", "completed", "error", "cancelled"): _sync_remote_status(jid, {"started": "running"}.get(event, event), args.registry_dir, dispatch_tokens.get(jid)) if event in TERMINAL_EVENTS: if jid in terminal: # Already finalised: ignore duplicates / late reorders. logger.info("ignoring duplicate terminal %s for %s", event, jid) continue terminal[jid] = event pending.discard(jid) finally: if connected: client.loop_stop() try: client.disconnect() except Exception: # pragma: no cover pass # All jobs reached a terminal state. error wins over completed. if any(state == "error" for state in terminal.values()): return RC_ERROR return RC_COMPLETED def main(argv=None) -> int: parser = argparse.ArgumentParser(description="Subscribe to Job events on MQTT") target = parser.add_mutually_exclusive_group(required=True) target.add_argument("--job", help="job id to watch") target.add_argument("--wait-any", action="store_true", help="watch every pending/running job in the registry") parser.add_argument("--timeout", type=float, default=None, help="wall-clock budget in seconds (default: job.timeout_sec or 3600)") parser.add_argument("--idle-timeout", type=float, default=None, help="max seconds with no new event (default: job.idle_timeout_sec or 120)") parser.add_argument("--ack-timeout", type=float, default=None, help="extra idle budget before the first event, for the injection to " "produce `started` (default: job.ack_timeout_sec; Issue #18)") parser.add_argument("--expect-retention", action="store_true", help="warn if no retained terminal event arrives promptly") parser.add_argument("--expect-dispatch-token", default=None, help="expected dispatch token for CAS verification") parser.add_argument("--attempt-id", type=int, default=None, help="minimum attempt id to accept events for") parser.add_argument("--registry-dir", default=default_registry_dir()) parser.add_argument("--remote", action="store_true", help="the agent shares no disk with this registry: mirror started/" "terminal events into the job record (Issue #18; automatic " "for jobs registered with --transport remote)") parser.add_argument("-v", "--verbose", action="store_true") args = parser.parse_args(argv) if not mqtt_common.load_env_or_report(): # rc 1, as the EnvFileError traceback exited before. The wrapper reads # the same file first and exits 1 itself, so it never sees this one. return RC_ERROR mqtt_common.setup_logging(logging.DEBUG if args.verbose else logging.WARNING) try: return _run_subscriber(args) except Exception as exc: logger.error("subscriber fatal error: %s", exc) return RC_BROKER_UNAVAILABLE if __name__ == "__main__": sys.exit(main())