|
| 1 | +"""Modules for communicating with specific Roborock devices over MQTT.""" |
| 2 | + |
| 3 | +import asyncio |
| 4 | +import logging |
| 5 | +from collections.abc import Callable |
| 6 | +from json import JSONDecodeError |
| 7 | + |
| 8 | +from roborock.containers import RRiot |
| 9 | +from roborock.exceptions import RoborockException |
| 10 | +from roborock.mqtt.session import MqttParams, MqttSession |
| 11 | +from roborock.protocol import create_mqtt_decoder, create_mqtt_encoder |
| 12 | +from roborock.roborock_message import RoborockMessage |
| 13 | + |
| 14 | +_LOGGER = logging.getLogger(__name__) |
| 15 | + |
| 16 | + |
| 17 | +class MqttChannel: |
| 18 | + """Simple RPC-style channel for communicating with a device over MQTT. |
| 19 | +
|
| 20 | + Handles request/response correlation and timeouts, but leaves message |
| 21 | + format most parsing to higher-level components. |
| 22 | + """ |
| 23 | + |
| 24 | + def __init__(self, mqtt_session: MqttSession, duid: str, local_key: str, rriot: RRiot, mqtt_params: MqttParams): |
| 25 | + self._mqtt_session = mqtt_session |
| 26 | + self._duid = duid |
| 27 | + self._local_key = local_key |
| 28 | + self._rriot = rriot |
| 29 | + self._mqtt_params = mqtt_params |
| 30 | + |
| 31 | + # RPC support |
| 32 | + self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {} |
| 33 | + self._decoder = create_mqtt_decoder(local_key) |
| 34 | + self._encoder = create_mqtt_encoder(local_key) |
| 35 | + self._queue_lock = asyncio.Lock() |
| 36 | + |
| 37 | + @property |
| 38 | + def _publish_topic(self) -> str: |
| 39 | + """Topic to send commands to the device.""" |
| 40 | + return f"rr/m/i/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}" |
| 41 | + |
| 42 | + @property |
| 43 | + def _subscribe_topic(self) -> str: |
| 44 | + """Topic to receive responses from the device.""" |
| 45 | + return f"rr/m/o/{self._rriot.u}/{self._mqtt_params.username}/{self._duid}" |
| 46 | + |
| 47 | + async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]: |
| 48 | + """Subscribe to the device's response topic. |
| 49 | +
|
| 50 | + The callback will be called with the message payload when a message is received. |
| 51 | +
|
| 52 | + All messages received will be processed through the provided callback, even |
| 53 | + those sent in response to the `send_command` command. |
| 54 | +
|
| 55 | + Returns a callable that can be used to unsubscribe from the topic. |
| 56 | + """ |
| 57 | + |
| 58 | + def message_handler(payload: bytes) -> None: |
| 59 | + if not (messages := self._decoder(payload)): |
| 60 | + _LOGGER.warning("Failed to decode MQTT message: %s", payload) |
| 61 | + return |
| 62 | + for message in messages: |
| 63 | + _LOGGER.debug("Received message: %s", message) |
| 64 | + asyncio.create_task(self._resolve_future_with_lock(message)) |
| 65 | + try: |
| 66 | + callback(message) |
| 67 | + except Exception as e: |
| 68 | + _LOGGER.exception("Uncaught error in message handler callback: %s", e) |
| 69 | + |
| 70 | + return await self._mqtt_session.subscribe(self._subscribe_topic, message_handler) |
| 71 | + |
| 72 | + async def _resolve_future_with_lock(self, message: RoborockMessage) -> None: |
| 73 | + """Resolve waiting future with proper locking.""" |
| 74 | + if (request_id := message.get_request_id()) is None: |
| 75 | + _LOGGER.debug("Received message with no request_id") |
| 76 | + return |
| 77 | + async with self._queue_lock: |
| 78 | + if (future := self._waiting_queue.pop(request_id, None)) is not None: |
| 79 | + future.set_result(message) |
| 80 | + else: |
| 81 | + _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id) |
| 82 | + |
| 83 | + async def send_command(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage: |
| 84 | + """Send a command message and wait for the response message. |
| 85 | +
|
| 86 | + Returns the raw response message - caller is responsible for parsing. |
| 87 | + """ |
| 88 | + try: |
| 89 | + if (request_id := message.get_request_id()) is None: |
| 90 | + raise RoborockException("Message must have a request_id for RPC calls") |
| 91 | + except (ValueError, JSONDecodeError) as err: |
| 92 | + _LOGGER.exception("Error getting request_id from message: %s", err) |
| 93 | + raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err |
| 94 | + |
| 95 | + future: asyncio.Future[RoborockMessage] = asyncio.Future() |
| 96 | + async with self._queue_lock: |
| 97 | + if request_id in self._waiting_queue: |
| 98 | + raise RoborockException(f"Request ID {request_id} already pending, cannot send command") |
| 99 | + self._waiting_queue[request_id] = future |
| 100 | + |
| 101 | + try: |
| 102 | + encoded_msg = self._encoder(message) |
| 103 | + await self._mqtt_session.publish(self._publish_topic, encoded_msg) |
| 104 | + |
| 105 | + return await asyncio.wait_for(future, timeout=timeout) |
| 106 | + |
| 107 | + except asyncio.TimeoutError as ex: |
| 108 | + async with self._queue_lock: |
| 109 | + self._waiting_queue.pop(request_id, None) |
| 110 | + raise RoborockException(f"Command timed out after {timeout}s") from ex |
| 111 | + except Exception: |
| 112 | + logging.exception("Uncaught error sending command") |
| 113 | + async with self._queue_lock: |
| 114 | + self._waiting_queue.pop(request_id, None) |
| 115 | + raise |
0 commit comments