feat: add MQTT QoS support and timestamp to A01 protocol payload (#891)

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.

Co-authored-by: NOisi-X <xcynoisi@gmail.com>

Author: NOisi-x

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,23 +40,32 @@ 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]: ...
 
 
 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)
 
     # For commands that set values: send the command and do not
     # 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

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,19 +89,23 @@ 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)
         except Exception as e:
             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

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.

roborock/mqtt/session.py

@@ -3,13 +3,32 @@
 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
 from roborock.mqtt.health_manager import HealthManager
 
 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

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,

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:

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

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_|

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