roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/devices/rpc/v1_channel.py

@@ -0,0 +1,457 @@
+"""V1 Channel for Roborock devices.
+
+This module provides a unified channel interface for V1 protocol devices,
+handling both MQTT and local connections with automatic fallback.
+"""
+
+import asyncio
+import datetime
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+from roborock_cli._vendor.roborock.data import HomeDataDevice, NetworkInfo, RoborockBase, UserData
+from roborock_cli._vendor.roborock.devices.cache import DeviceCache
+from roborock_cli._vendor.roborock.devices.transport.channel import Channel
+from roborock_cli._vendor.roborock.devices.transport.local_channel import LocalChannel, LocalSession, create_local_session
+from roborock_cli._vendor.roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+from roborock_cli._vendor.roborock.mqtt.health_manager import HealthManager
+from roborock_cli._vendor.roborock.mqtt.session import MqttParams, MqttSession
+from roborock_cli._vendor.roborock.protocols.v1_protocol import (
+    CommandType,
+    MapResponse,
+    ParamsType,
+    RequestMessage,
+    ResponseData,
+    ResponseMessage,
+    SecurityData,
+    V1RpcChannel,
+    create_map_response_decoder,
+    create_security_data,
+    decode_rpc_response,
+)
+from roborock_cli._vendor.roborock.roborock_message import RoborockMessage, RoborockMessageProtocol
+from roborock_cli._vendor.roborock.roborock_typing import RoborockCommand
+from roborock_cli._vendor.roborock.util import RoborockLoggerAdapter
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "create_v1_channel",
+]
+
+_T = TypeVar("_T", bound=RoborockBase)
+_TIMEOUT = 10.0
+
+
+# Exponential backoff parameters for reconnecting to local
+MIN_RECONNECT_INTERVAL = datetime.timedelta(minutes=1)
+MAX_RECONNECT_INTERVAL = datetime.timedelta(minutes=10)
+RECONNECT_MULTIPLIER = 1.5
+# After this many hours, the network info is refreshed
+NETWORK_INFO_REFRESH_INTERVAL = datetime.timedelta(hours=12)
+# Interval to check that the local connection is healthy
+LOCAL_CONNECTION_CHECK_INTERVAL = datetime.timedelta(seconds=15)
+
+
+@dataclass(frozen=True)
+class RpcStrategy:
+    """Strategy for encoding/sending/decoding RPC commands."""
+
+    name: str  # For debug logging
+    channel: LocalChannel | MqttChannel
+    encoder: Callable[[RequestMessage], RoborockMessage]
+    decoder: Callable[[RoborockMessage], ResponseMessage | MapResponse | None]
+    health_manager: HealthManager | None = None
+
+
+class RpcChannel(V1RpcChannel):
+    """Provides an RPC interface around a pub/sub transport channel."""
+
+    def __init__(self, rpc_strategies_cb: Callable[[], list[RpcStrategy]], logger: RoborockLoggerAdapter) -> None:
+        """Initialize the RpcChannel with an ordered list of strategies."""
+        self._rpc_strategies_cb = rpc_strategies_cb
+        self._logger = logger
+
+    async def send_command(
+        self,
+        method: CommandType,
+        *,
+        response_type: type[_T] | None = None,
+        params: ParamsType = None,
+    ) -> _T | Any:
+        """Send a command and return either a decoded or parsed response."""
+        request = RequestMessage(method, params=params)
+
+        # Try each channel in order until one succeeds
+        last_exception = None
+        for strategy in self._rpc_strategies_cb():
+            try:
+                decoded_response = await self._send_rpc(strategy, request, self._logger)
+            except RoborockException as e:
+                self._logger.debug("Command %s failed on %s channel: %s", method, strategy.name, e)
+                last_exception = e
+            except Exception as e:
+                self._logger.exception("Unexpected error sending command %s on %s channel", method, strategy.name)
+                last_exception = RoborockException(f"Unexpected error: {e}")
+            else:
+                if response_type is not None:
+                    if not isinstance(decoded_response, dict):
+                        raise RoborockException(
+                            f"Expected dict response to parse {response_type.__name__}, got {type(decoded_response)}"
+                        )
+                    return response_type.from_dict(decoded_response)
+                return decoded_response
+
+        raise last_exception or RoborockException("No available connection to send command")
+
+    @staticmethod
+    async def _send_rpc(
+        strategy: RpcStrategy, request: RequestMessage, logger: RoborockLoggerAdapter
+    ) -> ResponseData | bytes:
+        """Send a command and return a decoded response type.
+
+        This provides an RPC interface over a given channel strategy. The device
+        channel only supports publish and subscribe, so this function handles
+        associating requests with their corresponding responses.
+        """
+        future: asyncio.Future[ResponseData | bytes] = asyncio.Future()
+        logger.debug(
+            "Sending command (%s, request_id=%s): %s, params=%s",
+            strategy.name,
+            request.request_id,
+            request.method,
+            request.params,
+        )
+
+        message = strategy.encoder(request)
+
+        def find_response(response_message: RoborockMessage) -> None:
+            try:
+                decoded = strategy.decoder(response_message)
+            except RoborockException as ex:
+                logger.debug("Exception while decoding message (%s): %s", response_message, ex)
+                return
+            if decoded is None:
+                return
+            logger.debug("Received response (%s, request_id=%s)", strategy.name, decoded.request_id)
+            if decoded.request_id == request.request_id:
+                if isinstance(decoded, ResponseMessage) and decoded.api_error:
+                    future.set_exception(decoded.api_error)
+                else:
+                    future.set_result(decoded.data)
+
+        unsub = await strategy.channel.subscribe(find_response)
+        try:
+            await strategy.channel.publish(message)
+            result = await asyncio.wait_for(future, timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            if strategy.health_manager:
+                await strategy.health_manager.on_timeout()
+            future.cancel()
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+        finally:
+            unsub()
+        if strategy.health_manager:
+            await strategy.health_manager.on_success()
+        return result
+
+
+class V1Channel(Channel):
+    """Unified V1 protocol channel with automatic MQTT/local connection handling.
+
+    This channel abstracts away the complexity of choosing between MQTT and local
+    connections, and provides high-level V1 protocol methods. It automatically
+    handles connection setup, fallback logic, and protocol encoding/decoding.
+    """
+
+    def __init__(
+        self,
+        device_uid: str,
+        security_data: SecurityData,
+        mqtt_channel: MqttChannel,
+        local_session: LocalSession,
+        device_cache: DeviceCache,
+    ) -> None:
+        """Initialize the V1Channel."""
+        self._device_uid = device_uid
+        self._logger = RoborockLoggerAdapter(duid=device_uid, logger=_LOGGER)
+        self._security_data = security_data
+        self._mqtt_channel = mqtt_channel
+        self._local_session = local_session
+        self._local_channel: LocalChannel | None = None
+        self._mqtt_unsub: Callable[[], None] | None = None
+        self._local_unsub: Callable[[], None] | None = None
+        self._callback: Callable[[RoborockMessage], None] | None = None
+        self._device_cache = device_cache
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._last_network_info_refresh: datetime.datetime | None = None
+
+    @property
+    def is_connected(self) -> bool:
+        """Return whether any connection is available."""
+        return self.is_mqtt_connected or self.is_local_connected
+
+    @property
+    def is_local_connected(self) -> bool:
+        """Return whether local connection is available."""
+        return self._local_channel is not None and self._local_channel.is_connected
+
+    @property
+    def is_mqtt_connected(self) -> bool:
+        """Return whether MQTT connection is available."""
+        return self._mqtt_channel.is_connected
+
+    @property
+    def rpc_channel(self) -> V1RpcChannel:
+        """Return the combined RPC channel that prefers local with a fallback to MQTT.
+
+        The returned V1RpcChannel may be long lived and will respect the
+        current connection state of the underlying channels.
+        """
+
+        def rpc_strategies_cb() -> list[RpcStrategy]:
+            strategies = []
+            if local_rpc_strategy := self._create_local_rpc_strategy():
+                strategies.append(local_rpc_strategy)
+            strategies.append(self._create_mqtt_rpc_strategy())
+            return strategies
+
+        return RpcChannel(rpc_strategies_cb, self._logger)
+
+    @property
+    def mqtt_rpc_channel(self) -> V1RpcChannel:
+        """Return the MQTT-only RPC channel.
+
+        The returned V1RpcChannel may be long lived and will respect the
+        current connection state of the underlying channels.
+        """
+        return RpcChannel(lambda: [self._create_mqtt_rpc_strategy()], self._logger)
+
+    @property
+    def map_rpc_channel(self) -> V1RpcChannel:
+        """Return the map RPC channel used for fetching map content."""
+        decoder = create_map_response_decoder(security_data=self._security_data)
+        return RpcChannel(lambda: [self._create_mqtt_rpc_strategy(decoder)], self._logger)
+
+    def _create_local_rpc_strategy(self) -> RpcStrategy | None:
+        """Create the RPC strategy for local transport."""
+        if self._local_channel is None or not self.is_local_connected:
+            return None
+        return RpcStrategy(
+            name="local",
+            channel=self._local_channel,
+            encoder=self._local_encoder,
+            decoder=decode_rpc_response,
+        )
+
+    def _local_encoder(self, x: RequestMessage) -> RoborockMessage:
+        """Encode a request message for local transport.
+
+        This will read the current local channel's protocol version which
+        changes as the protocol version is discovered.
+        """
+        if self._local_channel is None:
+            raise ValueError("Local channel unavailable for encoding")
+        return x.encode_message(
+            RoborockMessageProtocol.GENERAL_REQUEST,
+            version=self._local_channel.protocol_version,
+        )
+
+    def _create_mqtt_rpc_strategy(self, decoder: Callable[[RoborockMessage], Any] = decode_rpc_response) -> RpcStrategy:
+        """Create the RPC strategy for MQTT transport with optional custom decoder."""
+        return RpcStrategy(
+            name="mqtt",
+            channel=self._mqtt_channel,
+            encoder=lambda x: x.encode_message(
+                RoborockMessageProtocol.RPC_REQUEST,
+                security_data=self._security_data,
+            ),
+            decoder=decoder,
+            health_manager=self._mqtt_channel.health_manager,
+        )
+
+    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
+        """Subscribe to all messages from the device.
+
+        This will first attempt to establish a local connection to the device
+        using cached network information if available. If that fails, it will
+        fall back to using the MQTT connection.
+
+        A background task will be started to monitor and maintain the local
+        connection, attempting to reconnect as needed.
+
+        Args:
+            callback: Callback to invoke for each received message.
+
+        Returns:
+            Unsubscribe function to stop receiving messages and clean up resources.
+        """
+        if self._callback is not None:
+            raise ValueError("Only one subscription allowed at a time")
+
+        # Make an initial, optimistic attempt to connect to local with the
+        # cache. The cache information will be refreshed by the background task.
+        try:
+            await self._local_connect(prefer_cache=True)
+        except RoborockException as err:
+            self._logger.debug("First local connection attempt failed, will retry: %s", err)
+
+        # Start a background task to manage the local connection health. This
+        # happens independent of whether we were able to connect locally now.
+        if self._reconnect_task is None:
+            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")
+
+        def unsub() -> None:
+            """Unsubscribe from all messages."""
+            if self._reconnect_task:
+                self._reconnect_task.cancel()
+                self._reconnect_task = None
+            if self._mqtt_unsub:
+                self._mqtt_unsub()
+                self._mqtt_unsub = None
+            if self._local_unsub:
+                self._local_unsub()
+                self._local_unsub = None
+            self._callback = None  # 修复：清除 callback，允许重新订阅
+            self._logger.debug("Unsubscribed from device")
+
+        self._callback = callback
+        return unsub
+
+    async def _get_networking_info(self, *, prefer_cache: bool = True) -> NetworkInfo:
+        """Retrieve networking information for the device.
+
+        This is a cloud only command used to get the local device's IP address.
+        """
+        device_cache_data = await self._device_cache.get()
+
+        if prefer_cache and device_cache_data.network_info:
+            self._logger.debug("Using cached network info")
+            return device_cache_data.network_info
+        try:
+            network_info = await self.mqtt_rpc_channel.send_command(
+                RoborockCommand.GET_NETWORK_INFO, response_type=NetworkInfo
+            )
+        except RoborockException as e:
+            self._logger.debug("Error fetching network info for device")
+            if device_cache_data.network_info:
+                self._logger.debug("Falling back to cached network info after error")
+                return device_cache_data.network_info
+            raise RoborockException(f"Network info failed for device {self._device_uid}") from e
+        self._logger.debug("Network info for device: %s", network_info)
+        self._last_network_info_refresh = datetime.datetime.now(datetime.UTC)
+
+        device_cache_data = await self._device_cache.get()
+        device_cache_data.network_info = network_info
+        await self._device_cache.set(device_cache_data)
+        return network_info
+
+    async def _local_connect(self, *, prefer_cache: bool = True) -> None:
+        """Set up local connection if possible."""
+        self._logger.debug("Attempting to connect to local channel (prefer_cache=%s)", prefer_cache)
+        networking_info = await self._get_networking_info(prefer_cache=prefer_cache)
+        host = networking_info.ip
+        self._logger.debug("Connecting to local channel at %s", host)
+        # Create a new local channel and connect
+        local_channel = self._local_session(host)
+        try:
+            await local_channel.connect()
+        except RoborockException as e:
+            raise RoborockException(f"Error connecting to local device {self._device_uid}: {e}") from e
+        # Wire up the new channel
+        self._local_channel = local_channel
+        self._local_unsub = await self._local_channel.subscribe(self._on_local_message)
+        self._logger.info("Connected to local channel successfully")
+
+    async def _background_reconnect(self) -> None:
+        """Task to run in the background to manage the local connection."""
+        self._logger.debug("Starting background task to manage local connection")
+        reconnect_backoff = MIN_RECONNECT_INTERVAL
+        local_connect_failures = 0
+
+        while True:
+            try:
+                if self.is_local_connected:
+                    await asyncio.sleep(LOCAL_CONNECTION_CHECK_INTERVAL.total_seconds())
+                    continue
+
+                # Not connected, so wait with backoff before trying to connect.
+                # The first time through, we don't sleep, we just try to connect.
+                local_connect_failures += 1
+                if local_connect_failures > 1:
+                    await asyncio.sleep(reconnect_backoff.total_seconds())
+                    reconnect_backoff = min(reconnect_backoff * RECONNECT_MULTIPLIER, MAX_RECONNECT_INTERVAL)
+
+                use_cache = self._should_use_cache(local_connect_failures)
+                await self._local_connect(prefer_cache=use_cache)
+                # Reset backoff and failures on success
+                reconnect_backoff = MIN_RECONNECT_INTERVAL
+                local_connect_failures = 0
+
+            except asyncio.CancelledError:
+                self._logger.debug("Background reconnect task cancelled")
+                if self._local_channel:
+                    self._local_channel.close()
+                return
+            except RoborockException as err:
+                self._logger.debug("Background reconnect failed: %s", err)
+            except Exception:
+                self._logger.exception("Unhandled exception in background reconnect task")
+
+    def _should_use_cache(self, local_connect_failures: int) -> bool:
+        """Determine whether to use cached network info on retries.
+
+        On the first retry we'll avoid the cache to handle the case where
+        the network ip may have recently changed. Otherwise, use the cache
+        if available then expire at some point.
+        """
+        if local_connect_failures == 1:
+            return False
+        elif self._last_network_info_refresh and (
+            datetime.datetime.now(datetime.UTC) - self._last_network_info_refresh > NETWORK_INFO_REFRESH_INTERVAL
+        ):
+            return False
+        return True
+
+    def _on_mqtt_message(self, message: RoborockMessage) -> None:
+        """Handle incoming MQTT messages."""
+        self._logger.debug("V1Channel received MQTT message: %s", message)
+        if self._callback:
+            self._callback(message)
+
+    def _on_local_message(self, message: RoborockMessage) -> None:
+        """Handle incoming local messages."""
+        self._logger.debug("V1Channel received local message: %s", message)
+        if self._callback:
+            self._callback(message)
+
+
+def create_v1_channel(
+    user_data: UserData,
+    mqtt_params: MqttParams,
+    mqtt_session: MqttSession,
+    device: HomeDataDevice,
+    device_cache: DeviceCache,
+) -> V1Channel:
+    """Create a V1Channel for the given device."""
+    security_data = create_security_data(user_data.rriot)
+    mqtt_channel = MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)
+    local_session = create_local_session(device.local_key, device.duid)
+    return V1Channel(
+        device.duid,
+        security_data,
+        mqtt_channel,
+        local_session=local_session,
+        device_cache=device_cache,
+    )

roborock_cli/_vendor/roborock/devices/traits/__init__.py

@@ -0,0 +1,28 @@
+"""Module for device traits.
+
+This package contains the trait definitions for different device protocols supported
+by Roborock devices.
+
+Submodules
+----------
+*   `v1`: Contains traits for standard Roborock vacuums (e.g., S-series, Q-series).
+    These devices use the V1 protocol and have rich feature sets split into
+    granular traits (e.g., `StatusTrait`, `ConsumableTrait`).
+*   `a01`: Contains APIs for A01 protocol devices, such as the Dyad (wet/dry vacuum)
+    and Zeo (washing machine). These devices use a different communication structure.
+*   `b01`: Contains APIs for B01 protocol devices.
+"""
+
+from abc import ABC
+
+__all__ = [
+    "Trait",
+    "traits_mixin",
+    "v1",
+    "a01",
+    "b01",
+]
+
+
+class Trait(ABC):
+    """Base class for all traits."""

roborock_cli/_vendor/roborock/devices/traits/a01/__init__.py

@@ -0,0 +1,191 @@
+"""Create traits for A01 devices.
+
+This module provides the API implementations for A01 protocol devices, which include
+Dyad (Wet/Dry Vacuums) and Zeo (Washing Machines).
+
+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`).
+
+The API provides two main 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.
+
+Note that these APIs fetch data directly from the device upon request and do not
+cache state internally.
+"""
+
+import json
+from collections.abc import Callable
+from datetime import time
+from typing import Any
+
+from roborock_cli._vendor.roborock.data import DyadProductInfo, DyadSndState, HomeDataProduct, RoborockCategory
+from roborock_cli._vendor.roborock.data.dyad.dyad_code_mappings import (
+    DyadBrushSpeed,
+    DyadCleanMode,
+    DyadError,
+    DyadSelfCleanLevel,
+    DyadSelfCleanMode,
+    DyadSuction,
+    DyadWarmLevel,
+    DyadWaterLevel,
+    RoborockDyadStateCode,
+)
+from roborock_cli._vendor.roborock.data.zeo.zeo_code_mappings import (
+    ZeoDetergentType,
+    ZeoDryingMode,
+    ZeoError,
+    ZeoMode,
+    ZeoProgram,
+    ZeoRinse,
+    ZeoSoftenerType,
+    ZeoSpin,
+    ZeoState,
+    ZeoTemperature,
+)
+from roborock_cli._vendor.roborock.devices.rpc.a01_channel import send_decoded_command
+from roborock_cli._vendor.roborock.devices.traits import Trait
+from roborock_cli._vendor.roborock.devices.transport.mqtt_channel import MqttChannel
+from roborock_cli._vendor.roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
+
+__init__ = [
+    "DyadApi",
+    "ZeoApi",
+]
+
+
+DYAD_PROTOCOL_ENTRIES: dict[RoborockDyadDataProtocol, Callable] = {
+    RoborockDyadDataProtocol.STATUS: lambda val: RoborockDyadStateCode(val).name,
+    RoborockDyadDataProtocol.SELF_CLEAN_MODE: lambda val: DyadSelfCleanMode(val).name,
+    RoborockDyadDataProtocol.SELF_CLEAN_LEVEL: lambda val: DyadSelfCleanLevel(val).name,
+    RoborockDyadDataProtocol.WARM_LEVEL: lambda val: DyadWarmLevel(val).name,
+    RoborockDyadDataProtocol.CLEAN_MODE: lambda val: DyadCleanMode(val).name,
+    RoborockDyadDataProtocol.SUCTION: lambda val: DyadSuction(val).name,
+    RoborockDyadDataProtocol.WATER_LEVEL: lambda val: DyadWaterLevel(val).name,
+    RoborockDyadDataProtocol.BRUSH_SPEED: lambda val: DyadBrushSpeed(val).name,
+    RoborockDyadDataProtocol.POWER: lambda val: int(val),
+    RoborockDyadDataProtocol.AUTO_DRY: lambda val: bool(val),
+    RoborockDyadDataProtocol.MESH_LEFT: lambda val: int(360000 - val * 60),
+    RoborockDyadDataProtocol.BRUSH_LEFT: lambda val: int(360000 - val * 60),
+    RoborockDyadDataProtocol.ERROR: lambda val: DyadError(val).name,
+    RoborockDyadDataProtocol.VOLUME_SET: lambda val: int(val),
+    RoborockDyadDataProtocol.STAND_LOCK_AUTO_RUN: lambda val: bool(val),
+    RoborockDyadDataProtocol.AUTO_DRY_MODE: lambda val: bool(val),
+    RoborockDyadDataProtocol.SILENT_DRY_DURATION: lambda val: int(val),  # in minutes
+    RoborockDyadDataProtocol.SILENT_MODE: lambda val: bool(val),
+    RoborockDyadDataProtocol.SILENT_MODE_START_TIME: lambda val: time(
+        hour=int(val / 60), minute=val % 60
+    ),  # in minutes since 00:00
+    RoborockDyadDataProtocol.SILENT_MODE_END_TIME: lambda val: time(
+        hour=int(val / 60), minute=val % 60
+    ),  # in minutes since 00:00
+    RoborockDyadDataProtocol.RECENT_RUN_TIME: lambda val: [
+        int(v) for v in val.split(",")
+    ],  # minutes of cleaning in past few days.
+    RoborockDyadDataProtocol.TOTAL_RUN_TIME: lambda val: int(val),
+    RoborockDyadDataProtocol.SND_STATE: lambda val: DyadSndState.from_dict(val),
+    RoborockDyadDataProtocol.PRODUCT_INFO: lambda val: DyadProductInfo.from_dict(val),
+}
+
+ZEO_PROTOCOL_ENTRIES: dict[RoborockZeoProtocol, Callable] = {
+    # read-only
+    RoborockZeoProtocol.STATE: lambda val: ZeoState(val).name,
+    RoborockZeoProtocol.COUNTDOWN: lambda val: int(val),
+    RoborockZeoProtocol.WASHING_LEFT: lambda val: int(val),
+    RoborockZeoProtocol.ERROR: lambda val: ZeoError(val).name,
+    RoborockZeoProtocol.TIMES_AFTER_CLEAN: lambda val: int(val),
+    RoborockZeoProtocol.DETERGENT_EMPTY: lambda val: bool(val),
+    RoborockZeoProtocol.SOFTENER_EMPTY: lambda val: bool(val),
+    # read-write
+    RoborockZeoProtocol.MODE: lambda val: ZeoMode(val).name,
+    RoborockZeoProtocol.PROGRAM: lambda val: ZeoProgram(val).name,
+    RoborockZeoProtocol.TEMP: lambda val: ZeoTemperature(val).name,
+    RoborockZeoProtocol.RINSE_TIMES: lambda val: ZeoRinse(val).name,
+    RoborockZeoProtocol.SPIN_LEVEL: lambda val: ZeoSpin(val).name,
+    RoborockZeoProtocol.DRYING_MODE: lambda val: ZeoDryingMode(val).name,
+    RoborockZeoProtocol.DETERGENT_TYPE: lambda val: ZeoDetergentType(val).name,
+    RoborockZeoProtocol.SOFTENER_TYPE: lambda val: ZeoSoftenerType(val).name,
+    RoborockZeoProtocol.SOUND_SET: lambda val: bool(val),
+}
+
+
+def convert_dyad_value(protocol_value: RoborockDyadDataProtocol, value: Any) -> Any:
+    """Convert a dyad protocol value to its corresponding type."""
+    if (converter := DYAD_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        try:
+            return converter(value)
+        except (ValueError, TypeError):
+            return None
+    return None
+
+
+def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
+    """Convert a zeo protocol value to its corresponding type."""
+    if (converter := ZEO_PROTOCOL_ENTRIES.get(protocol_value)) is not None:
+        try:
+            return converter(value)
+        except (ValueError, TypeError):
+            return None
+    return None
+
+
+class DyadApi(Trait):
+    """API for interacting with Dyad devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the Dyad API."""
+        self._channel = channel
+
+    async def query_values(self, protocols: list[RoborockDyadDataProtocol]) -> dict[RoborockDyadDataProtocol, Any]:
+        """Query the device for the values of the given Dyad protocols."""
+        response = await send_decoded_command(
+            self._channel,
+            {RoborockDyadDataProtocol.ID_QUERY: protocols},
+            value_encoder=json.dumps,
+        )
+        return {protocol: convert_dyad_value(protocol, response.get(protocol)) for protocol in protocols}
+
+    async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dict[RoborockDyadDataProtocol, Any]:
+        """Set a value for a specific protocol on the device."""
+        params = {protocol: value}
+        return await send_decoded_command(self._channel, params)
+
+
+class ZeoApi(Trait):
+    """API for interacting with Zeo devices."""
+
+    name = "zeo"
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the Zeo API."""
+        self._channel = channel
+
+    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}
+
+    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."""
+    match product.category:
+        case RoborockCategory.WET_DRY_VAC:
+            return DyadApi(mqtt_channel)
+        case RoborockCategory.WASHING_MACHINE:
+            return ZeoApi(mqtt_channel)
+        case _:
+            raise NotImplementedError(f"Unsupported category {product.category}")

roborock_cli/_vendor/roborock/devices/traits/b01/__init__.py

@@ -0,0 +1,12 @@
+"""Traits for B01 devices."""
+
+from . import q7, q10
+from .q7 import Q7PropertiesApi
+from .q10 import Q10PropertiesApi
+
+__all__ = [
+    "Q7PropertiesApi",
+    "Q10PropertiesApi",
+    "q7",
+    "q10",
+]