feat: Add a v1 protocol channel that can send messages across MQTT or Local connections, preferring local (#416)

* feat: Add a v1 protocol channel bridging across MQTT/Local channels

* chore: Remove whitespace

* feat: Fix tests referencing RoborockStateCode

* feat: Fix tests reverted by co-pilot

* fix: Update error message and add pydoc for exception handling on subscribe

* chore(deps): bump click from 8.1.8 to 8.2.1 (#401)

Bumps [click](https://github.com/pallets/click) from 8.1.8 to 8.2.1.
- [Release notes](https://github.com/pallets/click/releases)
- [Changelog](https://github.com/pallets/click/blob/main/CHANGES.rst)
- [Commits](pallets/click@8.1.8...8.2.1)

---
updated-dependencies:
- dependency-name: click
  dependency-version: 8.2.1
  dependency-type: direct:production
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* fix: Simplify local connection handling

* fix: Update pydoc for sending a raw command

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

Author: allenporter

roborock/cli.py

@@ -115,7 +115,16 @@ async def home_data_cache() -> HomeData:
     devices = await device_manager.get_devices()
     click.echo(f"Discovered devices: {', '.join([device.name for device in devices])}")
 
-    click.echo("MQTT session started. Listening for messages...")
+    click.echo("MQTT session started. Querying devices...")
+    for device in devices:
+        try:
+            status = await device.get_status()
+        except RoborockException as e:
+            click.echo(f"Failed to get status for {device.name}: {e}")
+        else:
+            click.echo(f"Device {device.name} status: {status.as_dict()}")
+
+    click.echo("Listening for messages.")
     await asyncio.sleep(duration)
 
     # Close the device manager (this will close all devices and MQTT session)

roborock/devices/device.py

@@ -9,10 +9,18 @@
 from collections.abc import Callable
 from functools import cached_property
 
-from roborock.containers import HomeDataDevice, HomeDataProduct, UserData
+from roborock.containers import (
+    HomeDataDevice,
+    HomeDataProduct,
+    ModelStatus,
+    S7MaxVStatus,
+    Status,
+    UserData,
+)
 from roborock.roborock_message import RoborockMessage
+from roborock.roborock_typing import RoborockCommand
 
-from .mqtt_channel import MqttChannel
+from .v1_channel import V1Channel
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -38,19 +46,18 @@ def __init__(
         user_data: UserData,
         device_info: HomeDataDevice,
         product_info: HomeDataProduct,
-        mqtt_channel: MqttChannel,
+        v1_channel: V1Channel,
     ) -> None:
         """Initialize the RoborockDevice.
 
-        The device takes ownership of the MQTT channel for communication with the device.
-        Use `connect()` to establish the connection, which will set up the MQTT channel
-        for receiving messages from the device. Use `close()` to unsubscribe from the MQTT
-        channel.
+        The device takes ownership of the V1 channel for communication with the device.
+        Use `connect()` to establish the connection, which will set up the appropriate
+        protocol channel. Use `close()` to clean up all connections.
         """
         self._user_data = user_data
         self._device_info = device_info
         self._product_info = product_info
-        self._mqtt_channel = mqtt_channel
+        self._v1_channel = v1_channel
         self._unsub: Callable[[], None] | None = None
 
     @property
@@ -82,27 +89,32 @@ def device_version(self) -> str:
         )
         return DeviceVersion.UNKNOWN
 
-    async def connect(self) -> None:
-        """Connect to the device using MQTT.
+    @property
+    def is_connected(self) -> bool:
+        """Return whether the device is connected."""
+        return self._v1_channel.is_mqtt_connected or self._v1_channel.is_local_connected
 
-        This method will set up the MQTT channel for communication with the device.
-        """
+    async def connect(self) -> None:
+        """Connect to the device using the appropriate protocol channel."""
         if self._unsub:
             raise ValueError("Already connected to the device")
-        self._unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
+        self._unsub = await self._v1_channel.subscribe(self._on_message)
+        _LOGGER.info("Connected to V1 device %s", self.name)
 
     async def close(self) -> None:
-        """Close the MQTT connection to the device.
-
-        This method will unsubscribe from the MQTT channel and clean up resources.
-        """
+        """Close all connections to the device."""
         if self._unsub:
             self._unsub()
             self._unsub = None
 
-    def _on_mqtt_message(self, message: RoborockMessage) -> None:
-        """Handle incoming MQTT messages from the device.
+    def _on_message(self, message: RoborockMessage) -> None:
+        """Handle incoming messages from the device."""
+        _LOGGER.debug("Received message from device: %s", message)
+
+    async def get_status(self) -> Status:
+        """Get the current status of the device.
 
-        This method should be overridden in subclasses to handle specific device messages.
+        This is a placeholder command and will likely be changed/moved in the future.
         """
-        _LOGGER.debug("Received message from device %s: %s", self.duid, message)
+        status_type: type[Status] = ModelStatus.get(self._product_info.model, S7MaxVStatus)
+        return await self._v1_channel.send_decoded_command(RoborockCommand.GET_STATUS, response_type=status_type)

roborock/devices/device_manager.py

@@ -10,13 +10,13 @@
     HomeDataProduct,
     UserData,
 )
-from roborock.devices.device import RoborockDevice
+from roborock.devices.device import DeviceVersion, RoborockDevice
 from roborock.mqtt.roborock_session import create_mqtt_session
 from roborock.mqtt.session import MqttSession
 from roborock.protocol import create_mqtt_params
 from roborock.web_api import RoborockApiClient
 
-from .mqtt_channel import MqttChannel
+from .v1_channel import create_v1_channel
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -114,8 +114,15 @@ async def create_device_manager(user_data: UserData, home_data_api: HomeDataApi)
     mqtt_session = await create_mqtt_session(mqtt_params)
 
     def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
-        mqtt_channel = MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)
-        return RoborockDevice(user_data, device, product, mqtt_channel)
+        # Check device version and only support V1 for now
+        if device.pv != DeviceVersion.V1.value:
+            raise NotImplementedError(
+                f"Device {device.name} has version {device.pv}, but only V1 devices "
+                f"are supported through the unified interface."
+            )
+        # Create V1 channel that handles both MQTT and local connections
+        v1_channel = create_v1_channel(user_data, mqtt_params, mqtt_session, device)
+        return RoborockDevice(user_data, device, product, v1_channel)
 
     manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session)
     await manager.discover_devices()

roborock/devices/local_channel.py

@@ -64,7 +64,7 @@ async def connect(self) -> None:
         except OSError as e:
             raise RoborockConnectionException(f"Failed to connect to {self._host}:{_PORT}") from e
 
-    async def close(self) -> None:
+    def close(self) -> None:
         """Disconnect from the device."""
         if self._transport:
             self._transport.close()
@@ -144,3 +144,25 @@ async def send_command(self, message: RoborockMessage, timeout: float = 10.0) ->
             async with self._queue_lock:
                 self._waiting_queue.pop(request_id, None)
             raise
+
+
+# This module provides a factory function to create LocalChannel instances.
+#
+# TODO: Make a separate LocalSession and use it to manage retries with the host,
+# similar to how MqttSession works. For now this is a simple factory function
+# for creating channels.
+LocalSession = Callable[[str], LocalChannel]
+
+
+def create_local_session(local_key: str) -> LocalSession:
+    """Creates a local session which can create local channels.
+
+    This plays a role similar to the MqttSession but is really just a factory
+    for creating LocalChannel instances with the same local key.
+    """
+
+    def create_local_channel(host: str) -> LocalChannel:
+        """Create a LocalChannel instance for the given host."""
+        return LocalChannel(host, local_key)
+
+    return create_local_channel

roborock/devices/v1_channel.py

@@ -0,0 +1,212 @@
+"""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 logging
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from roborock.containers import HomeDataDevice, NetworkInfo, RoborockBase, UserData
+from roborock.exceptions import RoborockException
+from roborock.mqtt.session import MqttParams, MqttSession
+from roborock.protocols.v1_protocol import (
+    CommandType,
+    ParamsType,
+    SecurityData,
+    create_mqtt_payload_encoder,
+    create_security_data,
+    decode_rpc_response,
+    encode_local_payload,
+)
+from roborock.roborock_message import RoborockMessage
+from roborock.roborock_typing import RoborockCommand
+
+from .local_channel import LocalChannel, LocalSession, create_local_session
+from .mqtt_channel import MqttChannel
+
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = [
+    "V1Channel",
+]
+
+_T = TypeVar("_T", bound=RoborockBase)
+
+
+class V1Channel:
+    """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,
+    ) -> None:
+        """Initialize the V1Channel.
+
+        Args:
+            mqtt_channel: MQTT channel for cloud communication
+            local_session: Factory that creates LocalChannels for a hostname.
+        """
+        self._device_uid = device_uid
+        self._mqtt_channel = mqtt_channel
+        self._mqtt_payload_encoder = create_mqtt_payload_encoder(security_data)
+        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._networking_info: NetworkInfo | None = None
+
+    @property
+    def is_local_connected(self) -> bool:
+        """Return whether local connection is available."""
+        return self._local_unsub is not None
+
+    @property
+    def is_mqtt_connected(self) -> bool:
+        """Return whether MQTT connection is available."""
+        return self._mqtt_unsub is not None
+
+    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
+        """Subscribe to all messages from the device.
+
+        This will establish MQTT connection first, and also attempt to set up
+        local connection if possible. Any failures to subscribe to MQTT will raise
+        a RoborockException. A local connection failure will not raise an exception,
+        since the local connection is optional.
+        """
+
+        if self._mqtt_unsub:
+            raise ValueError("Already connected to the device")
+        self._callback = callback
+
+        # First establish MQTT connection
+        self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
+        _LOGGER.debug("V1Channel connected to device %s via MQTT", self._device_uid)
+
+        # Try to establish an optional local connection as well.
+        try:
+            self._local_unsub = await self._local_connect()
+        except RoborockException as err:
+            _LOGGER.warning("Could not establish local connection for device %s: %s", self._device_uid, err)
+        else:
+            _LOGGER.debug("Local connection established for device %s", self._device_uid)
+
+        def unsub() -> None:
+            """Unsubscribe from all messages."""
+            if self._mqtt_unsub:
+                self._mqtt_unsub()
+                self._mqtt_unsub = None
+            if self._local_unsub:
+                self._local_unsub()
+                self._local_unsub = None
+            _LOGGER.debug("Unsubscribed from device %s", self._device_uid)
+
+        return unsub
+
+    async def _get_networking_info(self) -> NetworkInfo:
+        """Retrieve networking information for the device.
+
+        This is a cloud only command used to get the local device's IP address.
+        """
+        try:
+            return await self._send_mqtt_decoded_command(RoborockCommand.GET_NETWORK_INFO, response_type=NetworkInfo)
+        except RoborockException as e:
+            raise RoborockException(f"Network info failed for device {self._device_uid}") from e
+
+    async def _local_connect(self) -> Callable[[], None]:
+        """Set up local connection if possible."""
+        _LOGGER.debug("Attempting to connect to local channel for device %s", self._device_uid)
+        if self._networking_info is None:
+            self._networking_info = await self._get_networking_info()
+        host = self._networking_info.ip
+        _LOGGER.debug("Connecting to local channel at %s", host)
+        self._local_channel = self._local_session(host)
+        try:
+            await self._local_channel.connect()
+        except RoborockException as e:
+            self._local_channel = None
+            raise RoborockException(f"Error connecting to local device {self._device_uid}: {e}") from e
+
+        return await self._local_channel.subscribe(self._on_local_message)
+
+    async def send_decoded_command(
+        self,
+        method: CommandType,
+        *,
+        response_type: type[_T],
+        params: ParamsType = None,
+    ) -> _T:
+        """Send a command using the best available transport.
+
+        Will prefer local connection if available, falling back to MQTT.
+        """
+        connection = "local" if self.is_local_connected else "mqtt"
+        _LOGGER.debug("Sending command (%s): %s, params=%s", connection, method, params)
+        if self._local_channel:
+            return await self._send_local_decoded_command(method, response_type=response_type, params=params)
+        return await self._send_mqtt_decoded_command(method, response_type=response_type, params=params)
+
+    async def _send_mqtt_raw_command(self, method: CommandType, params: ParamsType | None = None) -> dict[str, Any]:
+        """Send a raw command and return a raw unparsed response."""
+        message = self._mqtt_payload_encoder(method, params)
+        _LOGGER.debug("Sending MQTT message for device %s: %s", self._device_uid, message)
+        response = await self._mqtt_channel.send_command(message)
+        return decode_rpc_response(response)
+
+    async def _send_mqtt_decoded_command(
+        self, method: CommandType, *, response_type: type[_T], params: ParamsType | None = None
+    ) -> _T:
+        """Send a command over MQTT and decode the response."""
+        decoded_response = await self._send_mqtt_raw_command(method, params)
+        return response_type.from_dict(decoded_response)
+
+    async def _send_local_raw_command(self, method: CommandType, params: ParamsType | None = None) -> dict[str, Any]:
+        """Send a raw command over local connection."""
+        if not self._local_channel:
+            raise RoborockException("Local channel is not connected")
+
+        message = encode_local_payload(method, params)
+        _LOGGER.debug("Sending local message for device %s: %s", self._device_uid, message)
+        response = await self._local_channel.send_command(message)
+        return decode_rpc_response(response)
+
+    async def _send_local_decoded_command(
+        self, method: CommandType, *, response_type: type[_T], params: ParamsType | None = None
+    ) -> _T:
+        """Send a command over local connection and decode the response."""
+        if not self._local_channel:
+            raise RoborockException("Local channel is not connected")
+        decoded_response = await self._send_local_raw_command(method, params)
+        return response_type.from_dict(decoded_response)
+
+    def _on_mqtt_message(self, message: RoborockMessage) -> None:
+        """Handle incoming MQTT messages."""
+        _LOGGER.debug("V1Channel received MQTT message from device %s: %s", self._device_uid, message)
+        if self._callback:
+            self._callback(message)
+
+    def _on_local_message(self, message: RoborockMessage) -> None:
+        """Handle incoming local messages."""
+        _LOGGER.debug("V1Channel received local message from device %s: %s", self._device_uid, message)
+        if self._callback:
+            self._callback(message)
+
+
+def create_v1_channel(
+    user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
+) -> 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)
+    return V1Channel(device.duid, security_data, mqtt_channel, local_session=local_session)