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 Optional, Tuple

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


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):
        client = self.client_factory("publisher", prepared.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(prepared.config.host, prepared.config.port, prepared.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"]))
            info = client.publish(prepared.topic, payload=prepared.body,
                                  qos=1, retain=prepared.retain)
            info.wait_for_publish(timeout=PUBLISH_ACK_TIMEOUT)
            if not info.is_published():
                raise TimeoutError("publish not acknowledged within timeout")
        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

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

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

    def open(self):
        return self

    def _ensure_client(self, config):
        if self._client is not None:
            if (config.host, config.port) != (self._config.host, self._config.port):
                raise ValueError(
                    "PersistentPublisher is bound to {}:{}; refusing an event for {}:{}".format(
                        self._config.host, self._config.port, config.host, config.port)
                )
            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
        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)
        info = client.publish(prepared.topic, payload=prepared.body,
                              qos=1, retain=prepared.retain)
        info.wait_for_publish(timeout=PUBLISH_ACK_TIMEOUT)
        if not info.is_published():
            # Surfaced to the caller's with_retry; paho reconnects underneath.
            raise TimeoutError("publish not acknowledged within timeout")

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