python-roborock 2.23.0__tar.gz → 2.25.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (38) hide show
  1. {python_roborock-2.23.0 → python_roborock-2.25.0}/PKG-INFO +1 -1
  2. {python_roborock-2.23.0 → python_roborock-2.25.0}/pyproject.toml +2 -1
  3. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/cli.py +15 -34
  4. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/devices/device.py +45 -2
  5. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/devices/device_manager.py +38 -7
  6. python_roborock-2.25.0/roborock/devices/local_channel.py +146 -0
  7. python_roborock-2.25.0/roborock/devices/mqtt_channel.py +115 -0
  8. python_roborock-2.25.0/roborock/protocols/v1_protocol.py +91 -0
  9. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_1_apis/roborock_client_v1.py +3 -33
  10. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_1_apis/roborock_local_client_v1.py +8 -17
  11. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_1_apis/roborock_mqtt_client_v1.py +7 -4
  12. {python_roborock-2.23.0 → python_roborock-2.25.0}/LICENSE +0 -0
  13. {python_roborock-2.23.0 → python_roborock-2.25.0}/README.md +0 -0
  14. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/__init__.py +0 -0
  15. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/api.py +0 -0
  16. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/cloud_api.py +0 -0
  17. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/code_mappings.py +0 -0
  18. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/command_cache.py +0 -0
  19. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/const.py +0 -0
  20. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/containers.py +0 -0
  21. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/devices/README.md +0 -0
  22. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/devices/__init__.py +0 -0
  23. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/exceptions.py +0 -0
  24. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/local_api.py +0 -0
  25. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/mqtt/__init__.py +0 -0
  26. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/mqtt/roborock_session.py +0 -0
  27. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/mqtt/session.py +0 -0
  28. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/protocol.py +0 -0
  29. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/py.typed +0 -0
  30. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/roborock_future.py +0 -0
  31. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/roborock_message.py +0 -0
  32. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/roborock_typing.py +0 -0
  33. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/util.py +0 -0
  34. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_1_apis/__init__.py +0 -0
  35. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_a01_apis/__init__.py +0 -0
  36. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_a01_apis/roborock_client_a01.py +0 -0
  37. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/version_a01_apis/roborock_mqtt_client_a01.py +0 -0
  38. {python_roborock-2.23.0 → python_roborock-2.25.0}/roborock/web_api.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: python-roborock
3
- Version: 2.23.0
3
+ Version: 2.25.0
4
4
  Summary: A package to control Roborock vacuums.
5
5
  Home-page: https://github.com/humbertogontijo/python-roborock
6
6
  License: GPL-3.0-only
@@ -1,6 +1,6 @@
1
1
  [tool.poetry]
2
2
  name = "python-roborock"
3
- version = "2.23.0"
3
+ version = "2.25.0"
4
4
  description = "A package to control Roborock vacuums."
5
5
  authors = ["humbertogontijo <humbertogontijo@users.noreply.github.com>"]
6
6
  license = "GPL-3.0-only"
@@ -76,3 +76,4 @@ select=["E", "F", "UP", "I"]
76
76
  asyncio_mode = "auto"
77
77
  asyncio_default_fixture_loop_scope = "function"
78
78
  timeout = 30
79
+ log_format = "%(asctime)s.%(msecs)03d %(levelname)s (%(threadName)s) [%(name)s] %(message)s"
@@ -12,9 +12,9 @@ from pyshark.capture.live_capture import LiveCapture, UnknownInterfaceException
12
12
  from pyshark.packet.packet import Packet # type: ignore
13
13
 
14
14
  from roborock import RoborockException
15
- from roborock.containers import DeviceData, HomeDataProduct, LoginData
16
- from roborock.mqtt.roborock_session import create_mqtt_session
17
- from roborock.protocol import MessageParser, create_mqtt_params
15
+ from roborock.containers import DeviceData, HomeData, HomeDataProduct, LoginData
16
+ from roborock.devices.device_manager import create_device_manager, create_home_data_api
17
+ from roborock.protocol import MessageParser
18
18
  from roborock.util import run_sync
19
19
  from roborock.version_1_apis.roborock_local_client_v1 import RoborockLocalClientV1
20
20
  from roborock.version_1_apis.roborock_mqtt_client_v1 import RoborockMqttClientV1
@@ -101,44 +101,25 @@ async def session(ctx, duration: int):
101
101
  context: RoborockContext = ctx.obj
102
102
  login_data = context.login_data()
103
103
 
104
- # Discovery devices if not already available
105
- if not login_data.home_data:
106
- await _discover(ctx)
107
- login_data = context.login_data()
108
- if not login_data.home_data or not login_data.home_data.devices:
109
- raise RoborockException("Unable to discover devices")
110
-
111
- all_devices = login_data.home_data.devices + login_data.home_data.received_devices
112
- click.echo(f"Discovered devices: {', '.join([device.name for device in all_devices])}")
113
-
114
- rriot = login_data.user_data.rriot
115
- params = create_mqtt_params(rriot)
116
-
117
- mqtt_session = await create_mqtt_session(params)
118
- click.echo("Starting MQTT session...")
119
- if not mqtt_session.connected:
120
- raise RoborockException("Failed to connect to MQTT broker")
104
+ home_data_api = create_home_data_api(login_data.email, login_data.user_data)
121
105
 
122
- def on_message(bytes: bytes):
123
- """Callback function to handle incoming MQTT messages."""
124
- # Decode the first 20 bytes of the message for display
125
- bytes = bytes[:20]
106
+ async def home_data_cache() -> HomeData:
107
+ if login_data.home_data is None:
108
+ login_data.home_data = await home_data_api()
109
+ context.update(login_data)
110
+ return login_data.home_data
126
111
 
127
- click.echo(f"Received message: {bytes}...")
112
+ # Create device manager
113
+ device_manager = await create_device_manager(login_data.user_data, home_data_cache)
128
114
 
129
- unsubs = []
130
- for device in all_devices:
131
- device_topic = f"rr/m/o/{rriot.u}/{params.username}/{device.duid}"
132
- unsub = await mqtt_session.subscribe(device_topic, on_message)
133
- unsubs.append(unsub)
115
+ devices = await device_manager.get_devices()
116
+ click.echo(f"Discovered devices: {', '.join([device.name for device in devices])}")
134
117
 
135
118
  click.echo("MQTT session started. Listening for messages...")
136
119
  await asyncio.sleep(duration)
137
120
 
138
- click.echo("Stopping MQTT session...")
139
- for unsub in unsubs:
140
- unsub()
141
- await mqtt_session.close()
121
+ # Close the device manager (this will close all devices and MQTT session)
122
+ await device_manager.close()
142
123
 
143
124
 
144
125
  async def _discover(ctx):
@@ -6,9 +6,13 @@ until the API is stable.
6
6
 
7
7
  import enum
8
8
  import logging
9
+ from collections.abc import Callable
9
10
  from functools import cached_property
10
11
 
11
12
  from roborock.containers import HomeDataDevice, HomeDataProduct, UserData
13
+ from roborock.roborock_message import RoborockMessage
14
+
15
+ from .mqtt_channel import MqttChannel
12
16
 
13
17
  _LOGGER = logging.getLogger(__name__)
14
18
 
@@ -29,11 +33,25 @@ class DeviceVersion(enum.StrEnum):
29
33
  class RoborockDevice:
30
34
  """Unified Roborock device class with automatic connection setup."""
31
35
 
32
- def __init__(self, user_data: UserData, device_info: HomeDataDevice, product_info: HomeDataProduct) -> None:
33
- """Initialize the RoborockDevice with device info, user data, and capabilities."""
36
+ def __init__(
37
+ self,
38
+ user_data: UserData,
39
+ device_info: HomeDataDevice,
40
+ product_info: HomeDataProduct,
41
+ mqtt_channel: MqttChannel,
42
+ ) -> None:
43
+ """Initialize the RoborockDevice.
44
+
45
+ The device takes ownership of the MQTT channel for communication with the device.
46
+ Use `connect()` to establish the connection, which will set up the MQTT channel
47
+ for receiving messages from the device. Use `close()` to unsubscribe from the MQTT
48
+ channel.
49
+ """
34
50
  self._user_data = user_data
35
51
  self._device_info = device_info
36
52
  self._product_info = product_info
53
+ self._mqtt_channel = mqtt_channel
54
+ self._unsub: Callable[[], None] | None = None
37
55
 
38
56
  @property
39
57
  def duid(self) -> str:
@@ -63,3 +81,28 @@ class RoborockDevice:
63
81
  self._device_info.name,
64
82
  )
65
83
  return DeviceVersion.UNKNOWN
84
+
85
+ async def connect(self) -> None:
86
+ """Connect to the device using MQTT.
87
+
88
+ This method will set up the MQTT channel for communication with the device.
89
+ """
90
+ if self._unsub:
91
+ raise ValueError("Already connected to the device")
92
+ self._unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
93
+
94
+ async def close(self) -> None:
95
+ """Close the MQTT connection to the device.
96
+
97
+ This method will unsubscribe from the MQTT channel and clean up resources.
98
+ """
99
+ if self._unsub:
100
+ self._unsub()
101
+ self._unsub = None
102
+
103
+ def _on_mqtt_message(self, message: RoborockMessage) -> None:
104
+ """Handle incoming MQTT messages from the device.
105
+
106
+ This method should be overridden in subclasses to handle specific device messages.
107
+ """
108
+ _LOGGER.debug("Received message from device %s: %s", self.duid, message)
@@ -1,5 +1,6 @@
1
1
  """Module for discovering Roborock devices."""
2
2
 
3
+ import asyncio
3
4
  import logging
4
5
  from collections.abc import Awaitable, Callable
5
6
 
@@ -10,8 +11,13 @@ from roborock.containers import (
10
11
  UserData,
11
12
  )
12
13
  from roborock.devices.device import RoborockDevice
14
+ from roborock.mqtt.roborock_session import create_mqtt_session
15
+ from roborock.mqtt.session import MqttSession
16
+ from roborock.protocol import create_mqtt_params
13
17
  from roborock.web_api import RoborockApiClient
14
18
 
19
+ from .mqtt_channel import MqttChannel
20
+
15
21
  _LOGGER = logging.getLogger(__name__)
16
22
 
17
23
  __all__ = [
@@ -34,11 +40,16 @@ class DeviceManager:
34
40
  self,
35
41
  home_data_api: HomeDataApi,
36
42
  device_creator: DeviceCreator,
43
+ mqtt_session: MqttSession,
37
44
  ) -> None:
38
- """Initialize the DeviceManager with user data and optional cache storage."""
45
+ """Initialize the DeviceManager with user data and optional cache storage.
46
+
47
+ This takes ownership of the MQTT session and will close it when the manager is closed.
48
+ """
39
49
  self._home_data_api = home_data_api
40
50
  self._device_creator = device_creator
41
51
  self._devices: dict[str, RoborockDevice] = {}
52
+ self._mqtt_session = mqtt_session
42
53
 
43
54
  async def discover_devices(self) -> list[RoborockDevice]:
44
55
  """Discover all devices for the logged-in user."""
@@ -46,9 +57,16 @@ class DeviceManager:
46
57
  device_products = home_data.device_products
47
58
  _LOGGER.debug("Discovered %d devices %s", len(device_products), home_data)
48
59
 
49
- self._devices = {
50
- duid: self._device_creator(device, product) for duid, (device, product) in device_products.items()
51
- }
60
+ # These are connected serially to avoid overwhelming the MQTT broker
61
+ new_devices = {}
62
+ for duid, (device, product) in device_products.items():
63
+ if duid in self._devices:
64
+ continue
65
+ new_device = self._device_creator(device, product)
66
+ await new_device.connect()
67
+ new_devices[duid] = new_device
68
+
69
+ self._devices.update(new_devices)
52
70
  return list(self._devices.values())
53
71
 
54
72
  async def get_device(self, duid: str) -> RoborockDevice | None:
@@ -59,6 +77,13 @@ class DeviceManager:
59
77
  """Get all discovered devices."""
60
78
  return list(self._devices.values())
61
79
 
80
+ async def close(self) -> None:
81
+ """Close all MQTT connections and clean up resources."""
82
+ tasks = [device.close() for device in self._devices.values()]
83
+ self._devices.clear()
84
+ tasks.append(self._mqtt_session.close())
85
+ await asyncio.gather(*tasks)
86
+
62
87
 
63
88
  def create_home_data_api(email: str, user_data: UserData) -> HomeDataApi:
64
89
  """Create a home data API wrapper.
@@ -67,7 +92,9 @@ def create_home_data_api(email: str, user_data: UserData) -> HomeDataApi:
67
92
  home data for the user.
68
93
  """
69
94
 
70
- client = RoborockApiClient(email, user_data)
95
+ # Note: This will auto discover the API base URL. This can be improved
96
+ # by caching this next to `UserData` if needed to avoid unnecessary API calls.
97
+ client = RoborockApiClient(email)
71
98
 
72
99
  async def home_data_api() -> HomeData:
73
100
  return await client.get_home_data(user_data)
@@ -83,9 +110,13 @@ async def create_device_manager(user_data: UserData, home_data_api: HomeDataApi)
83
110
  include caching or other optimizations.
84
111
  """
85
112
 
113
+ mqtt_params = create_mqtt_params(user_data.rriot)
114
+ mqtt_session = await create_mqtt_session(mqtt_params)
115
+
86
116
  def device_creator(device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:
87
- return RoborockDevice(user_data, device, product)
117
+ mqtt_channel = MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)
118
+ return RoborockDevice(user_data, device, product, mqtt_channel)
88
119
 
89
- manager = DeviceManager(home_data_api, device_creator)
120
+ manager = DeviceManager(home_data_api, device_creator, mqtt_session=mqtt_session)
90
121
  await manager.discover_devices()
91
122
  return manager
@@ -0,0 +1,146 @@
1
+ """Module for communicating with Roborock devices over a local network."""
2
+
3
+ import asyncio
4
+ import logging
5
+ from collections.abc import Callable
6
+ from dataclasses import dataclass
7
+ from json import JSONDecodeError
8
+
9
+ from roborock.exceptions import RoborockConnectionException, RoborockException
10
+ from roborock.protocol import Decoder, Encoder, create_local_decoder, create_local_encoder
11
+ from roborock.roborock_message import RoborockMessage
12
+
13
+ _LOGGER = logging.getLogger(__name__)
14
+ _PORT = 58867
15
+
16
+
17
+ @dataclass
18
+ class _LocalProtocol(asyncio.Protocol):
19
+ """Callbacks for the Roborock local client transport."""
20
+
21
+ messages_cb: Callable[[bytes], None]
22
+ connection_lost_cb: Callable[[Exception | None], None]
23
+
24
+ def data_received(self, data: bytes) -> None:
25
+ """Called when data is received from the transport."""
26
+ self.messages_cb(data)
27
+
28
+ def connection_lost(self, exc: Exception | None) -> None:
29
+ """Called when the transport connection is lost."""
30
+ self.connection_lost_cb(exc)
31
+
32
+
33
+ class LocalChannel:
34
+ """Simple RPC-style channel for communicating with a device over a local network.
35
+
36
+ Handles request/response correlation and timeouts, but leaves message
37
+ format most parsing to higher-level components.
38
+ """
39
+
40
+ def __init__(self, host: str, local_key: str):
41
+ self._host = host
42
+ self._transport: asyncio.Transport | None = None
43
+ self._protocol: _LocalProtocol | None = None
44
+ self._subscribers: list[Callable[[RoborockMessage], None]] = []
45
+ self._is_connected = False
46
+
47
+ # RPC support
48
+ self._waiting_queue: dict[int, asyncio.Future[RoborockMessage]] = {}
49
+ self._decoder: Decoder = create_local_decoder(local_key)
50
+ self._encoder: Encoder = create_local_encoder(local_key)
51
+ self._queue_lock = asyncio.Lock()
52
+
53
+ async def connect(self) -> None:
54
+ """Connect to the device."""
55
+ if self._is_connected:
56
+ _LOGGER.warning("Already connected")
57
+ return
58
+ _LOGGER.debug("Connecting to %s:%s", self._host, _PORT)
59
+ loop = asyncio.get_running_loop()
60
+ protocol = _LocalProtocol(self._data_received, self._connection_lost)
61
+ try:
62
+ self._transport, self._protocol = await loop.create_connection(lambda: protocol, self._host, _PORT)
63
+ self._is_connected = True
64
+ except OSError as e:
65
+ raise RoborockConnectionException(f"Failed to connect to {self._host}:{_PORT}") from e
66
+
67
+ async def close(self) -> None:
68
+ """Disconnect from the device."""
69
+ if self._transport:
70
+ self._transport.close()
71
+ else:
72
+ _LOGGER.warning("Close called but transport is already None")
73
+ self._transport = None
74
+ self._is_connected = False
75
+
76
+ def _data_received(self, data: bytes) -> None:
77
+ """Handle incoming data from the transport."""
78
+ if not (messages := self._decoder(data)):
79
+ _LOGGER.warning("Failed to decode local message: %s", data)
80
+ return
81
+ for message in messages:
82
+ _LOGGER.debug("Received message: %s", message)
83
+ asyncio.create_task(self._resolve_future_with_lock(message))
84
+ for callback in self._subscribers:
85
+ try:
86
+ callback(message)
87
+ except Exception as e:
88
+ _LOGGER.exception("Uncaught error in message handler callback: %s", e)
89
+
90
+ def _connection_lost(self, exc: Exception | None) -> None:
91
+ """Handle connection loss."""
92
+ _LOGGER.warning("Connection lost to %s", self._host, exc_info=exc)
93
+ self._transport = None
94
+ self._is_connected = False
95
+
96
+ async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
97
+ """Subscribe to all messages from the device."""
98
+ self._subscribers.append(callback)
99
+
100
+ def unsubscribe() -> None:
101
+ self._subscribers.remove(callback)
102
+
103
+ return unsubscribe
104
+
105
+ async def _resolve_future_with_lock(self, message: RoborockMessage) -> None:
106
+ """Resolve waiting future with proper locking."""
107
+ if (request_id := message.get_request_id()) is None:
108
+ _LOGGER.debug("Received message with no request_id")
109
+ return
110
+ async with self._queue_lock:
111
+ if (future := self._waiting_queue.pop(request_id, None)) is not None:
112
+ future.set_result(message)
113
+ else:
114
+ _LOGGER.debug("Received message with no waiting handler: request_id=%s", request_id)
115
+
116
+ async def send_command(self, message: RoborockMessage, timeout: float = 10.0) -> RoborockMessage:
117
+ """Send a command message and wait for the response message."""
118
+ if not self._transport or not self._is_connected:
119
+ raise RoborockConnectionException("Not connected to device")
120
+
121
+ try:
122
+ if (request_id := message.get_request_id()) is None:
123
+ raise RoborockException("Message must have a request_id for RPC calls")
124
+ except (ValueError, JSONDecodeError) as err:
125
+ _LOGGER.exception("Error getting request_id from message: %s", err)
126
+ raise RoborockException(f"Invalid message format, Message must have a request_id: {err}") from err
127
+
128
+ future: asyncio.Future[RoborockMessage] = asyncio.Future()
129
+ async with self._queue_lock:
130
+ if request_id in self._waiting_queue:
131
+ raise RoborockException(f"Request ID {request_id} already pending, cannot send command")
132
+ self._waiting_queue[request_id] = future
133
+
134
+ try:
135
+ encoded_msg = self._encoder(message)
136
+ self._transport.write(encoded_msg)
137
+ return await asyncio.wait_for(future, timeout=timeout)
138
+ except asyncio.TimeoutError as ex:
139
+ async with self._queue_lock:
140
+ self._waiting_queue.pop(request_id, None)
141
+ raise RoborockException(f"Command timed out after {timeout}s") from ex
142
+ except Exception:
143
+ logging.exception("Uncaught error sending command")
144
+ async with self._queue_lock:
145
+ self._waiting_queue.pop(request_id, None)
146
+ raise
@@ -0,0 +1,115 @@
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,0 +1,91 @@
1
+ """Roborock V1 Protocol Encoder."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import json
6
+ import math
7
+ import time
8
+ from collections.abc import Callable
9
+ from dataclasses import dataclass, field
10
+ from typing import Any
11
+
12
+ from roborock.roborock_message import MessageRetry, RoborockMessage, RoborockMessageProtocol
13
+ from roborock.roborock_typing import RoborockCommand
14
+ from roborock.util import get_next_int
15
+
16
+ CommandType = RoborockCommand | str
17
+ ParamsType = list | dict | int | None
18
+
19
+
20
+ @dataclass(frozen=True, kw_only=True)
21
+ class SecurityData:
22
+ """Security data included in the request for some V1 commands."""
23
+
24
+ endpoint: str
25
+ nonce: bytes
26
+
27
+ def to_dict(self) -> dict[str, Any]:
28
+ """Convert security data to a dictionary for sending in the payload."""
29
+ return {"security": {"endpoint": self.endpoint, "nonce": self.nonce.hex().lower()}}
30
+
31
+
32
+ @dataclass
33
+ class RequestMessage:
34
+ """Data structure for v1 RoborockMessage payloads."""
35
+
36
+ method: RoborockCommand | str
37
+ params: ParamsType
38
+ timestamp: int = field(default_factory=lambda: math.floor(time.time()))
39
+ request_id: int = field(default_factory=lambda: get_next_int(10000, 32767))
40
+
41
+ def as_payload(self, security_data: SecurityData | None) -> bytes:
42
+ """Convert the request arguments to a dictionary."""
43
+ inner = {
44
+ "id": self.request_id,
45
+ "method": self.method,
46
+ "params": self.params or [],
47
+ **(security_data.to_dict() if security_data else {}),
48
+ }
49
+ return bytes(
50
+ json.dumps(
51
+ {
52
+ "dps": {"101": json.dumps(inner, separators=(",", ":"))},
53
+ "t": self.timestamp,
54
+ },
55
+ separators=(",", ":"),
56
+ ).encode()
57
+ )
58
+
59
+
60
+ def create_mqtt_payload_encoder(security_data: SecurityData) -> Callable[[CommandType, ParamsType], RoborockMessage]:
61
+ """Create a payload encoder for V1 commands over MQTT."""
62
+
63
+ def _get_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
64
+ """Build the payload for a V1 command."""
65
+ request = RequestMessage(method=method, params=params)
66
+ payload = request.as_payload(security_data) # always secure
67
+ return RoborockMessage(
68
+ timestamp=request.timestamp,
69
+ protocol=RoborockMessageProtocol.RPC_REQUEST,
70
+ payload=payload,
71
+ )
72
+
73
+ return _get_payload
74
+
75
+
76
+ def encode_local_payload(method: CommandType, params: ParamsType) -> RoborockMessage:
77
+ """Encode payload for V1 commands over local connection."""
78
+
79
+ request = RequestMessage(method=method, params=params)
80
+ payload = request.as_payload(security_data=None)
81
+
82
+ message_retry: MessageRetry | None = None
83
+ if method == RoborockCommand.RETRY_REQUEST and isinstance(params, dict):
84
+ message_retry = MessageRetry(method=method, retry_id=params["retry_id"])
85
+
86
+ return RoborockMessage(
87
+ timestamp=request.timestamp,
88
+ protocol=RoborockMessageProtocol.GENERAL_REQUEST,
89
+ payload=payload,
90
+ message_retry=message_retry,
91
+ )
@@ -1,7 +1,6 @@
1
1
  import asyncio
2
2
  import dataclasses
3
3
  import json
4
- import math
5
4
  import struct
6
5
  import time
7
6
  from abc import ABC, abstractmethod
@@ -54,15 +53,15 @@ from roborock.roborock_message import (
54
53
  RoborockMessage,
55
54
  RoborockMessageProtocol,
56
55
  )
57
- from roborock.util import RepeatableTask, get_next_int, unpack_list
56
+ from roborock.util import RepeatableTask, unpack_list
57
+
58
+ CUSTOM_COMMANDS = {RoborockCommand.GET_MAP_CALIBRATION}
58
59
 
59
60
  COMMANDS_SECURED = {
60
61
  RoborockCommand.GET_MAP_V1,
61
62
  RoborockCommand.GET_MULTI_MAP,
62
63
  }
63
64
 
64
- CUSTOM_COMMANDS = {RoborockCommand.GET_MAP_CALIBRATION}
65
-
66
65
  CLOUD_REQUIRED = COMMANDS_SECURED.union(CUSTOM_COMMANDS)
67
66
 
68
67
  WASH_N_FILL_DOCK = [
@@ -340,35 +339,6 @@ class RoborockClientV1(RoborockClient, ABC):
340
339
  """Load the map into the vacuum's memory."""
341
340
  await self.send_command(RoborockCommand.LOAD_MULTI_MAP, [map_flag])
342
341
 
343
- def _get_payload(
344
- self,
345
- method: RoborockCommand | str,
346
- params: list | dict | int | None = None,
347
- secured=False,
348
- ):
349
- timestamp = math.floor(time.time())
350
- request_id = get_next_int(10000, 32767)
351
- inner = {
352
- "id": request_id,
353
- "method": method,
354
- "params": params or [],
355
- }
356
- if secured:
357
- inner["security"] = {
358
- "endpoint": self._endpoint,
359
- "nonce": self._nonce.hex().lower(),
360
- }
361
- payload = bytes(
362
- json.dumps(
363
- {
364
- "dps": {"101": json.dumps(inner, separators=(",", ":"))},
365
- "t": timestamp,
366
- },
367
- separators=(",", ":"),
368
- ).encode()
369
- )
370
- return request_id, timestamp, payload
371
-
372
342
  @abstractmethod
373
343
  async def _send_command(
374
344
  self,
@@ -4,9 +4,10 @@ from roborock.local_api import RoborockLocalClient
4
4
 
5
5
  from .. import CommandVacuumError, DeviceData, RoborockCommand, RoborockException
6
6
  from ..exceptions import VacuumError
7
- from ..roborock_message import MessageRetry, RoborockMessage, RoborockMessageProtocol
7
+ from ..protocols.v1_protocol import encode_local_payload
8
+ from ..roborock_message import RoborockMessage, RoborockMessageProtocol
8
9
  from ..util import RoborockLoggerAdapter
9
- from .roborock_client_v1 import COMMANDS_SECURED, RoborockClientV1
10
+ from .roborock_client_v1 import CLOUD_REQUIRED, RoborockClientV1
10
11
 
11
12
  _LOGGER = logging.getLogger(__name__)
12
13
 
@@ -21,26 +22,16 @@ class RoborockLocalClientV1(RoborockLocalClient, RoborockClientV1):
21
22
  self.queue_timeout = queue_timeout
22
23
  self._logger = RoborockLoggerAdapter(device_data.device.name, _LOGGER)
23
24
 
24
- def build_roborock_message(
25
- self, method: RoborockCommand | str, params: list | dict | int | None = None
26
- ) -> RoborockMessage:
27
- secured = True if method in COMMANDS_SECURED else False
28
- request_id, timestamp, payload = self._get_payload(method, params, secured)
29
- self._logger.debug("Building message id %s for method %s", request_id, method)
30
- request_protocol = RoborockMessageProtocol.GENERAL_REQUEST
31
- message_retry: MessageRetry | None = None
32
- if method == RoborockCommand.RETRY_REQUEST and isinstance(params, dict):
33
- message_retry = MessageRetry(method=params["method"], retry_id=params["retry_id"])
34
- return RoborockMessage(
35
- timestamp=timestamp, protocol=request_protocol, payload=payload, message_retry=message_retry
36
- )
37
-
38
25
  async def _send_command(
39
26
  self,
40
27
  method: RoborockCommand | str,
41
28
  params: list | dict | int | None = None,
42
29
  ):
43
- roborock_message = self.build_roborock_message(method, params)
30
+ if method in CLOUD_REQUIRED:
31
+ raise RoborockException(f"Method {method} is not supported over local connection")
32
+
33
+ roborock_message = encode_local_payload(method, params)
34
+ self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id(), method)
44
35
  return await self.send_message(roborock_message)
45
36
 
46
37
  async def send_message(self, roborock_message: RoborockMessage):
@@ -11,6 +11,7 @@ from roborock.cloud_api import RoborockMqttClient
11
11
  from ..containers import DeviceData, UserData
12
12
  from ..exceptions import CommandVacuumError, RoborockException, VacuumError
13
13
  from ..protocol import Utils
14
+ from ..protocols.v1_protocol import SecurityData, create_mqtt_payload_encoder
14
15
  from ..roborock_message import (
15
16
  RoborockMessage,
16
17
  RoborockMessageProtocol,
@@ -36,6 +37,9 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
36
37
  RoborockClientV1.__init__(self, device_info, endpoint)
37
38
  self.queue_timeout = queue_timeout
38
39
  self._logger = RoborockLoggerAdapter(device_info.device.name, _LOGGER)
40
+ self._payload_encoder = create_mqtt_payload_encoder(
41
+ SecurityData(endpoint=self._endpoint, nonce=self._nonce),
42
+ )
39
43
 
40
44
  async def send_message(self, roborock_message: RoborockMessage):
41
45
  await self.validate_connection()
@@ -78,10 +82,9 @@ class RoborockMqttClientV1(RoborockMqttClient, RoborockClientV1):
78
82
  if method in CUSTOM_COMMANDS:
79
83
  # When we have more custom commands do something more complicated here
80
84
  return await self._get_calibration_points()
81
- request_id, timestamp, payload = self._get_payload(method, params, True)
82
- self._logger.debug("Building message id %s for method %s", request_id, method)
83
- request_protocol = RoborockMessageProtocol.RPC_REQUEST
84
- roborock_message = RoborockMessage(timestamp=timestamp, protocol=request_protocol, payload=payload)
85
+
86
+ roborock_message = self._payload_encoder(method, params)
87
+ self._logger.debug("Building message id %s for method %s", roborock_message.get_request_id, method)
85
88
  return await self.send_message(roborock_message)
86
89
 
87
90
  async def _get_calibration_points(self):