roborock.devices.traits.a01

  1"""Create traits for A01 devices.
  2
  3This module provides the API implementations for A01 protocol devices, which include
  4Dyad (Wet/Dry Vacuums) and Zeo (Washing Machines).
  5
  6Using A01 APIs
  7--------------
  8A01 devices expose a single API object that handles all device interactions. This API is
  9available on the device instance (`device.dyad` or `device.zeo`).
 10
 11The API provides these methods:
 121.  **query_values(protocols)**: Fetches current state for specific data points.
 13    You must pass a list of protocol enums (e.g. `RoborockDyadDataProtocol` or
 14    `RoborockZeoProtocol`) to request specific data.
 152.  **set_value(protocol, value)**: Sends a command to the device to change a setting
 16    or perform an action.
 173.  **values**: The latest known state, merged from query responses and unsolicited
 18    pushes in arrival order.
 194.  **add_update_listener(callback)**: Registers a callback invoked whenever `values`
 20    changes; read `values` from the callback to get the updated state.
 21
 22The device pushes only the data points that changed, so `values` is the merged view
 23of everything seen so far. State tracking is active once the device is connected
 24(the device calls `start()` on the API, which subscribes to the MQTT topic).
 25"""
 26
 27import json
 28import logging
 29from abc import abstractmethod
 30from collections.abc import Callable
 31from datetime import UTC, datetime, time
 32from typing import Any, Generic, TypeVar
 33
 34from roborock.data import DyadProductInfo, DyadScheduleTask, DyadSndState, HomeDataProduct, RoborockCategory
 35from roborock.data.dyad.dyad_code_mappings import (
 36    DyadBrushSpeed,
 37    DyadCleanAssistPower,
 38    DyadCleanMode,
 39    DyadCleanserAmount,
 40    DyadDryingMode,
 41    DyadError,
 42    DyadPowerwheelDirection,
 43    DyadSelfCleanIntensity,
 44    DyadSelfCleanLevel,
 45    DyadSelfCleanMode,
 46    DyadSelfCleanWaterTemp,
 47    DyadSuction,
 48    DyadWarmLevel,
 49    DyadWaterLevel,
 50    RoborockDyadStateCode,
 51)
 52from roborock.data.zeo.zeo_code_mappings import (
 53    ZeoDetergentType,
 54    ZeoDryingMode,
 55    ZeoError,
 56    ZeoFeatureBits,
 57    ZeoMode,
 58    ZeoProgram,
 59    ZeoRinse,
 60    ZeoSoftenerType,
 61    ZeoSpin,
 62    ZeoState,
 63    ZeoTemperature,
 64)
 65from roborock.devices.rpc.a01_channel import send_decoded_command
 66from roborock.devices.traits import Trait
 67from roborock.devices.traits.a01.device_feature import (
 68    build_feature_dp_list,
 69    build_force_load_dp_list,
 70    supports_uv_light,
 71)
 72from roborock.devices.traits.common import TraitUpdateListener
 73from roborock.devices.transport.mqtt_channel import MqttChannel
 74from roborock.exceptions import RoborockException
 75from roborock.protocols.a01_protocol import decode_rpc_response
 76from roborock.roborock_message import (
 77    RoborockDyadDataProtocol,
 78    RoborockMessage,
 79    RoborockMessageProtocol,
 80    RoborockZeoProtocol,
 81)
 82
 83_LOGGER = logging.getLogger(__name__)
 84
 85__all__ = [
 86    "A01Api",
 87    "DyadApi",
 88    "ZeoApi",
 89]
 90
 91
 92def _decode_dyad_schedule(value: str) -> list[DyadScheduleTask]:
 93    """Decode the schedule data point, which is a JSON list or an empty string."""
 94    if not value:
 95        return []
 96    return [DyadScheduleTask.from_dict(task) for task in json.loads(value)]
 97
 98
 99DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
100    RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
101    RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
102    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
103    RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
104    RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
105    RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
106    RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
107    RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
108    RoborockDyadDataProtocol.POWER: lambda val: int(val),
109    RoborockDyadDataProtocol.AUTO_SELF_CLEAN_SET: lambda val: bool(val),  # self-clean after docking
110    RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
111    RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
112    RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
113    RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
114    RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
115    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
116    RoborockDyadDataProtocol.AUTO_SELF_CLEAN_SET_MODE: lambda val: DyadSelfCleanLevel(val).name,
117    RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
118    RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val),  # minutes, or DyadDryingMode on F25
119    RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
120    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
121        hour=int(val / 60), minute=val % 60
122    ),  # in minutes since 00:00
123    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
124        hour=int(val / 60), minute=val % 60
125    ),  # in minutes since 00:00
126    RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
127        int(v) for v in val.split(",")
128    ],  # minutes of cleaning in past few days.
129    RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
130    RoborockDyadDataProtocol.SET_CLEANSER_AMOUNT: lambda val: DyadCleanserAmount(val).name,
131    RoborockDyadDataProtocol.AUTO_SMART_SELF_CLEAN_SET: lambda val: bool(val),
132    RoborockDyadDataProtocol.LIGHT_SET: lambda val: bool(val),
133    RoborockDyadDataProtocol.CLEAN_ASSIST_POWER_SET: lambda val: DyadCleanAssistPower(val).name,
134    RoborockDyadDataProtocol.SELF_CLEAN_WATER_TEMP_LEVEL_SET: lambda val: DyadSelfCleanWaterTemp(val).name,
135    RoborockDyadDataProtocol.DOCK_TYPE: lambda val: int(val),
136    RoborockDyadDataProtocol.WATER_MODE_SUCTION: lambda val: DyadSuction(val).name,
137    RoborockDyadDataProtocol.WATER_MODE_BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
138    RoborockDyadDataProtocol.CHARGE_ERROR_STATUS: lambda val: bool(val),
139    RoborockDyadDataProtocol.DRYING_RATE: lambda val: int(val),  # percent
140    RoborockDyadDataProtocol.FLAT_STATUS: lambda val: bool(val),
141    RoborockDyadDataProtocol.POWERWHEEL_DIRECTION: lambda val: DyadPowerwheelDirection(val).name,
142    RoborockDyadDataProtocol.POWERWHEEL_SPEED: lambda val: int(val),
143    RoborockDyadDataProtocol.SELF_CLEAN_CUR_MODE: lambda val: DyadSelfCleanIntensity(val).name,
144    RoborockDyadDataProtocol.SELF_CLEAN_CUR_TEMP: lambda val: DyadSelfCleanWaterTemp(val).name,
145    RoborockDyadDataProtocol.DRY_CUR_MODE: lambda val: DyadDryingMode(val).name,
146    RoborockDyadDataProtocol.SCHEDULE_TASK: _decode_dyad_schedule,
147    RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
148    RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
149}
150
151ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
152    # read-only
153    RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
154    RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
155    RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
156    RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
157    RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
158    RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
159    RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
160    # read-write
161    RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
162    RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
163    RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
164    RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
165    RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
166    RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
167    RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
168    RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
169    RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
170    RoborockZeoProtocol.FEATURE_BITS: lambda val: int(val),
171}
172
173
174def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, value: Any) -> Any:
175    """Convert a dyad protocol value to its corresponding type."""
176    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
177        try:
178            return converter(value)
179        except (ValueError, TypeError):
180            return None
181    return None
182
183
184def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
185    """Convert a zeo protocol value to its corresponding type."""
186    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
187        try:
188            return converter(value)
189        except (ValueError, TypeError):
190            return None
191    return None
192
193
194_DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol)
195_ZEO_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockZeoProtocol)
196
197_P = TypeVar("_P", RoborockDyadDataProtocol, RoborockZeoProtocol)
198
199
200class A01Api(Trait, TraitUpdateListener, Generic[_P]):
201    """Base class for A01 device APIs with device state tracking.
202
203    Query responses and unsolicited pushes both arrive on the same MQTT topic,
204    so a single subscription merges every decoded message into `values` in
205    arrival order. Update listeners are notified whenever a value changes.
206    """
207
208    def __init__(self, channel: MqttChannel, initial_status: dict[int, Any] | None = None) -> None:
209        """Initialize the A01 API, optionally seeding `values` from a cloud status snapshot."""
210        TraitUpdateListener.__init__(self, _LOGGER)
211        self._channel = channel
212        self._values: dict[_P, Any] = {}
213        self._unsub: Callable[[], None] | None = None
214        self._last_message_time: datetime | None = None
215        if initial_status:
216            self._merge_values(self._decode_datapoints(initial_status))
217
218    @property
219    def values(self) -> dict[_P, Any]:
220        """Latest known device state, merged from query responses and pushes.
221
222        The device pushes only the data points that changed, so this is the
223        merged view of everything seen so far. A protocol the device has not
224        reported yet is absent from the dictionary.
225        """
226        return dict(self._values)
227
228    @property
229    def last_message_time(self) -> datetime | None:
230        """Time the last message was received from the device.
231
232        Updated on every decoded message, even when no value changed: idle
233        devices push an identical heartbeat, so this is the liveness signal
234        even when `values` stays the same and update listeners stay silent.
235        The initial cloud status snapshot does not count as a message.
236        """
237        return self._last_message_time
238
239    async def start(self) -> None:
240        """Subscribe to the device state topic and start tracking `values`."""
241        await self._ensure_subscribed()
242
243    def close(self) -> None:
244        """Unsubscribe from MQTT push and release resources."""
245        if self._unsub is not None:
246            self._unsub()
247            self._unsub = None
248
249    async def _ensure_subscribed(self) -> None:
250        """Subscribe to MQTT DPS push (idempotent)."""
251        if self._unsub is not None:
252            return
253        self._unsub = await self._channel.subscribe(self._on_message)
254
255    @abstractmethod
256    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[_P, Any]:
257        """Convert raw datapoints to typed values, skipping unknown codes."""
258
259    def _on_message(self, message: RoborockMessage) -> None:
260        """Handle a message on the device topic (query response or push)."""
261        if message.protocol != RoborockMessageProtocol.RPC_RESPONSE:
262            return
263        try:
264            datapoints = decode_rpc_response(message)
265        except RoborockException:
266            _LOGGER.debug("Dropped malformed push message", exc_info=True)
267            return
268        self._last_message_time = datetime.now(UTC)
269        self._merge_values(self._decode_datapoints(datapoints))
270
271    def _merge_query_response(self, values: dict[_P, Any]) -> None:
272        """Record a successful query response when there is no subscription.
273
274        When subscribed, the response was already merged in arrival order and
275        timestamped by `_on_message`; merging again here could overwrite a
276        push that arrived after it.
277        """
278        if self._unsub is not None:
279            return
280        self._last_message_time = datetime.now(UTC)
281        self._merge_values(values)
282
283    def _merge_values(self, values: dict[_P, Any]) -> None:
284        """Merge decoded values into the cache and notify on change."""
285        changed = False
286        for protocol, value in values.items():
287            if value is None:
288                continue
289            if protocol not in self._values or self._values[protocol] != value:
290                self._values[protocol] = value
291                changed = True
292        if changed:
293            self._notify_update()
294
295
296class DyadApi(A01Api[RoborockDyadDataProtocol]):
297    """API for interacting with Dyad devices."""
298
299    name = "dyad"
300
301    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockDyadDataProtocol, Any]:
302        """Convert raw datapoints to typed values, skipping unknown codes."""
303        values: dict[RoborockDyadDataProtocol, Any] = {}
304        for code, value in datapoints.items():
305            if code not in _DYAD_PROTOCOL_VALUES:
306                continue
307            protocol = RoborockDyadDataProtocol(code)
308            values[protocol] = convert_dyad_value(protocol, value)
309        return values
310
311    async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
312        """Query the device for the values of the given Dyad protocols."""
313        response = await send_decoded_command(
314            self._channel,
315            {RoborockDyadDataProtocol.ID_QUERY: protocols},
316            value_encoder=json.dumps,
317        )
318        values = {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
319        self._merge_query_response(values)
320        return values
321
322    async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
323        """Set a value for a specific protocol on the device."""
324        params = {protocol: value}
325        return await send_decoded_command(self._channel, params)
326
327    async def start_drying(self) -> None:
328        """Start drying with the drying mode currently set on the device.
329
330        The app sends the same command for fast and silent drying on the F25 Ultra.
331        """
332        await send_decoded_command(
333            self._channel,
334            {RoborockDyadDataProtocol.START: 1, RoborockDyadDataProtocol.WARM_LEVEL: 1},
335        )
336
337    async def stop_drying(self) -> None:
338        """Stop a running drying cycle."""
339        await send_decoded_command(self._channel, {RoborockDyadDataProtocol.START: 0})
340
341    async def reset_filter(self) -> None:
342        """Reset the filter usage counter."""
343        await send_decoded_command(self._channel, {RoborockDyadDataProtocol.MESH_RESET: 1})
344
345    async def reset_brush(self) -> None:
346        """Reset the brush roller usage counter."""
347        await send_decoded_command(self._channel, {RoborockDyadDataProtocol.BRUSH_RESET: 1})
348
349    async def add_listener(self, callback: Callable[[dict[RoborockDyadDataProtocol, Any]], None]) -> Callable[[], None]:
350        """Listen for state the device pushes on its own.
351
352        The callback is invoked with decoded values whenever the device sends a
353        message, including unsolicited pushes when its state changes. Only known
354        protocols are delivered. Returns a callable to remove the listener.
355
356        Prefer `add_update_listener` together with `values`, which handle the
357        merging of partial pushes for you.
358        """
359
360        def on_message(message: RoborockMessage) -> None:
361            try:
362                datapoints = decode_rpc_response(message)
363            except RoborockException:
364                return
365            if values := self._decode_datapoints(datapoints):
366                callback(values)
367
368        return await self._channel.subscribe(on_message)
369
370
371class ZeoApi(A01Api[RoborockZeoProtocol]):
372    """API for interacting with Zeo devices."""
373
374    name = "zeo"
375
376    def __init__(
377        self, channel: MqttChannel, model: str | None = None, initial_status: dict[int, Any] | None = None
378    ) -> None:
379        """Initialize the Zeo API."""
380        self._feature_bits: int = 0
381        self._model = model
382        super().__init__(channel, initial_status)
383
384    async def start(self) -> None:
385        """Subscribe to MQTT push and trigger a full state sync.
386
387        Subscribes to the DPS MQTT topic, then performs a two-stage
388        force-load: first the base DP list (including FEATURE_BITS),
389        then a second query for the DPs gated behind each enabled feature.
390        The device responds with a complete state dump;
391        subsequent changes arrive via incremental MQTT push.
392        """
393        await self._ensure_subscribed()
394        await self._force_load()
395        await self._load_feature_dps()
396
397    def _decode_datapoints(self, datapoints: dict[int, Any]) -> dict[RoborockZeoProtocol, Any]:
398        """Convert raw datapoints to typed values, skipping unknown codes."""
399        values: dict[RoborockZeoProtocol, Any] = {}
400        for code, value in datapoints.items():
401            if code not in _ZEO_PROTOCOL_VALUES:
402                continue
403            protocol = RoborockZeoProtocol(code)
404            values[protocol] = convert_zeo_value(protocol, value)
405        return values
406
407    async def _force_load(self) -> None:
408        """Send ID_QUERY with the base DP list to trigger a full state push.
409
410        For devices known to lack FEATURE_BITS, the DP is excluded
411        from the query list and ``_feature_bits`` stays at 0.
412        """
413        dp_list = build_force_load_dp_list(self._model)
414        result = await self.query_values(dp_list)
415        self._feature_bits = result.get(RoborockZeoProtocol.FEATURE_BITS, 0)
416
417    async def _load_feature_dps(self) -> None:
418        """Second-stage query for feature-gated DPs.
419
420        Called unconditionally after the first force-load; each DP is
421        independently gated:
422
423        - Feature-gated DPs are queried only when their feature bit is set
424          in FEATURE_BITS (DP 237).
425        - UV light (DP 228) is gated by :func:`supports_uv_light` (series
426          whitelist), independent of the feature bits.
427        """
428        feature_dps: list[RoborockZeoProtocol] = []
429        if self._feature_bits:
430            feature_dps.extend(build_feature_dp_list(self._feature_bits))
431        if supports_uv_light(self._model):
432            feature_dps.append(RoborockZeoProtocol.UV_LIGHT)
433        if not feature_dps:
434            return
435        try:
436            await self.query_values(feature_dps)
437        except RoborockException as exc:
438            _LOGGER.warning("Feature DPS load failed (non-fatal): %s", exc)
439
440    def supports(self, feature: ZeoFeatureBits) -> bool:
441        """Check whether the device supports a given feature bit."""
442        return bool(self._feature_bits & (1 << feature.value))
443
444    async def query_values(self, protocols: list[RoborockZeoProtocol]) -> dict[RoborockZeoProtocol, Any]:
445        """Query the device for the values of the given protocols."""
446        response = await send_decoded_command(
447            self._channel,
448            {RoborockZeoProtocol.ID_QUERY: protocols},
449            value_encoder=json.dumps,
450        )
451        values = {protocol: convert_zeo_value(protocol, response.get(protocol)) for protocol in protocols}
452        self._merge_query_response(values)
453        return values
454
455    async def set_value(self, protocol: RoborockZeoProtocol, value: Any) -> dict[RoborockZeoProtocol, Any]:
456        """Set a value for a specific protocol on the device."""
457        params = {protocol: value}
458        return await send_decoded_command(self._channel, params, value_encoder=lambda x: x)
459
460
461def _parse_device_status(device_status: dict | None) -> dict[int, Any] | None:
462    """Normalize the cloud home data status snapshot to integer datapoint codes."""
463    if not device_status:
464        return None
465    try:
466        return {int(code): value for code, value in device_status.items()}
467    except (TypeError, ValueError):
468        _LOGGER.debug("Ignoring malformed device status snapshot: %s", device_status)
469        return None
470
471
472def create(product: HomeDataProduct, mqtt_channel: MqttChannel, device_status: dict | None = None) -> DyadApi | ZeoApi:
473    """Create traits for A01 devices.
474
475    The optional `device_status` is the cloud home data status snapshot, used
476    to seed `values` so state is available before the first device round trip.
477    """
478    initial_status = _parse_device_status(device_status)
479    match product.category:
480        case RoborockCategory.WET_DRY_VAC:
481            return DyadApi(mqtt_channel, initial_status=initial_status)
482        case RoborockCategory.WASHING_MACHINE:
483            return ZeoApi(mqtt_channel, model=product.model, initial_status=initial_status)
484        case _:
485            raise NotImplementedError(f"Unsupported category {product.category}")