feat: Implement direct device trait updates from data protocol messages using dps metadata and add corresponding update listeners.

This uses the same dps converter patern used by q10, but does not share code explicitly.

Author: allenporter

roborock/data/v1/v1_containers.py

@@ -37,6 +37,7 @@
     ROBOROCK_G20S_Ultra,
 )
 from roborock.exceptions import RoborockException
+from roborock.roborock_message import RoborockDataProtocol
 
 from ..containers import NamedRoomMapping, RoborockBase, RoborockBaseTimer, _attr_repr
 from .v1_clean_modes import WashTowelModes
@@ -294,11 +295,11 @@ class StatusV2(RoborockBase):
 
     msg_ver: int | None = None
     msg_seq: int | None = None
-    state: RoborockStateCode | None = _requires_schema_code("state")
-    battery: int | None = _requires_schema_code("battery")
+    state: RoborockStateCode | None = field(default=None, metadata={"requires_schema_code": "state", "dps": RoborockDataProtocol.STATE})
+    battery: int | None = field(default=None, metadata={"requires_schema_code": "battery", "dps": RoborockDataProtocol.BATTERY})
     clean_time: int | None = None
     clean_area: int | None = None
-    error_code: RoborockErrorCode | None = _requires_schema_code("error_code")
+    error_code: RoborockErrorCode | None = field(default=None, metadata={"requires_schema_code": "error_code", "dps": RoborockDataProtocol.ERROR_CODE})
     map_present: int | None = None
     in_cleaning: RoborockInCleaning | None = None
     in_returning: int | None = None
@@ -308,12 +309,12 @@ class StatusV2(RoborockBase):
     back_type: int | None = None
     wash_phase: int | None = None
     wash_ready: int | None = None
-    fan_power: int | None = _requires_schema_code("fan_power")
+    fan_power: int | None = field(default=None, metadata={"requires_schema_code": "fan_power", "dps": RoborockDataProtocol.FAN_POWER})
     dnd_enabled: int | None = None
     map_status: int | None = None
     is_locating: int | None = None
     lock_status: int | None = None
-    water_box_mode: int | None = _requires_schema_code("water_box_mode")
+    water_box_mode: int | None = field(default=None, metadata={"requires_schema_code": "water_box_mode", "dps": RoborockDataProtocol.WATER_BOX_MODE})
     water_box_carriage_status: int | None = None
     mop_forbidden_enable: int | None = None
     camera_status: int | None = None
@@ -330,14 +331,14 @@ class StatusV2(RoborockBase):
     debug_mode: int | None = None
     collision_avoid_status: int | None = None
     switch_map_mode: int | None = None
-    dock_error_status: RoborockDockErrorCode | None = _requires_schema_code("dock_error_status")
-    charge_status: int | None = _requires_schema_code("charge_status")
+    dock_error_status: RoborockDockErrorCode | None = field(default=None, metadata={"requires_schema_code": "dock_error_status"})
+    charge_status: int | None = field(default=None, metadata={"requires_schema_code": "charge_status", "dps": RoborockDataProtocol.CHARGE_STATUS})
     unsave_map_reason: int | None = None
     unsave_map_flag: int | None = None
     wash_status: int | None = None
     distance_off: int | None = None
     in_warmup: int | None = None
-    dry_status: int | None = _requires_schema_code("drying_status")
+    dry_status: int | None = field(default=None, metadata={"requires_schema_code": "drying_status", "dps": RoborockDataProtocol.DRYING_STATUS})
     rdt: int | None = None
     clean_percent: int | None = None
     rss: int | None = None
@@ -643,9 +644,9 @@ class ConsumableField(FieldNameBase):
 
 @dataclass
 class Consumable(RoborockBase):
-    main_brush_work_time: int | None = field(metadata={"requires_schema_code": "main_brush_life"}, default=None)
-    side_brush_work_time: int | None = field(metadata={"requires_schema_code": "side_brush_life"}, default=None)
-    filter_work_time: int | None = field(metadata={"requires_schema_code": "filter_life"}, default=None)
+    main_brush_work_time: int | None = field(default=None, metadata={"requires_schema_code": "main_brush_life", "dps": RoborockDataProtocol.MAIN_BRUSH_WORK_TIME})
+    side_brush_work_time: int | None = field(default=None, metadata={"requires_schema_code": "side_brush_life", "dps": RoborockDataProtocol.SIDE_BRUSH_WORK_TIME})
+    filter_work_time: int | None = field(default=None, metadata={"requires_schema_code": "filter_life", "dps": RoborockDataProtocol.FILTER_WORK_TIME})
     filter_element_work_time: int | None = None
     sensor_dirty_time: int | None = None
     strainer_work_times: int | None = None

roborock/devices/device.py

@@ -199,7 +199,7 @@ async def connect(self) -> None:
         unsub = await self._channel.subscribe(self._on_message)
         try:
             if self.v1_properties is not None:
-                await self.v1_properties.discover_features()
+                await self.v1_properties.start()
             elif self.b01_q10_properties is not None:
                 await self.b01_q10_properties.start()
         except RoborockException:
@@ -216,6 +216,8 @@ async def close(self) -> None:
                 await self._connect_task
             except asyncio.CancelledError:
                 pass
+        if self.v1_properties is not None:
+            self.v1_properties.close()
         if self.b01_q10_properties is not None:
             await self.b01_q10_properties.close()
         if self._unsub:

roborock/devices/device_manager.py

@@ -236,6 +236,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                     channel.rpc_channel,
                     channel.mqtt_rpc_channel,
                     channel.map_rpc_channel,
+                    channel.add_dps_listener,
                     web_api,
                     device_cache=device_cache,
                     map_parser_config=map_parser_config,

roborock/devices/rpc/v1_channel.py

@@ -11,6 +11,7 @@
 from dataclasses import dataclass
 from typing import Any, TypeVar
 
+from roborock.callbacks import CallbackList
 from roborock.data import HomeDataDevice, NetworkInfo, RoborockBase, UserData
 from roborock.devices.cache import DeviceCache
 from roborock.devices.transport.channel import Channel
@@ -30,9 +31,10 @@
     V1RpcChannel,
     create_map_response_decoder,
     create_security_data,
+    decode_data_protocol_message,
     decode_rpc_response,
 )
-from roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from roborock.roborock_message import RoborockDataProtocol, RoborockMessage, RoborockMessageProtocol
 from roborock.roborock_typing import RoborockCommand
 from roborock.util import RoborockLoggerAdapter
 
@@ -188,6 +190,7 @@ def __init__(
         self._device_cache = device_cache
         self._reconnect_task: asyncio.Task[None] | None = None
         self._last_network_info_refresh: datetime.datetime | None = None
+        self._dps_listeners = CallbackList[dict[RoborockDataProtocol, Any]]()
 
     @property
     def is_connected(self) -> bool:
@@ -305,12 +308,14 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
             loop = asyncio.get_running_loop()
             self._reconnect_task = loop.create_task(self._background_reconnect())
 
-        if not self.is_local_connected:
-            # We were not able to connect locally, so fallback to MQTT and at least
-            # establish that connection explicitly. If this fails then raise an
-            # error and let the caller know we failed to subscribe.
-            self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
-            self._logger.debug("V1Channel connected to device via MQTT")
+        # Always subscribe to MQTT to receive protocol updates (data points)
+        # even if we have a local connection. Protocol updates only come via cloud/MQTT.
+        # Local connection is used for RPC commands, but push notifications come via MQTT.
+        self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
+        if self.is_local_connected:
+            self._logger.debug("V1Channel connected via local and MQTT (for protocol updates)")
+        else:
+            self._logger.debug("V1Channel connected via MQTT only")
 
         def unsub() -> None:
             """Unsubscribe from all messages."""
@@ -328,6 +333,16 @@ def unsub() -> None:
         self._callback = callback
         return unsub
 
+    def add_dps_listener(self, listener: Callable[[dict[RoborockDataProtocol, Any]], None]) -> Callable[[], None]:
+        """Add a listener for DPS updates.
+
+        This will attach a listener to the existing subscription, invoking
+        the listener whenever new DPS values arrive from the subscription.
+        This will only work if a subscription has already been setup, which is
+        handled by the device setup.
+        """
+        return self._dps_listeners.add_callback(listener)
+
     async def _get_networking_info(self, *, prefer_cache: bool = True) -> NetworkInfo:
         """Retrieve networking information for the device.
 
@@ -428,6 +443,11 @@ def _on_mqtt_message(self, message: RoborockMessage) -> None:
         self._logger.debug("V1Channel received MQTT message: %s", message)
         if self._callback:
             self._callback(message)
+        try:
+            if datapoints := decode_data_protocol_message(message):
+                self._dps_listeners(datapoints)
+        except RoborockException as e:
+            self._logger.debug("Error decoding data protocol message: %s", e)
 
     def _on_local_message(self, message: RoborockMessage) -> None:
         """Handle incoming local messages."""

roborock/devices/traits/v1/__init__.py

@@ -53,15 +53,18 @@
 """
 
 import logging
+from collections.abc import Callable
 from dataclasses import dataclass, field, fields
 from typing import Any, get_args
 
 from roborock.data.containers import HomeData, HomeDataProduct, RoborockBase
 from roborock.data.v1.v1_code_mappings import RoborockDockTypeCode
 from roborock.devices.cache import DeviceCache
 from roborock.devices.traits import Trait
+from roborock.exceptions import RoborockException
 from roborock.map.map_parser import MapParserConfig
-from roborock.protocols.v1_protocol import V1RpcChannel
+from roborock.protocols.v1_protocol import V1RpcChannel, decode_data_protocol_message
+from roborock.roborock_message import RoborockDataProtocol, RoborockMessage
 from roborock.web_api import UserWebApiClient
 
 from . import (
@@ -176,6 +179,7 @@ def __init__(
         rpc_channel: V1RpcChannel,
         mqtt_rpc_channel: V1RpcChannel,
         map_rpc_channel: V1RpcChannel,
+        add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
         web_api: UserWebApiClient,
         device_cache: DeviceCache,
         map_parser_config: MapParserConfig | None = None,
@@ -189,6 +193,8 @@ def __init__(
         self._web_api = web_api
         self._device_cache = device_cache
         self._region = region
+        self._unsub: Callable[[], None] | None = None
+        self._add_dps_listener = add_dps_listener
 
         self.device_features = DeviceFeaturesTrait(product, self._device_cache)
         self.status = StatusTrait(self.device_features, region=self._region)
@@ -227,6 +233,24 @@ def _get_rpc_channel(self, trait: V1TraitMixin) -> V1RpcChannel:
         else:
             return self._rpc_channel
 
+    async def start(self) -> None:
+        """Start the properties API and discover features."""
+        await self.discover_features()
+        self._unsub = self._add_dps_listener(self._on_dps_update)
+
+    def close(self) -> None:
+        if self._unsub:
+            self._unsub()
+
+    def _on_dps_update(self, dps: dict[RoborockDataProtocol, Any]) -> None:
+        """Handle incoming messages from the device.
+
+        This will notify all traits of the new values.
+        """
+        _LOGGER.debug("Received message from device: %s", dps)
+        self.status.update_from_dps(dps)
+        self.consumables.update_from_dps(dps)
+
     async def discover_features(self) -> None:
         """Populate any supported traits that were not initialized in __init__."""
         _LOGGER.debug("Starting optional trait discovery")
@@ -330,6 +354,7 @@ def create(
     rpc_channel: V1RpcChannel,
     mqtt_rpc_channel: V1RpcChannel,
     map_rpc_channel: V1RpcChannel,
+    add_dps_listener: Callable[[Callable[[dict[RoborockDataProtocol, Any]], None]], Callable[[], None]],
     web_api: UserWebApiClient,
     device_cache: DeviceCache,
     map_parser_config: MapParserConfig | None = None,
@@ -343,6 +368,7 @@ def create(
         rpc_channel,
         mqtt_rpc_channel,
         map_rpc_channel,
+        add_dps_listener,
         web_api,
         device_cache,
         map_parser_config,

roborock/devices/traits/v1/common.py

@@ -5,11 +5,14 @@
 
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from dataclasses import fields
-from typing import ClassVar
+from typing import Any, ClassVar
 
+from roborock.callbacks import CallbackList
 from roborock.data import RoborockBase
 from roborock.protocols.v1_protocol import V1RpcChannel
+from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
 
 _LOGGER = logging.getLogger(__name__)
@@ -173,3 +176,74 @@ def wrapper(*args, **kwargs):
 
     cls.map_rpc_channel = True  # type: ignore[attr-defined]
     return wrapper
+
+
+# TODO(allenporter): Merge with roborock.devices.traits.b01.q10.common.TraitUpdateListener
+class TraitUpdateListener(ABC):
+    """Trait update listener.
+
+    This is a base class for traits to support notifying listeners when they
+    have been updated. Clients may register callbacks to be notified when the
+    trait has been updated. When the listener callback is invoked, the client
+    should read the trait's properties to get the updated values.
+    """
+
+    def __init__(self, logger: logging.Logger) -> None:
+        """Initialize the trait update listener."""
+        self._update_callbacks: CallbackList[None] = CallbackList(logger=logger)
+
+    def add_update_listener(self, callback: Callable[[], None]) -> Callable[[], None]:
+        """Register a callback when the trait has been updated.
+
+        Returns a callable to remove the listener.
+        """
+        # We wrap the callback to ignore the value passed to it.
+        return self._update_callbacks.add_callback(lambda _: callback())
+
+    def _notify_update(self) -> None:
+        """Notify all update listeners."""
+        self._update_callbacks(None)
+
+
+class DpsDataConverter:
+    """Utility to handle the transformation and merging of DPS data into models.
+
+    This class pre-calculates the mapping between Data Point IDs and dataclass fields
+    to optimize repeated updates from device streams.
+    """
+
+    def __init__(self, dps_type_map: dict[RoborockDataProtocol, type], dps_field_map: dict[RoborockDataProtocol, str]):
+        """Initialize the converter for a specific RoborockBase-derived class."""
+        self._dps_type_map = dps_type_map
+        self._dps_field_map = dps_field_map
+
+    @classmethod
+    def from_dataclass(cls, dataclass_type: type[RoborockBase]):
+        """Initialize the converter for a specific RoborockBase-derived class."""
+        dps_type_map: dict[RoborockDataProtocol, type] = {}
+        dps_field_map: dict[RoborockDataProtocol, str] = {}
+        for field_obj in fields(dataclass_type):
+            if field_obj.metadata and "dps" in field_obj.metadata:
+                dps_id = field_obj.metadata["dps"]
+                dps_type_map[dps_id] = field_obj.type
+                dps_field_map[dps_id] = field_obj.name
+        return cls(dps_type_map, dps_field_map)
+
+    def update_from_dps(self, target: RoborockBase, decoded_dps: dict[RoborockDataProtocol, Any]) -> bool:
+        """Convert and merge raw DPS data into the target object.
+
+        Uses the pre-calculated type mapping to ensure values are converted to the
+        correct Python types before being updated on the target.
+
+        Args:
+            target: The target object to update.
+            decoded_dps: The decoded DPS data to convert.
+
+        Returns:
+            True if any values were updated, False otherwise.
+        """
+        conversions = RoborockBase.convert_dict(self._dps_type_map, decoded_dps)
+        for dps_id, value in conversions.items():
+            field_name = self._dps_field_map[dps_id]
+            setattr(target, field_name, value)
+        return bool(conversions)

roborock/devices/traits/v1/consumeable.py

@@ -5,16 +5,21 @@
 """
 
 from enum import StrEnum
-from typing import Self
+from typing import Any, Self
 
 from roborock.data import Consumable
 from roborock.devices.traits.v1 import common
+from roborock.roborock_message import RoborockDataProtocol
 from roborock.roborock_typing import RoborockCommand
 
+from .common import TraitUpdateListener
+
 __all__ = [
     "ConsumableTrait",
 ]
 
+_DPS_CONVERTER = common.DpsDataConverter.from_dataclass(Consumable)
+
 
 class ConsumableAttribute(StrEnum):
     """Enum for consumable attributes."""
@@ -35,7 +40,7 @@ def from_str(cls, value: str) -> Self:
         raise ValueError(f"Unknown ConsumableAttribute: {value}")
 
 
-class ConsumableTrait(Consumable, common.V1TraitMixin):
+class ConsumableTrait(Consumable, common.V1TraitMixin, TraitUpdateListener):
     """Trait for managing consumable attributes on Roborock devices.
 
     After the first refresh, you can tell what consumables are supported by
@@ -49,3 +54,12 @@ async def reset_consumable(self, consumable: ConsumableAttribute) -> None:
         """Reset a specific consumable attribute on the device."""
         await self.rpc_channel.send_command(RoborockCommand.RESET_CONSUMABLE, params=[consumable.value])
         await self.refresh()
+
+    def update_from_dps(self, decoded_dps: dict[RoborockDataProtocol, Any]) -> None:
+        """Update the trait from data protocol push message data.
+
+        This handles unsolicited status updates pushed by the device
+        via RoborockDataProtocol codes (e.g. STATE=121, BATTERY=122).
+        """
+        if _DPS_CONVERTER.update_from_dps(self, decoded_dps):
+            self._notify_update()