#!/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 timed out (wall-clock or idle) before all jobs finished """ from __future__ import annotations import argparse 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 ( DEFAULT_REGISTRY_DIR, SCHEMA_VERSION, broker_config_from_job, load_job, make_client, ) logger = logging.getLogger("delegate_job.job_subscriber") TERMINAL_EVENTS = ("completed", "error") 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 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")) except Exception: job_rec = {} if not disk_status or disk_status not in ("completed", "error", "cancelled"): try: disk_status = mqtt_common.read_logged_status(jid, mqtt_common.get_logs_dir()) except Exception: pass 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): self.events: "queue.Queue[Tuple[str, Dict[str, Any]]]" = queue.Queue() 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 = {} 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 # --- 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 self.last_seq[jid] = seq # 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. 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", ""), }) self.events.put((msg.topic, payload)) 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 2 if not jobs: return 2 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 watcher = _Watcher(expected_ids, tokens, seqs, dispatch_tokens, attempts) # 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)) # 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 subscribed_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") 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 for topic in subscribed_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 1 return 0 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 1 return 0 logger.error("broker connection failed: %s", exc) return 3 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 2 if not subscribed: last_event = now idle_left = idle_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 2 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 last_event = time.monotonic() retention_checked = True print(_format_line(topic, payload), flush=True) jid = payload["job_id"] event = payload.get("event") 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 1 return 0 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("--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("-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.WARNING) try: return _run_subscriber(args) except Exception as exc: logger.error("subscriber fatal error: %s", exc) return 3 if __name__ == "__main__": sys.exit(main())