feat: fix a01 and b01 response handling in new api

Author: allenporter

roborock/devices/a01_channel.py

@@ -1,19 +1,24 @@
 """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
 
 
 @overload
@@ -39,5 +44,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]
+
+    # We only block on a response for queries
+    param_values = {int(k): v for k, v in params.items()}
+    if not (
+        query_values := param_values.get(int(RoborockDyadDataProtocol.ID_QUERY))
+        or param_values.get(int(RoborockZeoProtocol.ID_QUERY))
+    ):
+        await mqtt_channel.publish(roborock_message)
+        return {}
+
+    # This can be simplified if we can assume a all results are returned in 
+    # single response. Otherwise, this will construct a result by merging in
+    # responses that contain the ids that were queried.
+    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:
+            return
+        for key, value in decoded.items():
+            if key in query_values:
+                result[key] = value
+        if len(result) != len(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/b01_channel.py

@@ -3,12 +3,10 @@
 from __future__ import annotations
 
 import logging
-from typing import Any
 
 from roborock.protocols.b01_protocol import (
     CommandType,
     ParamsType,
-    decode_rpc_response,
     encode_mqtt_payload,
 )
 
@@ -22,9 +20,8 @@ async def send_decoded_command(
     dps: int,
     command: CommandType,
     params: ParamsType,
-) -> dict[int, Any]:
+) -> None:
     """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)
-    response = await mqtt_channel.send_message(roborock_message)
-    return decode_rpc_response(response)  # type: ignore[return-value]
+    await mqtt_channel.publish(roborock_message)

roborock/devices/local_channel.py

@@ -4,14 +4,12 @@
 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
 from roborock.roborock_message import RoborockMessage
 
 from .channel import Channel
-from .pending import PendingRpcs
 
 _LOGGER = logging.getLogger(__name__)
 _PORT = 58867
@@ -47,8 +45,6 @@ def __init__(self, host: str, local_key: str):
         self._subscribers: list[Callable[[RoborockMessage], None]] = []
         self._is_connected = False
 
-        # RPC support
-        self._pending_rpcs: PendingRpcs[int, RoborockMessage] = PendingRpcs()
         self._decoder: Decoder = create_local_decoder(local_key)
         self._encoder: Encoder = create_local_encoder(local_key)
 
@@ -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,37 +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
-        await self._pending_rpcs.resolve(request_id, message)
+    async def publish(self, message: RoborockMessage) -> None:
+        """Send a command message.
 
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response 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] = await self._pending_rpcs.start(request_id)
         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:
-            await self._pending_rpcs.pop(request_id)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
+        except Exception as err:
             logging.exception("Uncaught error sending command")
-            await self._pending_rpcs.pop(request_id)
-            raise
+            raise RoborockException(f"Failed to send message: {message}") from err
 
 
 # This module provides a factory function to create LocalChannel instances.

roborock/devices/mqtt_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
 
@@ -69,7 +67,6 @@ def message_handler(payload: bytes) -> None:
                 return
             for message in messages:
                 _LOGGER.debug("Received message: %s", message)
-                asyncio.create_task(self._resolve_future_with_lock(message))
                 try:
                     callback(message)
                 except Exception as e:
@@ -84,40 +81,22 @@ def unsub_wrapper() -> 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
-        await self._pending_rpcs.resolve(request_id, message)
+    async def publish(self, message: RoborockMessage) -> None:
+        """Publish a command message.
 
-    async def send_message(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
-        """Send a command message and wait for the response 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] = await self._pending_rpcs.start(request_id)
-
         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:
-            await self._pending_rpcs.pop(request_id)
-            raise RoborockException(f"Command timed out after {timeout}s") from ex
-        except Exception:
-            logging.exception("Uncaught error sending command")
-            await self._pending_rpcs.pop(request_id)
-            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/v1_rpc_channel.py

@@ -6,11 +6,13 @@
 and local connections, preferring local when available.
 """
 
+import asyncio
 import logging
 from collections.abc import Callable
 from typing import Any, Protocol, TypeVar, overload
 
 from roborock.containers import RoborockBase
+from roborock.exceptions import RoborockException
 from roborock.protocols.v1_protocol import (
     CommandType,
     ParamsType,
@@ -24,6 +26,7 @@
 from .mqtt_channel import MqttChannel
 
 _LOGGER = logging.getLogger(__name__)
+_TIMEOUT = 10.0
 
 
 _T = TypeVar("_T", bound=RoborockBase)
@@ -132,8 +135,26 @@ async def _send_raw_command(
         _LOGGER.debug("Sending command (%s): %s, params=%s", self._name, method, params)
         request_message = RequestMessage(method, params=params)
         message = self._payload_encoder(request_message)
-        response = await self._channel.send_message(message)
-        return decode_rpc_response(response)
+
+        future: asyncio.Future[dict[str, Any]] = asyncio.Future()
+
+        def find_response(response_message: RoborockMessage) -> None:
+            try:
+                decoded = decode_rpc_response(response_message)
+            except RoborockException:
+                return
+            if decoded.request_id == request_message.request_id:
+                future.set_result(decoded.data)
+
+        unsub = await self._channel.subscribe(find_response)
+        try:
+            await self._channel.publish(message)
+            return await asyncio.wait_for(future, timeout=_TIMEOUT)
+        except TimeoutError as ex:
+            future.cancel()
+            raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
+        finally:
+            unsub()
 
 
 def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:

roborock/protocols/v1_protocol.py

@@ -95,7 +95,18 @@ def _as_payload(self, security_data: SecurityData | None) -> bytes:
         )
 
 
-def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
+@dataclass(kw_only=True, frozen=True)
+class ResponseMessage:
+    """Data structure for v1 RoborockMessage responses."""
+
+    request_id: int | None
+    """The request ID of the response."""
+
+    data: dict[str, Any]
+    """The data of the response."""
+
+
+def decode_rpc_response(message: RoborockMessage) -> ResponseMessage:
     """Decode a V1 RPC_RESPONSE message."""
     if not message.payload:
         raise RoborockException("Invalid V1 message format: missing payload")
@@ -109,14 +120,19 @@ def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
     if not isinstance(datapoints, dict):
         raise RoborockException(f"Invalid V1 message format: 'dps' should be a dictionary for {message.payload!r}")
 
-    if not (data_point := datapoints.get("102")):
-        raise RoborockException("Invalid V1 message format: missing '102' data point")
+    if not (data_point := datapoints.get(str(RoborockMessageProtocol.RPC_RESPONSE))):
+        raise RoborockException(
+            f"Invalid V1 message format: missing '{RoborockMessageProtocol.RPC_RESPONSE}' data point"
+        )
 
     try:
         data_point_response = json.loads(data_point)
     except (json.JSONDecodeError, TypeError) as e:
-        raise RoborockException(f"Invalid V1 message data point '102': {e} for {message.payload!r}") from e
+        raise RoborockException(
+            f"Invalid V1 message data point '{RoborockMessageProtocol.RPC_RESPONSE}': {e} for {message.payload!r}"
+        ) from e
 
+    request_id: int | None = data_point_response.get("id")
     if error := data_point_response.get("error"):
         raise RoborockException(f"Error in message: {error}")
 
@@ -127,7 +143,7 @@ def decode_rpc_response(message: RoborockMessage) -> dict[str, Any]:
         result = result[0]
     if not isinstance(result, dict):
         raise RoborockException(f"Invalid V1 message format: 'result' should be a dictionary for {message.payload!r}")
-    return result
+    return ResponseMessage(request_id=request_id, data=result)
 
 
 @dataclass

roborock/roborock_message.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import json
 import math
 import time
 from dataclasses import dataclass, field
@@ -247,12 +246,3 @@ class RoborockMessage:
     version: bytes = b"1.0"
     random: int = field(default_factory=lambda: get_next_int(10000, 99999))
     timestamp: int = field(default_factory=lambda: math.floor(time.time()))
-
-    def get_request_id(self) -> int | None:
-        if self.payload:
-            payload = json.loads(self.payload.decode())
-            for data_point_number, data_point in payload.get("dps").items():
-                if data_point_number in ["101", "102"]:
-                    data_point_response = json.loads(data_point)
-                    return data_point_response.get("id")
-        return None

roborock/version_1_apis/roborock_mqtt_client_v1.py

@@ -49,7 +49,7 @@ async def _send_command(
             RoborockMessageProtocol.RPC_REQUEST,
             security_data=self._security_data,
         )
-        self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id, method)
+        self._logger.debug("Building message id %s for method %s", request_message.request_id, method)
 
         await self.validate_connection()
         request_id = request_message.request_id