python-roborock 3.8.2__tar.gz → 3.8.4__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 (97) hide show
  1. {python_roborock-3.8.2 → python_roborock-3.8.4}/PKG-INFO +1 -1
  2. {python_roborock-3.8.2 → python_roborock-3.8.4}/pyproject.toml +1 -1
  3. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/cache.py +7 -1
  4. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/mqtt_channel.py +4 -0
  5. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/home.py +28 -6
  6. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/v1_channel.py +11 -7
  7. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/v1_rpc_channel.py +11 -1
  8. python_roborock-3.8.4/roborock/mqtt/health_manager.py +51 -0
  9. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/mqtt/roborock_session.py +77 -53
  10. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/mqtt/session.py +4 -0
  11. {python_roborock-3.8.2 → python_roborock-3.8.4}/.gitignore +0 -0
  12. {python_roborock-3.8.2 → python_roborock-3.8.4}/LICENSE +0 -0
  13. {python_roborock-3.8.2 → python_roborock-3.8.4}/README.md +0 -0
  14. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/__init__.py +0 -0
  15. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/api.py +0 -0
  16. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/broadcast_protocol.py +0 -0
  17. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/callbacks.py +0 -0
  18. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/cli.py +0 -0
  19. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/cloud_api.py +0 -0
  20. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/command_cache.py +0 -0
  21. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/const.py +0 -0
  22. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/__init__.py +0 -0
  23. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q10/__init__.py +0 -0
  24. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q10/b01_q10_code_mappings.py +0 -0
  25. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q10/b01_q10_containers.py +0 -0
  26. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q7/__init__.py +0 -0
  27. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q7/b01_q7_code_mappings.py +0 -0
  28. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/b01_q7/b01_q7_containers.py +0 -0
  29. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/code_mappings.py +0 -0
  30. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/containers.py +0 -0
  31. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/dyad/__init__.py +0 -0
  32. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/dyad/dyad_code_mappings.py +0 -0
  33. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/dyad/dyad_containers.py +0 -0
  34. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/v1/__init__.py +0 -0
  35. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/v1/v1_clean_modes.py +0 -0
  36. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/v1/v1_code_mappings.py +0 -0
  37. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/v1/v1_containers.py +0 -0
  38. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/zeo/__init__.py +0 -0
  39. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/zeo/zeo_code_mappings.py +0 -0
  40. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/data/zeo/zeo_containers.py +0 -0
  41. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/device_features.py +0 -0
  42. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/README.md +0 -0
  43. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/__init__.py +0 -0
  44. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/a01_channel.py +0 -0
  45. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/b01_channel.py +0 -0
  46. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/channel.py +0 -0
  47. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/device.py +0 -0
  48. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/device_manager.py +0 -0
  49. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/file_cache.py +0 -0
  50. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/local_channel.py +0 -0
  51. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/__init__.py +0 -0
  52. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/a01/__init__.py +0 -0
  53. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/b01/__init__.py +0 -0
  54. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/traits_mixin.py +0 -0
  55. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/__init__.py +0 -0
  56. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/child_lock.py +0 -0
  57. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/clean_summary.py +0 -0
  58. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/command.py +0 -0
  59. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/common.py +0 -0
  60. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/consumeable.py +0 -0
  61. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/device_features.py +0 -0
  62. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/do_not_disturb.py +0 -0
  63. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/dust_collection_mode.py +0 -0
  64. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/flow_led_status.py +0 -0
  65. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/led_status.py +0 -0
  66. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/map_content.py +0 -0
  67. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/maps.py +0 -0
  68. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/network_info.py +0 -0
  69. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/rooms.py +0 -0
  70. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/routines.py +0 -0
  71. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/smart_wash_params.py +0 -0
  72. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/status.py +0 -0
  73. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/valley_electricity_timer.py +0 -0
  74. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/volume.py +0 -0
  75. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/devices/traits/v1/wash_towel_mode.py +0 -0
  76. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/exceptions.py +0 -0
  77. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/map/__init__.py +0 -0
  78. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/map/map_parser.py +0 -0
  79. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/mqtt/__init__.py +0 -0
  80. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/protocol.py +0 -0
  81. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/protocols/__init__.py +0 -0
  82. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/protocols/a01_protocol.py +0 -0
  83. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/protocols/b01_protocol.py +0 -0
  84. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/protocols/v1_protocol.py +0 -0
  85. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/py.typed +0 -0
  86. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/roborock_future.py +0 -0
  87. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/roborock_message.py +0 -0
  88. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/roborock_typing.py +0 -0
  89. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/util.py +0 -0
  90. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_1_apis/__init__.py +0 -0
  91. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_1_apis/roborock_client_v1.py +0 -0
  92. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_1_apis/roborock_local_client_v1.py +0 -0
  93. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_1_apis/roborock_mqtt_client_v1.py +0 -0
  94. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_a01_apis/__init__.py +0 -0
  95. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_a01_apis/roborock_client_a01.py +0 -0
  96. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/version_a01_apis/roborock_mqtt_client_a01.py +0 -0
  97. {python_roborock-3.8.2 → python_roborock-3.8.4}/roborock/web_api.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: python-roborock
3
- Version: 3.8.2
3
+ Version: 3.8.4
4
4
  Summary: A package to control Roborock vacuums.
5
5
  Project-URL: Repository, https://github.com/humbertogontijo/python-roborock
6
6
  Project-URL: Documentation, https://python-roborock.readthedocs.io/
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "python-roborock"
3
- version = "3.8.2"
3
+ version = "3.8.4"
4
4
  description = "A package to control Roborock vacuums."
5
5
  authors = [{ name = "humbertogontijo", email = "humbertogontijo@users.noreply.github.com" }, {name="Lash-L"}, {name="allenporter"}]
6
6
  requires-python = ">=3.11, <4"
@@ -26,7 +26,13 @@ class CacheData:
26
26
  """Home map information indexed by map_flag."""
27
27
 
28
28
  home_map_content: dict[int, bytes] = field(default_factory=dict)
29
- """Home cache content for each map data indexed by map_flag."""
29
+ """Home cache content for each map data indexed by map_flag.
30
+
31
+ This is deprecated in favor of `home_map_content_base64`.
32
+ """
33
+
34
+ home_map_content_base64: dict[int, str] = field(default_factory=dict)
35
+ """Home cache content for each map data (encoded base64) indexed by map_flag."""
30
36
 
31
37
  device_features: DeviceFeatures | None = None
32
38
  """Device features information."""
@@ -82,6 +82,10 @@ class MqttChannel(Channel):
82
82
  _LOGGER.exception("Error publishing MQTT message: %s", e)
83
83
  raise RoborockException(f"Failed to publish MQTT message: {e}") from e
84
84
 
85
+ async def restart(self) -> None:
86
+ """Restart the underlying MQTT session."""
87
+ await self._mqtt_session.restart()
88
+
85
89
 
86
90
  def create_mqtt_channel(
87
91
  user_data: UserData, mqtt_params: MqttParams, mqtt_session: MqttSession, device: HomeDataDevice
@@ -16,6 +16,7 @@ the current map's information and room names as needed.
16
16
  """
17
17
 
18
18
  import asyncio
19
+ import base64
19
20
  import logging
20
21
  from typing import Self
21
22
 
@@ -86,14 +87,20 @@ class HomeTrait(RoborockBase, common.V1TraitMixin):
86
87
  After discovery, the home cache will be populated and can be accessed via the `home_map_info` property.
87
88
  """
88
89
  cache_data = await self._cache.get()
89
- if cache_data.home_map_info and cache_data.home_map_content:
90
+ if cache_data.home_map_info and (cache_data.home_map_content or cache_data.home_map_content_base64):
90
91
  _LOGGER.debug("Home cache already populated, skipping discovery")
91
92
  self._home_map_info = cache_data.home_map_info
92
93
  self._discovery_completed = True
93
94
  try:
94
- self._home_map_content = {
95
- k: self._map_content.parse_map_content(v) for k, v in cache_data.home_map_content.items()
96
- }
95
+ if cache_data.home_map_content_base64:
96
+ self._home_map_content = {
97
+ k: self._map_content.parse_map_content(base64.b64decode(v))
98
+ for k, v in cache_data.home_map_content_base64.items()
99
+ }
100
+ else:
101
+ self._home_map_content = {
102
+ k: self._map_content.parse_map_content(v) for k, v in cache_data.home_map_content.items()
103
+ }
97
104
  except (ValueError, RoborockException) as ex:
98
105
  _LOGGER.warning("Failed to parse cached home map content, will re-discover: %s", ex)
99
106
  self._home_map_content = {}
@@ -218,7 +225,12 @@ class HomeTrait(RoborockBase, common.V1TraitMixin):
218
225
  """Update the entire home cache with new map info and content."""
219
226
  cache_data = await self._cache.get()
220
227
  cache_data.home_map_info = home_map_info
221
- cache_data.home_map_content = {k: v.raw_api_response for k, v in home_map_content.items() if v.raw_api_response}
228
+ cache_data.home_map_content_base64 = {
229
+ k: base64.b64encode(v.raw_api_response).decode("utf-8")
230
+ for k, v in home_map_content.items()
231
+ if v.raw_api_response
232
+ }
233
+ cache_data.home_map_content = {}
222
234
  await self._cache.set(cache_data)
223
235
  self._home_map_info = home_map_info
224
236
  self._home_map_content = home_map_content
@@ -237,8 +249,18 @@ class HomeTrait(RoborockBase, common.V1TraitMixin):
237
249
  if update_cache:
238
250
  cache_data = await self._cache.get()
239
251
  cache_data.home_map_info[map_flag] = map_info
252
+ # Migrate existing cached content to base64 if needed
253
+ if cache_data.home_map_content and not cache_data.home_map_content_base64:
254
+ cache_data.home_map_content_base64 = {
255
+ k: base64.b64encode(v).decode("utf-8") for k, v in cache_data.home_map_content.items()
256
+ }
257
+ cache_data.home_map_content = {}
240
258
  if map_content.raw_api_response:
241
- cache_data.home_map_content[map_flag] = map_content.raw_api_response
259
+ if cache_data.home_map_content_base64 is None:
260
+ cache_data.home_map_content_base64 = {}
261
+ cache_data.home_map_content_base64[map_flag] = base64.b64encode(map_content.raw_api_response).decode(
262
+ "utf-8"
263
+ )
242
264
  await self._cache.set(cache_data)
243
265
 
244
266
  if self._home_map_info is None:
@@ -142,7 +142,7 @@ class V1Channel(Channel):
142
142
  # Make an initial, optimistic attempt to connect to local with the
143
143
  # cache. The cache information will be refreshed by the background task.
144
144
  try:
145
- await self._local_connect(use_cache=True)
145
+ await self._local_connect(prefer_cache=True)
146
146
  except RoborockException as err:
147
147
  _LOGGER.warning("Could not establish local connection for device %s: %s", self._device_uid, err)
148
148
 
@@ -175,13 +175,13 @@ class V1Channel(Channel):
175
175
  self._callback = callback
176
176
  return unsub
177
177
 
178
- async def _get_networking_info(self, *, use_cache: bool = True) -> NetworkInfo:
178
+ async def _get_networking_info(self, *, prefer_cache: bool = True) -> NetworkInfo:
179
179
  """Retrieve networking information for the device.
180
180
 
181
181
  This is a cloud only command used to get the local device's IP address.
182
182
  """
183
183
  cache_data = await self._cache.get()
184
- if use_cache and cache_data.network_info and (network_info := cache_data.network_info.get(self._device_uid)):
184
+ if prefer_cache and cache_data.network_info and (network_info := cache_data.network_info.get(self._device_uid)):
185
185
  _LOGGER.debug("Using cached network info for device %s", self._device_uid)
186
186
  return network_info
187
187
  try:
@@ -189,6 +189,10 @@ class V1Channel(Channel):
189
189
  RoborockCommand.GET_NETWORK_INFO, response_type=NetworkInfo
190
190
  )
191
191
  except RoborockException as e:
192
+ _LOGGER.debug("Error fetching network info for device %s", self._device_uid)
193
+ if cache_data.network_info and (network_info := cache_data.network_info.get(self._device_uid)):
194
+ _LOGGER.debug("Falling back to cached network info for device %s after error", self._device_uid)
195
+ return network_info
192
196
  raise RoborockException(f"Network info failed for device {self._device_uid}") from e
193
197
  _LOGGER.debug("Network info for device %s: %s", self._device_uid, network_info)
194
198
  self._last_network_info_refresh = datetime.datetime.now(datetime.UTC)
@@ -196,12 +200,12 @@ class V1Channel(Channel):
196
200
  await self._cache.set(cache_data)
197
201
  return network_info
198
202
 
199
- async def _local_connect(self, *, use_cache: bool = True) -> None:
203
+ async def _local_connect(self, *, prefer_cache: bool = True) -> None:
200
204
  """Set up local connection if possible."""
201
205
  _LOGGER.debug(
202
- "Attempting to connect to local channel for device %s (use_cache=%s)", self._device_uid, use_cache
206
+ "Attempting to connect to local channel for device %s (prefer_cache=%s)", self._device_uid, prefer_cache
203
207
  )
204
- networking_info = await self._get_networking_info(use_cache=use_cache)
208
+ networking_info = await self._get_networking_info(prefer_cache=prefer_cache)
205
209
  host = networking_info.ip
206
210
  _LOGGER.debug("Connecting to local channel at %s", host)
207
211
  # Create a new local channel and connect
@@ -236,7 +240,7 @@ class V1Channel(Channel):
236
240
  reconnect_backoff = min(reconnect_backoff * RECONNECT_MULTIPLIER, MAX_RECONNECT_INTERVAL)
237
241
 
238
242
  use_cache = self._should_use_cache(local_connect_failures)
239
- await self._local_connect(use_cache=use_cache)
243
+ await self._local_connect(prefer_cache=use_cache)
240
244
  # Reset backoff and failures on success
241
245
  reconnect_backoff = MIN_RECONNECT_INTERVAL
242
246
  local_connect_failures = 0
@@ -13,6 +13,7 @@ from typing import Any, Protocol, TypeVar, overload
13
13
 
14
14
  from roborock.data import RoborockBase
15
15
  from roborock.exceptions import RoborockException
16
+ from roborock.mqtt.health_manager import HealthManager
16
17
  from roborock.protocols.v1_protocol import (
17
18
  CommandType,
18
19
  MapResponse,
@@ -125,12 +126,14 @@ class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
125
126
  channel: MqttChannel | LocalChannel,
126
127
  payload_encoder: Callable[[RequestMessage], RoborockMessage],
127
128
  decoder: Callable[[RoborockMessage], ResponseMessage] | Callable[[RoborockMessage], MapResponse | None],
129
+ health_manager: HealthManager | None = None,
128
130
  ) -> None:
129
131
  """Initialize the channel with a raw channel and an encoder function."""
130
132
  self._name = name
131
133
  self._channel = channel
132
134
  self._payload_encoder = payload_encoder
133
135
  self._decoder = decoder
136
+ self._health_manager = health_manager
134
137
 
135
138
  async def _send_raw_command(
136
139
  self,
@@ -165,13 +168,19 @@ class PayloadEncodedV1RpcChannel(BaseV1RpcChannel):
165
168
  unsub = await self._channel.subscribe(find_response)
166
169
  try:
167
170
  await self._channel.publish(message)
168
- return await asyncio.wait_for(future, timeout=_TIMEOUT)
171
+ result = await asyncio.wait_for(future, timeout=_TIMEOUT)
169
172
  except TimeoutError as ex:
173
+ if self._health_manager:
174
+ await self._health_manager.on_timeout()
170
175
  future.cancel()
171
176
  raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
172
177
  finally:
173
178
  unsub()
174
179
 
180
+ if self._health_manager:
181
+ await self._health_manager.on_success()
182
+ return result
183
+
175
184
 
176
185
  def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityData) -> V1RpcChannel:
177
186
  """Create a V1 RPC channel using an MQTT channel."""
@@ -180,6 +189,7 @@ def create_mqtt_rpc_channel(mqtt_channel: MqttChannel, security_data: SecurityDa
180
189
  mqtt_channel,
181
190
  lambda x: x.encode_message(RoborockMessageProtocol.RPC_REQUEST, security_data=security_data),
182
191
  decode_rpc_response,
192
+ health_manager=HealthManager(mqtt_channel.restart),
183
193
  )
184
194
 
185
195
 
@@ -0,0 +1,51 @@
1
+ """A health manager for monitoring MQTT connections to Roborock devices.
2
+
3
+ We observe a problem where sometimes the MQTT connection appears to be alive but
4
+ no messages are being received. To mitigate this, we track consecutive timeouts
5
+ and restart the connection if too many timeouts occur in succession.
6
+ """
7
+
8
+ import datetime
9
+ from collections.abc import Awaitable, Callable
10
+
11
+ # Number of consecutive timeouts before considering the connection unhealthy.
12
+ TIMEOUT_THRESHOLD = 3
13
+
14
+ # We won't restart the session more often than this interval.
15
+ RESTART_COOLDOWN = datetime.timedelta(minutes=30)
16
+
17
+
18
+ class HealthManager:
19
+ """Manager for monitoring the health of MQTT connections.
20
+
21
+ This tracks communication timeouts and can trigger restarts of the MQTT
22
+ session if too many timeouts occur in succession.
23
+ """
24
+
25
+ def __init__(self, restart: Callable[[], Awaitable[None]]) -> None:
26
+ """Initialize the health manager.
27
+
28
+ Args:
29
+ restart: A callable to restart the MQTT session.
30
+ """
31
+ self._consecutive_timeouts = 0
32
+ self._restart = restart
33
+ self._last_restart: datetime.datetime | None = None
34
+
35
+ async def on_success(self) -> None:
36
+ """Record a successful communication event."""
37
+ self._consecutive_timeouts = 0
38
+
39
+ async def on_timeout(self) -> None:
40
+ """Record a timeout event.
41
+
42
+ This may trigger a restart of the MQTT session if too many timeouts
43
+ have occurred in succession.
44
+ """
45
+ self._consecutive_timeouts += 1
46
+ if self._consecutive_timeouts >= TIMEOUT_THRESHOLD:
47
+ now = datetime.datetime.now(datetime.UTC)
48
+ if self._last_restart is None or now - self._last_restart >= RESTART_COOLDOWN:
49
+ await self._restart()
50
+ self._last_restart = now
51
+ self._consecutive_timeouts = 0
@@ -49,13 +49,14 @@ class RoborockMqttSession(MqttSession):
49
49
 
50
50
  def __init__(self, params: MqttParams):
51
51
  self._params = params
52
- self._background_task: asyncio.Task[None] | None = None
52
+ self._reconnect_task: asyncio.Task[None] | None = None
53
53
  self._healthy = False
54
54
  self._stop = False
55
55
  self._backoff = MIN_BACKOFF_INTERVAL
56
56
  self._client: aiomqtt.Client | None = None
57
57
  self._client_lock = asyncio.Lock()
58
58
  self._listeners: CallbackMap[str, bytes] = CallbackMap(_LOGGER)
59
+ self._connection_task: asyncio.Task[None] | None = None
59
60
 
60
61
  @property
61
62
  def connected(self) -> bool:
@@ -72,7 +73,7 @@ class RoborockMqttSession(MqttSession):
72
73
  """
73
74
  start_future: asyncio.Future[None] = asyncio.Future()
74
75
  loop = asyncio.get_event_loop()
75
- self._background_task = loop.create_task(self._run_task(start_future))
76
+ self._reconnect_task = loop.create_task(self._run_reconnect_loop(start_future))
76
77
  try:
77
78
  await start_future
78
79
  except MqttError as err:
@@ -85,61 +86,47 @@ class RoborockMqttSession(MqttSession):
85
86
  async def close(self) -> None:
86
87
  """Cancels the MQTT loop and shutdown the client library."""
87
88
  self._stop = True
88
- if self._background_task:
89
- self._background_task.cancel()
90
- try:
91
- await self._background_task
92
- except asyncio.CancelledError:
93
- pass
94
- async with self._client_lock:
95
- if self._client:
96
- await self._client.close()
89
+ tasks = [task for task in [self._connection_task, self._reconnect_task] if task]
90
+ for task in tasks:
91
+ task.cancel()
92
+ try:
93
+ await asyncio.gather(*tasks)
94
+ except asyncio.CancelledError:
95
+ pass
97
96
 
98
97
  self._healthy = False
99
98
 
100
- async def _run_task(self, start_future: asyncio.Future[None] | None) -> None:
99
+ async def restart(self) -> None:
100
+ """Force the session to disconnect and reconnect.
101
+
102
+ The active connection task will be cancelled and restarted in the background, retried by
103
+ the reconnect loop. This is a no-op if there is no active connection.
104
+ """
105
+ _LOGGER.info("Forcing MQTT session restart")
106
+ if self._connection_task:
107
+ self._connection_task.cancel()
108
+ else:
109
+ _LOGGER.debug("No message loop task to cancel")
110
+
111
+ async def _run_reconnect_loop(self, start_future: asyncio.Future[None] | None) -> None:
101
112
  """Run the MQTT loop."""
102
113
  _LOGGER.info("Starting MQTT session")
103
114
  while True:
104
115
  try:
105
- async with self._mqtt_client(self._params) as client:
106
- # Reset backoff once we've successfully connected
107
- self._backoff = MIN_BACKOFF_INTERVAL
108
- self._healthy = True
109
- _LOGGER.info("MQTT Session connected.")
110
- if start_future:
111
- start_future.set_result(None)
112
- start_future = None
113
-
114
- await self._process_message_loop(client)
115
-
116
- except MqttError as err:
117
- if start_future:
118
- _LOGGER.info("MQTT error starting session: %s", err)
119
- start_future.set_exception(err)
120
- return
121
- _LOGGER.info("MQTT error: %s", err)
122
- except asyncio.CancelledError as err:
123
- if start_future:
124
- _LOGGER.debug("MQTT loop was cancelled while starting")
125
- start_future.set_exception(err)
126
- _LOGGER.debug("MQTT loop was cancelled")
127
- return
128
- # Catch exceptions to avoid crashing the loop
129
- # and to allow the loop to retry.
130
- except Exception as err:
131
- # This error is thrown when the MQTT loop is cancelled
132
- # and the generator is not stopped.
133
- if "generator didn't stop" in str(err) or "generator didn't yield" in str(err):
134
- _LOGGER.debug("MQTT loop was cancelled")
135
- return
136
- if start_future:
137
- _LOGGER.error("Uncaught error starting MQTT session: %s", err)
138
- start_future.set_exception(err)
116
+ self._connection_task = asyncio.create_task(self._run_connection(start_future))
117
+ await self._connection_task
118
+ except asyncio.CancelledError:
119
+ _LOGGER.debug("MQTT connection task cancelled")
120
+ except Exception:
121
+ # Exceptions are logged and handled in _run_connection.
122
+ # There is a special case for exceptions on startup where we return
123
+ # immediately. Otherwise, we let the reconnect loop retry with
124
+ # backoff when the reconnect loop is active.
125
+ if start_future and start_future.done() and start_future.exception():
139
126
  return
140
- _LOGGER.exception("Uncaught error during MQTT session: %s", err)
141
127
 
142
128
  self._healthy = False
129
+ start_future = None
143
130
  if self._stop:
144
131
  _LOGGER.debug("MQTT session closed, stopping retry loop")
145
132
  return
@@ -147,6 +134,45 @@ class RoborockMqttSession(MqttSession):
147
134
  await asyncio.sleep(self._backoff.total_seconds())
148
135
  self._backoff = min(self._backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_INTERVAL)
149
136
 
137
+ async def _run_connection(self, start_future: asyncio.Future[None] | None) -> None:
138
+ """Connect to the MQTT broker and listen for messages.
139
+
140
+ This is the primary connection loop for the MQTT session that is
141
+ long running and processes incoming messages. If the connection
142
+ is lost, this method will exit.
143
+ """
144
+ try:
145
+ async with self._mqtt_client(self._params) as client:
146
+ self._backoff = MIN_BACKOFF_INTERVAL
147
+ self._healthy = True
148
+ _LOGGER.info("MQTT Session connected.")
149
+ if start_future and not start_future.done():
150
+ start_future.set_result(None)
151
+
152
+ _LOGGER.debug("Processing MQTT messages")
153
+ async for message in client.messages:
154
+ _LOGGER.debug("Received message: %s", message)
155
+ self._listeners(message.topic.value, message.payload)
156
+ except MqttError as err:
157
+ if start_future and not start_future.done():
158
+ _LOGGER.info("MQTT error starting session: %s", err)
159
+ start_future.set_exception(err)
160
+ else:
161
+ _LOGGER.info("MQTT error: %s", err)
162
+ raise
163
+ except Exception as err:
164
+ # This error is thrown when the MQTT loop is cancelled
165
+ # and the generator is not stopped.
166
+ if "generator didn't stop" in str(err) or "generator didn't yield" in str(err):
167
+ _LOGGER.debug("MQTT loop was cancelled")
168
+ return
169
+ if start_future and not start_future.done():
170
+ _LOGGER.error("Uncaught error starting MQTT session: %s", err)
171
+ start_future.set_exception(err)
172
+ else:
173
+ _LOGGER.exception("Uncaught error during MQTT session: %s", err)
174
+ raise
175
+
150
176
  @asynccontextmanager
151
177
  async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
152
178
  """Connect to the MQTT broker and listen for messages."""
@@ -178,12 +204,6 @@ class RoborockMqttSession(MqttSession):
178
204
  async with self._client_lock:
179
205
  self._client = None
180
206
 
181
- async def _process_message_loop(self, client: aiomqtt.Client) -> None:
182
- _LOGGER.debug("Processing MQTT messages")
183
- async for message in client.messages:
184
- _LOGGER.debug("Received message: %s", message)
185
- self._listeners(message.topic.value, message.payload)
186
-
187
207
  async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
188
208
  """Subscribe to messages on the specified topic and invoke the callback for new messages.
189
209
 
@@ -271,6 +291,10 @@ class LazyMqttSession(MqttSession):
271
291
  """
272
292
  await self._session.close()
273
293
 
294
+ async def restart(self) -> None:
295
+ """Force the session to disconnect and reconnect."""
296
+ await self._session.restart()
297
+
274
298
 
275
299
  async def create_mqtt_session(params: MqttParams) -> MqttSession:
276
300
  """Create an MQTT session.
@@ -54,6 +54,10 @@ class MqttSession(ABC):
54
54
  This will raise an exception if the message could not be sent.
55
55
  """
56
56
 
57
+ @abstractmethod
58
+ async def restart(self) -> None:
59
+ """Force the session to disconnect and reconnect."""
60
+
57
61
  @abstractmethod
58
62
  async def close(self) -> None:
59
63
  """Cancels the mqtt loop"""
File without changes