multi-agent-mux/.agents/skills/lib_py/delegation/publish_event.py

#!/usr/bin/env python3
"""publish_event.py — the single entry point for emitting a Job event.

Loads the job record from the registry, resolves its broker, assigns the next
monotonic ``seq``, builds the schema-v1 JSON payload, and publishes it to
``<topic_prefix>/events`` over QoS 1 with exponential-backoff retry.

Silent by design: nothing is printed to stdout. Diagnostics go to stderr via
logging. Terminal events (``completed``/``error``) publish with retain=True so
a late subscriber still observes the final state (production hardening).

Exit codes:
  0  published successfully
  1  parameter / registry error (bad args, unknown job, no pending job)
  2  publish failed after retries (network / broker / ACK timeout)

Usage:
  publish_event.py --job <id> --event started [--detail "..."] [--data '{...}']
  publish_event.py --pick-pending --agent-session herdr:claude --event completed
  publish_event.py --job <id> --event completed --retained
"""
from __future__ import annotations

import argparse
import hashlib
import hmac
import json
import logging
import sys
from typing import Any, Dict, Optional

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,
    events_topic_for,
    load_job,
    next_seq,
    with_retry,
)

logger = logging.getLogger("delegate_job.publish_event")

VALID_EVENTS = ("started", "permission_required", "progress", "completed", "error", "cancelled")
TERMINAL_EVENTS = ("completed", "error", "cancelled")
# event -> registry status to sync as a best-effort side effect
EVENT_TO_STATUS = {
    "started": "running",
    "completed": "completed",
    "error": "error",
    "cancelled": "cancelled",
}

# Re-exported from the transport module, which owns the single implementation.
# Two copies of these numbers meant tuning one silently diverged from the other.
from lib_py.delegation.publisher import (  # noqa: F401  (re-exported for callers)
    CONNECT_ACK_TIMEOUT,
    PUBLISH_ACK_TIMEOUT,
    OneShotPublisher,
)


def build_payload(
    job_id: str,
    seq: int,
    event: str,
    detail: str,
    data: Optional[Dict[str, Any]],
    auth_token: Optional[str],
) -> Dict[str, Any]:
    payload: Dict[str, Any] = {
        "schema_version": SCHEMA_VERSION,
        "seq": seq,
        "job_id": job_id,
        "event": event,
        "timestamp": mqtt_common._utcnow(),
        "detail": detail,
        "data": dict(data) if data else {},
    }
    # Production: carry the per-job HMAC-SHA256 signature in `data.hmac_sig` so
    # the subscriber can verify the publisher without exposing the secret token.
    # The signature is calculated over the entire payload (with `data.hmac_sig` excluded).
    if auth_token:
        sign_payload = {k: v for k, v in payload.items() if k != "data"}
        sign_payload["data"] = {k: v for k, v in payload.get("data", {}).items() if k != "hmac_sig"}
        msg = json.dumps(sign_payload, sort_keys=True, separators=(",", ":")).encode()
        sig = hmac.new(auth_token.encode(), msg, hashlib.sha256).hexdigest()
        payload["data"]["hmac_sig"] = sig
    return payload


def _publish_once(config, topic: str, body: bytes, retain: bool) -> None:
    """Connect, publish one QoS-1 message, wait for the broker ACK, disconnect.

    Delegates to the transport module so the CLI path and the service path share
    one implementation. Keeping a second copy here is what let the persistent
    transport's handshake and not-connected bugs exist in only one of them.
    Raises on any failure so ``with_retry`` can re-run the whole sequence.
    """
    OneShotPublisher().publish_raw(config, topic, body, retain)


class PreparedEvent(object):
    """Everything decided before the wire: payload, topic, broker, retain flag.

    Splitting this out lets a different transport (one-shot vs persistent) be
    swapped in without touching the parts that must not vary: the seq is already
    consumed, the HMAC is already computed over the final payload, and the audit
    record is written from these exact fields.
    """

    __slots__ = ("job_id", "event", "detail", "seq", "topic", "retain",
                 "payload", "body", "config")

    def __init__(self, job_id, event, detail, seq, topic, retain, payload, body, config):
        self.job_id = job_id
        self.event = event
        self.detail = detail
        self.seq = seq
        self.topic = topic
        self.retain = retain
        self.payload = payload
        self.body = body
        self.config = config


def prepare_event(job_id, event, detail="", data=None,
                  registry_dir=DEFAULT_REGISTRY_DIR, retained=False) -> PreparedEvent:
    """Load the job, consume the next seq, and build the signed payload.

    Consuming the seq here (before any transport runs) is deliberate: a failed
    publish must still burn the number, so a replayed packet can never carry a
    seq the subscriber would accept.
    """
    job = load_job(job_id, registry_dir)
    config = broker_config_from_job(job)
    topic = job.get("topic_prefix")
    topic = f"{topic}/events" if topic else events_topic_for(job_id)
    seq = next_seq(job_id, registry_dir)
    payload = build_payload(
        job_id=job_id,
        seq=seq,
        event=event,
        detail=detail,
        data=data,
        auth_token=job.get("auth_token"),
    )
    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    retain = retained or event in TERMINAL_EVENTS
    return PreparedEvent(job_id, event, detail, seq, topic, retain, payload, body, config)


def record_outcome(prepared: PreparedEvent, registry_dir: str,
                   publish_ok: bool, publish_error: Optional[str]) -> None:
    """Persist the audit record, the registry event log and the status sync.

    All three are best effort and all three happen whether or not the publish
    succeeded — that is the contract G-1/G-2/G-4 pin down.
    """
    mqtt_common.append_event(prepared.job_id, {
        "event": "published",
        "source_event": prepared.event,
        "seq": prepared.seq,
        "topic": prepared.topic,
        "retain": prepared.retain,
        "timestamp": prepared.payload["timestamp"],
        "detail": prepared.detail,
        "payload": prepared.payload,
        "published": publish_ok,
        "publish_error": publish_error,
    })

    registry.append_event(prepared.job_id, registry_dir, prepared.payload)
    new_status = EVENT_TO_STATUS.get(prepared.event)
    if new_status:
        try:
            mqtt_common.update_job_status(prepared.job_id, registry_dir, status=new_status)
        except Exception as exc:  # pragma: no cover - best effort
            logger.warning("status sync failed: %s", exc)


def _resolve_job_id(args) -> Optional[str]:
    if args.pick_pending:
        return registry.pick_pending(args.agent_session, args.registry_dir)
    return args.job


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="Publish a Job event to MQTT")
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--job", help="job id to publish for")
    target.add_argument("--pick-pending", action="store_true",
                        help="auto-select a pending job for --agent-session")
    parser.add_argument("--agent-session", default="herdr:claude",
                        help="session label used with --pick-pending")
    parser.add_argument("--event", default="progress", choices=VALID_EVENTS)
    parser.add_argument("--detail", default="")
    parser.add_argument("--data", default=None, help="optional JSON object string")
    parser.add_argument("--retained", action="store_true",
                        help="force retain=True (auto for completed/error)")
    parser.add_argument("--registry-dir", default=DEFAULT_REGISTRY_DIR)
    parser.add_argument("--attempts", type=int, default=3)
    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)

    # --- parse optional data JSON (parameter error -> exit 1) ---
    data: Optional[Dict[str, Any]] = None
    if args.data:
        try:
            data = json.loads(args.data)
            if not isinstance(data, dict):
                raise ValueError("--data must be a JSON object")
        except (ValueError, json.JSONDecodeError) as exc:
            logger.error("invalid --data: %s", exc)
            return 1

    job_id = _resolve_job_id(args)
    if not job_id:
        logger.error("no job to publish for (unknown --job or no pending job)")
        return 1

    try:
        prepared = prepare_event(
            job_id=job_id,
            event=args.event,
            detail=args.detail,
            data=data,
            registry_dir=args.registry_dir,
            retained=args.retained,
        )
    except FileNotFoundError as exc:
        logger.error("%s", exc)
        return 1

    # Policy Note: seq consumption on failure is intentional (see prepare_event)
    # and the failure is documented explicitly in the audit record.
    publish = with_retry(
        _publish_once,
        attempts=args.attempts,
        exceptions=(OSError, TimeoutError, ConnectionError, ValueError),
    )
    publish_ok = True
    publish_error: Optional[str] = None
    try:
        publish(prepared.config, prepared.topic, prepared.body, prepared.retain)
    except Exception as exc:
        publish_ok = False
        publish_error = str(exc)
        logger.error("publish failed after %d attempts: %s", args.attempts, exc)

    record_outcome(prepared, args.registry_dir, publish_ok, publish_error)

    if not publish_ok:
        return 2

    logger.info("published %s seq=%d job=%s retain=%s",
                args.event, prepared.seq, job_id, prepared.retain)
    return 0


if __name__ == "__main__":
    sys.exit(main())