From a7553899c5ac97fe5e8b5890a2ed1b3ffd0beae4 Mon Sep 17 00:00:00 2001 From: Tewan Date: Fri, 18 Sep 2026 15:54:11 +0800 Subject: [PATCH 01/14] Add MQTT Reader enterprise workflow block Subscribes to an MQTT topic and exposes one received message per workflow run as a raw string and a parsed JSON payload. `read_mode` selects the newest message (`latest`) or the next unread one in arrival order (`sequential`); when nothing new arrived the previous outputs are repeated with `is_new=False`. Messages are buffered in the block instance between paho's network thread and the engine thread, the subscription is re-established on every reconnect, a retained copy re-sent after a reconnect is not reported as new, and connection parameters are fixed on the first run as in the MQTT Writer. The block lives in the `roboflow-workflows` package next to the MQTT Writer, and declares a HARD restriction for hosted serverless only, with dedicated deployments under the SOFT restriction, matching the Kafka consumer. The test fake broker now answers SUBSCRIBE, PINGREQ and the QoS handshakes, serves retained messages, can push PUBLISH packets to the client and re-accept a dropped connection, so subscriber blocks can be tested end to end without a real broker. Co-Authored-By: Claude Fable 5.1 --- .../integration_tests/execution/conftest.py | 221 ++++- .../test_workflow_with_mqtt_reader.py | 344 ++++++++ .../enterprise_blocks/loader.py | 2 + .../sinks/mqtt_reader/__init__.py | 0 .../enterprise_blocks/sinks/mqtt_reader/v1.py | 618 ++++++++++++++ .../core_steps/sinks/test_mqtt_reader.py | 791 ++++++++++++++++++ 6 files changed, 1963 insertions(+), 13 deletions(-) create mode 100644 tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py create mode 100644 workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/__init__.py create mode 100644 workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py create mode 100644 workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index 1b06e12d94..3b47810a04 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -2,6 +2,7 @@ import os.path import socket import tempfile +import threading from typing import Generator import cv2 @@ -148,8 +149,77 @@ def face_image() -> np.ndarray: # Below taken from https://github.com/eclipse-paho/paho.mqtt.python/blob/d45de3737879cfe7a6acc361631fa5cb1ef584bb/tests/testsupport/broker.py +# and extended with SUBSCRIBE/SUBACK, PINGREQ/PINGRESP, retained messages and +# broker-initiated PUBLISH so subscriber blocks can be tested as well. +MQTT_CONNECT = 0x10 +MQTT_PUBLISH = 0x30 +MQTT_PUBACK = 0x40 +MQTT_PUBREC = 0x50 +MQTT_PUBREL = 0x60 +MQTT_PUBCOMP = 0x70 +MQTT_SUBSCRIBE = 0x80 +MQTT_SUBACK = 0x90 +MQTT_PINGREQ = 0xC0 +MQTT_PINGRESP = 0xD0 +MQTT_DISCONNECT = 0xE0 +MQTT_SUBACK_FAILURE = 0x80 + + +def _encode_remaining_length(length: int) -> bytes: + encoded = b"" + while True: + byte, length = length % 128, length // 128 + if length > 0: + byte |= 0x80 + encoded += bytes([byte]) + if length == 0: + return encoded + + +def _split_mqtt_packets(buffer: bytes): + """Split a byte buffer into complete MQTT control packets; returns (packets, rest).""" + packets = [] + while buffer: + length, multiplier, index = 0, 1, 1 + while True: + if index >= len(buffer): + return packets, buffer + byte = buffer[index] + length += (byte & 0x7F) * multiplier + multiplier *= 128 + index += 1 + if not byte & 0x80: + break + end = index + length + if end > len(buffer): + return packets, buffer + packets.append(buffer[:end]) + buffer = buffer[end:] + return packets, buffer + + +def mqtt_topic_matches(topic_filter: str, topic: str) -> bool: + """MQTT topic-filter matching with `+` (one level) and `#` (rest) wildcards.""" + filter_levels = topic_filter.split("/") + topic_levels = topic.split("/") + for index, level in enumerate(filter_levels): + if level == "#": + return True + if index >= len(topic_levels): + return False + if level != "+" and level != topic_levels[index]: + return False + return len(filter_levels) == len(topic_levels) + + class FakeMQTTBroker: - def __init__(self, connack_reason_code: int = 0, listening: bool = True): + def __init__( + self, + connack_reason_code: int = 0, + listening: bool = True, + suback_reason_code: int = 0, + keep_serving: bool = False, + ): # Bind to "localhost" for maximum performance, as described in: # http://docs.python.org/howto/sockets.html#ipc sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) @@ -160,6 +230,16 @@ def __init__(self, connack_reason_code: int = 0, listening: bool = True): self.messages = [] self.messages_count_to_wait_for = 2 self.connack_reason_code = connack_reason_code + # granted QoS returned in SUBACK; 0x80 means the subscription is refused + self.suback_reason_code = suback_reason_code + # keep_serving: ignore recv timeouts and the message count, serve until + # finish() / drop_connection(); needed for subscriber tests + self.keep_serving = keep_serving + self.subscriptions = [] + self.retained = {} + self.connections_accepted = 0 + self._stop = False + self._send_lock = threading.Lock() sock.settimeout(5) @@ -181,20 +261,125 @@ def start(self): conn, address = self._sock.accept() conn.settimeout(1) self._conn = conn + self.connections_accepted += 1 + self._stop = False connack_sent = False - while len(self.messages) < self.messages_count_to_wait_for or not connack_sent: - packet = self.receive_packet(1000) - print(f"Received {packet}") - if not packet: + buffer = b"" + while not self._stop and ( + self.keep_serving + or len(self.messages) < self.messages_count_to_wait_for + or not connack_sent + ): + try: + chunk = self.receive_packet(1000) + except socket.timeout: + if self.keep_serving: + continue + raise + except (OSError, ValueError): + # connection closed, possibly from the test thread + break + print(f"Received {chunk}") + if not chunk: + break + buffer += chunk + packets, buffer = _split_mqtt_packets(buffer) + try: + for packet in packets: + if self._handle_packet(packet): + connack_sent = True + except (OSError, ValueError): break - if packet.startswith(b"\x10"): - print("sending CONNACK") - self._conn.send(b"\x20\x02\x00" + bytes([self.connack_reason_code])) - connack_sent = True - continue - self.messages.append(packet) + + def _handle_packet(self, packet: bytes) -> bool: + """Handle one control packet; returns True when a CONNACK was sent.""" + packet_type = packet[0] & 0xF0 + if packet_type == MQTT_CONNECT: + print("sending CONNACK") + self._send(b"\x20\x02\x00" + bytes([self.connack_reason_code])) + return True + if packet_type == MQTT_SUBSCRIBE: + self._handle_subscribe(packet) + return False + if packet_type == MQTT_PINGREQ: + self._send(bytes([MQTT_PINGRESP, 0])) + return False + if packet_type in (MQTT_PUBACK, MQTT_PUBREC, MQTT_PUBREL, MQTT_PUBCOMP): + # QoS 1/2 handshakes, not messages: a broker-initiated QoS 2 PUBLISH + # is answered with PUBREC and completed with PUBREL; a client's + # PUBREL (its own QoS 2 publish) is completed with PUBCOMP + if packet_type == MQTT_PUBREC: + self._send(bytes([MQTT_PUBREL | 0x02, 2]) + packet[2:4]) + if packet_type == MQTT_PUBREL: + self._send(bytes([MQTT_PUBCOMP, 2]) + packet[2:4]) + return False + if packet_type == MQTT_DISCONNECT: + return False + self.messages.append(packet) + return False + + def _handle_subscribe(self, packet: bytes) -> None: + # skip fixed header (type byte + remaining-length varint) + index = 1 + while packet[index] & 0x80: + index += 1 + index += 1 + packet_id = packet[index : index + 2] + index += 2 + granted = [] + while index < len(packet): + topic_length = int.from_bytes(packet[index : index + 2], "big") + index += 2 + topic_filter = packet[index : index + topic_length].decode("utf-8") + index += topic_length + requested_qos = packet[index] + index += 1 + self.subscriptions.append((topic_filter, requested_qos)) + if self.suback_reason_code == MQTT_SUBACK_FAILURE: + granted.append(MQTT_SUBACK_FAILURE) + else: + granted.append(min(requested_qos, self.suback_reason_code)) + suback = ( + bytes([MQTT_SUBACK]) + + _encode_remaining_length(2 + len(granted)) + + packet_id + + bytes(granted) + ) + self._send(suback) + if self.suback_reason_code == MQTT_SUBACK_FAILURE: + return + for topic_filter, _ in self.subscriptions[-len(granted) :]: + for topic, payload in self.retained.items(): + if mqtt_topic_matches(topic_filter, topic): + self.publish(topic, payload, retain=True) + + def publish( + self, topic: str, payload: bytes, retain: bool = False, qos: int = 0 + ) -> None: + """Send a PUBLISH packet from the broker to the connected client.""" + if isinstance(payload, str): + payload = payload.encode("utf-8") + topic_bytes = topic.encode("utf-8") + variable_header = len(topic_bytes).to_bytes(2, "big") + topic_bytes + if qos > 0: + variable_header += (1).to_bytes(2, "big") + body = variable_header + payload + header = bytes([MQTT_PUBLISH | (qos << 1) | int(retain)]) + self._send(header + _encode_remaining_length(len(body)) + body) + + def drop_connection(self) -> None: + """Close the client connection so a later start() can accept a new one.""" + self._stop = True + if self._conn is not None: + try: + self._conn.shutdown(socket.SHUT_RDWR) + except OSError: + pass + self._conn.close() + self._conn = None def finish(self): + self._stop = True if self._conn is not None: self._conn.close() self._conn = None @@ -204,12 +389,22 @@ def finish(self): self._sock = None def receive_packet(self, num_bytes): - if self._conn is None: + # capture locally: drop_connection()/finish() clear the attribute + # from another thread + conn = self._conn + if conn is None: raise ValueError("Connection is not open") - packet_in = self._conn.recv(num_bytes) + packet_in = conn.recv(num_bytes) return packet_in + def _send(self, data: bytes) -> None: + with self._send_lock: + conn = self._conn + if conn is None: + raise ValueError("Connection is not open") + conn.sendall(data) + @pytest.fixture(scope="function") def fake_mqtt_broker(): diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py new file mode 100644 index 0000000000..bce1bbf2b9 --- /dev/null +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -0,0 +1,344 @@ +import os +import threading +import time +from typing import Callable + +import pytest +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import MQTTReaderBlockV1 + +from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.conftest import FakeMQTTBroker + +RUNNING = b'{"state": "RUNNING"}' +PAUSED = b'{"state": "PAUSED"}' + + +def wait_until(predicate: Callable[[], bool], timeout: float = 5.0) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if predicate(): + return True + time.sleep(0.05) + return False + + +@pytest.fixture +def broker(): + broker = FakeMQTTBroker(keep_serving=True) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +def reader_kwargs(broker: FakeMQTTBroker, **overrides) -> dict: + kwargs = { + "host": broker.host, + "port": broker.port, + "topic": "plc/#", + "timeout": 2.0, + } + kwargs.update(overrides) + return kwargs + + +@pytest.mark.timeout(15) +def test_retained_message_returned_on_first_run(broker): + # given + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + + try: + # when + first = block.run(**reader_kwargs(broker)) + second = block.run(**reader_kwargs(broker)) + + # then + assert first == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": True, + "error_status": False, + "error_message": None, + } + assert second == {**first, "is_new": False} + assert broker.subscriptions == [("plc/#", 0)] + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_empty_first_run_then_latest_message_delivered(broker): + # given + block = MQTTReaderBlockV1() + + try: + # when + empty = block.run(**reader_kwargs(broker)) + for payload in (b"A", b"B", PAUSED): + broker.publish("plc/state", payload) + # the one-slot buffer holds a message as soon as A lands; wait for the + # newest one to have replaced it + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + result = block.run(**reader_kwargs(broker)) + + # then + assert empty["value"] is None + assert empty["is_new"] is False + assert empty["error_status"] is False + assert result["value"] == '{"state": "PAUSED"}' + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_sequential_mode_delivers_backlog_one_per_run(broker): + # given + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, read_mode="sequential") + + try: + # when + block.run(**kwargs) + broker.publish("plc/a", b"A") + broker.publish("plc/b", b"B") + broker.publish("plc/c", b"C") + assert wait_until(lambda: len(block._state.messages) == 3) + results = [block.run(**kwargs) for _ in range(4)] + + # then + assert [r["value"] for r in results] == ["A", "B", "C", "C"] + assert [r["topic"] for r in results] == ["plc/a", "plc/b", "plc/c", "plc/c"] + assert [r["is_new"] for r in results] == [True, True, True, False] + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_qos2_message_delivered_once(broker): + # given + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, qos=2, read_mode="sequential") + + try: + # when + block.run(**kwargs) + broker.publish("plc/state", RUNNING, qos=2) + assert wait_until(lambda: len(block._state.messages) == 1) + time.sleep(0.2) + first = block.run(**kwargs) + second = block.run(**kwargs) + + # then + assert broker.subscriptions == [("plc/#", 2)] + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + assert second["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_refused_subscription_reported_and_connection_kept(): + # given + broker = FakeMQTTBroker(keep_serving=True, suback_reason_code=0x80) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker)) + + # then + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._state.connected.is_set() + assert block._client is not None + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(20) +def test_subscription_restored_after_reconnect(broker): + # given - a retained state the broker re-sends on every subscription + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + + try: + first = block.run(**reader_kwargs(broker)) + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + + # when - the broker drops the connection and accepts a new one + broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + second_thread = threading.Thread(target=broker.start, daemon=True) + second_thread.start() + assert wait_until(lambda: len(broker.subscriptions) == 2, timeout=10) + assert wait_until(lambda: len(block._state.messages) == 1) + redelivered = block.run(**reader_kwargs(broker)) + broker.publish("plc/state", PAUSED) + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + updated = block.run(**reader_kwargs(broker)) + + # then - the subscription is live again, the retained copy is not new + assert broker.connections_accepted == 2 + assert redelivered == {**first, "is_new": False} + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_unreachable_broker_leaves_nothing_behind_and_later_run_recovers(): + # given - bound but not listening, so the TCP connect is refused + broker = FakeMQTTBroker(listening=False, keep_serving=True) + block = MQTTReaderBlockV1() + thread = None + + try: + # when + first = block.run(**reader_kwargs(broker, timeout=0.3)) + assert block._client is None + + broker.listen() + broker.retained["plc/state"] = RUNNING + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + second = block.run(**reader_kwargs(broker)) + + # then + assert first["error_status"] is True + assert second["error_status"] is False + assert second["payload"] == {"state": "RUNNING"} + finally: + block.close() + broker.finish() + if thread is not None: + thread.join(timeout=2) + + +GATED_WORKFLOW = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowParameter", "name": "host"}, + {"type": "WorkflowParameter", "name": "port"}, + ], + "steps": [ + { + "type": "roboflow_enterprise/mqtt_reader@v1", + "name": "reader", + "host": "$inputs.host", + "port": "$inputs.port", + "topic": "plc/state", + "timeout": 2.0, + }, + { + "type": "roboflow_core/continue_if@v1", + "name": "gate", + "condition_statement": { + "type": "StatementGroup", + "statements": [ + { + "type": "UnaryStatement", + "operand": { + "type": "DynamicOperand", + "operand_name": "raw", + "operations": [ + { + "type": "StringMatches", + "regex": '"state"\\s*:\\s*"RUNNING"', + } + ], + }, + "operator": {"type": "(Boolean) is True"}, + } + ], + }, + "evaluation_parameters": {"raw": "$steps.reader.value"}, + "next_steps": ["$steps.shout"], + }, + { + "type": "roboflow_core/property_definition@v1", + "name": "shout", + "data": "$steps.reader.value", + "operations": [{"type": "StringToUpperCase"}], + }, + ], + "outputs": [ + {"type": "JsonField", "name": "state", "selector": "$steps.reader.payload"}, + {"type": "JsonField", "name": "is_new", "selector": "$steps.reader.is_new"}, + {"type": "JsonField", "name": "shouted", "selector": "$steps.shout.output"}, + ], +} + + +@pytest.fixture +def enterprise_blocks_enabled(monkeypatch): + from inference.core.env import ENTERPRISE_BLOCKS_PLUGIN + from inference.core.workflows.execution_engine.introspection import blocks_loader + from inference.core.workflows.execution_engine.v1.compiler.core import ( + COMPILATION_CACHE, + ) + + plugins = [p for p in os.getenv("WORKFLOWS_PLUGINS", "").split(",") if p] + if ENTERPRISE_BLOCKS_PLUGIN not in plugins: # env.py may have expanded it already; + plugins.append(ENTERPRISE_BLOCKS_PLUGIN) # a duplicate trips the clash check + monkeypatch.setenv("WORKFLOWS_PLUGINS", ",".join(plugins)) + blocks_loader.clear_caches() + yield + blocks_loader.clear_caches() + # drop graphs compiled with enterprise blocks from the process-wide cache + with COMPILATION_CACHE._cache_lock: + COMPILATION_CACHE._cache.clear() + COMPILATION_CACHE._keys_buffer.clear() + + +@pytest.mark.timeout(20) +def test_workflow_gates_downstream_step_on_received_state( + broker, enterprise_blocks_enabled +): + # given - a retained RUNNING state, then a PAUSED update between frames + broker.retained["plc/state"] = RUNNING + execution_engine = ExecutionEngine.init( + workflow_definition=GATED_WORKFLOW, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + runtime_parameters = {"host": broker.host, "port": broker.port} + reader_block = execution_engine._engine._compiled_workflow.steps["reader"].step + + try: + # when + running_frame = execution_engine.run(runtime_parameters=runtime_parameters) + broker.publish("plc/state", PAUSED) + assert wait_until(lambda: len(reader_block._state.messages) == 1) + paused_frame = execution_engine.run(runtime_parameters=runtime_parameters) + repeated_frame = execution_engine.run(runtime_parameters=runtime_parameters) + + # then - one block instance served every frame + assert running_frame[0]["state"] == {"state": "RUNNING"} + assert running_frame[0]["is_new"] is True + assert running_frame[0]["shouted"] == '{"STATE": "RUNNING"}' + assert paused_frame[0]["state"] == {"state": "PAUSED"} + assert paused_frame[0]["is_new"] is True + assert paused_frame[0].get("shouted") is None + assert repeated_frame[0]["state"] == {"state": "PAUSED"} + assert repeated_frame[0]["is_new"] is False + assert broker.connections_accepted == 1 + finally: + reader_block.close() diff --git a/workflows/roboflow_workflows/enterprise_blocks/loader.py b/workflows/roboflow_workflows/enterprise_blocks/loader.py index 9981b8d31c..a9142f0a45 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/loader.py +++ b/workflows/roboflow_workflows/enterprise_blocks/loader.py @@ -12,6 +12,7 @@ from roboflow_workflows.enterprise_blocks.sinks.microsoft_sql_server.v1 import ( MicrosoftSQLServerSinkBlockV1, ) +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import MQTTReaderBlockV1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( MQTTWriterSinkBlockV1, ) @@ -43,6 +44,7 @@ def load_enterprise_blocks() -> List[Type[WorkflowBlock]]: return [ OPCWriterSinkBlockV1, MQTTWriterSinkBlockV1, + MQTTReaderBlockV1, PLCBlockV1, PLCReaderBlockV1, PLCWriterBlockV1, diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/__init__.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py new file mode 100644 index 0000000000..fa4293ebea --- /dev/null +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -0,0 +1,618 @@ +import json +import logging +import math +import threading +from collections import deque +from typing import Any, Deque, Dict, List, Literal, Optional, Tuple, Type, Union + +import paho.mqtt.client as mqtt +from pydantic import ConfigDict, Field +from typing_extensions import Annotated + +# Bind by literal name to the shared "inference" logger tree so server +# handlers / filters / propagation still apply when installed, without +# importing the server logger module. +logger = logging.getLogger("inference") +from roboflow_workflows.environment import GCP_SERVERLESS, LAMBDA +from roboflow_workflows.execution_engine.entities.base import OutputDefinition +from roboflow_workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + DICTIONARY_KIND, + FLOAT_KIND, + INTEGER_KIND, + SECRET_KIND, + STRING_KIND, + Selector, +) +from roboflow_workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeInputMode, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +MQTT Reader block for subscribing to an MQTT topic and exposing one received message +per workflow run as a raw string (`value`) and a parsed JSON object (`payload`). + +The first run connects synchronously: the TCP connect, the broker's session +acknowledgement and the subscription acknowledgement are each bounded by `timeout`, +and the first run then waits up to `timeout` for a retained message the broker +delivers on subscription, so a cold start may take up to four times `timeout` - +raise it for remote brokers. Afterwards a background network loop maintains the +connection, re-establishes the subscription after every reconnect and buffers +incoming messages between runs. + +`read_mode` selects what a run returns from that buffer: + +- `latest` (default): the newest message received since the previous run; anything + older is skipped. +- `sequential`: the next unread message, one per run, in arrival order. A backlog is + worked through one per run; when more than 1000 messages queue up between two runs + the oldest are dropped. + +When no new message arrived, the outputs of the previous run are repeated with +`is_new` set to False. Retained messages delivered by the broker on subscription are +treated like any other message, except that a retained copy the broker re-sends after +a reconnect is not reported as new when its topic and payload match the message +returned last. + +The subscription and the buffer live in the block instance only. An InferencePipeline +keeps them for its lifetime, a restart begins with an empty buffer, and over the HTTP +API every request builds a fresh instance, so only a retained message can be returned +there. There is no persistent session and no durable resume: messages published while +the block was not subscribed are not delivered. + +One block instance subscribes on a single broker connection: changing host, port, +credentials, timeout, topic, QoS or read mode between runs is rejected as a +configuration error. While the broker is unreachable every run waits up to `timeout` +for the background reconnect before reporting the failure. + +The block is not available on the Roboflow hosted platform (`GCP_SERVERLESS` or +`LAMBDA`): every run fails there without opening a connection. It is intended for +self-hosted inference servers and InferencePipelines that can reach the broker. + +Outputs: + - value (str): Raw message payload decoded as UTF-8, or None when nothing was + received yet. + - payload (dict): The payload parsed as a JSON object, or {} when the payload is + not a JSON object. + - topic (str): The topic the message arrived on (useful with wildcard filters). + - is_new (bool): True if this run surfaced a message the previous run did not. + - error_status (bool): True if the run failed (connection, subscription, invalid + configuration or an undecodable payload). + - error_message (str): Details of the failure, or None. + +Failures are returned in the outputs and logged; the workflow keeps running. +""" + +SUBSCRIPTION_REFUSED_QOS = 0x80 +LATEST_BUFFER_SIZE = 1 +SEQUENTIAL_BUFFER_SIZE = 1000 + +ReceivedMessage = Tuple[str, bytes, bool] + + +class MQTTReaderState: + """State shared between paho's network thread and the engine thread. + + Passed to paho as ``userdata`` so the callbacks never hold a reference to the + block itself (a bound-method callback would form a reference cycle and delay + the block's cleanup). + + Attributes: + connected: Set once the broker accepted the session; cleared on disconnect. + subscribed: Set once the broker granted the subscription; cleared on disconnect. + subscribe_failed: True when the broker refused the subscription. + topic: Topic filter to subscribe to on every (re)connect. + qos: Subscription QoS requested from the broker. + messages: Buffer of ``(topic, payload, retain)`` tuples appended by the + network thread and consumed by ``run()``; ``maxlen`` drops the oldest. + message_received: Set once the first message was buffered; lets the first + run wait for a retained message the broker sends right after SUBACK. + """ + + def __init__(self, topic: str, qos: int, buffer_size: int): + self.connected = threading.Event() + self.subscribed = threading.Event() + self.message_received = threading.Event() + self.subscribe_failed = False + self.topic = topic + self.qos = qos + self.messages: Deque[ReceivedMessage] = deque(maxlen=buffer_size) + + +def mqtt_on_connect( + client, state: MQTTReaderState, flags, reason_code, properties=None +): + # paho invokes on_connect for accepted and rejected CONNACK alike; + # only reason_code 0 means an established MQTT session + if reason_code != 0: + logger.error("MQTT connection refused, result code %s", reason_code) + state.connected.clear() + return + + logger.info("MQTT client connected") + state.connected.set() + # the broker forgets subscriptions on every reconnect (clean session), so + # the subscription is (re)established here rather than once after connect() + state.subscribe_failed = False + state.subscribed.clear() + try: + client.subscribe(state.topic, qos=state.qos) + except Exception as e: + # paho rejects e.g. an empty topic locally; report it through the + # refused path instead of leaving the run waiting for a SUBACK + logger.error("MQTT subscribe to topic %r failed: %s", state.topic, e) + state.subscribe_failed = True + + +def mqtt_on_connect_fail(client, state: MQTTReaderState): + # paho 1.6.1 invokes this callback with exactly (client, userdata) + logger.error("MQTT client failed to establish connection with broker") + state.connected.clear() + + +def mqtt_on_subscribe(client, state: MQTTReaderState, mid, granted_qos): + if any(qos == SUBSCRIPTION_REFUSED_QOS for qos in granted_qos): + logger.error("MQTT broker refused subscription to topic %s", state.topic) + state.subscribe_failed = True + state.subscribed.clear() + return + + logger.info("MQTT client subscribed to topic %s", state.topic) + state.subscribe_failed = False + state.subscribed.set() + + +def mqtt_on_message(client, state: MQTTReaderState, message): + # the only work done on the network thread: stash the raw message; + # decoding and parsing happen in run(), where a failure is reported + # instead of ending the network loop + state.messages.append((message.topic, message.payload, bool(message.retain))) + state.message_received.set() + + +def mqtt_on_disconnect(client, state: MQTTReaderState, reason_code, properties=None): + logger.info("MQTT client disconnected, result code %s", reason_code) + state.connected.clear() + state.subscribed.clear() + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "MQTT Reader", + "version": "v1", + "short_description": "Reads a message from an MQTT topic on every run: the newest one, or the next unread one.", + "long_description": LONG_DESCRIPTION, + "license": "Roboflow Enterprise License", + "block_type": "transformation", + "ui_manifest": { + "section": "industrial", + "icon": "fal fa-network-wired", + "blockPriority": 11, + "enterprise_only": True, + "local_only": True, + }, + } + ) + type: Literal["roboflow_enterprise/mqtt_reader@v1"] + host: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Host of the MQTT broker.", + examples=["localhost", "$inputs.mqtt_host"], + ) + port: Union[ + Annotated[int, Field(ge=1, le=65535)], + Selector(kind=[INTEGER_KIND]), + ] = Field( + description="Port of the MQTT broker (1-65535).", + examples=[1883, "$inputs.mqtt_port"], + ) + topic: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="MQTT topic filter to subscribe to. Single-level (`+`) and " + "multi-level (`#`) wildcards are allowed; the `topic` output tells which " + "topic each message arrived on.", + examples=["plc/state", "sensors/+/temperature", "$inputs.mqtt_topic"], + ) + read_mode: Literal["latest", "sequential"] = Field( + default="latest", + description="What a run returns from the messages received since the previous " + "run: `latest` returns the newest and skips older ones; `sequential` returns " + "the next unread message, one per run, in arrival order.", + examples=["latest", "sequential"], + json_schema_extra={ + "values_metadata": { + "latest": { + "name": "Latest", + "description": "Return the newest message; skip anything older.", + }, + "sequential": { + "name": "Sequential", + "description": "Return the next unread message, one per run, in arrival order.", + }, + }, + }, + ) + qos: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=0, + description="Quality of Service level requested for the subscription (0, 1 or 2).", + examples=[0, 1, 2], + ) + username: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Username for MQTT broker authentication.", + examples=["$inputs.mqtt_username"], + ) + password: Optional[Union[Selector(kind=[SECRET_KIND]), str]] = Field( + default=None, + description="Password for MQTT broker authentication. Requires username.", + examples=["$inputs.mqtt_password"], + json_schema_extra={"private": True}, + ) + timeout: Union[ + Annotated[float, Field(gt=0, allow_inf_nan=False)], + Selector(kind=[FLOAT_KIND]), + ] = Field( + default=0.5, + description="Timeout in seconds for connecting to the MQTT broker and for the " + "broker's connection and subscription acknowledgements. Must be a finite " + "number greater than 0.", + examples=[0.5], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="value", kind=[STRING_KIND]), + OutputDefinition(name="payload", kind=[DICTIONARY_KIND]), + OutputDefinition(name="topic", kind=[STRING_KIND]), + OutputDefinition(name="is_new", kind=[BOOLEAN_KIND]), + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="error_message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.4.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + RuntimeRestriction( + severity=Severity.HARD, + note=( + "On the Roboflow hosted platform every run returns " + "`error_status=true` with no message, and no MQTT connection is " + "opened from the hosted platform." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + ), + RuntimeRestriction( + severity=Severity.SOFT, + note=( + "The subscription and message buffer live in this block " + "instance. Over HTTP every request builds a fresh instance, so " + "only a retained message can be returned. Use an " + "InferencePipeline for a live subscription." + ), + applies_to_runtimes=[ + Runtime.SELF_HOSTED_CPU, + Runtime.SELF_HOSTED_GPU, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_input_modes=[RuntimeInputMode.IMAGE], + ), + ] + + +class MQTTReaderBlockV1(WorkflowBlock): + def __init__(self): + self._client: Optional[mqtt.Client] = None + self._state: Optional[MQTTReaderState] = None + self._connection_identity: Optional[Tuple] = None + # outputs of the last surfaced message, repeated with is_new=False + self._last: Optional[Dict[str, Any]] = None + # the first read after subscribing waits for a retained message, which + # the broker sends right after SUBACK as a separate packet + self._awaiting_retained = False + self._lifecycle_lock = threading.Lock() + + def close(self) -> None: + with self._lifecycle_lock: + client = self._client + state = self._state + if client is None: + return + self._client = None + self._state = None + self._connection_identity = None + self._last = None + self._awaiting_retained = False + try: + client.disconnect() + except Exception as e: + logger.error("Failed to disconnect MQTT client: %s", e) + finally: + try: + # loop_stop() joins the network thread without a timeout; + # accepted so the thread never outlives the block + client.loop_stop() + finally: + # only after the join can no callback re-set the events + if state is not None: + state.connected.clear() + state.subscribed.clear() + + def __del__(self): + try: + self.close() + except Exception as e: + logger.error("Failed to close MQTT client: %s", e) + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + @classmethod + def get_init_parameters(cls) -> List[str]: + return [] + + def run( + self, + host: str, + port: int, + topic: str, + read_mode: str = "latest", + qos: int = 0, + username: Optional[str] = None, + password: Optional[str] = None, + timeout: float = 0.5, + ) -> BlockResult: + if GCP_SERVERLESS or LAMBDA: + # never open an outbound broker connection from hosted workers + return self._handle_failure( + "MQTT Reader is not available on the Roboflow hosted platform" + ) + # selector-resolved values bypass manifest constraints, so validate here + try: + timeout_seconds = float(timeout) + except (TypeError, ValueError, OverflowError): + timeout_seconds = math.nan + if ( + not math.isfinite(timeout_seconds) + or not 0 < timeout_seconds <= threading.TIMEOUT_MAX + ): + return self._handle_failure( + f"Invalid timeout: {timeout!r}. Timeout must be a positive finite " + "number of seconds within the platform limit." + ) + timeout = timeout_seconds + # selector-supplied ports arrive uncoerced (numeric strings, floats + # from JSON), so coerce like the manifest would before range-checking; + # paho validates only port <= 0 and above 65535 the socket call raises + # OverflowError inside the network loop, silently killing it + try: + if isinstance(port, bool) or ( + isinstance(port, float) and not port.is_integer() + ): + raise ValueError + port_number = int(port) + except (TypeError, ValueError): + port_number = 0 + if not 1 <= port_number <= 65535: + return self._handle_failure( + f"Invalid port: {port!r}. Port must be an integer between 1 and 65535." + ) + port = port_number + try: + if isinstance(qos, bool) or ( + isinstance(qos, float) and not qos.is_integer() + ): + raise ValueError + qos_number = int(qos) + except (TypeError, ValueError): + qos_number = -1 + if qos_number not in (0, 1, 2): + return self._handle_failure(f"Invalid qos: {qos!r}. QoS must be 0, 1 or 2.") + qos = qos_number + if read_mode not in ("latest", "sequential"): + return self._handle_failure( + f"Invalid read_mode: {read_mode!r}. Must be 'latest' or 'sequential'." + ) + if not isinstance(topic, str) or not topic: + return self._handle_failure( + f"Invalid topic: {topic!r}. Topic must be a non-empty string." + ) + if password is not None and username is None: + return self._handle_failure( + "Password provided without username. Set username to enable MQTT authentication." + ) + with self._lifecycle_lock: + return self._connect_and_read( + host=host, + port=port, + topic=topic, + read_mode=read_mode, + qos=qos, + username=username, + password=password, + timeout=timeout, + ) + + def _connect_and_read( + self, + host: str, + port: int, + topic: str, + read_mode: str, + qos: int, + username: Optional[str], + password: Optional[str], + timeout: float, + ) -> BlockResult: + connection_identity = ( + host, + port, + username, + password, + timeout, + topic, + qos, + read_mode, + ) + if self._client is None: + buffer_size = ( + LATEST_BUFFER_SIZE if read_mode == "latest" else SEQUENTIAL_BUFFER_SIZE + ) + state = MQTTReaderState(topic=topic, qos=qos, buffer_size=buffer_size) + client = None + try: + client = mqtt.Client(userdata=state) + # a raising callback would end the network thread; log instead + client.suppress_exceptions = True + if username is not None: + client.username_pw_set(username, password) + client.on_connect = mqtt_on_connect + client.on_connect_fail = mqtt_on_connect_fail + client.on_subscribe = mqtt_on_subscribe + client.on_message = mqtt_on_message + client.on_disconnect = mqtt_on_disconnect + # min_delay stays below the readiness wait (= timeout) so the + # first reconnect attempt after a connection drop can finish + # within a single run's wait instead of starting as it expires + client.reconnect_delay_set(min_delay=timeout / 2, max_delay=2 * timeout) + # paho 1.6.1 has no public setter for its synchronous connect + # timeout; without this the TCP phase runs under paho's 5s + # default instead of the block's timeout + client._connect_timeout = timeout + # the TCP connect happens here, bounded by _connect_timeout + # (DNS resolution is not - it runs under the OS resolver + # timeout); the CONNACK wait below covers the handshake rest + client.connect(host, port) + client.loop_start() + except OSError as e: + # broker unreachable: nothing is kept and no loop was started, + # so a failed one-shot run leaves no background thread behind; + # the next run on this instance retries with a fresh client + return self._handle_failure( + f"MQTT broker not connected ({e}). Raise 'timeout' if the " + "broker needs longer to connect." + ) + except Exception as e: + if client is not None: + try: + client.loop_stop() + except Exception: + pass + return self._handle_failure(f"Failed to initialize MQTT client: {e}") + self._client = client + self._state = state + self._connection_identity = connection_identity + self._awaiting_retained = True + elif connection_identity != self._connection_identity: + return self._handle_failure( + "MQTT connection parameters (host, port, credentials, timeout, topic, " + "qos or read_mode) changed between runs; this block subscribes only " + "with the configuration of its first run." + ) + state = self._state + # the background loop owns (re)connecting and (re)subscribing; runs + # only wait for readiness + if not state.connected.wait(timeout=timeout): + return self._handle_failure( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." + ) + # a known refusal is permanent for this configuration: report it + # without paying the acknowledgement wait on every run + if state.subscribe_failed or ( + not state.subscribed.wait(timeout=timeout) and state.subscribe_failed + ): + return self._handle_failure( + f"MQTT broker refused subscription to topic {state.topic!r}." + ) + if not state.subscribed.is_set(): + return self._handle_failure( + "MQTT subscription was not acknowledged within timeout. Raise " + "'timeout' if the broker needs longer to respond." + ) + if self._awaiting_retained: + # one-off, bounded: a retained message (or nothing) arrives right + # after SUBACK; without this wait the first run would race it + self._awaiting_retained = False + state.message_received.wait(timeout=timeout) + try: + received_topic, payload, retained = state.messages.popleft() + except IndexError: + # nothing new since the previous run + return self._repeat_last() + try: + value = payload.decode("utf-8") + except UnicodeDecodeError: + return self._handle_failure( + f"Message received on topic {received_topic!r} is not valid UTF-8." + ) + parsed_payload: Dict[str, Any] = {} + try: + parsed = json.loads(value) + if isinstance(parsed, dict): + parsed_payload = parsed + except ValueError: + pass + if ( + retained + and self._last is not None + and self._last["topic"] == received_topic + and self._last["value"] == value + ): + # the broker re-sends the retained message on every resubscribe + # (reconnect); an unchanged copy is not a new message + return self._repeat_last() + self._last = { + "value": value, + "payload": parsed_payload, + "topic": received_topic, + } + return { + **self._last, + "is_new": True, + "error_status": False, + "error_message": None, + } + + def _repeat_last(self) -> BlockResult: + if self._last is None: + return self._empty_result() + return { + **self._last, + "is_new": False, + "error_status": False, + "error_message": None, + } + + @staticmethod + def _empty_result() -> BlockResult: + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": False, + "error_message": None, + } + + def _handle_failure(self, message: str) -> BlockResult: + logger.error("MQTT Reader failure: %s", message) + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": message, + } diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py new file mode 100644 index 0000000000..724d78b597 --- /dev/null +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -0,0 +1,791 @@ +import threading +import time +from types import SimpleNamespace +from typing import List, get_args +from unittest.mock import patch + +import pytest +from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader import v1 +from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import ( + LATEST_BUFFER_SIZE, + SEQUENTIAL_BUFFER_SIZE, + SUBSCRIPTION_REFUSED_QOS, + BlockManifest, + MQTTReaderBlockV1, + MQTTReaderState, + mqtt_on_connect, + mqtt_on_connect_fail, + mqtt_on_disconnect, + mqtt_on_message, + mqtt_on_subscribe, +) + +CLIENT_CLASS_PATH = ( + "roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1.mqtt.Client" +) + + +def manifest_payload(**overrides) -> dict: + payload = { + "type": "roboflow_enterprise/mqtt_reader@v1", + "name": "reader", + "host": "localhost", + "port": 1883, + "topic": "plc/state", + } + payload.update(overrides) + return payload + + +def run_kwargs(**overrides) -> dict: + kwargs = { + "host": "localhost", + "port": 1883, + "topic": "plc/state", + "timeout": 0.05, + } + kwargs.update(overrides) + return kwargs + + +class FakeClient: + """Stand-in for paho's Client that drives the block's real callbacks. + + connect() succeeds unless `connect_error` is set, loop_start() fires + on_connect with `connack_reason_code`, subscribe() answers with + `granted_qos` unless `ack_subscriptions` is False, and deliver() + pushes a message through on_message. + """ + + connect_error = None + connack_reason_code = 0 + granted_qos = 0 + ack_subscriptions = True + fire_on_connect = True + + def __init__(self, userdata=None): + self.userdata = userdata + self.on_connect = None + self.on_connect_fail = None + self.on_subscribe = None + self.on_message = None + self.on_disconnect = None + self.suppress_exceptions = False + self.credentials = None + self.connected_to = None + self.subscriptions: List[tuple] = [] + self.loop_started = False + self.loop_stopped = False + self.disconnected = False + self.reconnect_delays = None + + def username_pw_set(self, username, password=None): + self.credentials = (username, password) + + def reconnect_delay_set(self, min_delay, max_delay): + self.reconnect_delays = (min_delay, max_delay) + + def connect(self, host, port): + if self.connect_error is not None: + raise self.connect_error + self.connected_to = (host, port) + + def loop_start(self): + self.loop_started = True + if self.fire_on_connect: + self.on_connect(self, self.userdata, {}, self.connack_reason_code) + + def loop_stop(self): + self.loop_stopped = True + + def disconnect(self): + self.disconnected = True + + def subscribe(self, topic, qos=0): + self.subscriptions.append((topic, qos)) + if self.ack_subscriptions: + self.on_subscribe(self, self.userdata, 1, [self.granted_qos]) + + def deliver(self, payload, topic="plc/state", retain=False): + message = SimpleNamespace(topic=topic, payload=payload, retain=retain) + self.on_message(self, self.userdata, message) + + +@pytest.fixture +def clients() -> List[FakeClient]: + created: List[FakeClient] = [] + + def factory(userdata=None): + client = FakeClient(userdata=userdata) + created.append(client) + return client + + with patch(CLIENT_CLASS_PATH, side_effect=factory): + yield created + + +@pytest.fixture +def block() -> MQTTReaderBlockV1: + return MQTTReaderBlockV1() + + +class TestManifest: + def test_primary_identifier_is_namespaced(self): + identifiers = get_args(BlockManifest.model_fields["type"].annotation) + + assert identifiers == ("roboflow_enterprise/mqtt_reader@v1",) + + def test_defaults(self): + manifest = BlockManifest.model_validate(manifest_payload()) + + assert manifest.read_mode == "latest" + assert manifest.qos == 0 + assert manifest.timeout == 0.5 + assert manifest.username is None + assert manifest.password is None + + @pytest.mark.parametrize("read_mode", ["everything", "next", "", None]) + def test_read_mode_must_be_latest_or_sequential(self, read_mode): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(read_mode=read_mode)) + + @pytest.mark.parametrize("read_mode", ["latest", "sequential"]) + def test_valid_read_modes_accepted(self, read_mode): + manifest = BlockManifest.model_validate(manifest_payload(read_mode=read_mode)) + + assert manifest.read_mode == read_mode + + @pytest.mark.parametrize("port", [0, -1, 65536]) + def test_literal_port_must_be_within_tcp_range(self, port): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(port=port)) + + @pytest.mark.parametrize( + "timeout", [0, -1, float("nan"), float("inf"), float("-inf")] + ) + def test_literal_timeout_must_be_finite_and_positive(self, timeout): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(timeout=timeout)) + + def test_selectors_accepted_for_connection_fields(self): + manifest = BlockManifest.model_validate( + manifest_payload( + host="$inputs.host", + port="$inputs.port", + topic="$inputs.topic", + username="$inputs.user", + password="$inputs.secret", + timeout="$inputs.timeout", + qos="$inputs.qos", + ) + ) + + assert manifest.password == "$inputs.secret" + + def test_password_is_marked_private(self): + schema = BlockManifest.model_json_schema() + + assert schema["properties"]["password"]["private"] is True + + def test_describe_outputs_order(self): + names = [output.name for output in BlockManifest.describe_outputs()] + + assert names == [ + "value", + "payload", + "topic", + "is_new", + "error_status", + "error_message", + ] + + def test_restrictions_hard_for_hosted_serverless_only(self): + restrictions = BlockManifest.get_restrictions() + + assert [r.severity.value for r in restrictions] == ["hard", "soft"] + assert restrictions[0].applies_to_runtimes == [v1.Runtime.HOSTED_SERVERLESS] + assert restrictions[1].applies_to_runtimes == [ + v1.Runtime.SELF_HOSTED_CPU, + v1.Runtime.SELF_HOSTED_GPU, + v1.Runtime.DEDICATED_DEPLOYMENT, + ] + assert restrictions[1].applies_to_input_modes == [v1.RuntimeInputMode.IMAGE] + + def test_block_declares_no_init_parameters(self): + assert MQTTReaderBlockV1.get_init_parameters() == [] + + +class TestCallbacks: + def test_state_buffer_sizes_follow_read_mode_constants(self): + assert LATEST_BUFFER_SIZE == 1 + assert SEQUENTIAL_BUFFER_SIZE == 1000 + assert MQTTReaderState("t", 0, LATEST_BUFFER_SIZE).messages.maxlen == 1 + + def test_on_connect_subscribes_and_sets_event_for_accepted_connack(self): + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 0) + + assert client.subscriptions == [("plc/#", 1)] + assert state.connected.is_set() + + def test_on_connect_resubscribes_on_every_connect(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 0) + mqtt_on_disconnect(client, state, 1) + mqtt_on_connect(client, state, {}, 0) + + assert client.subscriptions == [("plc/#", 0), ("plc/#", 0)] + assert state.subscribed.is_set() + + @pytest.mark.parametrize("reason_code", [1, 2, 3, 4, 5]) + def test_on_connect_does_not_subscribe_for_rejected_connack(self, reason_code): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + state.connected.set() + + mqtt_on_connect(client, state, {}, reason_code) + + assert client.subscriptions == [] + assert not state.connected.is_set() + + def test_on_connect_fail_matches_paho_two_argument_signature(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + + mqtt_on_connect_fail(FakeClient(), state) + + assert not state.connected.is_set() + + def test_on_subscribe_flags_refused_subscription(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + + mqtt_on_subscribe(FakeClient(), state, 1, [SUBSCRIPTION_REFUSED_QOS]) + + assert state.subscribe_failed is True + assert not state.subscribed.is_set() + + @pytest.mark.parametrize("granted", [0, 1, 2]) + def test_on_subscribe_sets_event_for_granted_qos(self, granted): + state = MQTTReaderState(topic="t", qos=2, buffer_size=1) + + mqtt_on_subscribe(FakeClient(), state, 1, [granted]) + + assert state.subscribed.is_set() + assert state.subscribe_failed is False + + def test_on_message_appends_raw_tuple_only(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=10) + message = SimpleNamespace(topic="plc/a", payload=b"\xff\xfe", retain=1) + + mqtt_on_message(FakeClient(), state, message) + + assert list(state.messages) == [("plc/a", b"\xff\xfe", True)] + + def test_latest_buffer_keeps_only_newest(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=LATEST_BUFFER_SIZE) + for payload in (b"A", b"B", b"C"): + mqtt_on_message( + FakeClient(), + state, + SimpleNamespace(topic="t", payload=payload, retain=0), + ) + + assert list(state.messages) == [("t", b"C", False)] + + def test_on_disconnect_clears_both_events(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + state.subscribed.set() + + mqtt_on_disconnect(FakeClient(), state, 1) + + assert not state.connected.is_set() + assert not state.subscribed.is_set() + + def test_on_disconnect_accepts_mqtt5_properties_argument(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connected.set() + + mqtt_on_disconnect(FakeClient(), state, 0, None) + + assert not state.connected.is_set() + + +class TestInputValidation: + @pytest.mark.parametrize( + "timeout", [0, -1, float("nan"), float("inf"), "abc", None] + ) + def test_invalid_timeout_rejected_before_client_construction( + self, clients, block, timeout + ): + result = block.run(**run_kwargs(timeout=timeout)) + + assert result["error_status"] is True + assert "timeout" in result["error_message"].lower() + assert clients == [] + + @pytest.mark.parametrize("port", [0, 65536, "abc", 1.5, True, None]) + def test_invalid_port_rejected_before_client_construction( + self, clients, block, port + ): + result = block.run(**run_kwargs(port=port)) + + assert result["error_status"] is True + assert "port" in result["error_message"].lower() + assert clients == [] + + def test_selector_supplied_port_coerced_like_manifest(self, clients, block): + result = block.run(**run_kwargs(port="1883")) + + assert result["error_status"] is False + assert clients[0].connected_to == ("localhost", 1883) + + @pytest.mark.parametrize("qos", [-1, 3, "x", 1.5, True, None]) + def test_invalid_qos_rejected_before_client_construction(self, clients, block, qos): + result = block.run(**run_kwargs(qos=qos)) + + assert result["error_status"] is True + assert "qos" in result["error_message"].lower() + assert clients == [] + + def test_selector_supplied_qos_used_for_subscription(self, clients, block): + block.run(**run_kwargs(qos="2")) + + assert clients[0].subscriptions == [("plc/state", 2)] + + def test_invalid_read_mode_rejected(self, clients, block): + result = block.run(**run_kwargs(read_mode="everything")) + + assert result["error_status"] is True + assert "read_mode" in result["error_message"] + assert clients == [] + + def test_password_without_username_rejected(self, clients, block): + result = block.run(**run_kwargs(password="secret")) + + assert result["error_status"] is True + assert "username" in result["error_message"].lower() + assert clients == [] + + def test_username_and_password_configure_authentication(self, clients, block): + block.run(**run_kwargs(username="user", password="secret")) + + assert clients[0].credentials == ("user", "secret") + + @pytest.mark.parametrize("flag", ["GCP_SERVERLESS", "LAMBDA"]) + def test_hosted_platform_rejected_without_client(self, clients, block, flag): + with patch.object(v1, flag, True): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "hosted" in result["error_message"].lower() + assert clients == [] + + @pytest.mark.parametrize("topic", ["", None, 5]) + def test_invalid_topic_rejected_before_client_construction( + self, clients, block, topic + ): + result = block.run(**run_kwargs(topic=topic)) + + assert result["error_status"] is True + assert "topic" in result["error_message"].lower() + assert clients == [] + + def test_error_outputs_carry_every_declared_key(self, clients, block): + result = block.run(**run_kwargs(port=0)) + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": result["error_message"], + } + + +class TestConnectionLifecycle: + def test_first_run_connects_subscribes_and_starts_loop(self, clients, block): + result = block.run(**run_kwargs(username="u", password="p", timeout=0.25)) + + client = clients[0] + assert result["error_status"] is False + assert client.connected_to == ("localhost", 1883) + assert client.loop_started is True + assert client.suppress_exceptions is True + assert client.subscriptions == [("plc/state", 0)] + assert client.on_connect is mqtt_on_connect + assert client.on_connect_fail is mqtt_on_connect_fail + assert client.on_subscribe is mqtt_on_subscribe + assert client.on_message is mqtt_on_message + assert client.on_disconnect is mqtt_on_disconnect + assert client.reconnect_delays == (0.125, 0.5) + assert client._connect_timeout == 0.25 + + def test_client_reused_across_runs(self, clients, block): + block.run(**run_kwargs()) + block.run(**run_kwargs()) + + assert len(clients) == 1 + + def test_unreachable_broker_on_first_run_leaves_nothing_behind( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connect_error", OSError("connection refused")) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is None + assert clients[0].loop_started is False + + def test_next_run_after_failed_connect_retries_with_fresh_client( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connect_error", OSError("connection refused")) + block.run(**run_kwargs()) + + monkeypatch.setattr(FakeClient, "connect_error", None) + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert len(clients) == 2 + + def test_setup_failure_stops_loop_and_resets_client(self, clients, block): + with patch.object( + FakeClient, "reconnect_delay_set", side_effect=RuntimeError("boom") + ): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "initialize" in result["error_message"].lower() + assert block._client is None + assert clients[0].loop_stopped is True + + def test_connection_not_established_within_timeout_reported( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is clients[0] + + def test_rejected_connack_reported_as_not_connected( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert clients[0].subscriptions == [] + + def test_refused_subscription_reported_and_client_kept( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._client is clients[0] + assert block._state.connected.is_set() + + def test_refused_subscription_does_not_wait_for_timeout( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + block.run(**run_kwargs(timeout=1.0)) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=1.0)) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert time.monotonic() - started < 0.5 + + def test_subscribe_error_reported_through_refused_path( + self, clients, block, monkeypatch + ): + def failing_subscribe(client, topic, qos=0): + raise ValueError("Invalid topic.") + + monkeypatch.setattr(FakeClient, "subscribe", failing_subscribe) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "refused" in result["error_message"].lower() + assert block._state.connected.is_set() + + def test_unacknowledged_subscription_reported(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "ack_subscriptions", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "acknowledged" in result["error_message"].lower() + + @pytest.mark.parametrize( + "change", + [ + {"host": "other"}, + {"port": 8883}, + {"username": "u"}, + {"username": "u", "password": "p"}, + {"timeout": 0.1}, + {"topic": "plc/other"}, + {"qos": 1}, + {"read_mode": "sequential"}, + ], + ) + def test_changed_parameters_rejected_and_connection_kept( + self, clients, block, change + ): + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + assert clients[0].disconnected is False + + def test_close_disconnects_stops_loop_and_clears_state(self, clients, block): + block.run(**run_kwargs()) + state = block._state + + block.close() + + assert clients[0].disconnected is True + assert clients[0].loop_stopped is True + assert block._client is None + assert block._state is None + assert not state.connected.is_set() + assert not state.subscribed.is_set() + + def test_close_resets_last_outputs(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"x") + block.run(**run_kwargs()) + + block.close() + result = block.run(**run_kwargs()) + + assert len(clients) == 2 + assert result["value"] is None + assert result["is_new"] is False + + def test_close_without_client_is_noop(self, block): + block.close() + + assert block._client is None + + def test_close_when_disconnect_raises_still_stops_loop(self, clients, block): + block.run(**run_kwargs()) + with patch.object(FakeClient, "disconnect", side_effect=RuntimeError("boom")): + block.close() + + assert clients[0].loop_stopped is True + assert block._client is None + + +class TestReading: + def test_empty_first_run(self, clients, block): + result = block.run(**run_kwargs()) + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": False, + "error_message": None, + } + + def test_json_message_surfaced_with_parsed_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + + result = block.run(**run_kwargs()) + + assert result == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": True, + "error_status": False, + "error_message": None, + } + + def test_retained_message_delivered_before_first_run_is_surfaced( + self, clients, block + ): + original_loop_start = FakeClient.loop_start + + def loop_start_with_retained(client): + original_loop_start(client) + client.deliver(b'{"state": "RUNNING"}', retain=True) + + with patch.object(FakeClient, "loop_start", loop_start_with_retained): + result = block.run(**run_kwargs()) + + assert result["value"] == '{"state": "RUNNING"}' + assert result["is_new"] is True + + def test_retained_redelivery_of_same_message_is_not_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + first = block.run(**run_kwargs()) + # a reconnect resubscribes and the broker re-sends the retained copy + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + + second = block.run(**run_kwargs()) + + assert first["is_new"] is True + assert second == {**first, "is_new": False} + + def test_retained_message_with_new_payload_is_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', retain=True) + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "PAUSED"}', retain=True) + + result = block.run(**run_kwargs()) + + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + + def test_non_retained_duplicate_payload_is_new(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"tick") + block.run(**run_kwargs()) + clients[0].deliver(b"tick") + + result = block.run(**run_kwargs()) + + assert result["is_new"] is True + + def test_nothing_new_repeats_previous_outputs_with_is_new_false( + self, clients, block + ): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}', topic="plc/line1") + first = block.run(**run_kwargs()) + + second = block.run(**run_kwargs()) + + assert first["is_new"] is True + assert second == {**first, "is_new": False} + + def test_latest_mode_returns_newest_and_skips_older(self, clients, block): + block.run(**run_kwargs()) + for payload in (b"A", b"B", b"C"): + clients[0].deliver(payload) + + result = block.run(**run_kwargs()) + repeat = block.run(**run_kwargs()) + + assert result["value"] == "C" + assert result["is_new"] is True + assert repeat["value"] == "C" + assert repeat["is_new"] is False + + def test_sequential_mode_returns_messages_in_order_one_per_run( + self, clients, block + ): + block.run(**run_kwargs(read_mode="sequential")) + for payload in (b"A", b"B", b"C"): + clients[0].deliver(payload) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(4)] + + assert [r["value"] for r in results] == ["A", "B", "C", "C"] + assert [r["is_new"] for r in results] == [True, True, True, False] + + def test_sequential_mode_drops_oldest_beyond_buffer(self, clients, block): + block.run(**run_kwargs(read_mode="sequential")) + for index in range(SEQUENTIAL_BUFFER_SIZE + 1): + clients[0].deliver(str(index).encode()) + + result = block.run(**run_kwargs(read_mode="sequential")) + + assert result["value"] == "1" + assert len(block._state.messages) == SEQUENTIAL_BUFFER_SIZE - 1 + + def test_topic_output_reports_message_topic_for_wildcard_filters( + self, clients, block + ): + block.run(**run_kwargs(topic="plc/#")) + clients[0].deliver(b"1", topic="plc/line2/state") + + result = block.run(**run_kwargs(topic="plc/#")) + + assert result["topic"] == "plc/line2/state" + + def test_non_json_payload_gives_text_and_empty_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"hello") + + result = block.run(**run_kwargs()) + + assert result["value"] == "hello" + assert result["payload"] == {} + assert result["error_status"] is False + + def test_json_non_object_gives_empty_payload(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b"[1, 2]") + + result = block.run(**run_kwargs()) + + assert result["value"] == "[1, 2]" + assert result["payload"] == {} + + def test_non_utf8_payload_reported_as_error_without_previous_outputs( + self, clients, block + ): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + block.run(**run_kwargs()) + clients[0].deliver(b"\xff\xfe") + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "utf-8" in result["error_message"].lower() + assert result["value"] is None + assert result["payload"] == {} + assert result["is_new"] is False + + def test_message_received_on_network_thread_visible_to_run(self, clients, block): + block.run(**run_kwargs()) + delivered = threading.Event() + + def deliver(): + clients[0].deliver(b"from-thread") + delivered.set() + + threading.Thread(target=deliver).start() + assert delivered.wait(timeout=1) + + result = block.run(**run_kwargs()) + + assert result["value"] == "from-thread" + + def test_no_raw_prints_in_any_code_path(self, clients, block): + with patch("builtins.print") as mock_print: + block.run(**run_kwargs()) + clients[0].deliver(b"x") + block.run(**run_kwargs()) + block.run(**run_kwargs(port=0)) + block.close() + + mock_print.assert_not_called() From ef368086dd70c33c8cc7b2a336cd27cb2e0ceb15 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 13:04:42 +0800 Subject: [PATCH 02/14] Add operator broker policy and keepalive to the MQTT blocks Mirror the Kafka blocks' bootstrap-server policy for the MQTT Reader and MQTT Writer: MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST (default True) and MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS (host[:port] entries, no DNS, never echoed in errors) flow from env.py through the server configuration mapping into the roboflow_workflows engine configuration. When user-provided hosts are allowed the allowlist filters the workflow's host and port; when they are not, the first entry is the broker both blocks connect to, and the resolved address is the connection identity. A run the policy forbids is reported in the outputs (the writer honours fail_fast). Both blocks connect with a 15 s keepalive so a broker that disappears without closing the connection is noticed within about 25 s instead of paho's 60 s default. Co-Authored-By: Claude Fable 5.1 --- inference/core/env.py | 29 ++++ .../interfaces/workflows_configuration.py | 8 + .../test_workflows_configuration.py | 19 ++- .../core_steps/sinks/test_mqtt_env.py | 62 ++++++++ workflows/roboflow_workflows/configuration.py | 5 + .../enterprise_blocks/sinks/mqtt_common.py | 142 +++++++++++++++++ .../enterprise_blocks/sinks/mqtt_reader/v1.py | 22 ++- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 23 ++- workflows/roboflow_workflows/environment.py | 10 ++ .../core_steps/sinks/test_mqtt_common.py | 149 ++++++++++++++++++ .../core_steps/sinks/test_mqtt_reader.py | 76 ++++++++- .../core_steps/sinks/test_mqtt_writer.py | 88 ++++++++++- .../tests/unit_tests/test_configuration.py | 6 +- 13 files changed, 630 insertions(+), 9 deletions(-) create mode 100644 tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py create mode 100644 workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py create mode 100644 workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py diff --git a/inference/core/env.py b/inference/core/env.py index 7c8be62b08..764b583f93 100644 --- a/inference/core/env.py +++ b/inference/core/env.py @@ -157,6 +157,35 @@ for entry in KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS.split(",") if entry.strip() ] +# Operator policy for the MQTT Reader / MQTT Writer Workflow blocks, which open +# an outbound connection to the `host`/`port` taken from the workflow definition +# or its inputs. When True (default, preserves behaviour) that value is +# honoured, subject to the allowlist below. When False the workflow-provided +# value is ignored and the blocks connect to the first entry of +# MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS instead; with no entry configured +# there, the MQTT blocks are disabled (every run reports an error). +MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST = str2bool( + os.getenv("MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", True) +) +# Optional comma-separated list of `host[:port]` MQTT brokers the MQTT Workflow +# blocks may connect to. When set and a user-provided host is allowed, the +# workflow's `host:port` must match an entry or the run reports an error; an +# entry without a port matches that host on any port. When a user-provided host +# is not allowed, the first entry is what the blocks connect to (port 1883 when +# the entry has none). Hosts are compared after trimming whitespace, lowercasing +# and stripping IPv6 brackets; there is no DNS resolution. Empty entries are +# ignored. A variable that is set but holds no entries is an empty allowlist: +# nothing is permitted. Default None (no allowlist). +MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = os.getenv( + "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" +) +if MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS is not None: + # a list, not a set: the first entry is the operator's broker + MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = [ + entry.strip() + for entry in MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS.split(",") + if entry.strip() + ] # List of allowed origins ALLOW_ORIGINS = os.getenv("ALLOW_ORIGINS", "*") diff --git a/inference/core/interfaces/workflows_configuration.py b/inference/core/interfaces/workflows_configuration.py index 2e5a924855..0bb5716c60 100644 --- a/inference/core/interfaces/workflows_configuration.py +++ b/inference/core/interfaces/workflows_configuration.py @@ -79,6 +79,14 @@ def build_configuration_from_env() -> WorkflowsConfiguration: if env.KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS is None else tuple(env.KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS) ), + allow_mqtt_blocks_user_provided_host=env.MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, + # NOT sorted: the first entry is the operator's broker when + # user-provided hosts are not allowed. + mqtt_blocks_whitelisted_hosts=( + None + if env.MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS is None + else tuple(env.MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS) + ), ), tensor=TensorConfiguration( representation_enabled=env.ENABLE_TENSOR_DATA_REPRESENTATION, diff --git a/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py b/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py index 0022f2e08c..7fd87dde9e 100644 --- a/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py +++ b/tests/inference/unit_tests/core/interfaces/test_workflows_configuration.py @@ -99,6 +99,18 @@ else list(c.engine.kafka_sinks_whitelisted_bootstrap_servers) ), ), + ( + "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", + lambda c: c.engine.allow_mqtt_blocks_user_provided_host, + ), + ( + "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", + lambda c: ( + None + if c.engine.mqtt_blocks_whitelisted_hosts is None + else list(c.engine.mqtt_blocks_whitelisted_hosts) + ), + ), ("ENABLE_TENSOR_DATA_REPRESENTATION", lambda c: c.tensor.representation_enabled), ("WORKFLOWS_IMAGE_TENSOR_DEVICE", lambda c: c.tensor.image_tensor_device), ( @@ -226,7 +238,7 @@ def test_the_field_table_matches_the_facade_exports() -> None: "missing_from_table": sorted(exported - tabled), "missing_from_facade": sorted(tabled - exported), } - assert len(tabled) == 78, len(tabled) + assert len(tabled) == 80, len(tabled) def test_every_name_workflows_imports_from_the_facade_is_exported() -> None: @@ -291,6 +303,11 @@ def test_server_configuration_equals_env_field_by_field(name, reader) -> None: "KAFKA_WORKFLOWS_SINKS_WHITELISTED_BOOTSTRAP_SERVERS", ["kafka-2:9092", "kafka-1:9092"], ), + ("MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False), + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None), + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", []), + # order preserved, not sorted + ("MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker-2:1883", "broker-1"]), ("LAMBDA", True), ], ) diff --git a/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py b/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py new file mode 100644 index 0000000000..77bda17220 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/sinks/test_mqtt_env.py @@ -0,0 +1,62 @@ +"""Server-side parity: `inference.core.env` parsing of the MQTT blocks policy. + +Retained here (not in `workflows/tests`) because it imports the server package. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any, Dict + +ALLOW_USER_HOST = "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST" +WHITELISTED_HOSTS = "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" +_REPO_ROOT = str(Path(__file__).resolve().parents[5]) + + +def _import_env(extra_env: Dict[str, str]) -> Dict[str, Any]: + # env.py computes the values at import time: read them in a fresh interpreter + env = {**os.environ} + env.pop(ALLOW_USER_HOST, None) + env.pop(WHITELISTED_HOSTS, None) + env.update(extra_env) + env["PYTHONPATH"] = _REPO_ROOT + os.pathsep + env.get("PYTHONPATH", "") + code = ( + "import json; from inference.core import env; " + f"print('RESULT' + json.dumps([env.{ALLOW_USER_HOST}, " + f"env.{WHITELISTED_HOSTS}]))" + ) + result = subprocess.run( + [sys.executable, "-c", code], + env=env, + cwd=_REPO_ROOT, + capture_output=True, + text=True, + timeout=120, + ) + assert result.returncode == 0, result.stderr + line = [l for l in result.stdout.splitlines() if l.startswith("RESULT")][-1] + allow_user_provided, allowlist = json.loads(line[len("RESULT") :]) + return {"allow_user_provided": allow_user_provided, "allowlist": allowlist} + + +def test_env_defaults_allow_user_hosts_without_allowlist() -> None: + values = _import_env({}) + + assert values == {"allow_user_provided": True, "allowlist": None} + + +def test_env_parses_the_allowlist_in_order_without_empty_entries() -> None: + values = _import_env( + { + ALLOW_USER_HOST: "False", + WHITELISTED_HOSTS: " broker-2:1883, ,broker-1 ,broker-2:1883,", + } + ) + + # trimmed, blank entries dropped, otherwise exactly what the operator wrote + assert values == { + "allow_user_provided": False, + "allowlist": ["broker-2:1883", "broker-1", "broker-2:1883"], + } diff --git a/workflows/roboflow_workflows/configuration.py b/workflows/roboflow_workflows/configuration.py index a57be85f91..451ca8cee7 100644 --- a/workflows/roboflow_workflows/configuration.py +++ b/workflows/roboflow_workflows/configuration.py @@ -59,6 +59,11 @@ class EngineConfiguration: # None-vs-list distinction; the facade rebuilds the ORDERED list. allow_kafka_sinks_user_provided_bootstrap_servers: bool = True kafka_sinks_whitelisted_bootstrap_servers: Optional[Tuple[str, ...]] = None + # MQTT blocks broker-address policy, same shape as the Kafka pair: the + # facade rebuilds the ORDERED list and the first entry is the operator's + # broker when user-provided hosts are not allowed. + allow_mqtt_blocks_user_provided_host: bool = True + mqtt_blocks_whitelisted_hosts: Optional[Tuple[str, ...]] = None @dataclass(frozen=True) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py new file mode 100644 index 0000000000..da25fda589 --- /dev/null +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -0,0 +1,142 @@ +"""Shared pieces of the MQTT Reader and MQTT Writer blocks. + +Both blocks open an outbound connection to the broker named by the workflow. +`resolve_broker_address()` applies the operator's policy from +`MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` and +`MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (see `roboflow_workflows.environment`) +before either block builds a client. +""" + +import logging +from typing import List, Optional, Tuple + +from roboflow_workflows.environment import ( + MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, + MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS, +) + +logger = logging.getLogger("inference") + +DEFAULT_MQTT_PORT = 1883 +# paho defaults to 60 s; a silent link loss (no TCP close) is only noticed +# through the keepalive, within about 1.5 times this value +MQTT_KEEPALIVE_SECONDS = 15 + + +class ConfigurationError(ValueError): + """The operator's policy forbids the requested broker connection.""" + + +def _normalise_host(host: str) -> str: + host = host.strip().lower() + if host.startswith("[") and host.endswith("]"): + host = host[1:-1] + return host + + +def split_host_port(entry: str) -> Tuple[str, Optional[int]]: + """Split an allowlist entry into a normalised host and an optional port. + + Accepts ``host``, ``host:port``, ``[ipv6]``, ``[ipv6]:port`` and a bare + IPv6 address. There is no DNS resolution. + + Args: + entry: One operator-written allowlist entry. + + Returns: + ``(host, port)`` with the host lower-cased, whitespace trimmed and IPv6 + brackets stripped; ``port`` is ``None`` when the entry names none. + """ + entry = entry.strip() + if entry.startswith("["): + closing = entry.find("]") + if closing == -1: + return _normalise_host(entry), None + rest = entry[closing + 1 :] + port = int(rest[1:]) if rest.startswith(":") and rest[1:].isdigit() else None + return _normalise_host(entry[: closing + 1]), port + if entry.count(":") == 1: + host, port_text = entry.split(":") + if port_text.strip().isdigit(): + return _normalise_host(host), int(port_text) + return _normalise_host(entry), None + # no colon: a plain host; several colons: a bare IPv6 address + return _normalise_host(entry), None + + +def normalise_broker_address(host: str, port: int) -> str: + """Return the ``host:port`` form used when comparing addresses. + + Args: + host: Broker host as written in the workflow. + port: Broker port. + + Returns: + ``host:port`` with the host lower-cased and IPv6 brackets stripped. + """ + return f"{_normalise_host(str(host))}:{int(port)}" + + +def _operator_entries(allowlist: Optional[List[str]]) -> List[str]: + if allowlist is None: + return [] + return [str(entry).strip() for entry in allowlist if str(entry).strip()] + + +def resolve_broker_address( + host: str, port: int, log_override: bool = True +) -> Tuple[str, int]: + """Apply the operator's broker policy to a workflow-provided address. + + Args: + host: Broker host taken from the workflow. + port: Broker port taken from the workflow, already validated as an int. + log_override: Emit the warning when the workflow value is replaced by + the operator's broker; blocks pass ``False`` after their first run + so the warning is logged once per instance instead of once per frame. + + Returns: + The ``(host, port)`` the block must connect to: the workflow value when + permitted, or the operator's first allowlist entry when workflow-provided + hosts are not allowed. + + Raises: + ConfigurationError: When the policy forbids the connection. The message + never echoes the allowlist. + """ + # module attributes are read on every call so tests can patch them + user_provided_allowed = MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST + allowlist = MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS + operator_entries = _operator_entries(allowlist) + if not user_provided_allowed: + if not operator_entries: + raise ConfigurationError( + "MQTT blocks are disabled on this deployment: a workflow-provided " + "broker host is not allowed and the operator configured none." + ) + operator_host, operator_port = split_host_port(operator_entries[0]) + if operator_port is None: + operator_port = DEFAULT_MQTT_PORT + if log_override: + logger.warning( + "MQTT blocks: the workflow-provided host/port was replaced by the " + "operator-configured broker (%s:%s), because " + "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST is False.", + operator_host, + operator_port, + ) + return operator_host, operator_port + if allowlist is None: + return host, port + + wanted_host = _normalise_host(str(host)) + wanted_port = int(port) + for entry in operator_entries: + entry_host, entry_port = split_host_port(entry) + if entry_host == wanted_host and entry_port in (None, wanted_port): + return host, port + # the allowlist holds the operator's internal hostnames: never echo it + raise ConfigurationError( + f"Broker {host!r} port {wanted_port} is not permitted on this deployment: " + "the operator restricts the MQTT brokers Workflow blocks may connect to." + ) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index fa4293ebea..6739c459f7 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -13,6 +13,11 @@ # handlers / filters / propagation still apply when installed, without # importing the server logger module. logger = logging.getLogger("inference") +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, + ConfigurationError, + resolve_broker_address, +) from roboflow_workflows.environment import GCP_SERVERLESS, LAMBDA from roboflow_workflows.execution_engine.entities.base import OutputDefinition from roboflow_workflows.execution_engine.entities.types import ( @@ -75,6 +80,13 @@ `LAMBDA`): every run fails there without opening a connection. It is intended for self-hosted inference servers and InferencePipelines that can reach the broker. +The server operator may restrict which brokers the MQTT blocks connect to with +`MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` entries) and +`MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the workflow's host and +port are ignored and the first allowlist entry is used); a run the policy forbids is +reported in the outputs. The connection uses a 15 s keepalive, so a broker that +disappears without closing the connection is noticed within about 25 s. + Outputs: - value (str): Raw message payload decoded as UTF-8, or None when nothing was received yet. @@ -431,6 +443,14 @@ def run( return self._handle_failure( "Password provided without username. Set username to enable MQTT authentication." ) + try: + # the operator's broker policy; the RESOLVED address is what the + # block connects to and what the connection identity is built from + host, port = resolve_broker_address( + host, port, log_override=self._client is None + ) + except ConfigurationError as e: + return self._handle_failure(str(e)) with self._lifecycle_lock: return self._connect_and_read( host=host, @@ -492,7 +512,7 @@ def _connect_and_read( # the TCP connect happens here, bounded by _connect_timeout # (DNS resolution is not - it runs under the OS resolver # timeout); the CONNACK wait below covers the handshake rest - client.connect(host, port) + client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) client.loop_start() except OSError as e: # broker unreachable: nothing is kept and no loop was started, diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index e313a3704e..3d238650c1 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -12,6 +12,11 @@ # importing the server logger module. logger = logging.getLogger("inference") from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + MQTT_KEEPALIVE_SECONDS, + ConfigurationError, + resolve_broker_address, +) from roboflow_workflows.execution_engine.entities.base import OutputDefinition from roboflow_workflows.execution_engine.entities.types import ( BOOLEAN_KIND, @@ -44,6 +49,14 @@ connection: changing host, port, credentials or timeout between runs is rejected as a configuration error. +The server operator may restrict which brokers the MQTT blocks connect to +with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` +entries) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the +workflow's host and port are ignored and the first allowlist entry is used); +a run the policy forbids is reported like any other failure. The connection +uses a 15 s keepalive, so a broker that disappears without closing the +connection is noticed within about 25 s. + Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. True if there was an error, False if successful. @@ -278,6 +291,14 @@ def run( "Password provided without username. Set username to enable MQTT authentication.", fail_fast=fail_fast, ) + try: + # the operator's broker policy; the RESOLVED address is what the + # block connects to and what the connection identity is built from + host, port = resolve_broker_address( + host, port, log_override=self.mqtt_client is None + ) + except ConfigurationError as e: + return self._handle_failure(str(e), fail_fast=fail_fast) with self._lifecycle_lock: return self._connect_and_publish( host=host, @@ -326,7 +347,7 @@ def _connect_and_publish( # the TCP connect happens here, bounded by _connect_timeout # (DNS resolution is not - it runs under the OS resolver # timeout); the CONNACK wait below covers the handshake rest - client.connect(host, port) + client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) client.loop_start() except OSError as e: # broker unreachable: nothing is kept and no loop was started, diff --git a/workflows/roboflow_workflows/environment.py b/workflows/roboflow_workflows/environment.py index 4ac92fbe4d..51ffc8c440 100644 --- a/workflows/roboflow_workflows/environment.py +++ b/workflows/roboflow_workflows/environment.py @@ -68,6 +68,16 @@ if _CONFIGURATION.engine.kafka_sinks_whitelisted_bootstrap_servers is None else list(_CONFIGURATION.engine.kafka_sinks_whitelisted_bootstrap_servers) ) +MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST = ( + _CONFIGURATION.engine.allow_mqtt_blocks_user_provided_host +) +# A list, not a set: the first entry is the broker the MQTT blocks connect to +# when user-provided hosts are not allowed. +MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS = ( + None + if _CONFIGURATION.engine.mqtt_blocks_whitelisted_hosts is None + else list(_CONFIGURATION.engine.mqtt_blocks_whitelisted_hosts) +) # --- tensor representation --- ENABLE_TENSOR_DATA_REPRESENTATION = _CONFIGURATION.tensor.representation_enabled diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py new file mode 100644 index 0000000000..d865122adc --- /dev/null +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -0,0 +1,149 @@ +import logging + +import pytest +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common +from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + DEFAULT_MQTT_PORT, + MQTT_KEEPALIVE_SECONDS, + ConfigurationError, + normalise_broker_address, + resolve_broker_address, + split_host_port, +) + +ALLOW = "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST" +ALLOWLIST = "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS" + + +@pytest.fixture +def policy(monkeypatch): + def configure(allow: bool = True, allowlist=None): + monkeypatch.setattr(mqtt_common, ALLOW, allow) + monkeypatch.setattr(mqtt_common, ALLOWLIST, allowlist) + + configure() + return configure + + +@pytest.mark.parametrize( + "entry, expected", + [ + ("broker.local", ("broker.local", None)), + (" Broker.LOCAL ", ("broker.local", None)), + ("broker.local:1883", ("broker.local", 1883)), + ("Broker.LOCAL:8883 ", ("broker.local", 8883)), + ("10.0.0.5", ("10.0.0.5", None)), + ("[::1]", ("::1", None)), + ("[::1]:8883", ("::1", 8883)), + ("fe80::1", ("fe80::1", None)), + ("broker.local:notaport", ("broker.local:notaport", None)), + ], +) +def test_split_host_port(entry, expected): + assert split_host_port(entry) == expected + + +@pytest.mark.parametrize( + "host, port, expected", + [ + ("Broker.LOCAL", 1883, "broker.local:1883"), + ("[::1]", 8883, "::1:8883"), + (" 10.0.0.5 ", "1883", "10.0.0.5:1883"), + ], +) +def test_normalise_broker_address(host, port, expected): + assert normalise_broker_address(host, port) == expected + + +def test_keepalive_is_shorter_than_paho_default(): + assert 0 < MQTT_KEEPALIVE_SECONDS < 60 + + +def test_default_policy_passes_user_value_through_unchanged(policy): + assert resolve_broker_address("Broker.LOCAL", 1883) == ("Broker.LOCAL", 1883) + + +@pytest.mark.parametrize( + "host, port", + [ + ("broker.local", 1883), + ("Broker.LOCAL", 1883), + ("[::1]", 8883), + ("10.0.0.5", 1883), + ("10.0.0.5", 8883), + ], +) +def test_allowlist_accepts_listed_host_and_port_or_portless_host(policy, host, port): + policy(allowlist=["broker.local:1883", "[::1]:8883", "10.0.0.5"]) + + assert resolve_broker_address(host, port) == (host, port) + + +@pytest.mark.parametrize( + "host, port", + [ + ("other.local", 1883), + ("broker.local", 8883), + ("[::1]", 1883), + ], +) +def test_allowlist_rejects_unlisted_address_without_revealing_the_allowlist( + policy, host, port +): + policy(allowlist=["broker.local:1883", "[::1]:8883", "10.0.0.5"]) + + with pytest.raises(ConfigurationError) as error: + resolve_broker_address(host, port) + + message = str(error.value) + assert "not permitted" in message + assert "10.0.0.5" not in message + assert "broker.local" not in message or host == "broker.local" + + +def test_empty_allowlist_permits_nothing(policy): + policy(allowlist=[]) + + with pytest.raises(ConfigurationError): + resolve_broker_address("broker.local", 1883) + + +def test_allowlist_ignores_blank_entries(policy): + policy(allowlist=[" ", "broker.local", ""]) + + assert resolve_broker_address("broker.local", 1883) == ("broker.local", 1883) + + +def test_user_host_not_allowed_uses_first_operator_entry(policy, caplog): + policy(allow=False, allowlist=["Operator.Broker:8883", "other:1883"]) + + with caplog.at_level(logging.WARNING, logger="inference"): + resolved = resolve_broker_address("workflow.host", 1883) + + assert resolved == ("operator.broker", 8883) + assert "replaced by the operator-configured broker" in caplog.text + + +def test_user_host_not_allowed_defaults_port(policy): + policy(allow=False, allowlist=["operator.broker"]) + + assert resolve_broker_address("workflow.host", 1) == ( + "operator.broker", + DEFAULT_MQTT_PORT, + ) + + +def test_override_warning_can_be_silenced(policy, caplog): + policy(allow=False, allowlist=["operator.broker"]) + + with caplog.at_level(logging.WARNING, logger="inference"): + resolve_broker_address("workflow.host", 1883, log_override=False) + + assert caplog.text == "" + + +def test_user_host_not_allowed_without_operator_broker_is_disabled(policy): + policy(allow=False, allowlist=None) + + with pytest.raises(ConfigurationError, match="disabled"): + resolve_broker_address("workflow.host", 1883) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 724d78b597..351fd284c7 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -6,6 +6,7 @@ import pytest from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import ( LATEST_BUFFER_SIZE, @@ -86,10 +87,11 @@ def username_pw_set(self, username, password=None): def reconnect_delay_set(self, min_delay, max_delay): self.reconnect_delays = (min_delay, max_delay) - def connect(self, host, port): + def connect(self, host, port, keepalive=60): if self.connect_error is not None: raise self.connect_error self.connected_to = (host, port) + self.keepalive = keepalive def loop_start(self): self.loop_started = True @@ -428,6 +430,7 @@ def test_first_run_connects_subscribes_and_starts_loop(self, clients, block): assert client.on_disconnect is mqtt_on_disconnect assert client.reconnect_delays == (0.125, 0.5) assert client._connect_timeout == 0.25 + assert client.keepalive == 15 def test_client_reused_across_runs(self, clients, block): block.run(**run_kwargs()) @@ -789,3 +792,74 @@ def test_no_raw_prints_in_any_code_path(self, clients, block): block.close() mock_print.assert_not_called() + + +class TestBrokerPolicy: + def test_allowlisted_broker_is_used(self, clients, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["LocalHost"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert clients[0].connected_to == ("localhost", 1883) + + def test_unlisted_broker_rejected_before_client_construction( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not permitted" in result["error_message"] + assert "broker.internal" not in result["error_message"] + assert clients == [] + + def test_port_mismatch_rejected(self, clients, block, monkeypatch): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:8883"] + ) + + result = block.run(**run_kwargs(port=1883)) + + assert result["error_status"] is True + assert clients == [] + + def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["operator:8883"] + ) + + first = block.run(**run_kwargs(host="workflow-a")) + # identity is the resolved address: a different workflow host is not a change + second = block.run(**run_kwargs(host="workflow-b")) + + assert first["error_status"] is False + assert second["error_status"] is False + assert len(clients) == 1 + assert clients[0].connected_to == ("operator", 8883) + + def test_user_host_not_allowed_without_operator_broker_disables_block( + self, clients, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "disabled" in result["error_message"] + assert clients == [] diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index d7fd7df461..91308c32ba 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -4,6 +4,7 @@ import pytest from pydantic import ValidationError +from roboflow_workflows.enterprise_blocks.sinks import mqtt_common from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( BlockManifest, @@ -205,7 +206,9 @@ def test_selector_supplied_port_coerced_like_manifest( result = block.run(**run_kwargs(port=port)) assert result["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) + mock_client_cls.return_value.connect.assert_called_once_with( + "localhost", 1883, keepalive=15 + ) @pytest.mark.parametrize("qos", [-1, 3, "abc", 1.5, True, None]) def test_invalid_qos_rejected_before_client_construction( @@ -271,7 +274,7 @@ def test_client_connects_synchronously_before_background_loop( assert result["error_status"] is False mock_client_cls.assert_called_once_with(userdata=block._connected) - mock_client.connect.assert_called_once_with("localhost", 1883) + mock_client.connect.assert_called_once_with("localhost", 1883, keepalive=15) mock_client.connect_async.assert_not_called() called_methods = [call[0] for call in mock_client.method_calls] assert called_methods.index("connect") < called_methods.index("loop_start") @@ -423,7 +426,7 @@ def test_changed_connection_parameters_rejected(self, mock_client_cls, block): assert result["error_status"] is True assert "parameters" in result["message"].lower() - mock_client.connect.assert_called_once_with("broker-a", 1883) + mock_client.connect.assert_called_once_with("broker-a", 1883, keepalive=15) mock_client.publish.assert_called_once() def test_changed_credentials_rejected(self, mock_client_cls, block): @@ -627,3 +630,82 @@ def test_close_clears_readiness_only_after_loop_thread_joined( assert event_state_at_loop_stop == [True] assert not block._connected.is_set() + + +class TestBrokerPolicy: + def test_allowlisted_broker_is_used_and_keepalive_passed( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:1883"] + ) + block._connected.set() + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with( + "localhost", 1883, keepalive=15 + ) + + def test_unlisted_broker_rejected_before_client_construction( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not permitted" in result["message"] + assert "broker.internal" not in result["message"] + mock_client_cls.assert_not_called() + + def test_unlisted_broker_with_fail_fast_raises( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["broker.internal"] + ) + + with pytest.raises(RuntimeError, match="not permitted"): + block.run(**run_kwargs(fail_fast=True)) + + mock_client_cls.assert_not_called() + + def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["operator:8883"] + ) + block._connected.set() + + first = block.run(**run_kwargs(host="workflow-a")) + second = block.run(**run_kwargs(host="workflow-b")) + + assert first["error_status"] is False + assert second["error_status"] is False + mock_client_cls.return_value.connect.assert_called_once_with( + "operator", 8883, keepalive=15 + ) + + def test_user_host_not_allowed_without_operator_broker_disables_block( + self, mock_client_cls, block, monkeypatch + ): + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST", False + ) + monkeypatch.setattr( + mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", None + ) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "disabled" in result["message"] + mock_client_cls.assert_not_called() diff --git a/workflows/tests/unit_tests/test_configuration.py b/workflows/tests/unit_tests/test_configuration.py index b22744b341..8140b8a098 100644 --- a/workflows/tests/unit_tests/test_configuration.py +++ b/workflows/tests/unit_tests/test_configuration.py @@ -66,7 +66,7 @@ def test_configuration_is_frozen_in_every_group() -> None: total_fields += len(dataclasses.fields(value)) with pytest.raises(dataclasses.FrozenInstanceError): setattr(value, dataclasses.fields(value)[0].name, "mutated") - assert total_fields == 78, total_fields + assert total_fields == 80, total_fields def test_default_configuration_matches_env_pys_empty_environment_defaults() -> None: @@ -90,6 +90,8 @@ def test_default_configuration_matches_env_pys_empty_environment_defaults() -> N configuration.engine.allow_kafka_sinks_user_provided_bootstrap_servers is True ) assert configuration.engine.kafka_sinks_whitelisted_bootstrap_servers is None + assert configuration.engine.allow_mqtt_blocks_user_provided_host is True + assert configuration.engine.mqtt_blocks_whitelisted_hosts is None assert configuration.models.vlm_segmentation_max_polygon_vertices == 500 assert configuration.tensor.representation_enabled is False assert configuration.tensor.image_tensor_device is None @@ -412,7 +414,7 @@ def test_environment_facade_exports_every_owned_symbol() -> None: for name in vars(workflows_environment) if name.isupper() and not name.startswith("_") } - assert len(exported) == 78, sorted(exported) + assert len(exported) == 80, sorted(exported) assert isinstance(workflows_environment.WORKFLOW_DISABLED_BLOCK_TYPES, list) assert isinstance(workflows_environment.WORKFLOW_DISABLED_BLOCK_PATTERNS, list) assert isinstance(workflows_environment.ENABLE_TENSOR_DATA_REPRESENTATION, bool) From 568db616e1bb9d3671cd81259cb571e73c58681f Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 13:21:29 +0800 Subject: [PATCH 03/14] Add MQTT Reader and broker policy to the workflows changelog Co-Authored-By: Claude Fable 5.1 --- workflows/CHANGELOG.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 61913de8ed..21146792a2 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -2,6 +2,15 @@ ## Unreleased +### Added + +- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. +- Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. + +### Changed + +- MQTT Writer connects with a 15 s keepalive instead of paho's 60 s default, so a broker that disappears without closing the connection is reported within about 25 s. + --- ## `0.1.2` From 9de07eebf74b383bb3868edbe934787a75c27e10 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 13:24:37 +0800 Subject: [PATCH 04/14] Keep the MQTT Writer on paho's default keepalive The writer's connection behaviour is unchanged by this branch; only the broker policy call is added. Drops the keepalive argument, its test expectations and the changelog line. Co-Authored-By: Claude Fable 5.1 --- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index 3d238650c1..cbd68d3215 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -13,7 +13,6 @@ logger = logging.getLogger("inference") from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( - MQTT_KEEPALIVE_SECONDS, ConfigurationError, resolve_broker_address, ) @@ -53,9 +52,7 @@ with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` entries) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (when False, the workflow's host and port are ignored and the first allowlist entry is used); -a run the policy forbids is reported like any other failure. The connection -uses a 15 s keepalive, so a broker that disappears without closing the -connection is noticed within about 25 s. +a run the policy forbids is reported like any other failure. Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. @@ -347,7 +344,7 @@ def _connect_and_publish( # the TCP connect happens here, bounded by _connect_timeout # (DNS resolution is not - it runs under the OS resolver # timeout); the CONNACK wait below covers the handshake rest - client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) + client.connect(host, port) client.loop_start() except OSError as e: # broker unreachable: nothing is kept and no loop was started, From a0d32179a9cc18a14917d2a7c8a15021b452b2a0 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 13:25:04 +0800 Subject: [PATCH 05/14] Drop the writer keepalive expectations and changelog line Completes the previous commit: the writer tests no longer expect a keepalive argument and the changelog no longer lists the change. Co-Authored-By: Claude Fable 5.1 --- workflows/CHANGELOG.md | 4 ---- .../core_steps/sinks/test_mqtt_writer.py | 20 ++++++------------- 2 files changed, 6 insertions(+), 18 deletions(-) diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 21146792a2..bbceae38a9 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -7,10 +7,6 @@ - MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. -### Changed - -- MQTT Writer connects with a 15 s keepalive instead of paho's 60 s default, so a broker that disappears without closing the connection is reported within about 25 s. - --- ## `0.1.2` diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index 91308c32ba..b70e09179f 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -206,9 +206,7 @@ def test_selector_supplied_port_coerced_like_manifest( result = block.run(**run_kwargs(port=port)) assert result["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with( - "localhost", 1883, keepalive=15 - ) + mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) @pytest.mark.parametrize("qos", [-1, 3, "abc", 1.5, True, None]) def test_invalid_qos_rejected_before_client_construction( @@ -274,7 +272,7 @@ def test_client_connects_synchronously_before_background_loop( assert result["error_status"] is False mock_client_cls.assert_called_once_with(userdata=block._connected) - mock_client.connect.assert_called_once_with("localhost", 1883, keepalive=15) + mock_client.connect.assert_called_once_with("localhost", 1883) mock_client.connect_async.assert_not_called() called_methods = [call[0] for call in mock_client.method_calls] assert called_methods.index("connect") < called_methods.index("loop_start") @@ -426,7 +424,7 @@ def test_changed_connection_parameters_rejected(self, mock_client_cls, block): assert result["error_status"] is True assert "parameters" in result["message"].lower() - mock_client.connect.assert_called_once_with("broker-a", 1883, keepalive=15) + mock_client.connect.assert_called_once_with("broker-a", 1883) mock_client.publish.assert_called_once() def test_changed_credentials_rejected(self, mock_client_cls, block): @@ -633,9 +631,7 @@ def test_close_clears_readiness_only_after_loop_thread_joined( class TestBrokerPolicy: - def test_allowlisted_broker_is_used_and_keepalive_passed( - self, mock_client_cls, block, monkeypatch - ): + def test_allowlisted_broker_is_used(self, mock_client_cls, block, monkeypatch): monkeypatch.setattr( mqtt_common, "MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS", ["localhost:1883"] ) @@ -644,9 +640,7 @@ def test_allowlisted_broker_is_used_and_keepalive_passed( result = block.run(**run_kwargs()) assert result["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with( - "localhost", 1883, keepalive=15 - ) + mock_client_cls.return_value.connect.assert_called_once_with("localhost", 1883) def test_unlisted_broker_rejected_before_client_construction( self, mock_client_cls, block, monkeypatch @@ -690,9 +684,7 @@ def test_operator_broker_replaces_workflow_host_when_user_host_not_allowed( assert first["error_status"] is False assert second["error_status"] is False - mock_client_cls.return_value.connect.assert_called_once_with( - "operator", 8883, keepalive=15 - ) + mock_client_cls.return_value.connect.assert_called_once_with("operator", 8883) def test_user_host_not_allowed_without_operator_broker_disables_block( self, mock_client_cls, block, monkeypatch From 550aa48b186a905bec7e775073109801f7911241 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 13:28:56 +0800 Subject: [PATCH 06/14] Mention the MQTT Reader keepalive in its changelog entry Co-Authored-By: Claude Fable 5.1 --- workflows/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index bbceae38a9..27ed7d6577 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added -- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. +- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. --- From 2dbe9a9acb9f966832f639572ce92b21e291d114 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 16:06:06 +0800 Subject: [PATCH 07/14] Add TLS to the MQTT Reader and MQTT Writer blocks Both blocks gain an `encryption` dropdown (`none`, default, or `tls`) and a `ca_certificate_path` shown only for TLS. TLS wraps the connection before connect() and verifies the broker's certificate against the system trust store, or against the given PEM bundle for a private CA. The CA path is a server-side path chosen by the workflow, so it is gated by the engine's `allow_access_to_file_system` init parameter exactly as the Kafka blocks gate `ssl_ca_location`. Verification cannot be disabled and the port is not switched automatically. Both TLS inputs are part of the connection identity. A shared `configure_tls()` in mqtt_common carries the logic; the fake test broker can wrap accepted connections with a server SSL context, a fixture generates a throwaway CA and `localhost` certificate per test with `cryptography` (mirroring the mTLS enforcement test), and the broker drops a connection whose first byte is not an MQTT CONNECT, as a real broker does, so a TLS handshake against a plain broker fails fast in tests. Co-Authored-By: Claude Fable 5.1 --- .../integration_tests/execution/conftest.py | 90 ++++++++++- .../test_workflow_with_mqtt_reader.py | 75 +++++++++ .../test_workflow_with_mqtt_writer.py | 32 ++++ workflows/CHANGELOG.md | 1 + .../enterprise_blocks/sinks/mqtt_common.py | 49 +++++- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 72 ++++++++- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 86 +++++++++- .../core_steps/sinks/test_mqtt_common.py | 89 +++++++++++ .../core_steps/sinks/test_mqtt_reader.py | 144 ++++++++++++++++- .../core_steps/sinks/test_mqtt_writer.py | 150 ++++++++++++++++++ 10 files changed, 778 insertions(+), 10 deletions(-) diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index 3b47810a04..caffd35d93 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -1,9 +1,12 @@ import os import os.path import socket +import ssl import tempfile import threading -from typing import Generator +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from typing import Generator, Optional import cv2 import numpy as np @@ -219,6 +222,7 @@ def __init__( listening: bool = True, suback_reason_code: int = 0, keep_serving: bool = False, + tls_context: Optional[ssl.SSLContext] = None, ): # Bind to "localhost" for maximum performance, as described in: # http://docs.python.org/howto/sockets.html#ipc @@ -235,6 +239,10 @@ def __init__( # keep_serving: ignore recv timeouts and the message count, serve until # finish() / drop_connection(); needed for subscriber tests self.keep_serving = keep_serving + # tls_context: a server-side SSLContext; the accepted connection is + # wrapped with it before any MQTT packet is read + self.tls_context = tls_context + self.handshake_failures = 0 self.subscriptions = [] self.retained = {} self.connections_accepted = 0 @@ -260,6 +268,15 @@ def start(self): conn, address = self._sock.accept() conn.settimeout(1) + if self.tls_context is not None: + try: + conn = self.tls_context.wrap_socket(conn, server_side=True) + except (ssl.SSLError, OSError): + # the client refused our certificate or never spoke TLS: + # treat it like a closed connection + self.handshake_failures += 1 + conn.close() + return self._conn = conn self.connections_accepted += 1 self._stop = False @@ -282,6 +299,13 @@ def start(self): print(f"Received {chunk}") if not chunk: break + # exact byte: CONNECT's flags nibble is 0; a TLS ClientHello starts 0x16 + if not connack_sent and not buffer and chunk[0] != MQTT_CONNECT: + # not an MQTT CONNECT (e.g. a TLS ClientHello against this plain + # broker): a real broker drops the connection on the protocol + # error, which is what lets the client fail fast + self.drop_connection() + return buffer += chunk packets, buffer = _split_mqtt_packets(buffer) try: @@ -406,6 +430,70 @@ def _send(self, data: bytes) -> None: conn.sendall(data) +def _write_test_certificate(tmp_path, name, issuer=None, dns_names=()): + """Self-signed CA (no issuer) or a leaf signed by `issuer`, written as PEM. + + Mirrors `certificate()` in tests/inference/unit_tests/core/interfaces/http/ + test_mtls_enforcement.py; the leaf carries subjectAltName entries so paho's + hostname verification accepts `localhost`. + """ + from cryptography import x509 + from cryptography.hazmat.primitives import hashes, serialization + from cryptography.hazmat.primitives.asymmetric import ec + from cryptography.x509.oid import NameOID + + key = ec.generate_private_key(ec.SECP256R1()) + subject = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, name)]) + now = datetime.now(timezone.utc) + builder = ( + x509.CertificateBuilder() + .subject_name(subject) + .issuer_name(issuer[0].subject if issuer else subject) + .public_key(key.public_key()) + .serial_number(x509.random_serial_number()) + .not_valid_before(now - timedelta(minutes=1)) + .not_valid_after(now + timedelta(days=1)) + .add_extension( + x509.BasicConstraints(ca=issuer is None, path_length=None), critical=True + ) + ) + if dns_names: + builder = builder.add_extension( + x509.SubjectAlternativeName( + [x509.DNSName(dns_name) for dns_name in dns_names] + ), + critical=False, + ) + cert = builder.sign(issuer[1] if issuer else key, hashes.SHA256()) + cert_path, key_path = tmp_path / (name + ".pem"), tmp_path / (name + ".key") + cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM)) + key_path.write_bytes( + key.private_bytes( + serialization.Encoding.PEM, + serialization.PrivateFormat.PKCS8, + serialization.NoEncryption(), + ) + ) + return cert, key, cert_path, key_path + + +@pytest.fixture(scope="function") +def mqtt_test_certificates(tmp_path): + """A throwaway CA and a `localhost` server certificate for TLS broker tests. + + Yields `ca_path` (PEM the blocks are pointed at through `ca_certificate_path`) + and `server_context` (what `FakeMQTTBroker(tls_context=...)` wraps with). + """ + pytest.importorskip("cryptography") + ca = _write_test_certificate(tmp_path, "mqtt-test-ca") + server = _write_test_certificate( + tmp_path, "mqtt-test-server", issuer=ca, dns_names=("localhost",) + ) + server_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + server_context.load_cert_chain(certfile=str(server[2]), keyfile=str(server[3])) + return SimpleNamespace(ca_path=str(ca[2]), server_context=server_context) + + @pytest.fixture(scope="function") def fake_mqtt_broker(): print("Setup broker") diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index bce1bbf2b9..f813b82047 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -233,6 +233,81 @@ def test_unreachable_broker_leaves_nothing_behind_and_later_run_recovers(): thread.join(timeout=2) +@pytest.fixture +def tls_broker(mqtt_test_certificates): + broker = FakeMQTTBroker( + keep_serving=True, tls_context=mqtt_test_certificates.server_context + ) + thread = threading.Thread(target=broker.start, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_retained_message_over_tls_with_ca_certificate( + tls_broker, mqtt_test_certificates +): + # given + tls_broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + **reader_kwargs( + tls_broker, + encryption="tls", + ca_certificate_path=mqtt_test_certificates.ca_path, + ) + ) + + # then + assert result["payload"] == {"state": "RUNNING"} + assert result["is_new"] is True + assert tls_broker.connections_accepted == 1 + assert tls_broker.handshake_failures == 0 + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_tls_without_ca_certificate_is_rejected_by_verification(tls_broker): + # given - the test CA is not in the system trust store + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(tls_broker, encryption="tls")) + + # then - the client refused the certificate; nothing is kept + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is None + assert wait_until(lambda: tls_broker.handshake_failures == 1) + finally: + block.close() + + +@pytest.mark.timeout(15) +def test_tls_against_plain_broker_reports_not_connected(broker): + # given - the broker never speaks TLS, so the handshake waits for a + # ServerHello that never comes and paho's connect times out + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker, encryption="tls", timeout=1.0)) + + # then + assert result["error_status"] is True + assert "not connected" in result["error_message"].lower() + assert block._client is None + finally: + block.close() + + GATED_WORKFLOW = { "version": "1.0", "inputs": [ diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py index cf313ce021..a96d6b018a 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py @@ -248,6 +248,38 @@ def test_second_broker_is_rejected_instead_of_publishing_to_first(fake_mqtt_brok other_broker.finish() +@pytest.mark.timeout(15) +def test_publish_over_tls_with_ca_certificate(mqtt_test_certificates): + # given + broker = FakeMQTTBroker(tls_context=mqtt_test_certificates.server_context) + broker.messages_count_to_wait_for = 1 + broker_thread = threading.Thread(target=broker.start) + broker_thread.start() + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="encrypted payload", + timeout=2.0, + encryption="tls", + ca_certificate_path=mqtt_test_certificates.ca_path, + ) + broker_thread.join(timeout=2) + + # then + assert result["error_status"] is False + assert result["message"] == "Message published successfully" + assert b"encrypted payload" in broker.messages[-1] + assert broker.handshake_failures == 0 + finally: + block.close() + broker.finish() + + MQTT_SINK_WORKFLOW = { "version": "1.0", "inputs": [ diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 27ed7d6577..8682d0aa13 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -5,6 +5,7 @@ ### Added - MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. +- TLS for the MQTT Reader and MQTT Writer: an `encryption` dropdown (`none`, default, or `tls`) encrypts the connection and verifies the broker's certificate against the system trust store; `ca_certificate_path` (shown for TLS only) points at a PEM bundle for a private CA and requires `ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`. The port is not switched automatically and verification cannot be disabled. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. --- diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py index da25fda589..d6afa1ad82 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -8,7 +8,7 @@ """ import logging -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple from roboflow_workflows.environment import ( MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, @@ -24,7 +24,52 @@ class ConfigurationError(ValueError): - """The operator's policy forbids the requested broker connection.""" + """The operator's policy or the deployment's permissions forbid the connection.""" + + +def configure_tls( + client: Any, + *, + use_tls: bool, + ca_certificate_path: Optional[str], + allow_access_to_file_system: bool, +) -> None: + """Enable server-verified TLS on a paho client before it connects. + + Args: + client: The paho ``Client`` to configure; left untouched when ``use_tls`` + is False or when an error is raised. + use_tls: Whether the connection must be wrapped in TLS. + ca_certificate_path: Optional PEM bundle for a broker whose issuer is not in + the system trust store. Ignored when ``use_tls`` is False. + allow_access_to_file_system: The engine's file-system permission; a CA path + is a server-side path chosen by the workflow and is refused without it. + + Raises: + ConfigurationError: When the CA path is set but file-system access is + disabled, or when the CA bundle cannot be loaded. + """ + if not use_tls: + return + if ca_certificate_path and not allow_access_to_file_system: + # the path is chosen by the workflow and read by the server process + raise ConfigurationError( + "ca_certificate_path needs access to the local file system, which is " + "disabled on this deployment (ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=False)." + ) + try: + if ca_certificate_path: + client.tls_set(ca_certs=str(ca_certificate_path)) + else: + # system trust store; hostname verification stays on + client.tls_set() + except Exception as e: + raise ConfigurationError( + f"TLS could not be configured: could not load CA bundle " + f"{ca_certificate_path!r} ({e})." + if ca_certificate_path + else f"TLS could not be configured: {e}." + ) from e def _normalise_host(host: str) -> str: diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index 6739c459f7..dc9ab104c3 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -16,6 +16,7 @@ from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( MQTT_KEEPALIVE_SECONDS, ConfigurationError, + configure_tls, resolve_broker_address, ) from roboflow_workflows.environment import GCP_SERVERLESS, LAMBDA @@ -87,6 +88,14 @@ reported in the outputs. The connection uses a 15 s keepalive, so a broker that disappears without closing the connection is noticed within about 25 s. +Set `encryption` to `tls` to encrypt the connection and verify the broker's certificate +against the system trust store; the port is not switched automatically, so set it to the +broker's TLS port (usually 8883). A broker signed by a private CA needs +`ca_certificate_path`, a PEM bundle on the machine running inference, which requires +local file system access for Workflow blocks +(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). Certificate verification cannot be +disabled. + Outputs: - value (str): Raw message payload decoded as UTF-8, or None when nothing was received yet. @@ -102,6 +111,7 @@ """ SUBSCRIPTION_REFUSED_QOS = 0x80 +TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} LATEST_BUFFER_SIZE = 1 SEQUENTIAL_BUFFER_SIZE = 1000 @@ -275,6 +285,32 @@ class BlockManifest(WorkflowBlockManifest): "number greater than 0.", examples=[0.5], ) + encryption: Literal["none", "tls"] = Field( + default="none", + description="Transport security for the broker connection. `tls` encrypts the " + "connection and verifies the broker's certificate; set port to the broker's " + "TLS port (usually 8883).", + examples=["none", "tls"], + json_schema_extra={ + "values_metadata": { + "none": {"name": "None", "description": "Plain TCP (default)."}, + "tls": { + "name": "TLS", + "description": "Encrypted, broker certificate verified.", + }, + }, + }, + ) + ca_certificate_path: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Path to a PEM CA bundle on the machine running inference, for " + "brokers whose certificate issuer is not in the system trust store. Requires " + "local file system access for Workflow blocks " + "(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`); otherwise the run " + "reports an error. Used only when encryption is `tls`.", + examples=["/etc/ssl/certs/factory-ca.pem"], + json_schema_extra={"relevant_for": TLS_RELEVANT}, + ) @classmethod def describe_outputs(cls) -> List[OutputDefinition]: @@ -322,7 +358,10 @@ def get_restrictions(cls) -> List[RuntimeRestriction]: class MQTTReaderBlockV1(WorkflowBlock): - def __init__(self): + def __init__(self, allow_access_to_file_system: bool = False): + # gates `ca_certificate_path`, a server-side path chosen by the workflow; + # False by default so a hand-constructed block is safe + self._allow_access_to_file_system = allow_access_to_file_system self._client: Optional[mqtt.Client] = None self._state: Optional[MQTTReaderState] = None self._connection_identity: Optional[Tuple] = None @@ -371,7 +410,7 @@ def get_manifest(cls) -> Type[WorkflowBlockManifest]: @classmethod def get_init_parameters(cls) -> List[str]: - return [] + return ["allow_access_to_file_system"] def run( self, @@ -383,6 +422,8 @@ def run( username: Optional[str] = None, password: Optional[str] = None, timeout: float = 0.5, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, ) -> BlockResult: if GCP_SERVERLESS or LAMBDA: # never open an outbound broker connection from hosted workers @@ -443,6 +484,16 @@ def run( return self._handle_failure( "Password provided without username. Set username to enable MQTT authentication." ) + if encryption not in ("none", "tls"): + return self._handle_failure( + f"Invalid encryption: {encryption!r}. Must be 'none' or 'tls'." + ) + # an empty editor field is "unset"; anything but a string is a wiring error + if ca_certificate_path is not None and not isinstance(ca_certificate_path, str): + return self._handle_failure( + f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path." + ) + ca_certificate_path = ca_certificate_path or None try: # the operator's broker policy; the RESOLVED address is what the # block connects to and what the connection identity is built from @@ -461,6 +512,8 @@ def run( username=username, password=password, timeout=timeout, + encryption=encryption, + ca_certificate_path=ca_certificate_path, ) def _connect_and_read( @@ -473,6 +526,8 @@ def _connect_and_read( username: Optional[str], password: Optional[str], timeout: float, + encryption: str, + ca_certificate_path: Optional[str], ) -> BlockResult: connection_identity = ( host, @@ -483,6 +538,8 @@ def _connect_and_read( topic, qos, read_mode, + encryption, + ca_certificate_path, ) if self._client is None: buffer_size = ( @@ -496,6 +553,14 @@ def _connect_and_read( client.suppress_exceptions = True if username is not None: client.username_pw_set(username, password) + # TLS must be configured before connect(); the CA path is gated + # by the engine's file-system permission + configure_tls( + client, + use_tls=encryption == "tls", + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=self._allow_access_to_file_system, + ) client.on_connect = mqtt_on_connect client.on_connect_fail = mqtt_on_connect_fail client.on_subscribe = mqtt_on_subscribe @@ -514,6 +579,9 @@ def _connect_and_read( # timeout); the CONNACK wait below covers the handshake rest client.connect(host, port, keepalive=MQTT_KEEPALIVE_SECONDS) client.loop_start() + except ConfigurationError as e: + # TLS setup refused before any socket was opened: nothing kept + return self._handle_failure(str(e)) except OSError as e: # broker unreachable: nothing is kept and no loop was started, # so a failed one-shot run leaves no background thread behind; diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index cbd68d3215..a2e8b8e4bb 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -14,6 +14,7 @@ from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( ConfigurationError, + configure_tls, resolve_broker_address, ) from roboflow_workflows.execution_engine.entities.base import OutputDefinition @@ -54,6 +55,14 @@ workflow's host and port are ignored and the first allowlist entry is used); a run the policy forbids is reported like any other failure. +Set `encryption` to `tls` to encrypt the connection and verify the broker's +certificate against the system trust store; the port is not switched +automatically, so set it to the broker's TLS port (usually 8883). A broker +signed by a private CA needs `ca_certificate_path`, a PEM bundle on the +machine running inference, which requires local file system access for +Workflow blocks (`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). +Certificate verification cannot be disabled. + Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. True if there was an error, False if successful. @@ -70,6 +79,9 @@ """ +TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} + + def mqtt_on_connect(client, userdata, flags, reason_code, properties=None): # paho invokes on_connect for accepted and rejected CONNACK alike; # only reason_code 0 means an established MQTT session @@ -166,6 +178,32 @@ class BlockManifest(WorkflowBlockManifest): "requests; in streaming pipelines it stops processing.", examples=[False], ) + encryption: Literal["none", "tls"] = Field( + default="none", + description="Transport security for the broker connection. `tls` encrypts the " + "connection and verifies the broker's certificate; set port to the broker's " + "TLS port (usually 8883).", + examples=["none", "tls"], + json_schema_extra={ + "values_metadata": { + "none": {"name": "None", "description": "Plain TCP (default)."}, + "tls": { + "name": "TLS", + "description": "Encrypted, broker certificate verified.", + }, + }, + }, + ) + ca_certificate_path: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Path to a PEM CA bundle on the machine running inference, for " + "brokers whose certificate issuer is not in the system trust store. Requires " + "local file system access for Workflow blocks " + "(`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`); otherwise the run " + "reports an error. Used only when encryption is `tls`.", + examples=["/etc/ssl/certs/factory-ca.pem"], + json_schema_extra={"relevant_for": TLS_RELEVANT}, + ) @classmethod def describe_outputs(cls) -> List[OutputDefinition]: @@ -180,7 +218,12 @@ def get_execution_engine_compatibility(cls) -> Optional[str]: class MQTTWriterSinkBlockV1(WorkflowBlock): - def __init__(self, disable_sinks: bool = False): + def __init__( + self, disable_sinks: bool = False, allow_access_to_file_system: bool = False + ): + # gates `ca_certificate_path`, a server-side path chosen by the workflow; + # False by default so a hand-constructed block is safe + self._allow_access_to_file_system = allow_access_to_file_system self.mqtt_client: Optional[mqtt.Client] = None self._connected = threading.Event() self._connection_identity: Optional[Tuple] = None @@ -219,7 +262,7 @@ def get_manifest(cls) -> Type[WorkflowBlockManifest]: @classmethod def get_init_parameters(cls) -> List[str]: - return ["disable_sinks"] + return ["disable_sinks", "allow_access_to_file_system"] def run( self, @@ -233,6 +276,8 @@ def run( retain: bool = False, timeout: float = 0.5, fail_fast: bool = False, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, ) -> BlockResult: if self._disable_sinks: return disabled_sink_response() @@ -288,6 +333,18 @@ def run( "Password provided without username. Set username to enable MQTT authentication.", fail_fast=fail_fast, ) + if encryption not in ("none", "tls"): + return self._handle_failure( + f"Invalid encryption: {encryption!r}. Must be 'none' or 'tls'.", + fail_fast=fail_fast, + ) + # an empty editor field is "unset"; anything but a string is a wiring error + if ca_certificate_path is not None and not isinstance(ca_certificate_path, str): + return self._handle_failure( + f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path.", + fail_fast=fail_fast, + ) + ca_certificate_path = ca_certificate_path or None try: # the operator's broker policy; the RESOLVED address is what the # block connects to and what the connection identity is built from @@ -308,6 +365,8 @@ def run( retain=retain, timeout=timeout, fail_fast=fail_fast, + encryption=encryption, + ca_certificate_path=ca_certificate_path, ) def _connect_and_publish( @@ -322,14 +381,32 @@ def _connect_and_publish( retain: bool, timeout: float, fail_fast: bool, + encryption: str = "none", + ca_certificate_path: Optional[str] = None, ) -> BlockResult: - connection_identity = (host, port, username, password, timeout) + connection_identity = ( + host, + port, + username, + password, + timeout, + encryption, + ca_certificate_path, + ) if self.mqtt_client is None: client = None try: client = mqtt.Client(userdata=self._connected) if username is not None: client.username_pw_set(username, password) + # TLS must be configured before connect(); the CA path is gated + # by the engine's file-system permission + configure_tls( + client, + use_tls=encryption == "tls", + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=self._allow_access_to_file_system, + ) client.on_connect = mqtt_on_connect client.on_connect_fail = mqtt_on_connect_fail client.on_disconnect = mqtt_on_disconnect @@ -346,6 +423,9 @@ def _connect_and_publish( # timeout); the CONNACK wait below covers the handshake rest client.connect(host, port) client.loop_start() + except ConfigurationError as e: + # TLS setup refused before any socket was opened: nothing kept + return self._handle_failure(str(e), fail_fast=fail_fast) except OSError as e: # broker unreachable: nothing is kept and no loop was started, # so a failed one-shot run leaves no background thread behind; diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py index d865122adc..459fc87708 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -1,4 +1,5 @@ import logging +from unittest.mock import MagicMock import pytest from roboflow_workflows.enterprise_blocks.sinks import mqtt_common @@ -6,6 +7,7 @@ DEFAULT_MQTT_PORT, MQTT_KEEPALIVE_SECONDS, ConfigurationError, + configure_tls, normalise_broker_address, resolve_broker_address, split_host_port, @@ -147,3 +149,90 @@ def test_user_host_not_allowed_without_operator_broker_is_disabled(policy): with pytest.raises(ConfigurationError, match="disabled"): resolve_broker_address("workflow.host", 1883) + + +class TestConfigureTLS: + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + @pytest.mark.parametrize("allowed", [True, False]) + def test_disabled_tls_touches_nothing(self, ca_certificate_path, allowed): + client = MagicMock() + + configure_tls( + client, + use_tls=False, + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=allowed, + ) + + assert client.method_calls == [] + + @pytest.mark.parametrize("ca_certificate_path", [None, ""]) + def test_tls_without_ca_path_uses_system_trust_store(self, ca_certificate_path): + client = MagicMock() + + configure_tls( + client, + use_tls=True, + ca_certificate_path=ca_certificate_path, + allow_access_to_file_system=False, + ) + + client.tls_set.assert_called_once_with() + client.tls_insecure_set.assert_not_called() + + def test_ca_path_used_with_file_system_access(self): + client = MagicMock() + + configure_tls( + client, + use_tls=True, + ca_certificate_path="/etc/ssl/factory-ca.pem", + allow_access_to_file_system=True, + ) + + client.tls_set.assert_called_once_with(ca_certs="/etc/ssl/factory-ca.pem") + client.tls_insecure_set.assert_not_called() + + def test_ca_path_refused_without_file_system_access(self): + client = MagicMock() + + with pytest.raises(ConfigurationError) as error: + configure_tls( + client, + use_tls=True, + ca_certificate_path="/etc/passwd", + allow_access_to_file_system=False, + ) + + assert "ca_certificate_path" in str(error.value) + assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in str(error.value) + assert client.method_calls == [] + + def test_unloadable_ca_bundle_reported(self): + client = MagicMock() + client.tls_set.side_effect = FileNotFoundError("no such file") + + with pytest.raises( + ConfigurationError, match="could not load CA bundle" + ) as error: + configure_tls( + client, + use_tls=True, + ca_certificate_path="/missing/ca.pem", + allow_access_to_file_system=True, + ) + + assert "/missing/ca.pem" in str(error.value) + assert "no such file" in str(error.value) + + def test_system_store_failure_reported(self): + client = MagicMock() + client.tls_set.side_effect = ValueError("bad context") + + with pytest.raises(ConfigurationError, match="TLS could not be configured"): + configure_tls( + client, + use_tls=True, + ca_certificate_path=None, + allow_access_to_file_system=False, + ) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 351fd284c7..3cb59a3ba8 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -67,6 +67,8 @@ class FakeClient: def __init__(self, userdata=None): self.userdata = userdata + # ordered record of the calls that matter for TLS: must precede connect() + self.calls = [] self.on_connect = None self.on_connect_fail = None self.on_subscribe = None @@ -87,7 +89,14 @@ def username_pw_set(self, username, password=None): def reconnect_delay_set(self, min_delay, max_delay): self.reconnect_delays = (min_delay, max_delay) + def tls_set(self, ca_certs=None, **kwargs): + self.calls.append(("tls_set", ca_certs)) + + def tls_insecure_set(self, value): + self.calls.append(("tls_insecure_set", value)) + def connect(self, host, port, keepalive=60): + self.calls.append(("connect", host, port)) if self.connect_error is not None: raise self.connect_error self.connected_to = (host, port) @@ -214,8 +223,10 @@ def test_restrictions_hard_for_hosted_serverless_only(self): ] assert restrictions[1].applies_to_input_modes == [v1.RuntimeInputMode.IMAGE] - def test_block_declares_no_init_parameters(self): - assert MQTTReaderBlockV1.get_init_parameters() == [] + def test_block_declares_file_system_init_parameter_only(self): + assert MQTTReaderBlockV1.get_init_parameters() == [ + "allow_access_to_file_system" + ] class TestCallbacks: @@ -863,3 +874,132 @@ def test_user_host_not_allowed_without_operator_broker_disables_block( assert result["error_status"] is True assert "disabled" in result["error_message"] assert clients == [] + + +class TestTLS: + def test_manifest_defaults_and_dropdown_values(self): + manifest = BlockManifest.model_validate(manifest_payload()) + + assert manifest.encryption == "none" + assert manifest.ca_certificate_path is None + assert ( + BlockManifest.model_validate(manifest_payload(encryption="tls")).encryption + == "tls" + ) + + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None, ""]) + def test_manifest_rejects_other_encryption_values(self, encryption): + with pytest.raises(ValidationError): + BlockManifest.model_validate(manifest_payload(encryption=encryption)) + + def test_ca_certificate_path_is_shown_only_for_tls(self): + schema = BlockManifest.model_json_schema()["properties"] + + assert schema["ca_certificate_path"]["relevant_for"] == { + "encryption": {"values": ["tls"], "required": True} + } + + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + def test_default_encryption_never_configures_tls( + self, clients, block, ca_certificate_path + ): + result = block.run(**run_kwargs(ca_certificate_path=ca_certificate_path)) + + assert result["error_status"] is False + assert [name for name, *_ in clients[0].calls] == ["connect"] + + def test_tls_without_ca_path_uses_system_store_before_connect(self, clients, block): + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + assert clients[0].calls == [ + ("tls_set", None), + ("connect", "localhost", 1883), + ] + + def test_tls_with_ca_path_and_file_system_access(self, clients): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/ssl/ca.pem") + ) + + assert result["error_status"] is False + assert clients[0].calls == [ + ("tls_set", "/etc/ssl/ca.pem"), + ("connect", "localhost", 1883), + ] + + def test_tls_with_ca_path_refused_without_file_system_access(self, clients, block): + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd") + ) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["error_message"] + assert ( + "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in result["error_message"] + ) + assert block._client is None + assert clients[0].calls == [] + assert clients[0].loop_started is False + + def test_unloadable_ca_bundle_reported_and_nothing_kept(self, clients, monkeypatch): + def failing_tls_set(client, ca_certs=None, **kwargs): + raise FileNotFoundError("no such file") + + monkeypatch.setattr(FakeClient, "tls_set", failing_tls_set) + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/missing.pem") + ) + + assert result["error_status"] is True + assert "could not load CA bundle" in result["error_message"] + assert "/missing.pem" in result["error_message"] + assert block._client is None + assert clients[0].loop_started is False + + def test_next_run_after_tls_failure_retries_with_fresh_client(self, clients, block): + block.run(**run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd")) + + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + assert len(clients) == 2 + + @pytest.mark.parametrize( + "change", + [{"encryption": "tls"}, {"ca_certificate_path": "/other.pem"}], + ) + def test_changed_tls_parameters_rejected(self, clients, change): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + + def test_invalid_encryption_at_run_time_rejected(self, clients, block): + result = block.run(**run_kwargs(encryption="ssl")) + + assert result["error_status"] is True + assert "encryption" in result["error_message"] + assert clients == [] + + def test_non_string_ca_path_rejected(self, clients, block): + result = block.run(**run_kwargs(encryption="tls", ca_certificate_path=5)) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["error_message"] + assert clients == [] + + def test_tls_insecure_is_never_called(self, clients): + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + block.run(**run_kwargs(encryption="tls", ca_certificate_path="/ca.pem")) + + assert all(name != "tls_insecure_set" for name, *_ in clients[0].calls) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index b70e09179f..99630965f3 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -701,3 +701,153 @@ def test_user_host_not_allowed_without_operator_broker_disables_block( assert result["error_status"] is True assert "disabled" in result["message"] mock_client_cls.assert_not_called() + + +def _call_names(mock_client) -> list: + return [call[0] for call in mock_client.method_calls] + + +class TestTLS: + def test_manifest_defaults_and_relevant_for(self): + manifest = BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v1", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + } + ) + schema = BlockManifest.model_json_schema()["properties"] + + assert manifest.encryption == "none" + assert manifest.ca_certificate_path is None + assert schema["ca_certificate_path"]["relevant_for"] == { + "encryption": {"values": ["tls"], "required": True} + } + + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None]) + def test_manifest_rejects_other_encryption_values(self, encryption): + with pytest.raises(ValidationError): + BlockManifest.model_validate( + { + "type": "roboflow_enterprise/mqtt_writer_sink@v1", + "name": "mqtt", + "host": "localhost", + "port": 1883, + "topic": "test/topic", + "message": "Hello, MQTT!", + "encryption": encryption, + } + ) + + def test_init_parameters_include_file_system_access(self): + assert MQTTWriterSinkBlockV1.get_init_parameters() == [ + "disable_sinks", + "allow_access_to_file_system", + ] + + @pytest.mark.parametrize("ca_certificate_path", [None, "", "/ca.pem"]) + def test_default_encryption_never_configures_tls( + self, mock_client_cls, block, ca_certificate_path + ): + block._connected.set() + + result = block.run(**run_kwargs(ca_certificate_path=ca_certificate_path)) + + assert result["error_status"] is False + mock_client_cls.return_value.tls_set.assert_not_called() + mock_client_cls.return_value.tls_insecure_set.assert_not_called() + + def test_tls_without_ca_path_uses_system_store_before_connect( + self, mock_client_cls, block + ): + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run(**run_kwargs(encryption="tls")) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with() + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + mock_client.tls_insecure_set.assert_not_called() + + def test_tls_with_ca_path_and_file_system_access(self, mock_client_cls): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + block._connected.set() + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/ssl/ca.pem") + ) + + assert result["error_status"] is False + mock_client.tls_set.assert_called_once_with(ca_certs="/etc/ssl/ca.pem") + names = _call_names(mock_client) + assert names.index("tls_set") < names.index("connect") + + def test_tls_with_ca_path_refused_without_file_system_access( + self, mock_client_cls, block + ): + mock_client = mock_client_cls.return_value + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/etc/passwd") + ) + + assert result["error_status"] is True + assert "ca_certificate_path" in result["message"] + assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in result["message"] + assert block.mqtt_client is None + mock_client.tls_set.assert_not_called() + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + def test_tls_refused_with_fail_fast_raises(self, mock_client_cls, block): + with pytest.raises(RuntimeError, match="ca_certificate_path"): + block.run( + **run_kwargs( + encryption="tls", ca_certificate_path="/etc/passwd", fail_fast=True + ) + ) + + assert block.mqtt_client is None + + def test_unloadable_ca_bundle_reported_and_nothing_kept(self, mock_client_cls): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + mock_client = mock_client_cls.return_value + mock_client.tls_set.side_effect = FileNotFoundError("no such file") + + result = block.run( + **run_kwargs(encryption="tls", ca_certificate_path="/missing.pem") + ) + + assert result["error_status"] is True + assert "could not load CA bundle" in result["message"] + assert block.mqtt_client is None + mock_client.connect.assert_not_called() + mock_client.loop_start.assert_not_called() + + @pytest.mark.parametrize( + "change", + [{"encryption": "tls"}, {"ca_certificate_path": "/other.pem"}], + ) + def test_changed_tls_parameters_rejected(self, mock_client_cls, change): + block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) + block._connected.set() + block.run(**run_kwargs()) + + result = block.run(**run_kwargs(**change)) + + assert result["error_status"] is True + assert "changed between runs" in result["message"] + mock_client_cls.assert_called_once() + + def test_invalid_encryption_at_run_time_rejected(self, mock_client_cls, block): + result = block.run(**run_kwargs(encryption="ssl")) + + assert result["error_status"] is True + assert "encryption" in result["message"] + mock_client_cls.assert_not_called() From 016a9c6f14d292834e1bdd63cce9b94298458529 Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 16:27:14 +0800 Subject: [PATCH 08/14] Harden MQTT TLS after review configure_tls() refuses a CA path that is not a regular file before OpenSSL opens it: the loader blocks forever on a FIFO while the block holds its lifecycle lock. The refusal hint names the engine init parameter as well as the env flag. A synchronous connect failure with TLS enabled now points at the port and CA bundle instead of advising a longer timeout, and the changed-parameters message and both descriptions list the TLS inputs. The writer's description states the 60 s keepalive bound on a handshake against a silent peer. Tests use real temporary files for the CA path and one integration case feeds paho a real non-PEM bundle. Co-Authored-By: Claude Fable 5.1 --- .../test_workflow_with_mqtt_reader.py | 31 +++++++++++- .../enterprise_blocks/sinks/mqtt_common.py | 15 +++++- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 24 +++++++--- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 30 +++++++++--- .../core_steps/sinks/test_mqtt_common.py | 47 ++++++++++++++----- .../core_steps/sinks/test_mqtt_reader.py | 29 ++++++++---- .../core_steps/sinks/test_mqtt_writer.py | 26 +++++++--- 7 files changed, 158 insertions(+), 44 deletions(-) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index f813b82047..5dd49ae05f 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -284,16 +284,43 @@ def test_tls_without_ca_certificate_is_rejected_by_verification(tls_broker): # then - the client refused the certificate; nothing is kept assert result["error_status"] is True assert "not connected" in result["error_message"].lower() + assert "TLS is enabled" in result["error_message"] assert block._client is None assert wait_until(lambda: tls_broker.handshake_failures == 1) finally: block.close() +@pytest.mark.timeout(15) +def test_tls_with_unreadable_ca_bundle_is_reported_by_paho(tls_broker, tmp_path): + # given - a real file that is not a PEM bundle: this exercises paho's + # loader, not a mock + bad_bundle = tmp_path / "not-a-bundle.pem" + bad_bundle.write_text("this is not a certificate") + block = MQTTReaderBlockV1(allow_access_to_file_system=True) + + try: + # when + result = block.run( + **reader_kwargs( + tls_broker, encryption="tls", ca_certificate_path=str(bad_bundle) + ) + ) + + # then + assert result["error_status"] is True + assert "could not load CA bundle" in result["error_message"] + assert block._client is None + assert tls_broker.connections_accepted == 0 + finally: + block.close() + + @pytest.mark.timeout(15) def test_tls_against_plain_broker_reports_not_connected(broker): - # given - the broker never speaks TLS, so the handshake waits for a - # ServerHello that never comes and paho's connect times out + # given - the broker never speaks TLS; like a real broker it drops a + # connection whose first bytes are not an MQTT CONNECT, so the client's + # handshake fails at once block = MQTTReaderBlockV1() try: diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py index d6afa1ad82..a89122fbda 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -4,10 +4,13 @@ `resolve_broker_address()` applies the operator's policy from `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` and `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (see `roboflow_workflows.environment`) -before either block builds a client. +before either block builds a client, and `configure_tls()` enables +server-verified TLS on the client, gating a workflow-chosen CA path behind the +engine's file-system permission. """ import logging +import os from typing import Any, List, Optional, Tuple from roboflow_workflows.environment import ( @@ -55,7 +58,15 @@ def configure_tls( # the path is chosen by the workflow and read by the server process raise ConfigurationError( "ca_certificate_path needs access to the local file system, which is " - "disabled on this deployment (ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=False)." + "disabled on this deployment (ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=False " + "or the engine's allow_access_to_file_system init parameter)." + ) + if ca_certificate_path and not os.path.isfile(ca_certificate_path): + # OpenSSL's loader blocks forever on a FIFO and errors on a directory; + # refuse anything that is not a regular file before it is opened + raise ConfigurationError( + f"TLS could not be configured: ca_certificate_path {ca_certificate_path!r} " + "is not a readable file." ) try: if ca_certificate_path: diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index dc9ab104c3..b6221f136d 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -73,8 +73,8 @@ the block was not subscribed are not delivered. One block instance subscribes on a single broker connection: changing host, port, -credentials, timeout, topic, QoS or read mode between runs is rejected as a -configuration error. While the broker is unreachable every run waits up to `timeout` +credentials, timeout, topic, QoS, read mode or the TLS settings between runs is +rejected as a configuration error. While the broker is unreachable every run waits up to `timeout` for the background reconnect before reporting the failure. The block is not available on the Roboflow hosted platform (`GCP_SERVERLESS` or @@ -147,6 +147,19 @@ def __init__(self, topic: str, qos: int, buffer_size: int): self.messages: Deque[ReceivedMessage] = deque(maxlen=buffer_size) +def _not_connected_message(error: BaseException, encryption: str) -> str: + if encryption == "tls": + return ( + f"MQTT broker not connected ({error}). TLS is enabled: check that port is " + "the broker's TLS port and that ca_certificate_path matches the broker's " + "certificate authority." + ) + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " + "longer to connect." + ) + + def mqtt_on_connect( client, state: MQTTReaderState, flags, reason_code, properties=None ): @@ -587,8 +600,7 @@ def _connect_and_read( # so a failed one-shot run leaves no background thread behind; # the next run on this instance retries with a fresh client return self._handle_failure( - f"MQTT broker not connected ({e}). Raise 'timeout' if the " - "broker needs longer to connect." + _not_connected_message(e, encryption=encryption) ) except Exception as e: if client is not None: @@ -604,8 +616,8 @@ def _connect_and_read( elif connection_identity != self._connection_identity: return self._handle_failure( "MQTT connection parameters (host, port, credentials, timeout, topic, " - "qos or read_mode) changed between runs; this block subscribes only " - "with the configuration of its first run." + "qos, read_mode, encryption or ca_certificate_path) changed between " + "runs; this block subscribes only with the configuration of its first run." ) state = self._state # the background loop owns (re)connecting and (re)subscribing; runs diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index a2e8b8e4bb..a13c2bad68 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -46,8 +46,8 @@ background loop. Over the HTTP API every request builds a fresh block instance, so each request pays the bounded connect and a failed request is final for that request. One block instance publishes to a single broker -connection: changing host, port, credentials or timeout between runs is -rejected as a configuration error. +connection: changing host, port, credentials, timeout or the TLS settings +between runs is rejected as a configuration error. The server operator may restrict which brokers the MQTT blocks connect to with `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (an allowlist of `host[:port]` @@ -61,7 +61,10 @@ signed by a private CA needs `ca_certificate_path`, a PEM bundle on the machine running inference, which requires local file system access for Workflow blocks (`ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`). -Certificate verification cannot be disabled. +Certificate verification cannot be disabled. A peer that accepts the TCP +connection but never completes the TLS handshake is given up on after the +MQTT keepalive interval (60 s), not after `timeout`; real brokers close such +a connection immediately. Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. @@ -82,6 +85,19 @@ TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} +def _not_connected_message(error: BaseException, encryption: str) -> str: + if encryption == "tls": + return ( + f"MQTT broker not connected ({error}). TLS is enabled: check that port is " + "the broker's TLS port and that ca_certificate_path matches the broker's " + "certificate authority." + ) + return ( + f"MQTT broker not connected ({error}). Raise 'timeout' if the broker needs " + "longer to connect." + ) + + def mqtt_on_connect(client, userdata, flags, reason_code, properties=None): # paho invokes on_connect for accepted and rejected CONNACK alike; # only reason_code 0 means an established MQTT session @@ -431,8 +447,7 @@ def _connect_and_publish( # so a failed one-shot run leaves no background thread behind; # the next run on this instance retries with a fresh client return self._handle_failure( - f"MQTT broker not connected ({e}). Raise 'timeout' if the " - "broker needs longer to connect.", + _not_connected_message(e, encryption=encryption), fail_fast=fail_fast, ) except Exception as e: @@ -448,8 +463,9 @@ def _connect_and_publish( self._connection_identity = connection_identity elif connection_identity != self._connection_identity: return self._handle_failure( - "MQTT connection parameters (host, port, credentials or timeout) changed " - "between runs; this block publishes only to the connection configured " + "MQTT connection parameters (host, port, credentials, timeout, encryption " + "or ca_certificate_path) changed between runs; this block publishes only " + "to the connection configured on its first run." "on its first run.", fail_fast=fail_fast, ) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py index 459fc87708..adefec8814 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -1,4 +1,5 @@ import logging +import os from unittest.mock import MagicMock import pytest @@ -180,17 +181,19 @@ def test_tls_without_ca_path_uses_system_trust_store(self, ca_certificate_path): client.tls_set.assert_called_once_with() client.tls_insecure_set.assert_not_called() - def test_ca_path_used_with_file_system_access(self): + def test_ca_path_used_with_file_system_access(self, tmp_path): client = MagicMock() + ca_path = tmp_path / "factory-ca.pem" + ca_path.write_text("content is irrelevant: paho is faked here") configure_tls( client, use_tls=True, - ca_certificate_path="/etc/ssl/factory-ca.pem", + ca_certificate_path=str(ca_path), allow_access_to_file_system=True, ) - client.tls_set.assert_called_once_with(ca_certs="/etc/ssl/factory-ca.pem") + client.tls_set.assert_called_once_with(ca_certs=str(ca_path)) client.tls_insecure_set.assert_not_called() def test_ca_path_refused_without_file_system_access(self): @@ -208,22 +211,44 @@ def test_ca_path_refused_without_file_system_access(self): assert "ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE" in str(error.value) assert client.method_calls == [] - def test_unloadable_ca_bundle_reported(self): + @pytest.mark.parametrize("kind", ["missing", "directory", "fifo"]) + def test_ca_path_must_be_a_regular_file(self, tmp_path, kind): + # OpenSSL's loader blocks forever on a FIFO; the guard runs before it client = MagicMock() - client.tls_set.side_effect = FileNotFoundError("no such file") + if kind == "missing": + path = tmp_path / "missing.pem" + elif kind == "directory": + path = tmp_path + else: + path = tmp_path / "pipe" + os.mkfifo(path) + + with pytest.raises(ConfigurationError, match="is not a readable file"): + configure_tls( + client, + use_tls=True, + ca_certificate_path=str(path), + allow_access_to_file_system=True, + ) + + client.tls_set.assert_not_called() + + def test_unloadable_ca_bundle_reported(self, tmp_path): + client = MagicMock() + client.tls_set.side_effect = ValueError("no certificate or crl found") + ca_path = tmp_path / "not-a-pem.pem" + ca_path.write_text("garbage") - with pytest.raises( - ConfigurationError, match="could not load CA bundle" - ) as error: + with pytest.raises(ConfigurationError, match="could not load CA bundle") as e: configure_tls( client, use_tls=True, - ca_certificate_path="/missing/ca.pem", + ca_certificate_path=str(ca_path), allow_access_to_file_system=True, ) - assert "/missing/ca.pem" in str(error.value) - assert "no such file" in str(error.value) + assert str(ca_path) in str(e.value) + assert "no certificate or crl found" in str(e.value) def test_system_store_failure_reported(self): client = MagicMock() diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 3cb59a3ba8..69403b9788 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -876,6 +876,15 @@ def test_user_host_not_allowed_without_operator_broker_disables_block( assert clients == [] +@pytest.fixture +def ca_file(tmp_path): + # a real regular file: configure_tls() refuses anything else before paho + # (faked in these tests) would read it + path = tmp_path / "ca.pem" + path.write_text("content is irrelevant: paho is faked in these tests") + return path + + class TestTLS: def test_manifest_defaults_and_dropdown_values(self): manifest = BlockManifest.model_validate(manifest_payload()) @@ -917,16 +926,16 @@ def test_tls_without_ca_path_uses_system_store_before_connect(self, clients, blo ("connect", "localhost", 1883), ] - def test_tls_with_ca_path_and_file_system_access(self, clients): + def test_tls_with_ca_path_and_file_system_access(self, clients, ca_file): block = MQTTReaderBlockV1(allow_access_to_file_system=True) result = block.run( - **run_kwargs(encryption="tls", ca_certificate_path="/etc/ssl/ca.pem") + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) ) assert result["error_status"] is False assert clients[0].calls == [ - ("tls_set", "/etc/ssl/ca.pem"), + ("tls_set", str(ca_file)), ("connect", "localhost", 1883), ] @@ -944,7 +953,9 @@ def test_tls_with_ca_path_refused_without_file_system_access(self, clients, bloc assert clients[0].calls == [] assert clients[0].loop_started is False - def test_unloadable_ca_bundle_reported_and_nothing_kept(self, clients, monkeypatch): + def test_unloadable_ca_bundle_reported_and_nothing_kept( + self, clients, monkeypatch, ca_file + ): def failing_tls_set(client, ca_certs=None, **kwargs): raise FileNotFoundError("no such file") @@ -952,12 +963,12 @@ def failing_tls_set(client, ca_certs=None, **kwargs): block = MQTTReaderBlockV1(allow_access_to_file_system=True) result = block.run( - **run_kwargs(encryption="tls", ca_certificate_path="/missing.pem") + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) ) assert result["error_status"] is True assert "could not load CA bundle" in result["error_message"] - assert "/missing.pem" in result["error_message"] + assert str(ca_file) in result["error_message"] assert block._client is None assert clients[0].loop_started is False @@ -971,7 +982,7 @@ def test_next_run_after_tls_failure_retries_with_fresh_client(self, clients, blo @pytest.mark.parametrize( "change", - [{"encryption": "tls"}, {"ca_certificate_path": "/other.pem"}], + [{"encryption": "tls"}, {"ca_certificate_path": "other.pem"}], ) def test_changed_tls_parameters_rejected(self, clients, change): block = MQTTReaderBlockV1(allow_access_to_file_system=True) @@ -997,9 +1008,9 @@ def test_non_string_ca_path_rejected(self, clients, block): assert "ca_certificate_path" in result["error_message"] assert clients == [] - def test_tls_insecure_is_never_called(self, clients): + def test_tls_insecure_is_never_called(self, clients, ca_file): block = MQTTReaderBlockV1(allow_access_to_file_system=True) - block.run(**run_kwargs(encryption="tls", ca_certificate_path="/ca.pem")) + block.run(**run_kwargs(encryption="tls", ca_certificate_path=str(ca_file))) assert all(name != "tls_insecure_set" for name, *_ in clients[0].calls) diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index 99630965f3..d9cde3b188 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -707,6 +707,15 @@ def _call_names(mock_client) -> list: return [call[0] for call in mock_client.method_calls] +@pytest.fixture +def ca_file(tmp_path): + # a real regular file: configure_tls() refuses anything else before paho + # (faked in these tests) would read it + path = tmp_path / "ca.pem" + path.write_text("content is irrelevant: paho is faked in these tests") + return path + + class TestTLS: def test_manifest_defaults_and_relevant_for(self): manifest = BlockManifest.model_validate( @@ -727,7 +736,7 @@ def test_manifest_defaults_and_relevant_for(self): "encryption": {"values": ["tls"], "required": True} } - @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None]) + @pytest.mark.parametrize("encryption", ["ssl", "TLS", True, None, ""]) def test_manifest_rejects_other_encryption_values(self, encryption): with pytest.raises(ValidationError): BlockManifest.model_validate( @@ -774,17 +783,17 @@ def test_tls_without_ca_path_uses_system_store_before_connect( assert names.index("tls_set") < names.index("connect") mock_client.tls_insecure_set.assert_not_called() - def test_tls_with_ca_path_and_file_system_access(self, mock_client_cls): + def test_tls_with_ca_path_and_file_system_access(self, mock_client_cls, ca_file): block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) block._connected.set() mock_client = mock_client_cls.return_value result = block.run( - **run_kwargs(encryption="tls", ca_certificate_path="/etc/ssl/ca.pem") + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) ) assert result["error_status"] is False - mock_client.tls_set.assert_called_once_with(ca_certs="/etc/ssl/ca.pem") + mock_client.tls_set.assert_called_once_with(ca_certs=str(ca_file)) names = _call_names(mock_client) assert names.index("tls_set") < names.index("connect") @@ -815,24 +824,27 @@ def test_tls_refused_with_fail_fast_raises(self, mock_client_cls, block): assert block.mqtt_client is None - def test_unloadable_ca_bundle_reported_and_nothing_kept(self, mock_client_cls): + def test_unloadable_ca_bundle_reported_and_nothing_kept( + self, mock_client_cls, ca_file + ): block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) mock_client = mock_client_cls.return_value mock_client.tls_set.side_effect = FileNotFoundError("no such file") result = block.run( - **run_kwargs(encryption="tls", ca_certificate_path="/missing.pem") + **run_kwargs(encryption="tls", ca_certificate_path=str(ca_file)) ) assert result["error_status"] is True assert "could not load CA bundle" in result["message"] + assert str(ca_file) in result["message"] assert block.mqtt_client is None mock_client.connect.assert_not_called() mock_client.loop_start.assert_not_called() @pytest.mark.parametrize( "change", - [{"encryption": "tls"}, {"ca_certificate_path": "/other.pem"}], + [{"encryption": "tls"}, {"ca_certificate_path": "other.pem"}], ) def test_changed_tls_parameters_rejected(self, mock_client_cls, change): block = MQTTWriterSinkBlockV1(allow_access_to_file_system=True) From 5712a0212ed2ced06e5c2ca7b223faa69f2ae81c Mon Sep 17 00:00:00 2001 From: Tewan Date: Mon, 21 Sep 2026 16:49:17 +0800 Subject: [PATCH 09/14] Stabilise MQTT TLS integration tests on CI The TLS handshake, CONNACK, SUBACK and retained message each get a 5 s budget instead of 2 s, and the tests report the block's error message when a run fails. The fake broker exits quietly when finish() closes its listening socket while it is still in accept(), which a test whose block fails before connecting leaves behind. Co-Authored-By: Claude Fable 5.1 --- tests/workflows/integration_tests/execution/conftest.py | 7 ++++++- .../execution/test_workflow_with_mqtt_reader.py | 4 ++++ .../execution/test_workflow_with_mqtt_writer.py | 4 ++-- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index caffd35d93..0b6632e0fa 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -266,7 +266,12 @@ def start(self): if self._conn is not None: raise ValueError("Connection is already open") - conn, address = self._sock.accept() + try: + conn, address = self._sock.accept() + except OSError: + # finish() closed the listening socket while no client had + # connected (a test whose block fails before connecting) + return conn.settimeout(1) if self.tls_context is not None: try: diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index 5dd49ae05f..b320c590bb 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -260,10 +260,14 @@ def test_retained_message_over_tls_with_ca_certificate( tls_broker, encryption="tls", ca_certificate_path=mqtt_test_certificates.ca_path, + # the TLS handshake, CONNACK, SUBACK and the retained message + # each get this budget; CI runners are slower than a laptop + timeout=5.0, ) ) # then + assert result["error_status"] is False, result["error_message"] assert result["payload"] == {"state": "RUNNING"} assert result["is_new"] is True assert tls_broker.connections_accepted == 1 diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py index a96d6b018a..c8b8db6867 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py @@ -264,14 +264,14 @@ def test_publish_over_tls_with_ca_certificate(mqtt_test_certificates): port=broker.port, topic="RoboflowTopic", message="encrypted payload", - timeout=2.0, + timeout=5.0, encryption="tls", ca_certificate_path=mqtt_test_certificates.ca_path, ) broker_thread.join(timeout=2) # then - assert result["error_status"] is False + assert result["error_status"] is False, result["message"] assert result["message"] == "Message published successfully" assert b"encrypted payload" in broker.messages[-1] assert broker.handshake_failures == 0 From 087d151a91516b3c85d2b79e0446e2a03e9dd184 Mon Sep 17 00:00:00 2001 From: Tewan Date: Wed, 23 Sep 2026 16:16:50 +0800 Subject: [PATCH 10/14] Add persistent sessions to the MQTT Reader via an optional client_id When client_id is set the reader connects with that id and a non-clean session, so the broker keeps the subscription and queues QoS 1/2 messages while the block is away and delivers the backlog when the same id reconnects. Blank ids are unset, a persistent session needs qos 1 or 2, the id joins the connection identity and the connect log reports the broker's session-present flag. The fake broker used by the integration tests gains an opt-in session-aware mode: CONNECT parsing, per-id sessions with an offline queue, session-present CONNACK, same-id eviction, acknowledged and routed client publishes and multi-connection serving. The default mode is unchanged. Co-Authored-By: Claude Fable 5.1 --- .../integration_tests/execution/conftest.py | 455 ++++++++++++++---- .../test_workflow_with_mqtt_reader.py | 204 ++++++++ workflows/CHANGELOG.md | 1 + .../enterprise_blocks/sinks/mqtt_common.py | 30 +- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 80 ++- .../core_steps/sinks/test_mqtt_common.py | 22 + .../core_steps/sinks/test_mqtt_reader.py | 113 ++++- 7 files changed, 791 insertions(+), 114 deletions(-) diff --git a/tests/workflows/integration_tests/execution/conftest.py b/tests/workflows/integration_tests/execution/conftest.py index 0b6632e0fa..9f57e0021e 100644 --- a/tests/workflows/integration_tests/execution/conftest.py +++ b/tests/workflows/integration_tests/execution/conftest.py @@ -4,9 +4,10 @@ import ssl import tempfile import threading +from collections import deque from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Generator, Optional +from typing import Dict, Generator, Optional, Tuple import cv2 import numpy as np @@ -215,7 +216,70 @@ def mqtt_topic_matches(topic_filter: str, topic: str) -> bool: return len(filter_levels) == len(topic_levels) +class _Session: + """What a broker keeps for one client id: subscriptions, an offline queue and + the connection currently bound to it (None while the client is away).""" + + def __init__(self, client_id: str, clean: bool): + self.client_id = client_id + self.clean = clean + self.subscriptions: list = [] + self.queue: deque = deque() + self.connection: Optional["_Connection"] = None + + def matches(self, topic: str) -> Optional[int]: + """Highest granted QoS among the subscriptions matching `topic`, or None.""" + granted = [ + qos + for topic_filter, qos in self.subscriptions + if mqtt_topic_matches(topic_filter, topic) + ] + return max(granted) if granted else None + + +class _Connection: + """Per-connection state: the socket, its receive buffer and packet counter.""" + + def __init__(self, sock): + self.sock = sock + self.buffer = b"" + self.connack_sent = False + self.stop = False + self.next_packet_id = 1 + self.session: Optional[_Session] = None + self.send_lock = threading.Lock() + + def close(self) -> None: + self.stop = True + sock, self.sock = self.sock, None + if sock is None: + return + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + sock.close() + + class FakeMQTTBroker: + """A minimal MQTT 3.1.1 broker for tests. + + By default it serves one connection per `start()` call, answers CONNECT, + SUBSCRIBE and PINGREQ, re-sends retained messages on subscription and + records everything else (including a client's PUBLISH, which is never + acknowledged) in `messages`. + + `session_aware=True` adds what the persistent-session tests need: CONNECT is + parsed (client id, clean-session flag), sessions are kept per id in + `sessions`, a returning id gets CONNACK "session present" and its queued + messages, a second connection with a live id evicts the first, a client's + PUBLISH is acknowledged and routed to subscribers (queued for an offline + persistent session when published at QoS 1 or 2) and `serve()` accepts any + number of concurrent connections. Delivery is capped at the granted QoS, and + `suback_reason_code` is that cap (0 by default), so a persistent-session test + passes `suback_reason_code=2` for the broker to queue anything. + """ + def __init__( self, connack_reason_code: int = 0, @@ -223,6 +287,7 @@ def __init__( suback_reason_code: int = 0, keep_serving: bool = False, tls_context: Optional[ssl.SSLContext] = None, + session_aware: bool = False, ): # Bind to "localhost" for maximum performance, as described in: # http://docs.python.org/howto/sockets.html#ipc @@ -242,36 +307,43 @@ def __init__( # tls_context: a server-side SSLContext; the accepted connection is # wrapped with it before any MQTT packet is read self.tls_context = tls_context + self.session_aware = session_aware self.handshake_failures = 0 self.subscriptions = [] self.retained = {} self.connections_accepted = 0 - self._stop = False - self._send_lock = threading.Lock() + # session_aware only: (client_id, clean_session) per CONNECT, in order + self.connects = [] + # session_aware only: live and remembered sessions by client id + self.sessions: Dict[str, _Session] = {} + self._anonymous_sessions = 0 + self._connections: list = [] + self._closed = False + self._lock = threading.RLock() sock.settimeout(5) self._sock = sock - self._conn = None # bound but not listening: connection attempts are refused until listen() if listening: self.listen() def listen(self): - self._sock.listen(1) + self._sock.listen(5) - def start(self): - if self._sock is None: - raise ValueError("Socket is not open") - if self._conn is not None: - raise ValueError("Connection is already open") + # -- accepting ----------------------------------------------------------- + def _accept(self) -> Optional[_Connection]: + sock = self._sock + if sock is None: + raise ValueError("Socket is not open") try: - conn, address = self._sock.accept() + conn, address = sock.accept() except OSError: # finish() closed the listening socket while no client had - # connected (a test whose block fails before connecting) - return + # connected (a test whose block fails before connecting), or the + # accept timed out + return None conn.settimeout(1) if self.tls_context is not None: try: @@ -281,88 +353,185 @@ def start(self): # treat it like a closed connection self.handshake_failures += 1 conn.close() - return - self._conn = conn - self.connections_accepted += 1 - self._stop = False - connack_sent = False - buffer = b"" - while not self._stop and ( - self.keep_serving - or len(self.messages) < self.messages_count_to_wait_for - or not connack_sent - ): - try: - chunk = self.receive_packet(1000) - except socket.timeout: - if self.keep_serving: - continue - raise - except (OSError, ValueError): - # connection closed, possibly from the test thread - break - print(f"Received {chunk}") - if not chunk: - break - # exact byte: CONNECT's flags nibble is 0; a TLS ClientHello starts 0x16 - if not connack_sent and not buffer and chunk[0] != MQTT_CONNECT: - # not an MQTT CONNECT (e.g. a TLS ClientHello against this plain - # broker): a real broker drops the connection on the protocol - # error, which is what lets the client fail fast - self.drop_connection() - return - buffer += chunk - packets, buffer = _split_mqtt_packets(buffer) - try: - for packet in packets: - if self._handle_packet(packet): - connack_sent = True - except (OSError, ValueError): - break + return None + connection = _Connection(conn) + with self._lock: + self._connections.append(connection) + self.connections_accepted += 1 + return connection - def _handle_packet(self, packet: bytes) -> bool: + def start(self): + """Accept one connection and serve it in the calling thread.""" + if self._sock is None: + raise ValueError("Socket is not open") + if any(c.sock is not None for c in self._connections): + raise ValueError("Connection is already open") + connection = self._accept() + if connection is None: + return + self._serve_connection(connection) + + def serve(self): + """Accept connections until finish(), each served on its own thread.""" + while not self._closed: + connection = self._accept() + if connection is None: + continue + threading.Thread( + target=self._serve_connection, args=(connection,), daemon=True + ).start() + + def _serve_connection(self, connection: _Connection) -> None: + try: + while ( + not self._closed + and not connection.stop + and ( + self.keep_serving + or len(self.messages) < self.messages_count_to_wait_for + or not connection.connack_sent + ) + ): + try: + chunk = self._receive(connection, 1000) + except socket.timeout: + if self.keep_serving: + continue + raise + except (OSError, ValueError): + # connection closed, possibly from the test thread + break + print(f"Received {chunk}") + if not chunk: + break + # exact byte: CONNECT's flags nibble is 0; a TLS ClientHello starts 0x16 + if ( + not connection.connack_sent + and not connection.buffer + and chunk[0] != MQTT_CONNECT + ): + # not an MQTT CONNECT (e.g. a TLS ClientHello against this plain + # broker): a real broker drops the connection on the protocol + # error, which is what lets the client fail fast + break + connection.buffer += chunk + packets, connection.buffer = _split_mqtt_packets(connection.buffer) + try: + for packet in packets: + if self._handle_packet(connection, packet): + connection.connack_sent = True + except (OSError, ValueError): + break + finally: + self._release(connection) + + def _release(self, connection: _Connection) -> None: + with self._lock: + connection.close() + if connection in self._connections: + self._connections.remove(connection) + session = connection.session + if session is not None and session.connection is connection: + # a clean session ends with its connection; a persistent one + # stays, with its subscriptions and queue, until the id returns + session.connection = None + if session.clean: + self.sessions.pop(session.client_id, None) + + # -- packets ------------------------------------------------------------- + + def _handle_packet(self, connection: _Connection, packet: bytes) -> bool: """Handle one control packet; returns True when a CONNACK was sent.""" packet_type = packet[0] & 0xF0 if packet_type == MQTT_CONNECT: - print("sending CONNACK") - self._send(b"\x20\x02\x00" + bytes([self.connack_reason_code])) + self._handle_connect(connection, packet) return True if packet_type == MQTT_SUBSCRIBE: - self._handle_subscribe(packet) + self._handle_subscribe(connection, packet) return False if packet_type == MQTT_PINGREQ: - self._send(bytes([MQTT_PINGRESP, 0])) + self._send(connection, bytes([MQTT_PINGRESP, 0])) return False if packet_type in (MQTT_PUBACK, MQTT_PUBREC, MQTT_PUBREL, MQTT_PUBCOMP): # QoS 1/2 handshakes, not messages: a broker-initiated QoS 2 PUBLISH # is answered with PUBREC and completed with PUBREL; a client's # PUBREL (its own QoS 2 publish) is completed with PUBCOMP if packet_type == MQTT_PUBREC: - self._send(bytes([MQTT_PUBREL | 0x02, 2]) + packet[2:4]) + self._send(connection, bytes([MQTT_PUBREL | 0x02, 2]) + packet[2:4]) if packet_type == MQTT_PUBREL: - self._send(bytes([MQTT_PUBCOMP, 2]) + packet[2:4]) + self._send(connection, bytes([MQTT_PUBCOMP, 2]) + packet[2:4]) return False if packet_type == MQTT_DISCONNECT: return False self.messages.append(packet) + if packet_type == MQTT_PUBLISH and self.session_aware: + self._handle_publish(connection, packet) return False - def _handle_subscribe(self, packet: bytes) -> None: - # skip fixed header (type byte + remaining-length varint) + @staticmethod + def _body(packet: bytes) -> bytes: + """Strip the fixed header (type byte + remaining-length varint).""" index = 1 while packet[index] & 0x80: index += 1 - index += 1 - packet_id = packet[index : index + 2] + return packet[index + 1 :] + + @staticmethod + def _read_string(body: bytes, index: int) -> Tuple[str, int]: + length = int.from_bytes(body[index : index + 2], "big") index += 2 + return body[index : index + length].decode("utf-8"), index + length + + def _handle_connect(self, connection: _Connection, packet: bytes) -> None: + print("sending CONNACK") + if not self.session_aware: + self._send(connection, b"\x20\x02\x00" + bytes([self.connack_reason_code])) + return + body = self._body(packet) + _, index = self._read_string(body, 0) # protocol name + index += 1 # protocol level + clean_session = bool(body[index] & 0x02) + index += 3 # connect flags + keepalive + client_id, _ = self._read_string(body, index) + self.connects.append((client_id, clean_session)) + if self.connack_reason_code != 0: + self._send(connection, b"\x20\x02\x00" + bytes([self.connack_reason_code])) + return + with self._lock: + if not client_id: + # a broker assigns an id to an anonymous (clean) client + self._anonymous_sessions += 1 + client_id = f"anonymous-{self._anonymous_sessions}" + session = self.sessions.get(client_id) + if session is not None and session.connection is not None: + # a second connection with a live id: the broker drops the first + session.connection.close() + session.connection = None + session_present = session is not None and not clean_session + if not session_present: + session = _Session(client_id, clean=clean_session) + self.sessions[client_id] = session + session.connection = connection + connection.session = session + backlog = list(session.queue) + session.queue.clear() + self._send(connection, b"\x20\x02" + bytes([int(session_present), 0])) + # a real broker delivers the queued messages right after CONNACK, before + # the client's SUBSCRIBE arrives + for topic, payload, qos in backlog: + self._publish_to(connection, topic, payload, qos=qos) + + def _handle_subscribe(self, connection: _Connection, packet: bytes) -> None: + body = self._body(packet) + packet_id = body[0:2] + index = 2 granted = [] - while index < len(packet): - topic_length = int.from_bytes(packet[index : index + 2], "big") - index += 2 - topic_filter = packet[index : index + topic_length].decode("utf-8") - index += topic_length - requested_qos = packet[index] + requested = [] + while index < len(body): + topic_filter, index = self._read_string(body, index) + requested_qos = body[index] index += 1 + requested.append((topic_filter, requested_qos)) self.subscriptions.append((topic_filter, requested_qos)) if self.suback_reason_code == MQTT_SUBACK_FAILURE: granted.append(MQTT_SUBACK_FAILURE) @@ -374,65 +543,155 @@ def _handle_subscribe(self, packet: bytes) -> None: + packet_id + bytes(granted) ) - self._send(suback) + self._send(connection, suback) if self.suback_reason_code == MQTT_SUBACK_FAILURE: return - for topic_filter, _ in self.subscriptions[-len(granted) :]: - for topic, payload in self.retained.items(): + if connection.session is not None: + with self._lock: + for (topic_filter, _), granted_qos in zip(requested, granted): + connection.session.subscriptions = [ + s + for s in connection.session.subscriptions + if s[0] != topic_filter + ] + [(topic_filter, granted_qos)] + for topic_filter, _ in requested: + for topic, payload in list(self.retained.items()): if mqtt_topic_matches(topic_filter, topic): - self.publish(topic, payload, retain=True) + self._publish_to(connection, topic, payload, retain=True) + + def _handle_publish(self, connection: _Connection, packet: bytes) -> None: + """Acknowledge a client's PUBLISH and route it to the subscribers.""" + flags = packet[0] & 0x0F + qos = (flags >> 1) & 0x03 + retain = bool(flags & 0x01) + body = self._body(packet) + topic, index = self._read_string(body, 0) + packet_id = b"" + if qos > 0: + packet_id = body[index : index + 2] + index += 2 + payload = body[index:] + if qos == 1: + self._send(connection, bytes([MQTT_PUBACK, 2]) + packet_id) + elif qos == 2: + # completed with PUBCOMP when the client's PUBREL arrives + self._send(connection, bytes([MQTT_PUBREC, 2]) + packet_id) + if retain: + self.retained[topic] = payload + self._route(topic, payload, qos) + + # -- publishing ---------------------------------------------------------- def publish( self, topic: str, payload: bytes, retain: bool = False, qos: int = 0 ) -> None: - """Send a PUBLISH packet from the broker to the connected client.""" + """Publish from the broker side: to the connected client, or, when + session aware, to every matching subscriber (queued for an offline + persistent session when `qos` is 1 or 2).""" + if isinstance(payload, str): + payload = payload.encode("utf-8") + if self.session_aware: + if retain: + self.retained[topic] = payload + self._route(topic, payload, qos, retain=retain) + return + connection = self._current_connection() + if connection is None: + raise ValueError("Connection is not open") + self._publish_to(connection, topic, payload, retain=retain, qos=qos) + + def _route( + self, topic: str, payload: bytes, qos: int, retain: bool = False + ) -> None: + with self._lock: + targets = [] + for session in list(self.sessions.values()): + granted = session.matches(topic) + if granted is None: + continue + effective_qos = min(qos, granted) + if session.connection is not None: + targets.append((session.connection, effective_qos)) + elif not session.clean and effective_qos > 0: + # the broker queues only QoS 1 and 2 for an offline session + session.queue.append((topic, payload, effective_qos)) + for connection, effective_qos in targets: + try: + self._publish_to( + connection, topic, payload, retain=retain, qos=effective_qos + ) + except (OSError, ValueError): + pass + + def _publish_to( + self, + connection: _Connection, + topic: str, + payload: bytes, + retain: bool = False, + qos: int = 0, + ) -> None: + """Send a PUBLISH packet from the broker to one client.""" if isinstance(payload, str): payload = payload.encode("utf-8") topic_bytes = topic.encode("utf-8") variable_header = len(topic_bytes).to_bytes(2, "big") + topic_bytes if qos > 0: - variable_header += (1).to_bytes(2, "big") + with connection.send_lock: + packet_id = connection.next_packet_id + connection.next_packet_id = packet_id % 65535 + 1 + variable_header += packet_id.to_bytes(2, "big") body = variable_header + payload header = bytes([MQTT_PUBLISH | (qos << 1) | int(retain)]) - self._send(header + _encode_remaining_length(len(body)) + body) + self._send(connection, header + _encode_remaining_length(len(body)) + body) + + # -- lifecycle ----------------------------------------------------------- + + def _current_connection(self) -> Optional[_Connection]: + with self._lock: + live = [c for c in self._connections if c.sock is not None] + return live[-1] if live else None def drop_connection(self) -> None: - """Close the client connection so a later start() can accept a new one.""" - self._stop = True - if self._conn is not None: - try: - self._conn.shutdown(socket.SHUT_RDWR) - except OSError: - pass - self._conn.close() - self._conn = None + """Close the (most recent) client connection so a later start() or the + serving loop can accept a new one.""" + connection = self._current_connection() + if connection is not None: + connection.close() def finish(self): - self._stop = True - if self._conn is not None: - self._conn.close() - self._conn = None + self._closed = True + with self._lock: + connections = list(self._connections) + for connection in connections: + connection.close() if self._sock is not None: self._sock.close() self._sock = None def receive_packet(self, num_bytes): + connection = self._current_connection() + if connection is None: + raise ValueError("Connection is not open") + return self._receive(connection, num_bytes) + + @staticmethod + def _receive(connection: _Connection, num_bytes: int) -> bytes: # capture locally: drop_connection()/finish() clear the attribute # from another thread - conn = self._conn - if conn is None: + sock = connection.sock + if sock is None: raise ValueError("Connection is not open") + return sock.recv(num_bytes) - packet_in = conn.recv(num_bytes) - return packet_in - - def _send(self, data: bytes) -> None: - with self._send_lock: - conn = self._conn - if conn is None: + @staticmethod + def _send(connection: _Connection, data: bytes) -> None: + with connection.send_lock: + sock = connection.sock + if sock is None: raise ValueError("Connection is not open") - conn.sendall(data) + sock.sendall(data) def _write_test_certificate(tmp_path, name, issuer=None, dns_names=()): diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index b320c590bb..1eb28e48f4 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -339,6 +339,210 @@ def test_tls_against_plain_broker_reports_not_connected(broker): block.close() +# --- persistent sessions (client_id) ------------------------------------------ + + +@pytest.fixture +def session_broker(): + # granted QoS caps delivery, so the broker must grant up to QoS 2 to queue + # anything for an offline persistent session + broker = FakeMQTTBroker(keep_serving=True, session_aware=True, suback_reason_code=2) + thread = threading.Thread(target=broker.serve, daemon=True) + thread.start() + yield broker + broker.finish() + thread.join(timeout=2) + + +def persistent_kwargs(broker: FakeMQTTBroker, **overrides) -> dict: + kwargs = reader_kwargs(broker, client_id="line1", qos=1, read_mode="sequential") + kwargs.update(overrides) + return kwargs + + +def session_is_away(broker: FakeMQTTBroker, client_id: str = "line1") -> bool: + session = broker.sessions.get(client_id) + return session is not None and session.connection is None + + +def subscribe_and_leave(broker: FakeMQTTBroker, **overrides) -> None: + """A block instance with a persistent session subscribes, then goes away.""" + block = MQTTReaderBlockV1() + try: + result = block.run(**persistent_kwargs(broker, **overrides)) + assert result["error_status"] is False, result["error_message"] + finally: + block.close() + assert wait_until(lambda: session_is_away(broker)) + + +@pytest.mark.timeout(20) +def test_persistent_session_delivers_backlog_after_restart(session_broker, caplog): + # given - a subscriber with a fixed id went away and three QoS 1 messages + # were published meanwhile + with caplog.at_level("INFO", logger="inference"): + subscribe_and_leave(session_broker) + assert session_broker.connects == [("line1", False)] + assert "session present: False" in caplog.text + caplog.clear() + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=1) + block = MQTTReaderBlockV1() + + try: + # when - a new instance (a restarted pipeline) resumes the session + with caplog.at_level("INFO", logger="inference"): + results = [block.run(**persistent_kwargs(session_broker)) for _ in range(3)] + + # then - the broker resumed the session and delivered the backlog in order + assert session_broker.connects[-1] == ("line1", False) + assert "session present: True" in caplog.text + assert [r["value"] for r in results] == ["A", "B", "C"] + assert [r["is_new"] for r in results] == [True, True, True] + assert [r["error_status"] for r in results] == [False, False, False] + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_in_latest_mode_keeps_newest_of_backlog(session_broker): + # given + subscribe_and_leave(session_broker, read_mode="latest") + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=1) + block = MQTTReaderBlockV1() + + try: + # when + first = block.run(**persistent_kwargs(session_broker, read_mode="latest")) + second = block.run(**persistent_kwargs(session_broker, read_mode="latest")) + + # then - only the newest message of the burst survives + assert first["value"] == "C" + assert first["is_new"] is True + assert second["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_does_not_queue_qos0_messages(session_broker): + # given - the publisher used QoS 0, which a broker never queues + subscribe_and_leave(session_broker) + for payload in (b"A", b"B", b"C"): + session_broker.publish("plc/state", payload, qos=0) + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**persistent_kwargs(session_broker)) + + # then - the session resumed, but there was nothing to deliver + assert session_broker.connects[-1] == ("line1", False) + assert result["error_status"] is False + assert result["value"] is None + assert result["is_new"] is False + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_persistent_session_survives_reconnect_of_the_same_instance(session_broker): + # given - a live persistent subscription + block = MQTTReaderBlockV1() + + try: + first = block.run(**persistent_kwargs(session_broker)) + assert first["error_status"] is False + + # when - the connection drops, a message is published while the block is + # away, and the background loop reconnects + session_broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + assert wait_until(lambda: session_is_away(session_broker)) + session_broker.publish("plc/state", PAUSED, qos=1) + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + result = block.run(**persistent_kwargs(session_broker)) + + # then - the queued message reached the same instance after the reconnect + assert session_broker.connections_accepted == 2 + assert session_broker.connects == [("line1", False), ("line1", False)] + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + block.close() + + +@pytest.mark.timeout(20) +def test_second_connection_with_the_same_client_id_evicts_the_first(session_broker): + # given - two block instances sharing one client id on one broker + first = MQTTReaderBlockV1() + second = MQTTReaderBlockV1() + + try: + assert first.run(**persistent_kwargs(session_broker))["error_status"] is False + + # when - the second connects with the same id + result = second.run(**persistent_kwargs(session_broker)) + + # then - the broker drops the first; its reconnect would evict the second + # in turn, so the first is closed here (only the first eviction is tested) + assert wait_until(lambda: not first._state.connected.is_set()) + first.close() + assert result["error_status"] is False + assert session_broker.connects[:2] == [("line1", False), ("line1", False)] + assert wait_until( + lambda: second._state.connected.is_set() + and session_broker.sessions["line1"].connection is not None, + timeout=10, + ) + session_broker.publish("plc/state", PAUSED, qos=1) + assert wait_until(lambda: len(second._state.messages) == 1, timeout=10) + updated = second.run(**persistent_kwargs(session_broker)) + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + finally: + first.close() + second.close() + + +@pytest.mark.timeout(20) +def test_writer_without_client_id_delivers_to_an_away_persistent_subscriber( + session_broker, +): + # given - a persistent subscriber that went away, and the unchanged MQTT Writer + from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( + MQTTWriterSinkBlockV1, + ) + + subscribe_and_leave(session_broker) + writer = MQTTWriterSinkBlockV1() + reader = MQTTReaderBlockV1() + + try: + # when - the writer publishes at QoS 1 with a broker-generated id + published = writer.run( + host=session_broker.host, + port=session_broker.port, + topic="plc/state", + message=PAUSED.decode("utf-8"), + qos=1, + timeout=2.0, + ) + writer.close() + result = reader.run(**persistent_kwargs(session_broker)) + + # then - queued delivery depends only on the publish QoS, not on a writer id + assert published["error_status"] is False, published["message"] + assert published["message"] == "Message published successfully" + assert session_broker.connects[1] == ("", True) + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + finally: + writer.close() + reader.close() + + GATED_WORKFLOW = { "version": "1.0", "inputs": [ diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 8682d0aa13..1fa793f521 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -5,6 +5,7 @@ ### Added - MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. +- Persistent sessions for the MQTT Reader: an optional `client_id` (unique per broker, for example the pipeline name from a workflow input) connects with that id and a non-clean session, so the broker keeps the subscription and queues QoS 1/2 messages while the block is away and delivers the backlog when the same id reconnects, after a pipeline or process restart included. Requires `qos` 1 or 2 (a run with `client_id` and `qos` 0 reports an error); the publisher must also publish at QoS 1 or higher. Meant for `InferencePipeline`s; `sequential` works through the backlog one message per run. Left empty, the behaviour is unchanged. - TLS for the MQTT Reader and MQTT Writer: an `encryption` dropdown (`none`, default, or `tls`) encrypts the connection and verifies the broker's certificate against the system trust store; `ca_certificate_path` (shown for TLS only) points at a PEM bundle for a private CA and requires `ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`. The port is not switched automatically and verification cannot be disabled. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py index a89122fbda..777757b7fe 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -6,7 +6,8 @@ `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (see `roboflow_workflows.environment`) before either block builds a client, and `configure_tls()` enables server-verified TLS on the client, gating a workflow-chosen CA path behind the -engine's file-system permission. +engine's file-system permission. `normalise_client_id()` turns the reader's +optional `client_id` into either a usable id or "unset". """ import logging @@ -83,6 +84,33 @@ def configure_tls( ) from e +def normalise_client_id(value: Any) -> Optional[str]: + """Normalise the reader's optional ``client_id`` input. + + A persistent MQTT session exists only for a fixed client id, so the presence + of the id is the switch. An empty or whitespace-only value is "unset", so a + stray space never creates a persistent session. + + Args: + value: The workflow-provided ``client_id``: ``None``, a literal or a + selector-resolved value. + + Returns: + The id with surrounding whitespace removed, or ``None`` when unset. + + Raises: + ConfigurationError: When the value is not a string (a wiring error). + """ + if value is None: + return None + if not isinstance(value, str): + raise ConfigurationError( + f"Invalid client_id: {value!r}. Must be a string or left empty." + ) + value = value.strip() + return value or None + + def _normalise_host(host: str) -> str: host = host.strip().lower() if host.startswith("[") and host.endswith("]"): diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index b6221f136d..ec9259101d 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -17,6 +17,7 @@ MQTT_KEEPALIVE_SECONDS, ConfigurationError, configure_tls, + normalise_client_id, resolve_broker_address, ) from roboflow_workflows.environment import GCP_SERVERLESS, LAMBDA @@ -69,12 +70,26 @@ The subscription and the buffer live in the block instance only. An InferencePipeline keeps them for its lifetime, a restart begins with an empty buffer, and over the HTTP API every request builds a fresh instance, so only a retained message can be returned -there. There is no persistent session and no durable resume: messages published while -the block was not subscribed are not delivered. +there. By default the session is not persistent either: messages published while the +block was not subscribed are not delivered. + +Set `client_id` to a stable name that is unique on the broker (for example the camera +or pipeline name from a workflow input) to keep a persistent session instead. The +presence of the id is the switch: the block then connects with that id and asks the +broker to keep the session, so the broker remembers the subscription and queues QoS 1 +and 2 messages while the block is away and delivers the backlog when a block with the +same id reconnects, after a pipeline or process restart included. `qos` must then be 1 +or 2 and the publisher must also publish at QoS 1 or higher, because a broker queues +nothing for QoS 0. Use `sequential` to work through the backlog one message per run; +`latest` keeps only the newest of it. QoS 1 may redeliver a message whose +acknowledgement was lost, and `sequential` surfaces that duplicate; QoS 2 avoids it. +Two connections with the same id disconnect each other, so never share an id between +pipelines. Persistent sessions are meant for InferencePipelines: over the HTTP API each +request would resume the session, consume the whole backlog and return one message. One block instance subscribes on a single broker connection: changing host, port, -credentials, timeout, topic, QoS, read mode or the TLS settings between runs is -rejected as a configuration error. While the broker is unreachable every run waits up to `timeout` +credentials, timeout, topic, QoS, read mode, client id or the TLS settings between runs +is rejected as a configuration error. While the broker is unreachable every run waits up to `timeout` for the background reconnect before reporting the failure. The block is not available on the Roboflow hosted platform (`GCP_SERVERLESS` or @@ -170,10 +185,14 @@ def mqtt_on_connect( state.connected.clear() return - logger.info("MQTT client connected") + logger.info( + "MQTT client connected (session present: %s)", + bool(flags.get("session present")), + ) state.connected.set() - # the broker forgets subscriptions on every reconnect (clean session), so - # the subscription is (re)established here rather than once after connect() + # a clean-session broker forgets subscriptions on every reconnect, so the + # subscription is (re)established here rather than once after connect(); on + # a persistent session the repeated SUBSCRIBE is idempotent state.subscribe_failed = False state.subscribed.clear() try: @@ -274,9 +293,23 @@ class BlockManifest(WorkflowBlockManifest): ) qos: Union[int, Selector(kind=[INTEGER_KIND])] = Field( default=0, - description="Quality of Service level requested for the subscription (0, 1 or 2).", + description="Quality of Service level requested for the subscription (0, 1 or 2). " + "When client_id is set, qos must be 1 or 2: the broker only queues QoS 1 and 2 " + "messages for a persistent session, and a run with client_id and qos 0 reports " + "an error. The publisher must also publish at QoS 1 or higher for its messages " + "to be queued.", examples=[0, 1, 2], ) + client_id: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description="Leave empty for a fresh session on every connection (a " + "broker-generated id). Set a stable name that is unique on the broker, for " + "example the camera or pipeline name from a workflow input, to keep a " + "persistent session: the broker then remembers the subscription and queues " + "QoS 1/2 messages while this block is away and delivers them when it " + "reconnects. Two connections with the same id disconnect each other.", + examples=["$inputs.pipeline_name", "line1-camera3"], + ) username: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( default=None, description="Username for MQTT broker authentication.", @@ -437,6 +470,7 @@ def run( timeout: float = 0.5, encryption: str = "none", ca_certificate_path: Optional[str] = None, + client_id: Optional[str] = None, ) -> BlockResult: if GCP_SERVERLESS or LAMBDA: # never open an outbound broker connection from hosted workers @@ -507,6 +541,18 @@ def run( f"Invalid ca_certificate_path: {ca_certificate_path!r}. Must be a path." ) ca_certificate_path = ca_certificate_path or None + try: + # blank means "no persistent session"; the presence of the id is the switch + client_id = normalise_client_id(client_id) + except ConfigurationError as e: + return self._handle_failure(str(e)) + if client_id is not None and qos == 0: + # the broker queues only QoS 1 and 2 for an offline session: a + # persistent session at QoS 0 would look right and receive nothing + return self._handle_failure( + "client_id needs qos 1 or 2: the broker only queues QoS 1 and 2 " + "messages for a persistent session." + ) try: # the operator's broker policy; the RESOLVED address is what the # block connects to and what the connection identity is built from @@ -527,6 +573,7 @@ def run( timeout=timeout, encryption=encryption, ca_certificate_path=ca_certificate_path, + client_id=client_id, ) def _connect_and_read( @@ -541,6 +588,7 @@ def _connect_and_read( timeout: float, encryption: str, ca_certificate_path: Optional[str], + client_id: Optional[str], ) -> BlockResult: connection_identity = ( host, @@ -553,6 +601,7 @@ def _connect_and_read( read_mode, encryption, ca_certificate_path, + client_id, ) if self._client is None: buffer_size = ( @@ -561,7 +610,15 @@ def _connect_and_read( state = MQTTReaderState(topic=topic, qos=qos, buffer_size=buffer_size) client = None try: - client = mqtt.Client(userdata=state) + if client_id is None: + # broker-generated id, clean session: nothing outlives the connection + client = mqtt.Client(userdata=state) + else: + # a persistent session: the broker keeps the subscription and + # queues QoS 1/2 messages while this id is away + client = mqtt.Client( + client_id=client_id, clean_session=False, userdata=state + ) # a raising callback would end the network thread; log instead client.suppress_exceptions = True if username is not None: @@ -616,8 +673,9 @@ def _connect_and_read( elif connection_identity != self._connection_identity: return self._handle_failure( "MQTT connection parameters (host, port, credentials, timeout, topic, " - "qos, read_mode, encryption or ca_certificate_path) changed between " - "runs; this block subscribes only with the configuration of its first run." + "qos, read_mode, encryption, ca_certificate_path or client_id) changed " + "between runs; this block subscribes only with the configuration of its " + "first run." ) state = self._state # the background loop owns (re)connecting and (re)subscribing; runs diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py index adefec8814..fd9003d66e 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -10,6 +10,7 @@ ConfigurationError, configure_tls, normalise_broker_address, + normalise_client_id, resolve_broker_address, split_host_port, ) @@ -58,6 +59,27 @@ def test_normalise_broker_address(host, port, expected): assert normalise_broker_address(host, port) == expected +@pytest.mark.parametrize( + "value, expected", + [ + (None, None), + ("", None), + (" ", None), + ("\t\n", None), + ("line1", "line1"), + (" line1-camera3 ", "line1-camera3"), + ], +) +def test_normalise_client_id_treats_blank_as_unset_and_strips(value, expected): + assert normalise_client_id(value) == expected + + +@pytest.mark.parametrize("value", [5, 1.5, True, ["line1"], {"id": "line1"}]) +def test_normalise_client_id_rejects_non_strings_naming_the_field(value): + with pytest.raises(ConfigurationError, match="client_id"): + normalise_client_id(value) + + def test_keepalive_is_shorter_than_paho_default(): assert 0 < MQTT_KEEPALIVE_SECONDS < 60 diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 69403b9788..fb81dc1fad 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -64,8 +64,13 @@ class FakeClient: granted_qos = 0 ack_subscriptions = True fire_on_connect = True + # CONNACK flags handed to on_connect; a test sets {"session present": 1} + connect_flags: dict = {} - def __init__(self, userdata=None): + def __init__(self, client_id="", clean_session=None, userdata=None): + # paho 1.6.1's constructor arguments the block chooses + self.client_id = client_id + self.clean_session = clean_session self.userdata = userdata # ordered record of the calls that matter for TLS: must precede connect() self.calls = [] @@ -105,7 +110,9 @@ def connect(self, host, port, keepalive=60): def loop_start(self): self.loop_started = True if self.fire_on_connect: - self.on_connect(self, self.userdata, {}, self.connack_reason_code) + self.on_connect( + self, self.userdata, dict(self.connect_flags), self.connack_reason_code + ) def loop_stop(self): self.loop_stopped = True @@ -127,8 +134,10 @@ def deliver(self, payload, topic="plc/state", retain=False): def clients() -> List[FakeClient]: created: List[FakeClient] = [] - def factory(userdata=None): - client = FakeClient(userdata=userdata) + def factory(client_id="", clean_session=None, userdata=None): + client = FakeClient( + client_id=client_id, clean_session=clean_session, userdata=userdata + ) created.append(client) return client @@ -155,6 +164,19 @@ def test_defaults(self): assert manifest.timeout == 0.5 assert manifest.username is None assert manifest.password is None + assert manifest.client_id is None + + @pytest.mark.parametrize("client_id", ["line1-camera3", "$inputs.pipeline_name"]) + def test_client_id_accepts_literal_and_selector(self, client_id): + manifest = BlockManifest.model_validate(manifest_payload(client_id=client_id)) + + assert manifest.client_id == client_id + + def test_qos_description_states_the_persistent_session_rule(self): + schema = BlockManifest.model_json_schema() + + assert "client_id" in schema["properties"]["qos"]["description"] + assert "1 or 2" in schema["properties"]["qos"]["description"] @pytest.mark.parametrize("read_mode", ["everything", "next", "", None]) def test_read_mode_must_be_latest_or_sequential(self, read_mode): @@ -245,6 +267,24 @@ def test_on_connect_subscribes_and_sets_event_for_accepted_connack(self): assert client.subscriptions == [("plc/#", 1)] assert state.connected.is_set() + @pytest.mark.parametrize( + "flags, expected", + [ + ({}, "False"), + ({"session present": 0}, "False"), + ({"session present": 1}, "True"), + ], + ) + def test_on_connect_logs_session_present_flag(self, caplog, flags, expected): + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + with caplog.at_level("INFO", logger="inference"): + mqtt_on_connect(client, state, flags, 0) + + assert f"session present: {expected}" in caplog.text + def test_on_connect_resubscribes_on_every_connect(self): state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) client = FakeClient(userdata=state) @@ -411,6 +451,46 @@ def test_invalid_topic_rejected_before_client_construction( assert "topic" in result["error_message"].lower() assert clients == [] + @pytest.mark.parametrize("client_id", [None, "", " "]) + def test_blank_client_id_builds_a_clean_session_client( + self, clients, block, client_id + ): + result = block.run(**run_kwargs(client_id=client_id)) + + assert result["error_status"] is False + assert clients[0].client_id == "" + assert clients[0].clean_session is None + assert clients[0].userdata is block._state + + def test_client_id_builds_a_persistent_session_client(self, clients, block): + result = block.run(**run_kwargs(client_id=" line1 ", qos=1)) + + assert result["error_status"] is False + assert clients[0].client_id == "line1" + assert clients[0].clean_session is False + assert clients[0].userdata is block._state + assert clients[0].subscriptions == [("plc/state", 1)] + + def test_client_id_with_qos_zero_rejected_before_client_construction( + self, clients, block + ): + result = block.run(**run_kwargs(client_id="line1")) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert "qos 1 or 2" in result["error_message"] + assert clients == [] + + @pytest.mark.parametrize("client_id", [5, 1.5, True, ["line1"]]) + def test_non_string_client_id_rejected_before_client_construction( + self, clients, block, client_id + ): + result = block.run(**run_kwargs(client_id=client_id, qos=1)) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert clients == [] + def test_error_outputs_carry_every_declared_key(self, clients, block): result = block.run(**run_kwargs(port=0)) @@ -558,6 +638,7 @@ def test_unacknowledged_subscription_reported(self, clients, block, monkeypatch) {"topic": "plc/other"}, {"qos": 1}, {"read_mode": "sequential"}, + {"qos": 1, "client_id": "line1"}, ], ) def test_changed_parameters_rejected_and_connection_kept( @@ -572,6 +653,30 @@ def test_changed_parameters_rejected_and_connection_kept( assert len(clients) == 1 assert clients[0].disconnected is False + @pytest.mark.parametrize("later_client_id", [None, "", "line2"]) + def test_dropped_or_changed_client_id_rejected_and_connection_kept( + self, clients, block, later_client_id + ): + block.run(**run_kwargs(client_id="line1", qos=1)) + + result = block.run(**run_kwargs(client_id=later_client_id, qos=1)) + + assert result["error_status"] is True + assert "client_id" in result["error_message"] + assert "changed between runs" in result["error_message"] + assert len(clients) == 1 + assert clients[0].disconnected is False + + def test_same_client_id_with_different_whitespace_is_not_a_change( + self, clients, block + ): + block.run(**run_kwargs(client_id="line1", qos=1)) + + result = block.run(**run_kwargs(client_id=" line1 ", qos=1)) + + assert result["error_status"] is False + assert len(clients) == 1 + def test_close_disconnects_stops_loop_and_clears_state(self, clients, block): block.run(**run_kwargs()) state = block._state From 89a24ffca1d7c11e062e346e3e55ad6d4a21b576 Mon Sep 17 00:00:00 2001 From: Tewan Date: Wed, 23 Sep 2026 16:26:44 +0800 Subject: [PATCH 11/14] Mention persistent sessions in the MQTT Reader's HTTP restriction note The SOFT restriction surfaces at compile time; it only warned that HTTP requests see retained messages. With client_id set, each request resumes the session, consumes the backlog and returns one message, and concurrent requests with the same id disconnect each other. Co-Authored-By: Claude Fable 5.1 --- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 7 +++++-- .../tests/unit_tests/core_steps/sinks/test_mqtt_reader.py | 3 +++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index ec9259101d..89d05fad80 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -390,8 +390,11 @@ def get_restrictions(cls) -> List[RuntimeRestriction]: note=( "The subscription and message buffer live in this block " "instance. Over HTTP every request builds a fresh instance, so " - "only a retained message can be returned. Use an " - "InferencePipeline for a live subscription." + "only a retained message can be returned; with client_id set, " + "each request resumes the persistent session, consumes its whole " + "backlog and returns one message, and concurrent requests with " + "the same id disconnect each other. Use an InferencePipeline " + "for a live subscription or a persistent session." ), applies_to_runtimes=[ Runtime.SELF_HOSTED_CPU, diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index fb81dc1fad..28a0aa1f8e 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -244,6 +244,9 @@ def test_restrictions_hard_for_hosted_serverless_only(self): v1.Runtime.DEDICATED_DEPLOYMENT, ] assert restrictions[1].applies_to_input_modes == [v1.RuntimeInputMode.IMAGE] + # the HTTP note must warn about persistent sessions, not only retained messages + assert "client_id" in restrictions[1].note + assert "InferencePipeline" in restrictions[1].note def test_block_declares_file_system_init_parameter_only(self): assert MQTTReaderBlockV1.get_init_parameters() == [ From ed91baee2915fe37520171cef8c8102b21436af1 Mon Sep 17 00:00:00 2001 From: Tewan Date: Thu, 24 Sep 2026 01:29:42 +0800 Subject: [PATCH 12/14] Stop the MQTT blocks from retrying a refused connection When the broker refused the CONNECT (bad credentials, not authorised, rejected client id, protocol version), paho reconnected with the same credentials about once a second for the life of the block, and the run reported "connection was not established within timeout ... raise timeout". Brokers with authentication rate limiting banned the host. Both blocks now record the CONNACK reason code, report it in the outputs with paho's reason text as soon as the broker answers, and stop the network loop for a permanent refusal so every later run repeats the refusal without touching the broker. "Broker unavailable" keeps retrying. The writer hands paho a small state object instead of a bare Event and drops a duplicated fragment from its changed-parameters message. Co-Authored-By: Claude Fable 5.1 --- .../test_workflow_with_mqtt_reader.py | 65 +++++ .../test_workflow_with_mqtt_writer.py | 44 ++++ workflows/CHANGELOG.md | 6 +- .../enterprise_blocks/sinks/mqtt_common.py | 39 ++- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 71 +++++- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 134 ++++++++-- .../core_steps/sinks/test_mqtt_common.py | 44 ++++ .../core_steps/sinks/test_mqtt_reader.py | 151 +++++++++++- .../core_steps/sinks/test_mqtt_writer.py | 230 ++++++++++++++++-- 9 files changed, 733 insertions(+), 51 deletions(-) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index 1eb28e48f4..c8b5c2edb6 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -233,6 +233,71 @@ def test_unreachable_broker_leaves_nothing_behind_and_later_run_recovers(): thread.join(timeout=2) +def refusing_broker(reason_code: int): + """A broker that answers every CONNECT with `reason_code`, accepting any + number of connections so retries, if any, are counted.""" + broker = FakeMQTTBroker(connack_reason_code=reason_code, keep_serving=True) + thread = threading.Thread(target=broker.serve, daemon=True) + thread.start() + return broker, thread + + +@pytest.mark.timeout(15) +def test_refused_connection_is_reported_once_and_not_retried(): + # given - a broker that refuses the credentials + broker, thread = refusing_broker(5) + block = MQTTReaderBlockV1() + + try: + # when - the first run hits the refusal, then the block lives on past + # several reconnect intervals (max delay is 2 * timeout) + started = time.monotonic() + first = block.run(**reader_kwargs(broker, timeout=0.5)) + first_elapsed = time.monotonic() - started + time.sleep(2.5) + started = time.monotonic() + second = block.run(**reader_kwargs(broker, timeout=0.5)) + second_elapsed = time.monotonic() - started + + # then - the refusal is named at once, the loop stopped, no retry happened + assert first["error_status"] is True + assert "not authorised" in first["error_message"] + assert "code 5" in first["error_message"] + assert "client_id" in first["error_message"] + assert "Raise 'timeout'" not in first["error_message"] + assert first_elapsed < 0.4 + assert second == first + assert second_elapsed < 0.1 + assert broker.connections_accepted == 1 + assert broker.subscriptions == [] + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + +@pytest.mark.timeout(15) +def test_unavailable_broker_keeps_retrying_and_is_named(): + # given - a broker that answers "unavailable", a passing condition + broker, thread = refusing_broker(3) + block = MQTTReaderBlockV1() + + try: + # when + result = block.run(**reader_kwargs(broker, timeout=0.5)) + time.sleep(2.5) + + # then - the reason is reported and the client kept reconnecting + assert result["error_status"] is True + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert broker.connections_accepted >= 2 + finally: + block.close() + broker.finish() + thread.join(timeout=2) + + @pytest.fixture def tls_broker(mqtt_test_certificates): broker = FakeMQTTBroker( diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py index c8b8db6867..6590fe4b22 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_writer.py @@ -88,6 +88,8 @@ def test_rejected_connack_is_not_treated_as_connected(): # then assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] assert not block._connected.is_set() assert broker.messages == [] finally: @@ -95,6 +97,48 @@ def test_rejected_connack_is_not_treated_as_connected(): broker.finish() +@pytest.mark.timeout(15) +def test_refused_connection_is_reported_once_and_not_retried(): + # given - a broker that refuses the credentials and would accept any retry + broker = FakeMQTTBroker(connack_reason_code=5, keep_serving=True) + broker_thread = threading.Thread(target=broker.serve, daemon=True) + broker_thread.start() + block = MQTTWriterSinkBlockV1() + kwargs = dict( + host=broker.host, + port=broker.port, + topic="RoboflowTopic", + message="should not be delivered", + timeout=0.5, + ) + + try: + # when - the first run hits the refusal, then the block lives on past + # several reconnect intervals (max delay is 2 * timeout) + started = time.monotonic() + first = block.run(**kwargs) + first_elapsed = time.monotonic() - started + time.sleep(2.5) + started = time.monotonic() + second = block.run(**kwargs) + second_elapsed = time.monotonic() - started + + # then - the refusal is named at once, the loop stopped, nothing published + assert first["error_status"] is True + assert "not authorised" in first["message"] + assert "Check username and password" in first["message"] + assert "Raise 'timeout'" not in first["message"] + assert first_elapsed < 0.4 + assert second == first + assert second_elapsed < 0.1 + assert broker.connections_accepted == 1 + assert broker.messages == [] + finally: + block.close() + broker.finish() + broker_thread.join(timeout=2) + + @pytest.mark.timeout(15) def test_block_recovers_when_broker_becomes_available_after_first_run(): # given - the port is bound but refuses connections, so the first TCP diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 860a853015..8a1809dfe4 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -2,9 +2,13 @@ ## Unreleased +### Fixed + +- MQTT Writer: a broker that refuses the connection (bad user name or password, not authorised, unacceptable protocol version) is now reported in the outputs with the broker's reason, and the client stops reconnecting instead of retrying the same credentials about once a second, which tripped brokers' authentication rate limiting. The refusal is reported as soon as the broker answers rather than after `timeout`, and every later run repeats it without touching the broker; fix the configuration and restart the pipeline. A broker answering "unavailable" keeps retrying in the background. + ### Added -- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. +- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. A broker that refuses the connection (bad user name or password, not authorised, rejected client id) is reported with the broker's reason as soon as it answers, and the client stops reconnecting; a broker answering "unavailable" keeps retrying. - Persistent sessions for the MQTT Reader: an optional `client_id` (unique per broker, for example the pipeline name from a workflow input) connects with that id and a non-clean session, so the broker keeps the subscription and queues QoS 1/2 messages while the block is away and delivers the backlog when the same id reconnects, after a pipeline or process restart included. Requires `qos` 1 or 2 (a run with `client_id` and `qos` 0 reports an error); the publisher must also publish at QoS 1 or higher. Meant for `InferencePipeline`s; `sequential` works through the backlog one message per run. Left empty, the behaviour is unchanged. - TLS for the MQTT Reader and MQTT Writer: an `encryption` dropdown (`none`, default, or `tls`) encrypts the connection and verifies the broker's certificate against the system trust store; `ca_certificate_path` (shown for TLS only) points at a PEM bundle for a private CA and requires `ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`. The port is not switched automatically and verification cannot be disabled. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py index 777757b7fe..ded0da20c0 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_common.py @@ -7,13 +7,17 @@ before either block builds a client, and `configure_tls()` enables server-verified TLS on the client, gating a workflow-chosen CA path behind the engine's file-system permission. `normalise_client_id()` turns the reader's -optional `client_id` into either a usable id or "unset". +optional `client_id` into either a usable id or "unset", and +`connection_refused_message()` with `PERMANENT_CONNACK_CODES` turn a broker's +CONNACK refusal into the outputs' error text and the decision whether paho may +keep reconnecting. """ import logging import os from typing import Any, List, Optional, Tuple +import paho.mqtt.client as mqtt from roboflow_workflows.environment import ( MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST, MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS, @@ -111,6 +115,39 @@ def normalise_client_id(value: Any) -> Optional[str]: return value or None +# CONNACK refusals a retry cannot change: unacceptable protocol version (1), +# identifier rejected (2), bad user name or password (4), not authorised (5). +# 3 (broker unavailable) is transient and keeps paho's reconnect. +PERMANENT_CONNACK_CODES = frozenset({1, 2, 4, 5}) +TRANSIENT_CONNACK_CODE = 3 + + +def connection_refused_message(reason_code: int, *, block_inputs: str) -> str: + """Return the outputs' error text for a broker that refused the CONNECT. + + Args: + reason_code: The CONNACK return code, 1 or higher. + block_inputs: The block's inputs the user should check, for example + ``"username and password"``. + + Returns: + A message naming the broker's reason. For a permanent refusal it says the + client does not retry; for "broker unavailable" it says the client keeps + retrying in the background. + """ + reason = mqtt.connack_string(reason_code) + if reason_code == TRANSIENT_CONNACK_CODE: + return ( + f"MQTT broker refused the connection: {reason} (code {reason_code}); " + "the client keeps retrying in the background." + ) + return ( + f"MQTT broker refused the connection: {reason} (code {reason_code}). " + f"Check {block_inputs}; the client does not retry until the block is closed " + "or the pipeline restarted." + ) + + def _normalise_host(host: str) -> str: host = host.strip().lower() if host.startswith("[") and host.endswith("]"): diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index 89d05fad80..d560ad2559 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -15,8 +15,11 @@ logger = logging.getLogger("inference") from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( MQTT_KEEPALIVE_SECONDS, + PERMANENT_CONNACK_CODES, + TRANSIENT_CONNACK_CODE, ConfigurationError, configure_tls, + connection_refused_message, normalise_client_id, resolve_broker_address, ) @@ -103,6 +106,13 @@ reported in the outputs. The connection uses a 15 s keepalive, so a broker that disappears without closing the connection is noticed within about 25 s. +A broker that refuses the connection (bad user name or password, not authorised, a +rejected client id or an unacceptable protocol version) is reported in the outputs with +the broker's reason, and the block stops reconnecting: retrying with the same credentials +cannot succeed and would only trip the broker's authentication rate limiting. Fix the +configuration and restart the pipeline. A broker answering "unavailable" is a passing +condition, so the client keeps retrying in the background. + Set `encryption` to `tls` to encrypt the connection and verify the broker's certificate against the system trust store; the port is not switched automatically, so set it to the broker's TLS port (usually 8883). A broker signed by a private CA needs @@ -126,6 +136,12 @@ """ SUBSCRIPTION_REFUSED_QOS = 0x80 +READER_CONNECTION_INPUTS = "username, password and client_id" +NOT_CONNECTED_WITHIN_TIMEOUT = ( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." +) TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} LATEST_BUFFER_SIZE = 1 SEQUENTIAL_BUFFER_SIZE = 1000 @@ -142,6 +158,12 @@ class MQTTReaderState: Attributes: connected: Set once the broker accepted the session; cleared on disconnect. + connack: Set on every CONNACK, accepted or refused, so the first run can + stop waiting as soon as the broker answered; cleared on a transport + disconnect, kept while a refusal is recorded. + refused_code: The CONNACK return code of the last refusal, or None after an + accepted CONNACK. A permanent code (see ``PERMANENT_CONNACK_CODES``) means + the network loop was stopped and every run reports the refusal. subscribed: Set once the broker granted the subscription; cleared on disconnect. subscribe_failed: True when the broker refused the subscription. topic: Topic filter to subscribe to on every (re)connect. @@ -154,6 +176,8 @@ class MQTTReaderState: def __init__(self, topic: str, qos: int, buffer_size: int): self.connected = threading.Event() + self.connack = threading.Event() + self.refused_code: Optional[int] = None self.subscribed = threading.Event() self.message_received = threading.Event() self.subscribe_failed = False @@ -180,11 +204,23 @@ def mqtt_on_connect( ): # paho invokes on_connect for accepted and rejected CONNACK alike; # only reason_code 0 means an established MQTT session + state.connack.set() if reason_code != 0: - logger.error("MQTT connection refused, result code %s", reason_code) + state.refused_code = reason_code + logger.error( + "MQTT connection refused: %s (code %s)", + mqtt.connack_string(reason_code), + reason_code, + ) state.connected.clear() + if reason_code in PERMANENT_CONNACK_CODES: + # paho would otherwise reconnect with the same credentials forever; + # disconnect() puts it in the disconnecting state, so its loop ends + # after this CONNACK instead of retrying + client.disconnect() return + state.refused_code = None logger.info( "MQTT client connected (session present: %s)", bool(flags.get("session present")), @@ -234,6 +270,10 @@ def mqtt_on_disconnect(client, state: MQTTReaderState, reason_code, properties=N logger.info("MQTT client disconnected, result code %s", reason_code) state.connected.clear() state.subscribed.clear() + if state.refused_code is None: + # a transport drop: the next run waits for the reconnect's CONNACK; after + # a refusal the answer stays visible so run() reports it without waiting + state.connack.clear() class BlockManifest(WorkflowBlockManifest): @@ -681,14 +721,33 @@ def _connect_and_read( "first run." ) state = self._state + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) # the background loop owns (re)connecting and (re)subscribing; runs - # only wait for readiness - if not state.connected.wait(timeout=timeout): + # only wait for readiness. The wait ends on any CONNACK, so a refusal is + # reported as soon as the broker answers instead of after the timeout + if not state.connack.wait(timeout=timeout): + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) + if state.refused_code in PERMANENT_CONNACK_CODES: return self._handle_failure( - "MQTT broker not connected (connection was not established within timeout); " - "the client keeps retrying in the background. Raise 'timeout' if the " - "broker needs longer to connect." + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) # a known refusal is permanent for this configuration: report it # without paying the acknowledgement wait on every run if state.subscribe_failed or ( diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index a13c2bad68..176dead7fb 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -13,8 +13,11 @@ logger = logging.getLogger("inference") from roboflow_workflows.core_steps.sinks.noop import disabled_sink_response from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( + PERMANENT_CONNACK_CODES, + TRANSIENT_CONNACK_CODE, ConfigurationError, configure_tls, + connection_refused_message, resolve_broker_address, ) from roboflow_workflows.execution_engine.entities.base import OutputDefinition @@ -66,6 +69,14 @@ MQTT keepalive interval (60 s), not after `timeout`; real brokers close such a connection immediately. +A broker that refuses the connection (bad user name or password, not +authorised, an unacceptable protocol version) is reported in the outputs with +the broker's reason, and the block stops reconnecting: retrying with the same +credentials cannot succeed and would only trip the broker's authentication +rate limiting. Fix the configuration and restart the pipeline. A broker +answering "unavailable" is a passing condition, so the client keeps retrying +in the background. + Outputs: - error_status (bool): Indicates if an error occurred during the MQTT publishing process. True if there was an error, False if successful. @@ -98,26 +109,79 @@ def _not_connected_message(error: BaseException, encryption: str) -> str: ) -def mqtt_on_connect(client, userdata, flags, reason_code, properties=None): +WRITER_CONNECTION_INPUTS = "username and password" +NOT_CONNECTED_WITHIN_TIMEOUT = ( + "MQTT broker not connected (connection was not established within timeout); " + "the client keeps retrying in the background. Raise 'timeout' if the " + "broker needs longer to connect." +) + + +class MQTTWriterState: + """Connection state shared between paho's network thread and the engine thread. + + Passed to paho as ``userdata`` so the callbacks never hold a reference to the + block itself. + + Attributes: + connected: Set once the broker accepted the session; cleared on disconnect. + connack: Set on every CONNACK, accepted or refused, so a run can stop + waiting as soon as the broker answered; cleared on a transport + disconnect, kept while a refusal is recorded. + refused_code: The CONNACK return code of the last refusal, or None after an + accepted CONNACK. A permanent code (see ``PERMANENT_CONNACK_CODES``) means + the network loop was stopped and every run reports the refusal. + """ + + def __init__(self): + self.connected = threading.Event() + self.connack = threading.Event() + self.refused_code: Optional[int] = None + + def reset(self) -> None: + self.connected.clear() + self.connack.clear() + self.refused_code = None + + +def mqtt_on_connect( + client, state: MQTTWriterState, flags, reason_code, properties=None +): # paho invokes on_connect for accepted and rejected CONNACK alike; # only reason_code 0 means an established MQTT session + state.connack.set() if reason_code == 0: + state.refused_code = None logger.info("MQTT client connected") - userdata.set() - else: - logger.error("MQTT connection refused, result code %s", reason_code) - userdata.clear() + state.connected.set() + return + state.refused_code = reason_code + logger.error( + "MQTT connection refused: %s (code %s)", + mqtt.connack_string(reason_code), + reason_code, + ) + state.connected.clear() + if reason_code in PERMANENT_CONNACK_CODES: + # paho would otherwise reconnect with the same credentials forever; + # disconnect() puts it in the disconnecting state, so its loop ends + # after this CONNACK instead of retrying + client.disconnect() -def mqtt_on_connect_fail(client, userdata): +def mqtt_on_connect_fail(client, state: MQTTWriterState): # paho 1.6.1 invokes this callback with exactly (client, userdata) logger.error("MQTT client failed to establish connection with broker") - userdata.clear() + state.connected.clear() -def mqtt_on_disconnect(client, userdata, reason_code, properties=None): +def mqtt_on_disconnect(client, state: MQTTWriterState, reason_code, properties=None): logger.info("MQTT client disconnected, result code %s", reason_code) - userdata.clear() + state.connected.clear() + if state.refused_code is None: + # a transport drop: the next run waits for the reconnect's CONNACK; after + # a refusal the answer stays visible so run() reports it without waiting + state.connack.clear() class BlockManifest(WorkflowBlockManifest): @@ -241,7 +305,9 @@ def __init__( # False by default so a hand-constructed block is safe self._allow_access_to_file_system = allow_access_to_file_system self.mqtt_client: Optional[mqtt.Client] = None - self._connected = threading.Event() + self._connection = MQTTWriterState() + # the same Event object as the state's, kept under its historical name + self._connected = self._connection.connected self._connection_identity: Optional[Tuple] = None self._lifecycle_lock = threading.Lock() self._disable_sinks = disable_sinks @@ -263,8 +329,9 @@ def close(self) -> None: # accepted so the thread never outlives the block client.loop_stop() finally: - # only after the join can no callback re-set the event - self._connected.clear() + # only after the join can no callback re-set the events; a + # client rebuilt by a later run starts without a stale refusal + self._connection.reset() def __del__(self): try: @@ -412,7 +479,7 @@ def _connect_and_publish( if self.mqtt_client is None: client = None try: - client = mqtt.Client(userdata=self._connected) + client = mqtt.Client(userdata=self._connection) if username is not None: client.username_pw_set(username, password) # TLS must be configured before connect(); the CA path is gated @@ -465,18 +532,45 @@ def _connect_and_publish( return self._handle_failure( "MQTT connection parameters (host, port, credentials, timeout, encryption " "or ca_certificate_path) changed between runs; this block publishes only " - "to the connection configured on its first run." - "on its first run.", + "to the connection configured on its first run.", fail_fast=fail_fast, ) - # the background loop owns (re)connecting; runs only wait for readiness - if not self._connected.wait(timeout=timeout): + state = self._connection + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: return self._handle_failure( - "MQTT broker not connected (connection was not established within timeout); " - "the client keeps retrying in the background. Raise 'timeout' if the " - "broker needs longer to connect.", + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), fail_fast=fail_fast, ) + # the background loop owns (re)connecting; runs only wait for readiness. + # The wait ends on any CONNACK, so a refusal is reported as soon as the + # broker answers instead of after the timeout + if not state.connected.is_set(): + if not state.connack.wait(timeout=timeout): + return self._handle_failure( + NOT_CONNECTED_WITHIN_TIMEOUT, fail_fast=fail_fast + ) + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), + fail_fast=fail_fast, + ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=WRITER_CONNECTION_INPUTS + ), + fail_fast=fail_fast, + ) + return self._handle_failure( + NOT_CONNECTED_WITHIN_TIMEOUT, fail_fast=fail_fast + ) try: res: mqtt.MQTTMessageInfo = self.mqtt_client.publish( topic, message, qos=qos, retain=retain diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py index fd9003d66e..f11a59fd3e 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_common.py @@ -7,8 +7,10 @@ from roboflow_workflows.enterprise_blocks.sinks.mqtt_common import ( DEFAULT_MQTT_PORT, MQTT_KEEPALIVE_SECONDS, + PERMANENT_CONNACK_CODES, ConfigurationError, configure_tls, + connection_refused_message, normalise_broker_address, normalise_client_id, resolve_broker_address, @@ -80,6 +82,48 @@ def test_normalise_client_id_rejects_non_strings_naming_the_field(value): normalise_client_id(value) +def test_permanent_connack_codes_exclude_success_and_broker_unavailable(): + assert PERMANENT_CONNACK_CODES == {1, 2, 4, 5} + assert 0 not in PERMANENT_CONNACK_CODES + assert 3 not in PERMANENT_CONNACK_CODES + + +@pytest.mark.parametrize( + "code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + (5, "not authorised"), + ], +) +def test_permanent_refusal_message_names_reason_inputs_and_no_retry(code, reason): + message = connection_refused_message(code, block_inputs="username and password") + + assert reason in message + assert f"code {code}" in message + assert "Check username and password" in message + assert "does not retry" in message + assert "Raise 'timeout'" not in message + + +def test_broker_unavailable_message_says_retrying(): + message = connection_refused_message(3, block_inputs="username and password") + + assert "broker unavailable" in message + assert "code 3" in message + assert "retrying in the background" in message + assert "does not retry" not in message + + +def test_unknown_refusal_code_uses_permanent_wording(): + message = connection_refused_message(7, block_inputs="username and password") + + assert "unknown reason" in message + assert "code 7" in message + assert "does not retry" in message + + def test_keepalive_is_shorter_than_paho_default(): assert 0 < MQTT_KEEPALIVE_SECONDS < 60 diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 28a0aa1f8e..2af605e522 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -10,6 +10,7 @@ from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import ( LATEST_BUFFER_SIZE, + NOT_CONNECTED_WITHIN_TIMEOUT, SEQUENTIAL_BUFFER_SIZE, SUBSCRIPTION_REFUSED_QOS, BlockManifest, @@ -310,6 +311,47 @@ def test_on_connect_does_not_subscribe_for_rejected_connack(self, reason_code): assert client.subscriptions == [] assert not state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code == reason_code + + @pytest.mark.parametrize("reason_code", [1, 2, 4, 5]) + def test_on_connect_stops_the_loop_for_a_permanent_refusal(self, reason_code): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + + mqtt_on_connect(client, state, {}, reason_code) + + assert client.disconnected is True + + def test_on_connect_keeps_the_loop_for_broker_unavailable(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + + mqtt_on_connect(client, state, {}, 3) + + assert client.disconnected is False + assert state.refused_code == 3 + + def test_on_connect_logs_the_refusal_reason(self, caplog): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + + with caplog.at_level("ERROR", logger="inference"): + mqtt_on_connect(FakeClient(userdata=state), state, {}, 5) + + assert "not authorised" in caplog.text + assert "code 5" in caplog.text + + def test_accepted_connack_after_broker_unavailable_resets_refusal(self): + state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, 3) + mqtt_on_connect(client, state, {}, 0) + + assert state.refused_code is None + assert state.connected.is_set() + assert state.connack.is_set() def test_on_connect_fail_matches_paho_two_argument_signature(self): state = MQTTReaderState(topic="t", qos=0, buffer_size=1) @@ -365,6 +407,24 @@ def test_on_disconnect_clears_both_events(self): assert not state.connected.is_set() assert not state.subscribed.is_set() + def test_on_disconnect_after_transport_drop_clears_connack(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + state.connack.set() + + mqtt_on_disconnect(FakeClient(), state, 1) + + assert not state.connack.is_set() + + def test_on_disconnect_after_refusal_keeps_connack(self): + state = MQTTReaderState(topic="t", qos=0, buffer_size=1) + client = FakeClient(userdata=state) + mqtt_on_connect(client, state, {}, 5) + + mqtt_on_disconnect(client, state, 0) + + assert state.connack.is_set() + assert state.refused_code == 5 + def test_on_disconnect_accepts_mqtt5_properties_argument(self): state = MQTTReaderState(topic="t", qos=0, buffer_size=1) state.connected.set() @@ -576,14 +636,101 @@ def test_connection_not_established_within_timeout_reported( assert "not connected" in result["error_message"].lower() assert block._client is clients[0] - def test_rejected_connack_reported_as_not_connected( + def test_rejected_connack_reported_as_a_refusal(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not authorised" in result["error_message"] + assert "code 5" in result["error_message"] + assert "client_id" in result["error_message"] + assert "Raise 'timeout'" not in result["error_message"] + assert clients[0].subscriptions == [] + assert clients[0].disconnected is True + + @pytest.mark.parametrize( + "reason_code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + ], + ) + def test_other_permanent_refusals_name_their_reason( + self, clients, block, monkeypatch, reason_code, reason + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", reason_code) + result = block.run(**run_kwargs()) + + assert reason in result["error_message"] + assert clients[0].disconnected is True + + def test_refusal_is_reported_without_waiting_for_timeout( self, clients, block, monkeypatch ): monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert result["error_status"] is True + assert time.monotonic() - started < 0.5 + + def test_next_run_after_refusal_repeats_it_without_a_new_connection( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + first = block.run(**run_kwargs(timeout=2.0)) + + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0)) + + assert second == first + assert time.monotonic() - started < 0.5 + assert len(clients) == 1 + assert [call for call in clients[0].calls if call[0] == "connect"] == [ + ("connect", "localhost", 1883) + ] + + def test_broker_unavailable_keeps_retrying_and_names_the_reason( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 3) result = block.run(**run_kwargs()) assert result["error_status"] is True - assert clients[0].subscriptions == [] + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert clients[0].disconnected is False + + def test_run_succeeds_once_an_unavailable_broker_accepts( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "connack_reason_code", 3) + block.run(**run_kwargs()) + client = clients[0] + client.on_connect(client, client.userdata, {}, 0) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert block._state.refused_code is None + + def test_no_connack_at_all_reports_not_connected(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert result["error_message"] == NOT_CONNECTED_WITHIN_TIMEOUT + + def test_close_after_refusal_completes(self, clients, block, monkeypatch): + monkeypatch.setattr(FakeClient, "connack_reason_code", 5) + block.run(**run_kwargs()) + + block.close() + + assert block._client is None + assert clients[0].loop_stopped is True def test_refused_subscription_reported_and_client_kept( self, clients, block, monkeypatch diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index d9cde3b188..2334e8a804 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -1,4 +1,4 @@ -import threading +import time from typing import get_args from unittest.mock import MagicMock, patch @@ -7,8 +7,10 @@ from roboflow_workflows.enterprise_blocks.sinks import mqtt_common from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_writer.v1 import ( + NOT_CONNECTED_WITHIN_TIMEOUT, BlockManifest, MQTTWriterSinkBlockV1, + MQTTWriterState, mqtt_on_connect, mqtt_on_connect_fail, mqtt_on_disconnect, @@ -128,45 +130,231 @@ def test_valid_ports_accepted(self, port): class TestCallbacks: def test_on_connect_sets_event_only_for_accepted_connack(self): - event = threading.Event() + state = MQTTWriterState() - mqtt_on_connect(MagicMock(), event, {}, 0) + mqtt_on_connect(MagicMock(), state, {}, 0) - assert event.is_set() + assert state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code is None @pytest.mark.parametrize("reason_code", [1, 2, 3, 4, 5]) def test_on_connect_clears_event_for_rejected_connack(self, reason_code): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() - mqtt_on_connect(MagicMock(), event, {}, reason_code) + mqtt_on_connect(MagicMock(), state, {}, reason_code) - assert not event.is_set() + assert not state.connected.is_set() + assert state.connack.is_set() + assert state.refused_code == reason_code + + @pytest.mark.parametrize("reason_code", [1, 2, 4, 5]) + def test_on_connect_stops_the_loop_for_a_permanent_refusal(self, reason_code): + client = MagicMock() + + mqtt_on_connect(client, MQTTWriterState(), {}, reason_code) + + client.disconnect.assert_called_once_with() + + def test_on_connect_keeps_the_loop_for_broker_unavailable(self): + client = MagicMock() + state = MQTTWriterState() + + mqtt_on_connect(client, state, {}, 3) + + client.disconnect.assert_not_called() + assert state.refused_code == 3 + + def test_on_connect_logs_the_refusal_reason(self, caplog): + with caplog.at_level("ERROR", logger="inference"): + mqtt_on_connect(MagicMock(), MQTTWriterState(), {}, 4) + + assert "bad user name or password" in caplog.text + assert "code 4" in caplog.text + + def test_accepted_connack_after_broker_unavailable_resets_refusal(self): + state = MQTTWriterState() + + mqtt_on_connect(MagicMock(), state, {}, 3) + mqtt_on_connect(MagicMock(), state, {}, 0) + + assert state.refused_code is None + assert state.connected.is_set() def test_on_connect_fail_matches_paho_two_argument_signature(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() # paho 1.6.1 invokes on_connect_fail with exactly (client, userdata) - mqtt_on_connect_fail(MagicMock(), event) + mqtt_on_connect_fail(MagicMock(), state) - assert not event.is_set() + assert not state.connected.is_set() def test_on_disconnect_clears_event(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1) + + assert not state.connected.is_set() + + def test_on_disconnect_after_transport_drop_clears_connack(self): + state = MQTTWriterState() + state.connack.set() + + mqtt_on_disconnect(MagicMock(), state, 1) - mqtt_on_disconnect(MagicMock(), event, 1) + assert not state.connack.is_set() - assert not event.is_set() + def test_on_disconnect_after_refusal_keeps_connack(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + + mqtt_on_disconnect(MagicMock(), state, 0) + + assert state.connack.is_set() + assert state.refused_code == 5 def test_on_disconnect_accepts_mqtt5_properties_argument(self): - event = threading.Event() - event.set() + state = MQTTWriterState() + state.connected.set() + + mqtt_on_disconnect(MagicMock(), state, 1, properties=None) + + assert not state.connected.is_set() + + def test_reset_clears_everything(self): + state = MQTTWriterState() + mqtt_on_connect(MagicMock(), state, {}, 5) + + state.reset() + + assert not state.connected.is_set() + assert not state.connack.is_set() + assert state.refused_code is None + + +def _answer_connect_with(mock_client_cls, reason_code: int) -> None: + """Make the mocked client's loop_start deliver a CONNACK with `reason_code` + to the state the block handed paho as userdata.""" + mock_client = mock_client_cls.return_value + + def fire_connack(): + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client, state, {}, reason_code) + + mock_client.loop_start.side_effect = fire_connack + + +class TestRefusedConnection: + def test_refusal_reported_with_reason_and_inputs(self, mock_client_cls, block): + _answer_connect_with(mock_client_cls, 5) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert "not authorised" in result["message"] + assert "code 5" in result["message"] + assert "Check username and password" in result["message"] + assert "Raise 'timeout'" not in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + mock_client_cls.return_value.publish.assert_not_called() + + @pytest.mark.parametrize( + "reason_code, reason", + [ + (1, "unacceptable protocol version"), + (2, "identifier rejected"), + (4, "bad user name or password"), + ], + ) + def test_other_permanent_refusals_name_their_reason( + self, mock_client_cls, block, reason_code, reason + ): + _answer_connect_with(mock_client_cls, reason_code) + + result = block.run(**run_kwargs()) + + assert reason in result["message"] + mock_client_cls.return_value.disconnect.assert_called_once_with() + + def test_refusal_is_reported_without_waiting_for_timeout( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert result["error_status"] is True + assert time.monotonic() - started < 0.5 + + def test_next_run_after_refusal_repeats_it_without_a_new_connection( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + first = block.run(**run_kwargs(timeout=2.0)) + + started = time.monotonic() + second = block.run(**run_kwargs(timeout=2.0)) + + assert second == first + assert time.monotonic() - started < 0.5 + assert mock_client_cls.call_count == 1 + mock_client_cls.return_value.connect.assert_called_once() + + def test_refusal_with_fail_fast_raises_with_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + + with pytest.raises(Exception, match="not authorised"): + block.run(**run_kwargs(fail_fast=True)) + + def test_broker_unavailable_keeps_retrying_and_names_the_reason( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + + result = block.run(**run_kwargs()) - mqtt_on_disconnect(MagicMock(), event, 1, properties=None) + assert result["error_status"] is True + assert "broker unavailable" in result["message"] + assert "retrying in the background" in result["message"] + mock_client_cls.return_value.disconnect.assert_not_called() + + def test_run_succeeds_once_an_unavailable_broker_accepts( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 3) + block.run(**run_kwargs()) + state = mock_client_cls.call_args.kwargs["userdata"] + mqtt_on_connect(mock_client_cls.return_value, state, {}, 0) + + result = block.run(**run_kwargs()) + + assert result["error_status"] is False + assert state.refused_code is None + + def test_no_connack_at_all_reports_not_connected(self, mock_client_cls, block): + result = block.run(**run_kwargs()) + + assert result["error_status"] is True + assert result["message"] == NOT_CONNECTED_WITHIN_TIMEOUT + + def test_close_after_refusal_completes_and_resets_state( + self, mock_client_cls, block + ): + _answer_connect_with(mock_client_cls, 5) + block.run(**run_kwargs()) - assert not event.is_set() + block.close() + + assert block.mqtt_client is None + assert block._connection.refused_code is None + assert not block._connection.connack.is_set() + mock_client_cls.return_value.loop_stop.assert_called_once() class TestRunValidation: @@ -271,7 +459,7 @@ def test_client_connects_synchronously_before_background_loop( result = block.run(**run_kwargs()) assert result["error_status"] is False - mock_client_cls.assert_called_once_with(userdata=block._connected) + mock_client_cls.assert_called_once_with(userdata=block._connection) mock_client.connect.assert_called_once_with("localhost", 1883) mock_client.connect_async.assert_not_called() called_methods = [call[0] for call in mock_client.method_calls] From 2d8528ecb2be84ff9669223538980db9e21a2708 Mon Sep 17 00:00:00 2001 From: Tewan Date: Thu, 24 Sep 2026 02:20:57 +0800 Subject: [PATCH 13/14] Repeat the MQTT Reader's last message while the broker is down While the broker connection was down, every run waited the full timeout for the reconnect and then returned an empty message, although the block still held the last message it surfaced. Readiness waits are now a cold-start concern: once the first run reached SUBACK, later runs never block. They drain the buffer first and, when it is empty and no fresh data can arrive (not connected, not resubscribed, broker unavailable, or a connection refused on the reconnect), return the last message with error_status set and a message naming the cause. An outage is logged once at ERROR when first observed and once at INFO with its duration and run count when a live message arrives again. The same-id eviction integration test gains a retained state so both first runs return at once; its second run could otherwise be evicted by the first's reconnect mid-wait, which the new behaviour reports honestly. Co-Authored-By: Claude Fable 5.1 --- .../test_workflow_with_mqtt_reader.py | 121 ++++++- workflows/CHANGELOG.md | 2 +- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 222 +++++++++--- .../core_steps/sinks/test_mqtt_reader.py | 321 ++++++++++++++++++ 4 files changed, 618 insertions(+), 48 deletions(-) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index c8b5c2edb6..a1bdf22353 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -204,6 +204,76 @@ def test_subscription_restored_after_reconnect(broker): block.close() +@pytest.mark.timeout(20) +def test_outage_repeats_last_message_without_waiting_then_recovers(broker, caplog): + # given - a warm reader holding RUNNING + broker.retained["plc/state"] = RUNNING + block = MQTTReaderBlockV1() + kwargs = reader_kwargs(broker, timeout=2.0) + second_thread = None + + try: + first = block.run(**kwargs) + assert first["payload"] == {"state": "RUNNING"} + + # when - the connection drops and three frames arrive while it is down + broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + with caplog.at_level("INFO", logger="inference"): + during = [] + for _ in range(3): + started = time.monotonic() + during.append(block.run(**kwargs)) + assert time.monotonic() - started < 0.2 + + # and the broker comes back: the retained copy is re-sent on the + # resubscribe (the buffer holds one message in latest mode), then a + # new message is published + second_thread = threading.Thread(target=broker.start, daemon=True) + second_thread.start() + assert wait_until(lambda: len(broker.subscriptions) == 2, timeout=10) + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + after_reconnect = block.run(**kwargs) + broker.publish("plc/state", PAUSED) + assert wait_until( + lambda: bool(block._state.messages) + and block._state.messages[-1][1] == PAUSED + ) + updated = block.run(**kwargs) + + # then - every run during the outage repeated RUNNING with the error set + for result in during: + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": result["error_message"], + } + assert "not connected" in result["error_message"] + errors = [ + r + for r in caplog.records + if r.levelname == "ERROR" and r.getMessage().startswith("MQTT Reader ") + ] + assert len(errors) == 1 + # the retained copy re-sent on resubscribe is not new, but it is live: + # it closes the outage; the next run returns the new message + assert after_reconnect == {**first, "is_new": False} + assert updated["payload"] == {"state": "PAUSED"} + assert updated["is_new"] is True + assert updated["error_status"] is False + recoveries = [ + r.getMessage() for r in caplog.records if "restored after" in r.getMessage() + ] + assert len(recoveries) == 1 + assert "3 runs returned the last message" in recoveries[0] + assert broker.connections_accepted == 2 + finally: + block.close() + if second_thread is not None: + second_thread.join(timeout=2) + + @pytest.mark.timeout(15) def test_unreachable_broker_leaves_nothing_behind_and_later_run_recovers(): # given - bound but not listening, so the TCP connect is refused @@ -538,9 +608,58 @@ def test_persistent_session_survives_reconnect_of_the_same_instance(session_brok block.close() +@pytest.mark.timeout(20) +def test_persistent_session_outage_repeats_last_then_delivers_queued(session_broker): + # given - a warm persistent subscriber holding RUNNING + block = MQTTReaderBlockV1() + kwargs = persistent_kwargs(session_broker, timeout=2.0) + + try: + block.run(**kwargs) + session_broker.publish("plc/state", RUNNING, qos=1) + assert wait_until(lambda: len(block._state.messages) == 1) + first = block.run(**kwargs) + assert first["payload"] == {"state": "RUNNING"} + assert first["is_new"] is True + + # when - the connection drops, a QoS 1 message is queued while it is away, + # and a frame arrives during the outage + session_broker.drop_connection() + assert wait_until(lambda: not block._state.connected.is_set()) + assert wait_until(lambda: session_is_away(session_broker)) + session_broker.publish("plc/state", PAUSED, qos=1) + started = time.monotonic() + during = block.run(**kwargs) + assert time.monotonic() - started < 0.2 + + # and paho reconnects on its own + assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) + after = block.run(**kwargs) + + # then - the outage frame repeated RUNNING with the error set, and the + # queued message came through afterwards + assert during == { + **first, + "is_new": False, + "error_status": True, + "error_message": during["error_message"], + } + assert "not connected" in during["error_message"] + assert after["payload"] == {"state": "PAUSED"} + assert after["is_new"] is True + assert after["error_status"] is False + assert session_broker.connections_accepted == 2 + finally: + block.close() + + @pytest.mark.timeout(20) def test_second_connection_with_the_same_client_id_evicts_the_first(session_broker): - # given - two block instances sharing one client id on one broker + # given - two block instances sharing one client id on one broker; a retained + # state makes each first run return at once instead of waiting `timeout` for + # a retained message, so the evicted first cannot reconnect and evict the + # second while the second's first run is still in progress + session_broker.retained["plc/state"] = RUNNING first = MQTTReaderBlockV1() second = MQTTReaderBlockV1() diff --git a/workflows/CHANGELOG.md b/workflows/CHANGELOG.md index 8a1809dfe4..a77b679e76 100644 --- a/workflows/CHANGELOG.md +++ b/workflows/CHANGELOG.md @@ -8,7 +8,7 @@ ### Added -- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. A broker that refuses the connection (bad user name or password, not authorised, rejected client id) is reported with the broker's reason as soon as it answers, and the client stops reconnecting; a broker answering "unavailable" keeps retrying. +- MQTT Reader enterprise block (`roboflow_enterprise/mqtt_reader@v1`): subscribes to an MQTT topic and returns one message per run as raw text and parsed JSON, with `read_mode` selecting the newest message (`latest`) or the next unread one (`sequential`). Retained messages are returned on the first run; the subscription is restored after a reconnect, and a 15 s keepalive reports a silently lost broker within about 25 s. Live subscriptions need an `InferencePipeline`; over the HTTP API each request only sees retained messages. A broker that refuses the connection (bad user name or password, not authorised, rejected client id) is reported with the broker's reason as soon as it answers, and the client stops reconnecting; a broker answering "unavailable" keeps retrying. After the first subscription a run never waits for the broker: while the connection is down, runs return the last message with `error_status` set instead of an empty message, buffered messages first; the outage is logged once at its start and once at recovery. - Persistent sessions for the MQTT Reader: an optional `client_id` (unique per broker, for example the pipeline name from a workflow input) connects with that id and a non-clean session, so the broker keeps the subscription and queues QoS 1/2 messages while the block is away and delivers the backlog when the same id reconnects, after a pipeline or process restart included. Requires `qos` 1 or 2 (a run with `client_id` and `qos` 0 reports an error); the publisher must also publish at QoS 1 or higher. Meant for `InferencePipeline`s; `sequential` works through the backlog one message per run. Left empty, the behaviour is unchanged. - TLS for the MQTT Reader and MQTT Writer: an `encryption` dropdown (`none`, default, or `tls`) encrypts the connection and verifies the broker's certificate against the system trust store; `ca_certificate_path` (shown for TLS only) points at a PEM bundle for a private CA and requires `ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE=True`. The port is not switched automatically and verification cannot be disabled. - Operator policy for the MQTT Reader and MQTT Writer broker address: `MQTT_WORKFLOWS_BLOCKS_WHITELISTED_HOSTS` (comma-separated `host[:port]` allowlist, no DNS) and `MQTT_WORKFLOWS_BLOCKS_ALLOW_USER_PROVIDED_HOST` (default `True`; when `False` the workflow's host and port are ignored and the first allowlist entry is used). Defaults preserve existing behaviour. diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index d560ad2559..6cccd93fb5 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -2,6 +2,7 @@ import logging import math import threading +import time from collections import deque from typing import Any, Deque, Dict, List, Literal, Optional, Tuple, Type, Union @@ -90,6 +91,15 @@ pipelines. Persistent sessions are meant for InferencePipelines: over the HTTP API each request would resume the session, consume the whole backlog and return one message. +After the first subscription a run never waits for the broker. While the connection is +down the outputs repeat the last message received with `error_status` True and `is_new` +False, so a downstream step can gate on `error_status`; messages buffered before the drop +are still returned first, and with `client_id` the messages the broker queued meanwhile +follow after the reconnect. A connection refused on a reconnect is reported the same way, +with the broker's reason, until the pipeline is restarted. The log carries one ERROR line +when an outage starts and one INFO line with its duration when a live message arrives +again. + One block instance subscribes on a single broker connection: changing host, port, credentials, timeout, topic, QoS, read mode, client id or the TLS settings between runs is rejected as a configuration error. While the broker is unreachable every run waits up to `timeout` @@ -128,8 +138,11 @@ not a JSON object. - topic (str): The topic the message arrived on (useful with wildcard filters). - is_new (bool): True if this run surfaced a message the previous run did not. - - error_status (bool): True if the run failed (connection, subscription, invalid - configuration or an undecodable payload). + - error_status (bool): True if the run failed (invalid configuration, refused + subscription, undecodable payload) or the broker connection + is down or was refused after a healthy start; in the latter + cases value, payload and topic repeat the last message + received. - error_message (str): Details of the failure, or None. Failures are returned in the outputs and logged; the workflow keeps running. @@ -142,6 +155,14 @@ "the client keeps retrying in the background. Raise 'timeout' if the " "broker needs longer to connect." ) +NOT_CONNECTED_REPEATING_LAST = ( + "MQTT broker not connected; the client keeps reconnecting in the background and " + "the outputs repeat the last message received." +) +NOT_RESUBSCRIBED_REPEATING_LAST = ( + "MQTT subscription not re-established yet after a reconnect; the outputs repeat " + "the last message received." +) TLS_RELEVANT = {"encryption": {"values": ["tls"], "required": True}} LATEST_BUFFER_SIZE = 1 SEQUENTIAL_BUFFER_SIZE = 1000 @@ -366,9 +387,11 @@ class BlockManifest(WorkflowBlockManifest): Selector(kind=[FLOAT_KIND]), ] = Field( default=0.5, - description="Timeout in seconds for connecting to the MQTT broker and for the " - "broker's connection and subscription acknowledgements. Must be a finite " - "number greater than 0.", + description="Timeout in seconds for the first connection: the TCP connect, the " + "broker's connection acknowledgement (accepted or refused) and the subscription " + "acknowledgement. Later runs never wait: while the connection is down they return " + "the last message with error_status set, and the client reconnects in the " + "background. Must be a finite number greater than 0.", examples=[0.5], ) encryption: Literal["none", "tls"] = Field( @@ -459,6 +482,16 @@ def __init__(self, allow_access_to_file_system: bool = False): # the first read after subscribing waits for a retained message, which # the broker sends right after SUBACK as a separate packet self._awaiting_retained = False + # readiness waits are a cold-start concern: once the first run reached + # SUBACK, later runs never block on the broker + self._subscribed_once = False + # an outage is bracketed in the log: one ERROR line when a run first + # observes it, one INFO line with duration and run count on recovery + self._outage_started_at: Optional[float] = None + self._outage_runs = 0 + # the resolved broker address, kept for the log lines + self._host: Optional[str] = None + self._port: Optional[int] = None self._lifecycle_lock = threading.Lock() def close(self) -> None: @@ -472,6 +505,11 @@ def close(self) -> None: self._connection_identity = None self._last = None self._awaiting_retained = False + self._subscribed_once = False + self._outage_started_at = None + self._outage_runs = 0 + self._host = None + self._port = None try: client.disconnect() except Exception as e: @@ -713,6 +751,7 @@ def _connect_and_read( self._state = state self._connection_identity = connection_identity self._awaiting_retained = True + self._host, self._port = host, port elif connection_identity != self._connection_identity: return self._handle_failure( "MQTT connection parameters (host, port, credentials, timeout, topic, " @@ -721,56 +760,50 @@ def _connect_and_read( "first run." ) state = self._state - # a permanent refusal stopped the network loop: report it on every run - # without waiting and without touching the broker again - if state.refused_code in PERMANENT_CONNACK_CODES: + if not self._subscribed_once: + # cold start: bounded waits for the broker's answers, so the first + # run of a pipeline gets a message or a precise failure + failure = self._wait_until_subscribed(state, timeout) + if failure is not None: + return failure + self._subscribed_once = True + if self._awaiting_retained: + # one-off, bounded: a retained message (or nothing) arrives right + # after SUBACK; without this wait the first run would race it + self._awaiting_retained = False + state.message_received.wait(timeout=timeout) + elif state.subscribe_failed: + # a known refusal is permanent for this configuration return self._handle_failure( - connection_refused_message( - state.refused_code, block_inputs=READER_CONNECTION_INPUTS - ) + f"MQTT broker refused subscription to topic {state.topic!r}." ) - # the background loop owns (re)connecting and (re)subscribing; runs - # only wait for readiness. The wait ends on any CONNACK, so a refusal is - # reported as soon as the broker answers instead of after the timeout - if not state.connack.wait(timeout=timeout): - return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) - if state.refused_code in PERMANENT_CONNACK_CODES: - return self._handle_failure( - connection_refused_message( - state.refused_code, block_inputs=READER_CONNECTION_INPUTS + # never wait once subscribed; drain first, because messages received + # before a drop are real whatever happened to the connection afterwards + try: + received_topic, payload, retained = state.messages.popleft() + except IndexError: + if state.refused_code in PERMANENT_CONNACK_CODES: + # a reconnect refused with 1/2/4/5 stopped the loop: no fresh + # data can arrive until the pipeline is restarted + return self._repeat_last_with_error( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) ) - ) - if not state.connected.wait(timeout=timeout): - if state.refused_code == TRANSIENT_CONNACK_CODE: - return self._handle_failure( + if not state.connected.is_set(): + return self._repeat_last_with_error( connection_refused_message( state.refused_code, block_inputs=READER_CONNECTION_INPUTS ) + if state.refused_code == TRANSIENT_CONNACK_CODE + else NOT_CONNECTED_REPEATING_LAST ) - return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) - # a known refusal is permanent for this configuration: report it - # without paying the acknowledgement wait on every run - if state.subscribe_failed or ( - not state.subscribed.wait(timeout=timeout) and state.subscribe_failed - ): - return self._handle_failure( - f"MQTT broker refused subscription to topic {state.topic!r}." - ) - if not state.subscribed.is_set(): - return self._handle_failure( - "MQTT subscription was not acknowledged within timeout. Raise " - "'timeout' if the broker needs longer to respond." - ) - if self._awaiting_retained: - # one-off, bounded: a retained message (or nothing) arrives right - # after SUBACK; without this wait the first run would race it - self._awaiting_retained = False - state.message_received.wait(timeout=timeout) - try: - received_topic, payload, retained = state.messages.popleft() - except IndexError: + if not state.subscribed.is_set(): + return self._repeat_last_with_error(NOT_RESUBSCRIBED_REPEATING_LAST) # nothing new since the previous run return self._repeat_last() + # a live message ends an outage, a retained duplicate included + self._log_recovery_if_needed() try: value = payload.decode("utf-8") except UnicodeDecodeError: @@ -805,6 +838,52 @@ def _connect_and_read( "error_message": None, } + def _wait_until_subscribed( + self, state: MQTTReaderState, timeout: float + ) -> Optional[BlockResult]: + """The first run's readiness waits; returns the failure result or None.""" + # a permanent refusal stopped the network loop: report it on every run + # without waiting and without touching the broker again + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + # the background loop owns (re)connecting and (re)subscribing; runs + # only wait for readiness. The wait ends on any CONNACK, so a refusal is + # reported as soon as the broker answers instead of after the timeout + if not state.connack.wait(timeout=timeout): + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) + if state.refused_code in PERMANENT_CONNACK_CODES: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + if not state.connected.wait(timeout=timeout): + if state.refused_code == TRANSIENT_CONNACK_CODE: + return self._handle_failure( + connection_refused_message( + state.refused_code, block_inputs=READER_CONNECTION_INPUTS + ) + ) + return self._handle_failure(NOT_CONNECTED_WITHIN_TIMEOUT) + # a known refusal is permanent for this configuration: report it + # without paying the acknowledgement wait on every run + if state.subscribe_failed or ( + not state.subscribed.wait(timeout=timeout) and state.subscribe_failed + ): + return self._handle_failure( + f"MQTT broker refused subscription to topic {state.topic!r}." + ) + if not state.subscribed.is_set(): + return self._handle_failure( + "MQTT subscription was not acknowledged within timeout. Raise " + "'timeout' if the broker needs longer to respond." + ) + return None + def _repeat_last(self) -> BlockResult: if self._last is None: return self._empty_result() @@ -815,6 +894,57 @@ def _repeat_last(self) -> BlockResult: "error_message": None, } + def _repeat_last_with_error(self, message: str) -> BlockResult: + """The outputs while no fresh data can arrive: the last message with the + error fields set, or the empty failure result when nothing was received yet. + + The first call of an outage logs one ERROR line; later calls log nothing + (a video pipeline would otherwise write one line per frame). + """ + if self._outage_started_at is None: + self._outage_started_at = time.monotonic() + self._outage_runs = 0 + # host, port and topic identify the block instance in a shared log + logger.error( + "MQTT Reader %s:%s %r: %s", + self._host, + self._port, + self._state.topic if self._state is not None else None, + message, + ) + self._outage_runs += 1 + if self._last is None: + return { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": message, + } + return { + **self._last, + "is_new": False, + "error_status": True, + "error_message": message, + } + + def _log_recovery_if_needed(self) -> None: + """Close the outage bracket in the log once a live message arrived again.""" + if self._outage_started_at is None: + return + logger.info( + "MQTT Reader %s:%s %r: connection restored after %.1f s; %d runs returned " + "the last message received", + self._host, + self._port, + self._state.topic if self._state is not None else None, + time.monotonic() - self._outage_started_at, + self._outage_runs, + ) + self._outage_started_at = None + self._outage_runs = 0 + @staticmethod def _empty_result() -> BlockResult: return { diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 2af605e522..72a1619ee4 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -10,7 +10,9 @@ from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader import v1 from roboflow_workflows.enterprise_blocks.sinks.mqtt_reader.v1 import ( LATEST_BUFFER_SIZE, + NOT_CONNECTED_REPEATING_LAST, NOT_CONNECTED_WITHIN_TIMEOUT, + NOT_RESUBSCRIBED_REPEATING_LAST, SEQUENTIAL_BUFFER_SIZE, SUBSCRIPTION_REFUSED_QOS, BlockManifest, @@ -173,6 +175,20 @@ def test_client_id_accepts_literal_and_selector(self, client_id): assert manifest.client_id == client_id + def test_timeout_description_scopes_the_wait_to_the_first_connection(self): + schema = BlockManifest.model_json_schema() + + description = schema["properties"]["timeout"]["description"] + assert "first connection" in description + assert "Later runs never wait" in description + + def test_outputs_unchanged_by_outage_handling(self): + names = [output.name for output in BlockManifest.describe_outputs()] + + assert "error_status" in names + assert "error_message" in names + assert len(names) == 6 + def test_qos_description_states_the_persistent_session_rule(self): schema = BlockManifest.model_json_schema() @@ -1060,6 +1076,311 @@ def test_no_raw_prints_in_any_code_path(self, clients, block): mock_print.assert_not_called() +class TestOutageBookkeeping: + """The two helpers behind the outage outputs, driven directly.""" + + def warm_block_with_last(self, clients, block): + block.run(**run_kwargs()) + clients[0].deliver(b'{"state": "RUNNING"}') + block.run(**run_kwargs()) + return block + + def test_repeat_with_error_returns_last_message_with_error_fields( + self, clients, block + ): + self.warm_block_with_last(clients, block) + + result = block._repeat_last_with_error("boom") + + assert result == { + "value": '{"state": "RUNNING"}', + "payload": {"state": "RUNNING"}, + "topic": "plc/state", + "is_new": False, + "error_status": True, + "error_message": "boom", + } + + def test_repeat_with_error_without_last_message_is_the_empty_failure( + self, clients, block + ): + block.run(**run_kwargs()) + + result = block._repeat_last_with_error("boom") + + assert result == { + "value": None, + "payload": {}, + "topic": None, + "is_new": False, + "error_status": True, + "error_message": "boom", + } + + def test_outage_logs_one_error_line_naming_the_instance( + self, clients, block, caplog + ): + self.warm_block_with_last(clients, block) + + with caplog.at_level("ERROR", logger="inference"): + block._repeat_last_with_error("boom") + block._repeat_last_with_error("boom") + + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 1 + assert "MQTT Reader localhost:1883 'plc/state': boom" in errors[0].getMessage() + assert block._outage_runs == 2 + + def test_recovery_logs_one_info_line_with_duration_and_count_and_resets( + self, clients, block, caplog + ): + self.warm_block_with_last(clients, block) + block._repeat_last_with_error("boom") + block._repeat_last_with_error("boom") + + with caplog.at_level("INFO", logger="inference"): + block._log_recovery_if_needed() + block._log_recovery_if_needed() + + infos = [ + r + for r in caplog.records + if r.levelname == "INFO" and "restored after" in r.getMessage() + ] + assert len(infos) == 1 + assert "2 runs returned the last message" in infos[0].getMessage() + assert "localhost:1883 'plc/state'" in infos[0].getMessage() + assert block._outage_started_at is None + assert block._outage_runs == 0 + + def test_recovery_while_healthy_logs_nothing(self, clients, block, caplog): + self.warm_block_with_last(clients, block) + + with caplog.at_level("INFO", logger="inference"): + block._log_recovery_if_needed() + + assert "restored after" not in caplog.text + + def test_first_run_keeps_the_resolved_address_for_the_log(self, clients, block): + block.run(**run_kwargs()) + + assert (block._host, block._port) == ("localhost", 1883) + + def test_close_resets_the_bookkeeping(self, clients, block): + self.warm_block_with_last(clients, block) + block._repeat_last_with_error("boom") + block._subscribed_once = True + + block.close() + + assert block._subscribed_once is False + assert block._outage_started_at is None + assert block._outage_runs == 0 + assert (block._host, block._port) == (None, None) + + +def drop(client: FakeClient) -> None: + client.on_disconnect(client, client.userdata, 1) + + +def reconnect(client: FakeClient, reason_code: int = 0) -> None: + client.on_connect(client, client.userdata, {}, reason_code) + + +RUNNING = b'{"state": "RUNNING"}' +PAUSED = b'{"state": "PAUSED"}' + + +class TestOutage: + """A warm instance (first run reached SUBACK) while the broker is away.""" + + def warm(self, clients, block, **kwargs): + block.run(**run_kwargs(**kwargs)) + clients[0].deliver(RUNNING) + first = block.run(**run_kwargs(**kwargs)) + assert first["is_new"] is True + return clients[0], first + + def test_drop_repeats_last_message_with_error_without_waiting(self, clients, block): + client, first = self.warm(clients, block, timeout=2.0) + drop(client) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert time.monotonic() - started < 0.5 + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": NOT_CONNECTED_REPEATING_LAST, + } + assert "not connected" in result["error_message"] + + def test_repeated_runs_during_outage_log_one_error_line( + self, clients, block, caplog + ): + client, _ = self.warm(clients, block) + drop(client) + + with caplog.at_level("ERROR", logger="inference"): + results = [block.run(**run_kwargs()) for _ in range(4)] + + assert len({str(r) for r in results}) == 1 + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 1 + assert "localhost:1883 'plc/state'" in errors[0].getMessage() + + def test_buffered_messages_drain_before_the_outage_is_reported( + self, clients, block + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + client.deliver(b"B") + drop(client) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(3)] + + assert [r["value"] for r in results] == ["A", "B", "B"] + assert [r["is_new"] for r in results] == [True, True, False] + assert [r["error_status"] for r in results] == [False, False, True] + + def test_drop_before_any_message_returns_empty_failure_without_waiting( + self, clients, block + ): + block.run(**run_kwargs(timeout=2.0)) + drop(clients[0]) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=2.0)) + + assert time.monotonic() - started < 0.5 + assert result["payload"] == {} + assert result["value"] is None + assert result["error_status"] is True + assert result["error_message"] == NOT_CONNECTED_REPEATING_LAST + + def test_recovery_returns_live_message_and_logs_once(self, clients, block, caplog): + client, _ = self.warm(clients, block) + drop(client) + for _ in range(3): + block.run(**run_kwargs()) + + reconnect(client) + client.deliver(PAUSED) + with caplog.at_level("INFO", logger="inference"): + result = block.run(**run_kwargs()) + block.run(**run_kwargs()) + + assert result["payload"] == {"state": "PAUSED"} + assert result["is_new"] is True + assert result["error_status"] is False + infos = [ + r.getMessage() for r in caplog.records if "restored after" in r.getMessage() + ] + assert len(infos) == 1 + assert "3 runs returned the last message" in infos[0] + + def test_second_outage_logs_error_again(self, clients, block, caplog): + client, _ = self.warm(clients, block) + with caplog.at_level("ERROR", logger="inference"): + drop(client) + block.run(**run_kwargs()) + reconnect(client) + client.deliver(PAUSED) + block.run(**run_kwargs()) + drop(client) + block.run(**run_kwargs()) + + errors = [r for r in caplog.records if r.levelname == "ERROR"] + assert len(errors) == 2 + + def test_reconnected_with_nothing_new_repeats_without_error_and_no_recovery_line( + self, clients, block, caplog + ): + client, first = self.warm(clients, block) + drop(client) + block.run(**run_kwargs()) + reconnect(client) + + with caplog.at_level("INFO", logger="inference"): + result = block.run(**run_kwargs()) + + assert result == {**first, "is_new": False} + assert "restored after" not in caplog.text + + def test_connected_but_not_resubscribed_repeats_with_resubscribe_message( + self, clients, block + ): + client, first = self.warm(clients, block) + drop(client) + client.ack_subscriptions = False + reconnect(client) + + result = block.run(**run_kwargs()) + + assert result == { + **first, + "is_new": False, + "error_status": True, + "error_message": NOT_RESUBSCRIBED_REPEATING_LAST, + } + + def test_broker_unavailable_on_reconnect_repeats_with_its_reason( + self, clients, block + ): + client, first = self.warm(clients, block) + drop(client) + reconnect(client, 3) + + result = block.run(**run_kwargs()) + + assert result["value"] == first["value"] + assert result["error_status"] is True + assert "broker unavailable" in result["error_message"] + assert "retrying in the background" in result["error_message"] + assert client.disconnected is False + + def test_refused_reconnect_drains_then_repeats_with_refusal( + self, clients, block, caplog + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + drop(client) + reconnect(client, 5) + + with caplog.at_level("ERROR", logger="inference"): + results = [ + block.run(**run_kwargs(read_mode="sequential")) for _ in range(3) + ] + + assert [r["value"] for r in results] == ["A", "A", "A"] + assert [r["error_status"] for r in results] == [False, True, True] + assert "not authorised" in results[1]["error_message"] + assert "does not retry" in results[1]["error_message"] + assert results[2] == results[1] + assert client.disconnected is True + run_path_errors = [ + r + for r in caplog.records + if r.levelname == "ERROR" and r.getMessage().startswith("MQTT Reader ") + ] + assert len(run_path_errors) == 1 + assert "restored after" not in caplog.text + + def test_cold_instance_still_waits_for_the_broker( + self, clients, block, monkeypatch + ): + monkeypatch.setattr(FakeClient, "fire_on_connect", False) + + started = time.monotonic() + result = block.run(**run_kwargs(timeout=0.3)) + + assert time.monotonic() - started >= 0.3 + assert result["error_message"] == NOT_CONNECTED_WITHIN_TIMEOUT + assert block._subscribed_once is False + + class TestBrokerPolicy: def test_allowlisted_broker_is_used(self, clients, block, monkeypatch): monkeypatch.setattr( From 9175f6534c52e533275c1bc052b6d84afc816324 Mon Sep 17 00:00:00 2001 From: Tewan Date: Thu, 24 Sep 2026 02:38:04 +0800 Subject: [PATCH 14/14] Address the local review of the MQTT refusal and outage changes The connect callbacks set the CONNACK event before recording the outcome, so a first run woken by it could read no refusal, wait the full timeout and report the generic not-connected text. Both blocks now record the reason code and the connected state before the event fires; a deterministic test observes the state at the moment the event is set. The reader now always drains its buffer before returning an empty result: a refused subscription on a warm instance, or a cold start whose persistent-session backlog landed before a refused SUBACK, returns the buffered messages first. Two timing bounds get slack: the cold-start lower bound tolerates timer granularity and the fast-path upper bound in the integration tests allows a loaded runner. Co-Authored-By: Claude Fable 5.1 --- .../test_workflow_with_mqtt_reader.py | 4 +- .../enterprise_blocks/sinks/mqtt_reader/v1.py | 41 +++++++----- .../enterprise_blocks/sinks/mqtt_writer/v1.py | 11 ++-- .../core_steps/sinks/test_mqtt_reader.py | 65 ++++++++++++++++++- .../core_steps/sinks/test_mqtt_writer.py | 22 +++++++ 5 files changed, 119 insertions(+), 24 deletions(-) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py index a1bdf22353..85128b6ea8 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_mqtt_reader.py @@ -224,7 +224,7 @@ def test_outage_repeats_last_message_without_waiting_then_recovers(broker, caplo for _ in range(3): started = time.monotonic() during.append(block.run(**kwargs)) - assert time.monotonic() - started < 0.2 + assert time.monotonic() - started < 0.5 # and the broker comes back: the retained copy is re-sent on the # resubscribe (the buffer holds one message in latest mode), then a @@ -630,7 +630,7 @@ def test_persistent_session_outage_repeats_last_then_delivers_queued(session_bro session_broker.publish("plc/state", PAUSED, qos=1) started = time.monotonic() during = block.run(**kwargs) - assert time.monotonic() - started < 0.2 + assert time.monotonic() - started < 0.5 # and paho reconnects on its own assert wait_until(lambda: len(block._state.messages) == 1, timeout=10) diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py index 6cccd93fb5..d24a1a4ae7 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_reader/v1.py @@ -224,16 +224,18 @@ def mqtt_on_connect( client, state: MQTTReaderState, flags, reason_code, properties=None ): # paho invokes on_connect for accepted and rejected CONNACK alike; - # only reason_code 0 means an established MQTT session - state.connack.set() + # only reason_code 0 means an established MQTT session. The outcome is + # recorded before `connack` wakes a waiting run, so the run never reads a + # half-updated state if reason_code != 0: state.refused_code = reason_code + state.connected.clear() + state.connack.set() logger.error( "MQTT connection refused: %s (code %s)", mqtt.connack_string(reason_code), reason_code, ) - state.connected.clear() if reason_code in PERMANENT_CONNACK_CODES: # paho would otherwise reconnect with the same credentials forever; # disconnect() puts it in the disconnecting state, so its loop ends @@ -247,6 +249,7 @@ def mqtt_on_connect( bool(flags.get("session present")), ) state.connected.set() + state.connack.set() # a clean-session broker forgets subscriptions on every reconnect, so the # subscription is (re)established here rather than once after connect(); on # a persistent session the repeated SUBSCRIBE is idempotent @@ -764,24 +767,28 @@ def _connect_and_read( # cold start: bounded waits for the broker's answers, so the first # run of a pipeline gets a message or a precise failure failure = self._wait_until_subscribed(state, timeout) - if failure is not None: + if failure is not None and not state.messages: return failure - self._subscribed_once = True - if self._awaiting_retained: - # one-off, bounded: a retained message (or nothing) arrives right - # after SUBACK; without this wait the first run would race it - self._awaiting_retained = False - state.message_received.wait(timeout=timeout) - elif state.subscribe_failed: - # a known refusal is permanent for this configuration - return self._handle_failure( - f"MQTT broker refused subscription to topic {state.topic!r}." - ) - # never wait once subscribed; drain first, because messages received - # before a drop are real whatever happened to the connection afterwards + if failure is None: + self._subscribed_once = True + if self._awaiting_retained: + # one-off, bounded: a retained message (or nothing) arrives + # right after SUBACK; without this wait the first run would + # race it + self._awaiting_retained = False + state.message_received.wait(timeout=timeout) + # never wait once subscribed; always drain first, because messages + # received before a drop or a refusal are real whatever happened to the + # connection afterwards (a persistent session's backlog even lands + # before the SUBACK that may refuse the subscription) try: received_topic, payload, retained = state.messages.popleft() except IndexError: + if state.subscribe_failed: + # a known refusal is permanent for this configuration + return self._handle_failure( + f"MQTT broker refused subscription to topic {state.topic!r}." + ) if state.refused_code in PERMANENT_CONNACK_CODES: # a reconnect refused with 1/2/4/5 stopped the loop: no fresh # data can arrive until the pipeline is restarted diff --git a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py index 176dead7fb..b7a69afec3 100644 --- a/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py +++ b/workflows/roboflow_workflows/enterprise_blocks/sinks/mqtt_writer/v1.py @@ -148,20 +148,23 @@ def mqtt_on_connect( client, state: MQTTWriterState, flags, reason_code, properties=None ): # paho invokes on_connect for accepted and rejected CONNACK alike; - # only reason_code 0 means an established MQTT session - state.connack.set() + # only reason_code 0 means an established MQTT session. The outcome is + # recorded before `connack` wakes a waiting run, so the run never reads a + # half-updated state if reason_code == 0: state.refused_code = None - logger.info("MQTT client connected") state.connected.set() + state.connack.set() + logger.info("MQTT client connected") return state.refused_code = reason_code + state.connected.clear() + state.connack.set() logger.error( "MQTT connection refused: %s (code %s)", mqtt.connack_string(reason_code), reason_code, ) - state.connected.clear() if reason_code in PERMANENT_CONNACK_CODES: # paho would otherwise reconnect with the same credentials forever; # disconnect() puts it in the disconnecting state, so its loop ends diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py index 72a1619ee4..8417fc8dda 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_reader.py @@ -348,6 +348,29 @@ def test_on_connect_keeps_the_loop_for_broker_unavailable(self): assert client.disconnected is False assert state.refused_code == 3 + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_completes_the_state_before_waking_a_waiter(self, reason_code): + # a run wakes on `connack`; what it then reads must already be final + state = MQTTReaderState(topic="plc/#", qos=1, buffer_size=1) + seen = {} + + class RecordingEvent(threading.Event): + def set(self): + seen["refused_code"] = state.refused_code + seen["connected"] = state.connected.is_set() + super().set() + + state.connack = RecordingEvent() + client = FakeClient(userdata=state) + client.on_subscribe = mqtt_on_subscribe + + mqtt_on_connect(client, state, {}, reason_code) + + assert seen == { + "refused_code": reason_code or None, + "connected": reason_code == 0, + } + def test_on_connect_logs_the_refusal_reason(self, caplog): state = MQTTReaderState(topic="plc/#", qos=0, buffer_size=1) @@ -1368,6 +1391,45 @@ def test_refused_reconnect_drains_then_repeats_with_refusal( assert len(run_path_errors) == 1 assert "restored after" not in caplog.text + def test_refused_resubscription_drains_then_returns_empty_failure( + self, clients, block + ): + client, _ = self.warm(clients, block, read_mode="sequential") + client.deliver(b"A") + drop(client) + client.granted_qos = SUBSCRIPTION_REFUSED_QOS + reconnect(client) + + results = [block.run(**run_kwargs(read_mode="sequential")) for _ in range(2)] + + assert results[0]["value"] == "A" + assert results[0]["error_status"] is False + assert results[1]["payload"] == {} + assert results[1]["error_status"] is True + assert "refused subscription" in results[1]["error_message"] + + def test_cold_refused_subscription_drains_the_backlog_first( + self, clients, block, monkeypatch + ): + # a persistent session's backlog lands right after CONNACK, before the + # SUBACK that refuses the subscription: the messages are still real + monkeypatch.setattr(FakeClient, "granted_qos", SUBSCRIPTION_REFUSED_QOS) + original_loop_start = FakeClient.loop_start + + def loop_start_with_backlog(client): + original_loop_start(client) + client.deliver(b"A") + + with patch.object(FakeClient, "loop_start", loop_start_with_backlog): + first = block.run(**run_kwargs(read_mode="sequential")) + second = block.run(**run_kwargs(read_mode="sequential")) + + assert first["value"] == "A" + assert first["error_status"] is False + assert block._subscribed_once is False + assert second["payload"] == {} + assert "refused subscription" in second["error_message"] + def test_cold_instance_still_waits_for_the_broker( self, clients, block, monkeypatch ): @@ -1376,7 +1438,8 @@ def test_cold_instance_still_waits_for_the_broker( started = time.monotonic() result = block.run(**run_kwargs(timeout=0.3)) - assert time.monotonic() - started >= 0.3 + # a hair of timer slack is allowed on the lower bound + assert time.monotonic() - started >= 0.25 assert result["error_message"] == NOT_CONNECTED_WITHIN_TIMEOUT assert block._subscribed_once is False diff --git a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py index 2334e8a804..7b1034f1a7 100644 --- a/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py +++ b/workflows/tests/unit_tests/core_steps/sinks/test_mqtt_writer.py @@ -1,3 +1,4 @@ +import threading import time from typing import get_args from unittest.mock import MagicMock, patch @@ -166,6 +167,27 @@ def test_on_connect_keeps_the_loop_for_broker_unavailable(self): client.disconnect.assert_not_called() assert state.refused_code == 3 + @pytest.mark.parametrize("reason_code", [0, 3, 5]) + def test_on_connect_completes_the_state_before_waking_a_waiter(self, reason_code): + # a run wakes on `connack`; what it then reads must already be final + state = MQTTWriterState() + seen = {} + + class RecordingEvent(threading.Event): + def set(self): + seen["refused_code"] = state.refused_code + seen["connected"] = state.connected.is_set() + super().set() + + state.connack = RecordingEvent() + + mqtt_on_connect(MagicMock(), state, {}, reason_code) + + assert seen == { + "refused_code": reason_code or None, + "connected": reason_code == 0, + } + def test_on_connect_logs_the_refusal_reason(self, caplog): with caplog.at_level("ERROR", logger="inference"): mqtt_on_connect(MagicMock(), MQTTWriterState(), {}, 4)