From 0f9bad2c946492d8da02e84f8f925045d9f73fe5 Mon Sep 17 00:00:00 2001 From: NOisi-X Date: Mon, 20 Jul 2026 11:39:56 +0800 Subject: [PATCH] feat: add MQTT QoS support and timestamp to A01 protocol payload Add MqttQos enum (AT_MOST_ONCE=0, AT_LEAST_ONCE=1, EXACTLY_ONCE=2) and thread a qos parameter through the publish chain (MqttSession -> MqttChannel -> send_decoded_command). All existing callers keep default AT_MOST_ONCE (backward compatible). Also add a unix timestamp field to A01 encode_mqtt_payload, required by Zeo/Dyad devices for command acceptance. --- roborock/devices/rpc/a01_channel.py | 15 +++++++++-- roborock/devices/transport/mqtt_channel.py | 10 ++++--- roborock/mqtt/roborock_session.py | 18 ++++++++----- roborock/mqtt/session.py | 26 ++++++++++++++++++- roborock/protocols/a01_protocol.py | 6 ++++- roborock/testing/channel.py | 7 ++++- tests/devices/traits/a01/test_init.py | 12 ++++++--- .../__snapshots__/test_device_manager.ambr | 9 ++++--- tests/fixtures/logging_fixtures.py | 1 + 9 files changed, 82 insertions(+), 22 deletions(-) diff --git a/roborock/devices/rpc/a01_channel.py b/roborock/devices/rpc/a01_channel.py index 2e2f9ceac..66b8b27d7 100644 --- a/roborock/devices/rpc/a01_channel.py +++ b/roborock/devices/rpc/a01_channel.py @@ -7,6 +7,7 @@ from roborock.devices.transport.mqtt_channel import MqttChannel from roborock.exceptions import RoborockException +from roborock.mqtt.session import MqttQos from roborock.protocols.a01_protocol import ( decode_rpc_response, encode_mqtt_payload, @@ -30,6 +31,7 @@ async def send_decoded_command( mqtt_channel: MqttChannel, params: dict[RoborockDyadDataProtocol, Any], value_encoder: Callable[[Any], Any] | None = None, + qos: MqttQos = MqttQos.AT_MOST_ONCE, ) -> dict[RoborockDyadDataProtocol, Any]: ... @@ -38,6 +40,7 @@ async def send_decoded_command( mqtt_channel: MqttChannel, params: dict[RoborockZeoProtocol, Any], value_encoder: Callable[[Any], Any] | None = None, + qos: MqttQos = MqttQos.AT_MOST_ONCE, ) -> dict[RoborockZeoProtocol, Any]: ... @@ -45,8 +48,16 @@ async def send_decoded_command( mqtt_channel: MqttChannel, params: dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any], value_encoder: Callable[[Any], Any] | None = None, + qos: MqttQos = MqttQos.AT_MOST_ONCE, ) -> dict[RoborockDyadDataProtocol, Any] | dict[RoborockZeoProtocol, Any]: - """Send a command on the MQTT channel and get a decoded response.""" + """Send a command on the MQTT channel and get a decoded response. + + Args: + mqtt_channel: The MQTT channel to send the command on. + params: The parameters to send. + value_encoder: A function to encode the values of the dictionary. + qos: The MQTT QoS level. Defaults to AT_MOST_ONCE. + """ _LOGGER.debug("Sending MQTT command: %s", params) roborock_message = encode_mqtt_payload(params, value_encoder) @@ -54,7 +65,7 @@ async def send_decoded_command( # block waiting for a response. Queries are handled below. param_values = {int(k): v for k, v in params.items()} if not (query_values := param_values.get(_ID_QUERY)): - await mqtt_channel.publish(roborock_message) + await mqtt_channel.publish(roborock_message, qos=qos) return {} # Merge any results together than contain the requested data. This diff --git a/roborock/devices/transport/mqtt_channel.py b/roborock/devices/transport/mqtt_channel.py index 5ff0ab085..1a0f81e64 100644 --- a/roborock/devices/transport/mqtt_channel.py +++ b/roborock/devices/transport/mqtt_channel.py @@ -8,7 +8,7 @@ from roborock.data import HomeDataDevice, RRiot, UserData from roborock.exceptions import RoborockException from roborock.mqtt.health_manager import HealthManager -from roborock.mqtt.session import MqttParams, MqttSession, MqttSessionException +from roborock.mqtt.session import MqttParams, MqttQos, MqttSession, MqttSessionException from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder from roborock.roborock_message import RoborockMessage from roborock.util import RoborockLoggerAdapter @@ -89,11 +89,15 @@ async def subscribe_stream(self) -> AsyncGenerator[RoborockMessage, None]: finally: unsub() - async def publish(self, message: RoborockMessage) -> None: + async def publish(self, message: RoborockMessage, qos: MqttQos = MqttQos.AT_MOST_ONCE) -> None: """Publish a command message. The caller is responsible for handling any responses and associating them with the incoming request. + + Args: + message: The message to publish. + qos: The MQTT QoS level. Defaults to AT_MOST_ONCE. """ try: encoded_msg = self._encoder(message) @@ -101,7 +105,7 @@ async def publish(self, message: RoborockMessage) -> None: self._logger.exception("Error encoding MQTT message: %s", e) raise RoborockException(f"Failed to encode MQTT message: {e}") from e try: - return await self._mqtt_session.publish(self._publish_topic, encoded_msg) + return await self._mqtt_session.publish(self._publish_topic, encoded_msg, qos=qos) except MqttSessionException as e: self._logger.debug("Error publishing MQTT message: %s", e) raise RoborockException(f"Failed to publish MQTT message: {e}") from e diff --git a/roborock/mqtt/roborock_session.py b/roborock/mqtt/roborock_session.py index ec6e5aa71..15202372e 100644 --- a/roborock/mqtt/roborock_session.py +++ b/roborock/mqtt/roborock_session.py @@ -22,7 +22,7 @@ from roborock.diagnostics import Diagnostics, redact_topic_name from .health_manager import HealthManager -from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized +from .session import MqttParams, MqttQos, MqttSession, MqttSessionException, MqttSessionUnauthorized _LOGGER = logging.getLogger(__name__) _MQTT_LOGGER = logging.getLogger(f"{__name__}.aiomqtt") @@ -361,8 +361,14 @@ def delayed_unsub(): return delayed_unsub - async def publish(self, topic: str, message: bytes) -> None: - """Publish a message on the topic.""" + async def publish(self, topic: str, message: bytes, qos: MqttQos = MqttQos.AT_MOST_ONCE) -> None: + """Publish a message on the topic. + + Args: + topic: The MQTT topic to publish to. + message: The message payload. + qos: The MQTT QoS level. Defaults to AT_MOST_ONCE. + """ _LOGGER.debug("Sending message to topic %s: %s", topic, message) client: aiomqtt.Client async with self._client_lock: @@ -371,7 +377,7 @@ async def publish(self, topic: str, message: bytes) -> None: client = self._client try: with self._diagnostics.timer("publish"): - await client.publish(topic, message) + await client.publish(topic, message, qos=qos) except MqttError as err: raise MqttSessionException(f"Error publishing message: {err}") from err @@ -417,13 +423,13 @@ async def subscribe(self, device_id: str, callback: Callable[[bytes], None]) -> await self._maybe_start() return await self._session.subscribe(device_id, callback) - async def publish(self, topic: str, message: bytes) -> None: + async def publish(self, topic: str, message: bytes, qos: MqttQos = MqttQos.AT_MOST_ONCE) -> None: """Publish a message on the specified topic. This will raise an exception if the message could not be sent. """ await self._maybe_start() - return await self._session.publish(topic, message) + return await self._session.publish(topic, message, qos=qos) async def close(self) -> None: """Cancels the mqtt loop. diff --git a/roborock/mqtt/session.py b/roborock/mqtt/session.py index 9e77b4a86..319615fc4 100644 --- a/roborock/mqtt/session.py +++ b/roborock/mqtt/session.py @@ -3,6 +3,7 @@ from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass, field +from enum import IntEnum from roborock.diagnostics import Diagnostics from roborock.exceptions import RoborockException @@ -10,6 +11,24 @@ DEFAULT_TIMEOUT = 30.0 + +class MqttQos(IntEnum): + """MQTT Quality of Service levels. + + A01 devices (Zeo, Dyad) require ``AT_LEAST_ONCE`` for DP200 (start) + commands. Other protocol versions use ``AT_MOST_ONCE``. + """ + + AT_MOST_ONCE = 0 + """Fire-and-forget. No acknowledgment required.""" + + AT_LEAST_ONCE = 1 + """Guaranteed delivery with possible duplicates. Broker sends PUBACK.""" + + EXACTLY_ONCE = 2 + """Guaranteed delivery with no duplicates. Broker sends PUBREC/PUBREL/PUBCOMP.""" + + SessionUnauthorizedHook = Callable[[], None] @@ -76,10 +95,15 @@ async def subscribe(self, device_id: str, callback: Callable[[bytes], None]) -> """ @abstractmethod - async def publish(self, topic: str, message: bytes) -> None: + async def publish(self, topic: str, message: bytes, qos: MqttQos = MqttQos.AT_MOST_ONCE) -> None: """Publish a message on the specified topic. This will raise an exception if the message could not be sent. + + Args: + topic: The MQTT topic to publish to. + message: The message payload. + qos: The MQTT QoS level. Defaults to AT_MOST_ONCE. """ @abstractmethod diff --git a/roborock/protocols/a01_protocol.py b/roborock/protocols/a01_protocol.py index f3166de87..46db6ef4a 100644 --- a/roborock/protocols/a01_protocol.py +++ b/roborock/protocols/a01_protocol.py @@ -2,6 +2,7 @@ import json import logging +import time from collections.abc import Callable from typing import Any @@ -42,7 +43,10 @@ def encode_mqtt_payload( """ if value_encoder is None: value_encoder = _no_encode - dps_data = {"dps": {key: value_encoder(value) for key, value in data.items()}} + dps_data = { + "dps": {key: value_encoder(value) for key, value in data.items()}, + "t": int(time.time()), + } payload = pad(json.dumps(dps_data).encode("utf-8"), AES.block_size) return RoborockMessage( protocol=RoborockMessageProtocol.RPC_REQUEST, diff --git a/roborock/testing/channel.py b/roborock/testing/channel.py index 022d449bd..65a36bc79 100644 --- a/roborock/testing/channel.py +++ b/roborock/testing/channel.py @@ -12,6 +12,7 @@ from roborock.devices.transport.channel import Channel from roborock.mqtt.health_manager import HealthManager +from roborock.mqtt.session import MqttQos from roborock.protocols.v1_protocol import LocalProtocolVersion from roborock.roborock_message import RoborockMessage @@ -83,6 +84,7 @@ def __init__(self, is_local: bool = False): self.close = MagicMock(side_effect=self._close) self.protocol_version = LocalProtocolVersion.V1 + self.restart = AsyncMock() self.health_manager = HealthManager(self.restart) @@ -102,10 +104,13 @@ def is_local_connected(self) -> bool: """Return true if locally connected.""" return self._is_connected and self._is_local - async def _publish(self, message: RoborockMessage) -> None: + async def _publish(self, message: RoborockMessage, qos: MqttQos = MqttQos.AT_MOST_ONCE) -> None: """Default publish implementation. Records the message in ``published_messages`` and executes ``publish_handler``. + + The ``qos`` parameter is accepted for compatibility with + ``MqttChannel.publish`` but not simulated by the fake channel. """ self.published_messages.append(message) if self.publish_side_effect: diff --git a/tests/devices/traits/a01/test_init.py b/tests/devices/traits/a01/test_init.py index 8e1cb7dd8..7e2fd6a69 100644 --- a/tests/devices/traits/a01/test_init.py +++ b/tests/devices/traits/a01/test_init.py @@ -77,7 +77,8 @@ async def test_dyad_api_query_values(dyad_api: DyadApi, fake_channel: FakeChanne assert message.protocol == RoborockMessageProtocol.RPC_REQUEST assert message.version == b"A01" payload_data = json.loads(unpad(message.payload, AES.block_size)) - assert payload_data == {"dps": {"10000": "[209, 201, 207, 214, 215, 227, 229, 230, 222, 224]"}} + assert payload_data["dps"] == {"10000": "[209, 201, 207, 214, 215, 227, 229, 230, 222, 224]"} + assert "t" in payload_data @pytest.mark.parametrize( @@ -174,7 +175,8 @@ async def test_zeo_api_query_values(zeo_api: ZeoApi, fake_channel: FakeChannel): assert message.protocol == RoborockMessageProtocol.RPC_REQUEST assert message.version == b"A01" payload_data = json.loads(unpad(message.payload, AES.block_size)) - assert payload_data == {"dps": {"10000": "[203, 207, 226, 227, 224, 218]"}} + assert payload_data["dps"] == {"10000": "[203, 207, 226, 227, 224, 218]"} + assert "t" in payload_data @pytest.mark.parametrize( @@ -245,7 +247,8 @@ async def test_dyad_api_set_value(dyad_api: DyadApi, fake_channel: FakeChannel): # decode the payload to verify contents payload_data = json.loads(unpad(message.payload, AES.block_size)) # A01 protocol expects values to be strings in the dps dict - assert payload_data == {"dps": {"209": 1}} + assert payload_data["dps"] == {"209": 1} + assert "t" in payload_data async def test_zeo_api_set_value(zeo_api: ZeoApi, fake_channel: FakeChannel): @@ -261,4 +264,5 @@ async def test_zeo_api_set_value(zeo_api: ZeoApi, fake_channel: FakeChannel): # decode the payload to verify contents payload_data = json.loads(unpad(message.payload, AES.block_size)) # A01 protocol expects values to be strings in the dps dict - assert payload_data == {"dps": {"204": "standard"}} + assert payload_data["dps"] == {"204": "standard"} + assert "t" in payload_data diff --git a/tests/e2e/__snapshots__/test_device_manager.ambr b/tests/e2e/__snapshots__/test_device_manager.ambr index 90f2398dc..157f89c38 100644 --- a/tests/e2e/__snapshots__/test_device_manager.ambr +++ b/tests/e2e/__snapshots__/test_device_manager.ambr @@ -13,12 +13,13 @@ [mqtt <] 00000000 90 04 00 01 00 00 |......| [mqtt >] - 00000000 30 5a 00 20 72 72 2f 6d 2f 69 2f 75 73 65 72 31 |0Z. rr/m/i/user1| + 00000000 30 6a 00 20 72 72 2f 6d 2f 69 2f 75 73 65 72 31 |0j. rr/m/i/user1| 00000010 32 33 2f 31 39 36 34 38 66 39 34 2f 7a 65 6f 5f |23/19648f94/zeo_| 00000020 64 75 69 64 00 41 30 31 00 00 23 82 00 00 23 83 |duid.A01..#...#.| - 00000030 68 a6 a2 24 00 65 00 20 c5 de 2b f6 a9 ba 32 7e |h..$.e. ..+...2~| - 00000040 6b 73 82 bb d8 67 d4 db 7e cd 61 aa 8c 38 56 53 |ks...g..~.a..8VS| - 00000050 ca 4e 15 0d b1 b7 80 a2 0f 16 58 36 |.N........X6| + 00000030 68 a6 a2 24 00 65 00 30 c5 de 2b f6 a9 ba 32 7e |h..$.e.0..+...2~| + 00000040 6b 73 82 bb d8 67 d4 db 7d 80 60 67 80 96 b8 a1 |ks...g..}.`g....| + 00000050 c6 bc 9e d2 da 07 fb d3 79 f5 6f 6d 04 9c 71 00 |........y.om..q.| + 00000060 48 66 3d 7e 5d fe d3 df 18 e4 26 38 |Hf=~].....&8| [mqtt <] 00000000 30 5e 00 20 72 72 2f 6d 2f 6f 2f 75 73 65 72 31 |0^. rr/m/o/user1| 00000010 32 33 2f 31 39 36 34 38 66 39 34 2f 7a 65 6f 5f |23/19648f94/zeo_| diff --git a/tests/fixtures/logging_fixtures.py b/tests/fixtures/logging_fixtures.py index e267f0488..136326983 100644 --- a/tests/fixtures/logging_fixtures.py +++ b/tests/fixtures/logging_fixtures.py @@ -52,6 +52,7 @@ def get_token_bytes(n: int) -> bytes: with ( patch("roborock.devices.transport.local_channel.get_next_int", side_effect=get_next_int), + patch("roborock.protocols.a01_protocol.time.time", return_value=1755750947.0), patch("roborock.protocols.b01_q7_protocol.get_next_int", side_effect=get_next_int), patch("roborock.protocols.v1_protocol.get_next_int", side_effect=get_next_int), patch("roborock.protocols.v1_protocol.get_timestamp", side_effect=get_timestamp),