Merge branch 'main' into b01_improvements

Author: Lash-L

poetry.lock

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "python-roborock"
-version = "2.35.0"
+version = "2.37.0"
 description = "A package to control Roborock vacuums."
 authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
 license = "GPL-3.0-only"

pyproject.toml

@@ -733,6 +733,29 @@ class NetworkInfo(RoborockBase):
     rssi: int | None = None
 
 
+@dataclass
+class AppInitStatusLocalInfo(RoborockBase):
+    location: str
+    bom: str | None = None
+    featureset: int | None = None
+    language: str | None = None
+    logserver: str | None = None
+    wifiplan: str | None = None
+    timezone: str | None = None
+    name: str | None = None
+
+
+@dataclass
+class AppInitStatus(RoborockBase):
+    local_info: AppInitStatusLocalInfo
+    feature_info: list[int]
+    new_feature_info: int
+    new_feature_info_str: str
+    new_feature_info_2: int | None = None
+    carriage_type: int | None = None
+    dsp_version: int | None = None
+
+
 @dataclass
 class DeviceData(RoborockBase):
     device: HomeDataDevice

roborock/containers.py

@@ -1,19 +1,28 @@
 """Thin wrapper around the MQTT channel for Roborock A01 devices."""
 
-from __future__ import annotations
-
+import asyncio
 import logging
 from typing import Any, overload
 
+from roborock.exceptions import RoborockException
 from roborock.protocols.a01_protocol import (
     decode_rpc_response,
     encode_mqtt_payload,
 )
-from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
+from roborock.roborock_message import (
+    RoborockDyadDataProtocol,
+    RoborockMessage,
+    RoborockZeoProtocol,
+)
 
 from .mqtt_channel import MqttChannel
 
 _LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
+
+# Both RoborockDyadDataProtocol and RoborockZeoProtocol have the same
+# value for ID_QUERY
+_ID_QUERY = int(RoborockDyadDataProtocol.ID_QUERY)
 
 
 @overload
@@ -39,5 +48,46 @@ async def send_decoded_command(
     """Send a command on the MQTT channel and get a decoded response."""
     _LOGGER.debug("Sending MQTT command: %s", params)
     roborock_message = encode_mqtt_payload(params)
-    response = await mqtt_channel.send_message(roborock_message)
-    return decode_rpc_response(response)  # type: ignore[return-value]
+
+    # For commands that set values: send the command and do not
+    # block waiting for a response. Queries are handled below.
+    param_values = {int(k): v for k, v in params.items()}
+    if not (query_values := param_values.get(_ID_QUERY)):
+        await mqtt_channel.publish(roborock_message)
+        return {}
+
+    # Merge any results together than contain the requested data. This
+    # does not use a future since it needs to merge results across responses.
+    # This could be simplified if we can assume there is a single response.
+    finished = asyncio.Event()
+    result: dict[int, Any] = {}
+
+    def find_response(response_message: RoborockMessage) -> None:
+        """Handle incoming messages and resolve the future."""
+        try:
+            decoded = decode_rpc_response(response_message)
+        except RoborockException as ex:
+            _LOGGER.info("Failed to decode a01 message: %s: %s", response_message, ex)
+            return
+        for key, value in decoded.items():
+            if key in query_values:
+                result[key] = value
+        if len(result) != len(query_values):
+            _LOGGER.debug("Incomplete query response: %s != %s", result, query_values)
+            return
+        _LOGGER.debug("Received query response: %s", result)
+        if not finished.is_set():
+            finished.set()
+
+    unsub = await mqtt_channel.subscribe(find_response)
+
+    try:
+        await mqtt_channel.publish(roborock_message)
+        try:
+            await asyncio.wait_for(finished.wait(), timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+    finally:
+        unsub()
+
+    return result  # type: ignore[return-value]

roborock/devices/a01_channel.py

@@ -24,4 +24,4 @@ async def send_decoded_command(
     """Send a command on the MQTT channel and get a decoded response."""
     _LOGGER.debug("Sending MQTT command: %s", params)
     roborock_message = encode_mqtt_payload(dps, command, params)
-    await mqtt_channel.send_message_no_wait(roborock_message)
+    await mqtt_channel.publish(roborock_message)

roborock/devices/b01_channel.py

@@ -4,7 +4,6 @@
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass
-from json import JSONDecodeError
 
 from roborock.exceptions import RoborockConnectionException, RoborockException
 from roborock.protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
@@ -46,11 +45,8 @@ def __init__(self, host: str, local_key: str):
         self._subscribers: list[Callable[[RoborockMessage], None]] = []
         self._is_connected = False
 
-        # RPC support
-        self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {}
         self._decoder: Decoder = create_local_decoder(local_key)
         self._encoder: Encoder = create_local_encoder(local_key)
-        self._queue_lock = asyncio.Lock()
 
     @property
     def is_connected(self) -> bool:
@@ -87,7 +83,6 @@ def _data_received(self, data: bytes) -> None:
             return
         for message in messages:
             _LOGGER.debug("Received message: %s", message)
-            asyncio.create_task(self._resolve_future_with_lock(message))
             for callback in self._subscribers:
                 try:
                     callback(message)
@@ -109,48 +104,24 @@ def unsubscribe() -> None:
 
         return unsubscribe
 
-    async def _resolve_future_with_lock(self, message: RoborockMessage) -> None:
-        """Resolve waiting future with proper locking."""
-        if (request_id := message.get_request_id()) is None:
-            _LOGGER.debug("Received message with no request_id")
-            return
-        async with self._queue_lock:
-            if (future := self._waiting_queue.pop(request_id, None)) is not None:
-                future.set_result(message)
-            else:
-                _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id)
-
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response message."""
+    async def publish(self, message: RoborockMessage) -> None:
+        """Send a command message.
+
+        The caller is responsible for associating the message with its response.
+        """
         if not self._transport or not self._is_connected:
             raise RoborockConnectionException("Not connected to device")
 
-        try:
-            if (request_id := message.get_request_id()) is None:
-                raise RoborockException("Message must have a request_id for RPC calls")
-        except (ValueError, JSONDecodeError) as err:
-            _LOGGER.exception("Error getting request_id from message: %s", err)
-            raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err
-
-        future: asyncio.Future[RoborockMessage] = asyncio.Future()
-        async with self._queue_lock:
-            if request_id in self._waiting_queue:
-                raise RoborockException(f"Request ID {request_id} already pending, cannot send command")
-            self._waiting_queue[request_id] = future
-
         try:
             encoded_msg = self._encoder(message)
+        except Exception as err:
+            _LOGGER.exception("Error encoding MQTT message: %s", err)
+            raise RoborockException(f"Failed to encode MQTT message: {err}") from err
+        try:
             self._transport.write(encoded_msg)
-            return await asyncio.wait_for(future, timeout=timeout)
-        except asyncio.TimeoutError as ex:
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
+        except Exception as err:
             logging.exception("Uncaught error sending command")
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise
+            raise RoborockException(f"Failed to send message: {message}") from err
 
 
 # This module provides a factory function to create LocalChannel instances.

roborock/devices/local_channel.py

@@ -1,13 +1,11 @@
 """Modules for communicating with specific Roborock devices over MQTT."""
 
-import asyncio
 import logging
 from collections.abc import Callable
-from json import JSONDecodeError
 
 from roborock.containers import HomeDataDevice, RRiot, UserData
 from roborock.exceptions import RoborockException
-from roborock.mqtt.session import MqttParams, MqttSession
+from roborock.mqtt.session import MqttParams, MqttSession, MqttSessionException
 from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder
 from roborock.roborock_message import RoborockMessage
 
@@ -30,17 +28,16 @@ def __init__(self, mqtt_session: MqttSession, duid: str, local_key: str, rriot:
         self._rriot = rriot
         self._mqtt_params = mqtt_params
 
-        # RPC support
-        self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {}
         self._decoder = create_mqtt_decoder(local_key)
         self._encoder = create_mqtt_encoder(local_key)
-        self._queue_lock = asyncio.Lock()
-        self._mqtt_unsub: Callable[[], None] | None = None
 
     @property
     def is_connected(self) -> bool:
-        """Return true if the channel is connected."""
-        return (self._mqtt_unsub is not None) and self._mqtt_session.connected
+        """Return true if the channel is connected.
+
+        This passes through the underlying MQTT session's connected state.
+        """
+        return self._mqtt_session.connected
 
     @property
     def _publish_topic(self) -> str:
@@ -57,9 +54,6 @@ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callab
 
         The callback will be called with the message payload when a message is received.
 
-        All messages received will be processed through the provided callback, even
-        those sent in response to the `send_command` command.
-
         Returns a callable that can be used to unsubscribe from the topic.
         """
 
@@ -69,75 +63,29 @@ def message_handler(payload: bytes) -> None:
                 return
             for message in messages:
                 _LOGGER.debug("Received message: %s", message)
-                if message.version != b"B01":
-                    asyncio.create_task(self._resolve_future_with_lock(message))
                 try:
                     callback(message)
                 except Exception as e:
                     _LOGGER.exception("Uncaught error in message handler callback: %s", e)
 
-        self._mqtt_unsub = await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
-
-        def unsub_wrapper() -> None:
-            if self._mqtt_unsub is not None:
-                self._mqtt_unsub()
-                self._mqtt_unsub = None
-
-        return unsub_wrapper
-
-    async def _resolve_future_with_lock(self, message: RoborockMessage) -> None:
-        """Resolve waiting future with proper locking."""
-        if (request_id := message.get_request_id()) is None:
-            _LOGGER.debug("Received message with no request_id")
-            return
-        async with self._queue_lock:
-            if (future := self._waiting_queue.pop(request_id, None)) is not None:
-                future.set_result(message)
-            else:
-                _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id)
+        return await self._mqtt_session.subscribe(self._subscribe_topic, message_handler)
 
-    async def send_message_no_wait(self, message: RoborockMessage) -> None:
-        """Send a command message without waiting for a response."""
-        try:
-            encoded_msg = self._encoder(message)
-            await self._mqtt_session.publish(self._publish_topic, encoded_msg)
-        except Exception:
-            logging.exception("Uncaught error sending command")
-            raise
-
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response message.
+    async def publish(self, message: RoborockMessage) -> None:
+        """Publish a command message.
 
-        Returns the raw response message - caller is responsible for parsing.
+        The caller is responsible for handling any responses and associating them
+        with the incoming request.
         """
-        try:
-            if (request_id := message.get_request_id()) is None:
-                raise RoborockException("Message must have a request_id for RPC calls")
-        except (ValueError, JSONDecodeError) as err:
-            _LOGGER.exception("Error getting request_id from message: %s", err)
-            raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err
-
-        future: asyncio.Future[RoborockMessage] = asyncio.Future()
-        async with self._queue_lock:
-            if request_id in self._waiting_queue:
-                raise RoborockException(f"Request ID {request_id} already pending, cannot send command")
-            self._waiting_queue[request_id] = future
-
         try:
             encoded_msg = self._encoder(message)
-            await self._mqtt_session.publish(self._publish_topic, encoded_msg)
-
-            return await asyncio.wait_for(future, timeout=timeout)
-
-        except asyncio.TimeoutError as ex:
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
-            logging.exception("Uncaught error sending command")
-            async with self._queue_lock:
-                self._waiting_queue.pop(request_id, None)
-            raise
+        except Exception as e:
+            _LOGGER.exception("Error encoding MQTT message: %s", e)
+            raise RoborockException(f"Failed to encode MQTT message: {e}") from e
+        try:
+            return await self._mqtt_session.publish(self._publish_topic, encoded_msg)
+        except MqttSessionException as e:
+            _LOGGER.exception("Error publishing MQTT message: %s", e)
+            raise RoborockException(f"Failed to publish MQTT message: {e}") from e
 
 
 def create_mqtt_channel(

roborock/devices/mqtt_channel.py

@@ -27,6 +27,6 @@ def __init__(self, channel: MqttChannel) -> None:
 
     async def query_values(self, props: list[RoborockB01Props]) -> None:
         """Query the device for the values of the given Dyad protocols."""
-        return await send_decoded_command(
+        await send_decoded_command(
             self._channel, dps=10000, command=RoborockB01Methods.GET_PROP, params={"property": props}
         )

roborock/devices/traits/b01/props.py

@@ -79,7 +79,7 @@ def is_local_connected(self) -> bool:
     @property
     def is_mqtt_connected(self) -> bool:
         """Return whether MQTT connection is available."""
-        return self._mqtt_unsub is not None
+        return self._mqtt_unsub is not None and self._mqtt_channel.is_connected
 
     @property
     def rpc_channel(self) -> V1RpcChannel: