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

"""Publish transports: one connection per event, or one connection reused.

Phase 1 ships the persistent transport as a prototype only. Nothing in
production uses it: the CLI and the shell helpers exit after a single publish,
so reusing a connection buys them nothing. The consumer is the phase 2a engine,
which is long-running and pays the CONNACK cost on every event today.

Both transports raise on failure so the caller's `with_retry` policy stays the
single place that decides how many attempts an event gets.
"""
from __future__ import annotations

import logging
import time
from typing import Any, Optional, Tuple

import paho.mqtt.client as mqtt

from lib_py.delegation import mqtt_common

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

CONNECT_ACK_TIMEOUT = 10  # seconds to wait for CONNACK
PUBLISH_ACK_TIMEOUT = 5   # seconds to wait for QoS-1 PUBACK


def _send(client, topic: str, body: bytes, retain: bool) -> None:
    """Publish one QoS-1 message and wait for the broker ACK.

    paho queues a publish issued while disconnected and returns
    ``MQTT_ERR_NO_CONN``; ``wait_for_publish`` then raises ``RuntimeError``,
    which is deliberately NOT in the retry tuple (a RuntimeError normally means a
    programming error and must not be retried blindly). Left alone, that turned
    the one state a persistent connection exists to survive -- the gap between a
    drop and paho's reconnect -- into a single-attempt failure. Detect that one
    condition explicitly and raise a retryable error instead.
    """
    info = client.publish(topic, payload=body, qos=1, retain=retain)
    if getattr(info, "rc", mqtt.MQTT_ERR_SUCCESS) == mqtt.MQTT_ERR_NO_CONN:
        raise ConnectionError("client is not connected (paho is reconnecting)")
    info.wait_for_publish(timeout=PUBLISH_ACK_TIMEOUT)
    if not info.is_published():
        raise TimeoutError("publish not acknowledged within timeout")


class _BasePublisher(object):
    """Shared retry wrapper so both transports report outcomes identically."""

    def publish(self, prepared):  # pragma: no cover - interface
        raise NotImplementedError

    def publish_with_retry(self, prepared, attempts=3) -> Tuple[bool, Optional[str]]:
        runner = mqtt_common.with_retry(
            self.publish,
            attempts=attempts,
            exceptions=(OSError, TimeoutError, ConnectionError, ValueError),
        )
        try:
            runner(prepared)
        except Exception as exc:
            logger.error("publish failed after %d attempts: %s", attempts, exc)
            return False, str(exc)
        return True, None

    def open(self):
        return self

    def close(self):
        return None


class OneShotPublisher(_BasePublisher):
    """Connect, publish, wait for the ACK, disconnect. The current behaviour."""

    def __init__(self, client_factory=None):
        self.client_factory = client_factory or mqtt_common.make_client
        self.connect_count = 0

    def publish(self, prepared):
        self.publish_raw(prepared.config, prepared.topic, prepared.body, prepared.retain)

    def publish_raw(self, config, topic: str, body: bytes, retain: bool, extra=()) -> None:
        """Connect, publish, wait for the ACK, disconnect.

        The single implementation of one-shot publishing: ``publish_event``'s
        CLI path calls straight into this, so the two entry points cannot drift.
        ``extra`` messages (``(topic, body, retain)`` tuples, Issue #18's retained
        job state) go out over the same connection, each awaited like the first.
        """
        client = self.client_factory("publisher", config)
        self.connect_count += 1
        connected = {"rc": None}

        def on_connect(_c, _u, _flags, reason_code, _props):
            connected["rc"] = reason_code

        client.on_connect = on_connect
        client.connect(config.host, config.port, config.keepalive)
        client.loop_start()
        try:
            deadline = time.monotonic() + CONNECT_ACK_TIMEOUT
            while connected["rc"] is None and time.monotonic() < deadline:
                time.sleep(0.05)
            if connected["rc"] is None:
                raise TimeoutError("no CONNACK from broker")
            if mqtt_common.reason_code_value(connected["rc"]) != 0:
                raise ConnectionError("broker refused connection: rc={}".format(connected["rc"]))
            _send(client, topic, body, retain)
            for x_topic, x_body, x_retain in (extra or ()):
                _send(client, x_topic, x_body, x_retain)
        finally:
            client.loop_stop()
            try:
                client.disconnect()
            except Exception:  # pragma: no cover - disconnect best effort
                pass


class PersistentPublisher(_BasePublisher):
    """One connection reused across events; paho handles reconnect underneath.

    Bound to a single broker: the connection is established from the first
    event's config, so an event for a job on a different broker is refused
    rather than silently published to the wrong place.
    """

    def __init__(self, client_factory=None):
        self.client_factory = client_factory or mqtt_common.make_client
        self.connect_count = 0
        self._client = None
        self._config = None
        self._fingerprint_bound = None

    def __enter__(self):
        return self.open()

    def __exit__(self, exc_type, exc, tb):
        self.close()
        return False

    def open(self):
        return self

    @staticmethod
    def _fingerprint(config) -> tuple:
        """An immutable snapshot of everything that decides *which connection*.

        Endpoint alone is not enough: two jobs can share a host and port while
        carrying different credentials or TLS material, and reusing the first
        connection would publish under the wrong principal or silently drop a
        requested TLS setting. Snapshotting also means later mutation of the
        original config object cannot change what the guard believes is bound.
        """
        fields = ("host", "port", "username", "password", "tls", "ca_certs",
                  "certfile", "keyfile", "keepalive", "client_id_prefix")
        return tuple((name, getattr(config, name, None)) for name in fields)

    @staticmethod
    def _describe_mismatch(bound: tuple, incoming: tuple) -> str:
        """Name the differing fields without ever echoing their values."""
        differing = [name for (name, old), (_n, new) in zip(bound, incoming) if old != new]
        return ", ".join(differing) or "connection settings"

    def _ensure_client(self, config):
        if self._client is not None:
            incoming = self._fingerprint(config)
            if incoming != self._fingerprint_bound:
                raise ValueError(
                    "PersistentPublisher is bound to {}:{}; refusing an event whose "
                    "{} differ".format(
                        self._config.host, self._config.port,
                        self._describe_mismatch(self._fingerprint_bound, incoming))
                )
            return self._client

        client = self.client_factory("publisher", config)
        self.connect_count += 1

        # The handshake result belongs to THIS client, so it lives in a local the
        # callback closes over. Holding it on self would let a refused attempt's
        # reason code be read by the next attempt: the retry would either reject a
        # healthy connection or, worse, accept one whose CONNACK never arrived.
        connected = {"rc": None}

        def on_connect(_c, _u, _flags, reason_code, _props):
            connected["rc"] = reason_code

        client.on_connect = on_connect
        client.connect(config.host, config.port, config.keepalive)
        client.loop_start()
        reconnect_delay_set = getattr(client, "reconnect_delay_set", None)
        if reconnect_delay_set is not None:
            reconnect_delay_set(min_delay=1, max_delay=16)
        try:
            deadline = time.monotonic() + CONNECT_ACK_TIMEOUT
            while connected["rc"] is None and time.monotonic() < deadline:
                time.sleep(0.05)
            if connected["rc"] is None:
                raise TimeoutError("no CONNACK from broker")
            if mqtt_common.reason_code_value(connected["rc"]) != 0:
                raise ConnectionError("broker refused connection: rc={}".format(connected["rc"]))
        except Exception:
            self._discard(client)   # never leak the half-open client on the way out
            raise
        self._client = client
        self._config = config
        self._fingerprint_bound = self._fingerprint(config)
        return client

    @staticmethod
    def _discard(client):
        try:
            client.loop_stop()
            client.disconnect()
        except Exception:  # pragma: no cover - best effort
            pass

    def publish(self, prepared):
        client = self._ensure_client(prepared.config)
        # Both failure shapes (queued-while-disconnected and no PUBACK) are raised
        # as retryable errors by _send, so the caller's with_retry can ride out a
        # reconnect instead of recording a one-attempt failure.
        _send(client, prepared.topic, prepared.body, prepared.retain)

    def close(self):
        if self._client is None:
            return
        try:
            self._discard(self._client)
        finally:
            self._client = None
            self._config = None
            self._fingerprint_bound = None