feat: Add initial Q10 support

Author: allenporter

roborock/cli.py

@@ -43,6 +43,7 @@
 
 from roborock import SHORT_MODEL_TO_ENUM, RoborockCommand
 from roborock.data import DeviceData, RoborockBase, UserData
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
 from roborock.device_features import DeviceFeatures
 from roborock.devices.cache import Cache, CacheData
 from roborock.devices.device import RoborockDevice
@@ -91,7 +92,12 @@ def wrapper(*args, **kwargs):
         context: RoborockContext = ctx.obj
 
         async def run():
-            return await func(*args, **kwargs)
+            try:
+                await func(*args, **kwargs)
+            except Exception:
+                _LOGGER.exception("Uncaught exception in command")
+                click.echo(f"Error: {sys.exc_info()[1]}", err=True)
+            await context.cleanup()
 
         if context.is_session_mode():
             # Session mode - run in the persistent loop
@@ -739,6 +745,16 @@ async def network_info(ctx, device_id: str):
     await _display_v1_trait(context, device_id, lambda v1: v1.network_info)
 
 
+def _parse_b01_q10_command(cmd: str) -> B01_Q10_DP | None:
+    """Parse B01_Q10 command from either enum name or value."""
+    for func in (B01_Q10_DP.from_code, B01_Q10_DP.from_name, B01_Q10_DP.from_value):
+        try:
+            return func(cmd)
+        except ValueError:
+            continue
+    return None
+
+
 @click.command()
 @click.option("--device_id", required=True)
 @click.option("--cmd", required=True)
@@ -749,12 +765,20 @@ async def command(ctx, cmd, device_id, params):
     context: RoborockContext = ctx.obj
     device_manager = await context.get_device_manager()
     device = await device_manager.get_device(device_id)
-    if device.v1_properties is None:
-        raise RoborockException(f"Device {device.name} does not support V1 protocol")
-    command_trait: Trait = device.v1_properties.command
-    result = await command_trait.send(cmd, json.loads(params) if params is not None else None)
-    if result:
-        click.echo(dump_json(result))
+    if device.v1_properties is not None:
+        command_trait: Trait = device.v1_properties.command
+        result = await command_trait.send(cmd, json.loads(params) if params is not None else {})
+        if result:
+            click.echo(dump_json(result))
+    elif device.b01_q10_properties is not None:
+        # Parse B01_Q10_DP from either enum name or the value
+        if (cmd_value := _parse_b01_q10_command(cmd)) is None:
+            raise RoborockException(f"Invalid command {cmd} for B01_Q10 device")
+        await device.b01_q10_properties.send(cmd_value, json.loads(params) if params is not None else {})
+        # B10 Commands don't have a specific time to respond, so wait a bit
+        await asyncio.sleep(5)
+    else:
+        raise RoborockException(f"Device {device.name} does not support sending raw commands")
 
 
 @click.command()

roborock/data/code_mappings.py

@@ -65,11 +65,28 @@ def __new__(cls, value: str, code: int) -> RoborockModeEnum:
 
     @classmethod
     def from_code(cls, code: int):
+        """Find enum member by code."""
         for member in cls:
             if member.code == code:
                 return member
         raise ValueError(f"{code} is not a valid code for {cls.__name__}")
 
+    @classmethod
+    def from_name(cls, name: str):
+        """Find enum member by name (case-insensitive)."""
+        for member in cls:
+            if member.name.lower() == name.lower():
+                return member
+        raise ValueError(f"{name} is not a valid name for {cls.__name__}")
+
+    @classmethod
+    def from_value(cls, value: str):
+        """Find enum member by value (case-insensitive)."""
+        for member in cls:
+            if member.value.lower() == value.lower():
+                return member
+        raise ValueError(f"{value} is not a valid value for {cls.__name__}")
+
     @classmethod
     def keys(cls) -> list[str]:
         """Returns a list of all member values."""

roborock/devices/b01_q10_channel.py

@@ -0,0 +1,63 @@
+"""Thin wrapper around the MQTT channel for Roborock B01 devices."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.exceptions import RoborockException
+from roborock.protocols.b01_q10_protocol import (
+    ParamsType,
+    decode_rpc_response,
+    encode_mqtt_payload,
+)
+
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
+
+
+async def send_command(
+    mqtt_channel: MqttChannel,
+    command: B01_Q10_DP,
+    params: ParamsType,
+) -> None:
+    """Send a command on the MQTT channel, without waiting for a response"""
+    _LOGGER.debug(
+        "Sending B01 MQTT command: cmd=%s params=%s",
+        command,
+        params,
+    )
+    roborock_message = encode_mqtt_payload(command, params)
+    _LOGGER.debug("Sending MQTT message: %s", roborock_message)
+    try:
+        await mqtt_channel.publish(roborock_message)
+    except RoborockException as ex:
+        _LOGGER.debug(
+            "Error sending B01 decoded command (method=%s params=%s): %s",
+            command,
+            params,
+            ex,
+        )
+        raise
+
+
+async def stream_decoded_responses(
+    mqtt_channel: MqttChannel,
+) -> AsyncGenerator[dict[B01_Q10_DP, Any], None]:
+    """Stream decoded DPS messages received via MQTT."""
+
+    async for response_message in mqtt_channel.subscribe_stream():
+        try:
+            decoded_dps = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.debug(
+                "Failed to decode B01 RPC response: %s: %s",
+                response_message,
+                ex,
+            )
+            continue
+        yield decoded_dps

roborock/devices/b01_channel.py roborock/devices/b01_q7_channel.pyroborock/devices/b01_channel.py renamed to roborock/devices/b01_q7_channel.py

@@ -8,7 +8,7 @@
 from typing import Any
 
 from roborock.exceptions import RoborockException
-from roborock.protocols.b01_protocol import (
+from roborock.protocols.b01_q7_protocol import (
     CommandType,
     ParamsType,
     decode_rpc_response,

roborock/devices/device.py

@@ -195,12 +195,14 @@ async def connect(self) -> None:
         if self._unsub:
             raise ValueError("Already connected to the device")
         unsub = await self._channel.subscribe(self._on_message)
-        if self.v1_properties is not None:
-            try:
+        try:
+            if self.v1_properties is not None:
                 await self.v1_properties.discover_features()
-            except RoborockException:
-                unsub()
-                raise
+            elif self.b01_q10_properties is not None:
+                await self.b01_q10_properties.start()
+        except RoborockException:
+            unsub()
+            raise
         self._logger.info("Connected to device")
         self._unsub = unsub
 
@@ -212,6 +214,8 @@ async def close(self) -> None:
                 await self._connect_task
             except asyncio.CancelledError:
                 pass
+        if self.b01_q10_properties is not None:
+            await self.b01_q10_properties.close()
         if self._unsub:
             self._unsub()
             self._unsub = None

roborock/devices/device_manager.py

@@ -240,9 +240,7 @@ def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDat
                 channel = create_mqtt_channel(user_data, mqtt_params, mqtt_session, device)
                 model_part = product.model.split(".")[-1]
                 if "ss" in model_part:
-                    raise UnsupportedDeviceError(
-                        f"Device {device.name} has unsupported version B01 product model {product.model}"
-                    )
+                    trait = b01.q10.create(channel)
                 elif "sc" in model_part:
                     # Q7 devices start with 'sc' in their model naming.
                     trait = b01.q7.create(channel)

roborock/devices/mqtt_channel.py

@@ -1,7 +1,8 @@
 """Modules for communicating with specific Roborock devices over MQTT."""
 
+import asyncio
 import logging
-from collections.abc import Callable
+from collections.abc import AsyncGenerator, Callable
 
 from roborock.callbacks import decoder_callback
 from roborock.data import HomeDataDevice, RRiot, UserData
@@ -73,6 +74,21 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
         dispatch = decoder_callback(self._decoder, callback, _LOGGER)
         return await self._mqtt_session.subscribe(self._subscribe_topic, dispatch)
 
+    async def subscribe_stream(self) -> AsyncGenerator[RoborockMessage, None]:
+        """Subscribe to the device's message stream.
+        
+        This is useful for processing all incoming messages in an async for loop,
+        when they are not necessarily associated with a specific request.
+        """
+        message_queue: asyncio.Queue[RoborockMessage] = asyncio.Queue()
+        unsub = await self.subscribe(message_queue.put_nowait)
+        try:
+            while True:
+                message = await message_queue.get()
+                yield message
+        finally:
+            unsub()
+
     async def publish(self, message: RoborockMessage) -> None:
         """Publish a command message.
 

roborock/devices/traits/b01/__init__.py

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

roborock/devices/traits/b01/q10/__init__.py

@@ -1 +1,104 @@
-"""Q10"""
+"""Traits for Q10 B01 devices."""
+
+import asyncio
+import logging
+from typing import Any
+
+from roborock import B01Props
+from roborock.data.b01_q10.b01_q10_code_mappings import B01_Q10_DP
+from roborock.devices.b01_q10_channel import ParamsType, send_command, stream_decoded_responses
+from roborock.devices.mqtt_channel import MqttChannel
+from roborock.devices.traits import Trait
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "Q10PropertiesApi",
+]
+
+
+class Q10PropertiesApi(Trait):
+    """API for interacting with B01 devices."""
+
+    def __init__(self, channel: MqttChannel) -> None:
+        """Initialize the B01Props API."""
+        self._channel = channel
+        self._task: asyncio.Task | None = None
+
+    async def start(self) -> None:
+        """Start any necessary subscriptions for the trait."""
+        self._task = asyncio.create_task(self._run_loop())
+
+    async def close(self) -> None:
+        """Close any resources held by the trait."""
+        if self._task is not None:
+            self._task.cancel()
+            try:
+                await self._task
+            except asyncio.CancelledError:
+                pass
+            self._task = None
+
+    async def start_clean(self) -> None:
+        """Start cleaning."""
+        await self.send(
+            command=B01_Q10_DP.START_CLEAN,
+            # TODO: figure out other commands
+            # 1 = start cleaning
+            # 2 = electoral clean, also has "clean_paramters"
+            # 4 = fast create map
+            params={"cmd": 1},
+        )
+
+    async def pause_clean(self) -> None:
+        """Pause cleaning."""
+        await self.send(
+            command=B01_Q10_DP.PAUSE,
+            params={},
+        )
+
+    async def resume_clean(self) -> None:
+        """Pause cleaning."""
+        await self.send(
+            command=B01_Q10_DP.RESUME,
+            params={},
+        )
+
+    async def stop_clean(self) -> None:
+        """Stop cleaning."""
+        await self.send(
+            command=B01_Q10_DP.STOP,
+            params={},
+        )
+
+    async def return_to_dock(self) -> None:
+        """Return to dock."""
+        await self.send(
+            command=B01_Q10_DP.START_DOCK_TASK,
+            params={},
+        )
+
+    async def send(self, command: B01_Q10_DP, params: ParamsType) -> None:
+        """Send a command to the device."""
+        await send_command(
+            self._channel,
+            command=command,
+            params=params,
+        )
+
+    async def _run_loop(self) -> None:
+        """Run the main loop for processing incoming messages."""
+        async for decoded_dps in stream_decoded_responses(self._channel):
+            _LOGGER.debug("Received B01 Q10 decoded DPS: %s", decoded_dps)
+
+            # Temporary debugging: Log all common values
+            if B01_Q10_DP.COMMON not in decoded_dps:
+                continue
+            common_values = decoded_dps[B01_Q10_DP.COMMON]
+            for key, value in common_values.items():
+                _LOGGER.debug("%s: %s", key, value)
+
+
+def create(channel: MqttChannel) -> Q10PropertiesApi:
+    """Create traits for B01 devices."""
+    return Q10PropertiesApi(channel)

roborock/devices/traits/b01/q7/__init__.py

@@ -10,7 +10,7 @@
     SCWindMapping,
     WaterLevelMapping,
 )
-from roborock.devices.b01_channel import CommandType, ParamsType, send_decoded_command
+from roborock.devices.b01_q7_channel import CommandType, ParamsType, send_decoded_command
 from roborock.devices.mqtt_channel import MqttChannel
 from roborock.devices.traits import Trait
 from roborock.roborock_message import RoborockB01Props