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,
                  dispatch_token=None, attempt_id=None) -> 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_data = dict(data) if data else {}
    eff_dispatch_token = dispatch_token
    if eff_dispatch_token is not None and "dispatch_token" not in payload_data:
        payload_data["dispatch_token"] = eff_dispatch_token
    eff_attempt_id = attempt_id
    if eff_attempt_id is not None and "attempt_id" not in payload_data:
        payload_data["attempt_id"] = eff_attempt_id

    payload = build_payload(
        job_id=job_id,
        seq=seq,
        event=event,
        detail=detail,
        data=payload_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],
                   expected_dispatch_token: Optional[str] = None,
                   expected_attempt: Optional[int] = None) -> 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:
            exp_token = expected_dispatch_token
            if exp_token is None and isinstance(prepared.payload.get("data"), dict):
                exp_token = prepared.payload["data"].get("dispatch_token")
            exp_att = expected_attempt
            if exp_att is None and isinstance(prepared.payload.get("data"), dict):
                exp_att = prepared.payload["data"].get("attempt_id")

            # Agent authority on purpose: an agent may never move a terminal
            # record, which is what the operator-only requests are gated for.
            mqtt_common.update_job_status(prepared.job_id, registry_dir,
                                          _source="publish_event", status=new_status,
                                          expected_dispatch_token=exp_token,
                                          expected_attempt=exp_att)
        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("--finalize-report", default=None, metavar="PATH",
                        help="atomically finalize this report and declare its "
                             "completion manifest before publishing")
    parser.add_argument("--side-effect", action="append", default=[], dest="side_effects",
                        metavar="KIND=REFERENCE",
                        help="evidence for an external side effect this job performed")
    parser.add_argument("--attempt-id", type=int, default=None,
                        help="attempt this finalization belongs to (default: the "
                             "job's current iteration)")
    parser.add_argument("--expect-dispatch-token", default=None,
                        help="expected dispatch token for CAS verification")
    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

    # Finalize before publishing: a `completed` event whose report is still being
    # written would invite exactly the partial-report adoption the contract
    # forbids. Failing here fails the publish, on purpose.
    if args.finalize_report:
        from lib_py.delegation import finalize  # same package: no sibling hop (see __init__)

        try:
            side_effects = finalize.parse_side_effects(args.side_effects)
        except ValueError as exc:
            logger.error("%s", exc)
            return 1
        attempt_id = args.attempt_id
        if attempt_id is None:
            try:
                rec = load_job(job_id, args.registry_dir)
                attempt_id = rec.get("attempt_id", rec.get("iteration", 1))
            except FileNotFoundError:
                attempt_id = 1
        try:
            manifest = finalize.finalize_report(
                args.finalize_report, job_id, attempt_id,
                side_effects=side_effects, declared_by="agent")
        except (OSError, FileNotFoundError) as exc:
            logger.error("could not finalize report: %s", exc)
            return 1
        logger.info("finalized %s (sha256 %s...)", manifest["report_path"],
                    manifest["sha256"][:12])

    check_token = args.expect_dispatch_token
    if check_token is None and isinstance(data, dict) and "dispatch_token" in data:
        check_token = data.get("dispatch_token")
    try:
        rec = load_job(job_id, args.registry_dir)
        actual_tok = rec.get("dispatch_token")
        rec_attempt = rec.get("attempt_id", rec.get("iteration", 1))
    except FileNotFoundError:
        rec = None
        actual_tok = None
        rec_attempt = 1

    check_attempt = args.attempt_id
    if check_attempt is None and isinstance(data, dict) and "attempt_id" in data:
        check_attempt = data.get("attempt_id")

    if rec is not None and actual_tok is not None:
        if check_token is None and check_attempt is None:
            if args.event in ("completed", "error", "cancelled") and rec_attempt > 1:
                logger.error("dispatch token required: redispatched job %s carries a dispatch token "
                             "but none was provided (aborting terminal publish)", job_id)
                return 4
        elif check_token is not None and actual_tok != check_token:
            logger.error("dispatch token mismatch: expected %s, got %s (aborting publish)",
                         check_token, actual_tok)
            return 4
        elif check_attempt is not None and rec_attempt != check_attempt:
            logger.error("attempt mismatch: expected %s, got %s (aborting publish)",
                         check_attempt, rec_attempt)
            return 4
    elif check_token is not None and actual_tok != check_token:
        logger.error("dispatch token mismatch: expected %s, got %s (aborting publish)",
                     check_token, actual_tok)
        return 4

    try:
        prepared = prepare_event(
            job_id=job_id,
            event=args.event,
            detail=args.detail,
            data=data,
            registry_dir=args.registry_dir,
            retained=args.retained,
            dispatch_token=args.expect_dispatch_token,
            attempt_id=args.attempt_id,
        )
    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)

    sync_expected_token = args.expect_dispatch_token
    if sync_expected_token is None and isinstance(data, dict):
        sync_expected_token = data.get("dispatch_token")
    sync_expected_attempt = args.attempt_id
    if sync_expected_attempt is None and isinstance(data, dict):
        sync_expected_attempt = data.get("attempt_id")

    record_outcome(prepared, args.registry_dir, publish_ok, publish_error,
                   expected_dispatch_token=sync_expected_token,
                   expected_attempt=sync_expected_attempt)

    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())