feat: track A01 device state from push updates (#941)

* feat: track A01 device state from push updates

* fix: update last_message_time on unsubscribed query responses

Author: piitaya

docs/DEVICES.md

@@ -12,9 +12,11 @@ Cloud and Network.
     *   **Vacuums (V1)**: Use `device.v1_properties` to access traits like `status` or `consumables`.
         *   Call `await trait.refresh()` to update state.
         *   Use `device.v1_properties.command.send()` for raw commands (start/stop).
-    *   **Washers (A01)**: Use `device.a01_properties` for Dyad/Zeo devices.
-        *   Use `await device.a01_properties.query_values([...])` to get state.
-        *   Use `await device.a01_properties.set_value(protocol, value)` to control.
+    *   **Washers (A01)**: Use `device.dyad` or `device.zeo` for Dyad/Zeo devices.
+        *   Read `api.values` for the latest known state, merged from query responses and unsolicited pushes.
+        *   Use `api.add_update_listener(callback)` to be notified when `values` changes.
+        *   Use `await api.query_values([...])` to poll specific data points.
+        *   Use `await api.set_value(protocol, value)` to control.
     *   **Vacuums (B01 Q10)**: Use `device.b01_q10_properties` for Q10 series devices.
         *   Use `device.b01_q10_properties.vacuum` to access vacuum commands (start, pause, stop, dock, empty dustbin, set clean mode, set fan level).
         *   Use `device.b01_q10_properties.command.send()` for raw DP commands.

roborock/devices/device.py

@@ -206,6 +206,8 @@ async def connect(self) -> None:
                 await self.b01_q7_properties.start()
             elif self.zeo is not None:
                 await self.zeo.start()
+            elif self.dyad is not None:
+                await self.dyad.start()
         except RoborockException:
             # Expected: start() can fail transiently. Unsubscribe before propagating
             # so the retry by connect_loop() gets a clean channel.
@@ -238,6 +240,8 @@ async def close(self) -> None:
             await self.b01_q7_properties.close()
         if self.zeo is not None:
             self.zeo.close()
+        if self.dyad is not None:
+            self.dyad.close()
         if self._unsub:
             self._unsub()
             self._unsub = None

roborock/devices/device_manager.py

@@ -255,7 +255,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                 )
             case DeviceVersion.A01:
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
-                trait = a01.create(product, channel)
+                trait = a01.create(product, channel, device_status=device.device_status)
             case DeviceVersion.B01:
                 mqtt_channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 model_part = product.model.split(".")[-1]

roborock/devices/traits/a01/__init__.py

@@ -6,26 +6,30 @@
 Using A01 APIs
 --------------
 A01 devices expose a single API object that handles all device interactions. This API is
-available on the device instance (typically via `device.a01_properties`).
+available on the device instance (`device.dyad` or `device.zeo`).
 
 The API provides these methods:
 1.  **query_values(protocols)**: Fetches current state for specific data points.
     You must pass a list of protocol enums (e.g. `RoborockDyadDataProtocol` or
     `RoborockZeoProtocol`) to request specific data.
 2.  **set_value(protocol, value)**: Sends a command to the device to change a setting
     or perform an action.
-3.  **add_listener(callback)**: Subscribes to state the device pushes on its own (for
-    example when its state changes), invoking the callback with decoded values.
-
-Note that these APIs fetch data directly from the device upon request and do not
-cache state internally.
+3.  **values**: The latest known state, merged from query responses and unsolicited
+    pushes in arrival order.
+4.  **add_update_listener(callback)**: Registers a callback invoked whenever `values`
+    changes; read `values` from the callback to get the updated state.
+
+The device pushes only the data points that changed, so `values` is the merged view
+of everything seen so far. State tracking is active once the device is connected
+(the device calls `start()` on the API, which subscribes to the MQTT topic).
 """
 
 import json
 import logging
+from abc import abstractmethod
 from collections.abc import Callable
-from datetime import time
-from typing import Any
+from datetime import UTC, datetime, time
+from typing import Any, Generic, TypeVar
 
 from roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
 from roborock.data.dyad.dyad_code_mappings import (
@@ -73,6 +77,7 @@
 _LOGGER = logging.getLogger(__name__)
 
 __all__ = [
+    "A01Api",
     "DyadApi",
     "ZeoApi",
 ]
@@ -155,14 +160,121 @@ def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
 
 
 _DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol)
+_ZEO_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockZeoProtocol)
 
+_P = TypeVar("_P", RoborockDyadDataProtocol, RoborockZeoProtocol)
 
-class DyadApi(Trait):
-    """API for interacting with Dyad devices."""
 
-    def __init__(self, channel: MqttChannel) -> None:
-        """Initialize the Dyad API."""
+class A01Api(Trait, TraitUpdateListener, Generic[_P]):
+    """Base class for A01 device APIs with device state tracking.
+
+    Query responses and unsolicited pushes both arrive on the same MQTT topic,
+    so a single subscription merges every decoded message into `values` in
+    arrival order. Update listeners are notified whenever a value changes.
+    """
+
+    def __init__(self, channel: MqttChannel, initial_status: dict[int, Any] | None = None) -> None:
+        """Initialize the A01 API, optionally seeding `values` from a cloud status snapshot."""
+        TraitUpdateListener.__init__(self, _LOGGER)
         self._channel = channel
+        self._values: dict[_P, Any] = {}
+        self._unsub: Callable[[], None] | None = None
+        self._last_message_time: datetime | None = None
+        if initial_status:
+            self._merge_values(self._decode_datapoints(initial_status))
+
+    @property
+    def values(self) -> dict[_P, Any]:
+        """Latest known device state, merged from query responses and pushes.
+
+        The device pushes only the data points that changed, so this is the
+        merged view of everything seen so far. A protocol the device has not
+        reported yet is absent from the dictionary.
+        """
+        return dict(self._values)
+
+    @property
+    def last_message_time(self) -> datetime | None:
+        """Time the last message was received from the device.
+
+        Updated on every decoded message, even when no value changed: idle
+        devices push an identical heartbeat, so this is the liveness signal
+        even when `values` stays the same and update listeners stay silent.
+        The initial cloud status snapshot does not count as a message.
+        """
+        return self._last_message_time
+
+    async def start(self) -> None:
+        """Subscribe to the device state topic and start tracking `values`."""
+        await self._ensure_subscribed()
+
+    def close(self) -> None:
+        """Unsubscribe from MQTT push and release resources."""
+        if self._unsub is not None:
+            self._unsub()
+            self._unsub = None
+
+    async def _ensure_subscribed(self) -> None:
+        """Subscribe to MQTT DPS push (idempotent)."""
+        if self._unsub is not None:
+            return
+        self._unsub = await self._channel.subscribe(self._on_message)
+
+    @abstractmethod
+    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[_P, Any]:
+        """Convert raw datapoints to typed values, skipping unknown codes."""
+
+    def _on_message(self, message: RoborockMessage) -> None:
+        """Handle a message on the device topic (query response or push)."""
+        if message.protocol != RoborockMessageProtocol.RPC_RESPONSE:
+            return
+        try:
+            datapoints = decode_rpc_response(message)
+        except RoborockException:
+            _LOGGER.debug("Dropped malformed push message", exc_info=True)
+            return
+        self._last_message_time = datetime.now(UTC)
+        self._merge_values(self._decode_datapoints(datapoints))
+
+    def _merge_query_response(self, values: dict[_P, Any]) -> None:
+        """Record a successful query response when there is no subscription.
+
+        When subscribed, the response was already merged in arrival order and
+        timestamped by `_on_message`; merging again here could overwrite a
+        push that arrived after it.
+        """
+        if self._unsub is not None:
+            return
+        self._last_message_time = datetime.now(UTC)
+        self._merge_values(values)
+
+    def _merge_values(self, values: dict[_P, Any]) -> None:
+        """Merge decoded values into the cache and notify on change."""
+        changed = False
+        for protocol, value in values.items():
+            if value is None:
+                continue
+            if protocol not in self._values or self._values[protocol] != value:
+                self._values[protocol] = value
+                changed = True
+        if changed:
+            self._notify_update()
+
+
+class DyadApi(A01Api[RoborockDyadDataProtocol]):
+    """API for interacting with Dyad devices."""
+
+    name = "dyad"
+
+    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockDyadDataProtocol, Any]:
+        """Convert raw datapoints to typed values, skipping unknown codes."""
+        values: dict[RoborockDyadDataProtocol, Any] = {}
+        for code, value in datapoints.items():
+            if code not in _DYAD_PROTOCOL_VALUES:
+                continue
+            protocol = RoborockDyadDataProtocol(code)
+            values[protocol] = convert_dyad_value(protocol, value)
+        return values
 
     async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
         """Query the device for the values of the given Dyad protocols."""
@@ -171,7 +283,9 @@ async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[
             {RoborockDyadDataProtocol.ID_QUERY: protocols},
             value_encoder=json.dumps,
         )
-        return {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
+        values = {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
+        self._merge_query_response(values)
+        return values
 
     async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
         """Set a value for a specific protocol on the device."""
@@ -184,38 +298,34 @@ async def add_listener(self, callback: Callable[[dict[RoborockDyadDataProtocol,
         The callback is invoked with decoded values whenever the device sends a
         message, including unsolicited pushes when its state changes. Only known
         protocols are delivered. Returns a callable to remove the listener.
+
+        Prefer `add_update_listener` together with `values`, which handle the
+        merging of partial pushes for you.
         """
 
         def on_message(message: RoborockMessage) -> None:
             try:
                 datapoints = decode_rpc_response(message)
             except RoborockException:
                 return
-            values: dict[RoborockDyadDataProtocol, Any] = {}
-            for code, value in datapoints.items():
-                if code not in _DYAD_PROTOCOL_VALUES:
-                    continue
-                protocol = RoborockDyadDataProtocol(code)
-                values[protocol] = convert_dyad_value(protocol, value)
-            if values:
+            if values := self._decode_datapoints(datapoints):
                 callback(values)
 
         return await self._channel.subscribe(on_message)
 
 
-class ZeoApi(Trait, TraitUpdateListener):
+class ZeoApi(A01Api[RoborockZeoProtocol]):
     """API for interacting with Zeo devices."""
 
     name = "zeo"
 
-    def __init__(self, channel: MqttChannel, model: str | None = None) -> None:
+    def __init__(
+        self, channel: MqttChannel, model: str | None = None, initial_status: dict[int, Any] | None = None
+    ) -> None:
         """Initialize the Zeo API."""
-        TraitUpdateListener.__init__(self, _LOGGER)
-        self._channel = channel
-        self._dps_cache: dict[int, Any] = {}
-        self._dps_unsub: Callable[[], None] | None = None
         self._feature_bits: int = 0
         self._model = model
+        super().__init__(channel, initial_status)
 
     async def start(self) -> None:
         """Subscribe to MQTT push and trigger a full state sync.
@@ -230,17 +340,15 @@ async def start(self) -> None:
         await self._force_load()
         await self._load_feature_dps()
 
-    def close(self) -> None:
-        """Unsubscribe from MQTT push and release resources."""
-        if self._dps_unsub is not None:
-            self._dps_unsub()
-            self._dps_unsub = None
-
-    async def _ensure_subscribed(self) -> None:
-        """Subscribe to MQTT DPS push (idempotent)."""
-        if self._dps_unsub is not None:
-            return
-        self._dps_unsub = await self._channel.subscribe(self._on_dps_message)
+    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockZeoProtocol, Any]:
+        """Convert raw datapoints to typed values, skipping unknown codes."""
+        values: dict[RoborockZeoProtocol, Any] = {}
+        for code, value in datapoints.items():
+            if code not in _ZEO_PROTOCOL_VALUES:
+                continue
+            protocol = RoborockZeoProtocol(code)
+            values[protocol] = convert_zeo_value(protocol, value)
+        return values
 
     async def _force_load(self) -> None:
         """Send ID_QUERY with the base DP list to trigger a full state push.
@@ -279,39 +387,45 @@ def supports(self, feature: ZeoFeatureBits) -> bool:
         """Check whether the device supports a given feature bit."""
         return bool(self._feature_bits & (1 << feature.value))
 
-    def _on_dps_message(self, message: RoborockMessage) -> None:
-        """Handle unsolicited MQTT push (protocol 102 — RPC_RESPONSE)."""
-        if message.protocol != RoborockMessageProtocol.RPC_RESPONSE:
-            return
-        try:
-            decoded = decode_rpc_response(message)
-        except RoborockException:
-            _LOGGER.debug("Dropped malformed push message", exc_info=True)
-            return
-        self._dps_cache.update(decoded)
-        self._notify_update()
-
     async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
         """Query the device for the values of the given protocols."""
         response = await send_decoded_command(
             self._channel,
             {RoborockZeoProtocol.ID_QUERY: protocols},
             value_encoder=json.dumps,
         )
-        return {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
+        values = {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
+        self._merge_query_response(values)
+        return values
 
     async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
         """Set a value for a specific protocol on the device."""
         params = {protocol: value}
         return await send_decoded_command(self._channel, params, value_encoder=lambda x: x)
 
 
-def create(product: HomeDataProduct, mqtt_channel: MqttChannel) -> DyadApi | ZeoApi:
-    """Create traits for A01 devices."""
+def _parse_device_status(device_status: dict | None) -> dict[int, Any] | None:
+    """Normalize the cloud home data status snapshot to integer datapoint codes."""
+    if not device_status:
+        return None
+    try:
+        return {int(code): value for code, value in device_status.items()}
+    except (TypeError, ValueError):
+        _LOGGER.debug("Ignoring malformed device status snapshot: %s", device_status)
+        return None
+
+
+def create(product: HomeDataProduct, mqtt_channel: MqttChannel, device_status: dict | None = None) -> DyadApi | ZeoApi:
+    """Create traits for A01 devices.
+
+    The optional `device_status` is the cloud home data status snapshot, used
+    to seed `values` so state is available before the first device round trip.
+    """
+    initial_status = _parse_device_status(device_status)
     match product.category:
         case RoborockCategory.WET_DRY_VAC:
-            return DyadApi(mqtt_channel)
+            return DyadApi(mqtt_channel, initial_status=initial_status)
         case RoborockCategory.WASHING_MACHINE:
-            return ZeoApi(mqtt_channel, model=product.model)
+            return ZeoApi(mqtt_channel, model=product.model, initial_status=initial_status)
         case _:
             raise NotImplementedError(f"Unsupported category {product.category}")