Add a hook for handling background rate limit errors

When a rate limit error is received the background loop will jump to the maximum backoff (now 6 hours) and will also invoke a callback so that the caller can decide to re-authenticate or stop harder.

The new backoff follows this trajectory, where the change in behavior is introduced after 15 attempts:
- attempt 1: wait 10 seconds
- attempt 5: waits 50 seconds
- attempt 7: waits 2 minutes
- attempt 10: waits 6 minutes
- attempt 15: waits 32 minutes
- attempt 17: waits 2 hours
- attempt 20: waits 6 hours

Author: allenporter

roborock/devices/device_manager.py

@@ -20,7 +20,7 @@
 from roborock.exceptions import RoborockException
 from roborock.map.map_parser import MapParserConfig
 from roborock.mqtt.roborock_session import create_lazy_mqtt_session
-from roborock.mqtt.session import MqttSession
+from roborock.mqtt.session import MqttSession, SessionUnauthorizedHook
 from roborock.protocol import create_mqtt_params
 from roborock.web_api import RoborockApiClient, UserWebApiClient
 
@@ -173,6 +173,7 @@ async def create_device_manager(
     map_parser_config: MapParserConfig | None = None,
     session: aiohttp.ClientSession | None = None,
     ready_callback: DeviceReadyCallback | None = None,
+    mqtt_session_unauthorized_hook: SessionUnauthorizedHook | None = None,
 ) -> DeviceManager:
     """Convenience function to create and initialize a DeviceManager.
 
@@ -196,6 +197,7 @@ async def create_device_manager(
 
     mqtt_params = create_mqtt_params(user_data.rriot)
     mqtt_params.diagnostics = diagnostics.subkey("mqtt_session")
+    mqtt_params.unauthorized_hook = mqtt_session_unauthorized_hook
     mqtt_session = await create_lazy_mqtt_session(mqtt_params)
 
     def device_creator(home_data: HomeData, device: HomeDataDevice, product: HomeDataProduct) -> RoborockDevice:

roborock/mqtt/roborock_session.py

@@ -31,7 +31,7 @@
 
 # Exponential backoff parameters
 MIN_BACKOFF_INTERVAL = datetime.timedelta(seconds=10)
-MAX_BACKOFF_INTERVAL = datetime.timedelta(minutes=30)
+MAX_BACKOFF_INTERVAL = datetime.timedelta(hours=6)
 BACKOFF_MULTIPLIER = 1.5
 
 
@@ -79,6 +79,7 @@ def __init__(
         self._idle_timers: dict[str, asyncio.Task[None]] = {}
         self._diagnostics = params.diagnostics
         self._health_manager = HealthManager(self.restart)
+        self._unauthorized_hook = params.unauthorized_hook
 
     @property
     def connected(self) -> bool:
@@ -199,14 +200,28 @@ async def _run_connection(self, start_future: asyncio.Future[None] | None) -> No
                         _LOGGER.debug("Received message: %s", message)
                         with self._diagnostics.timer("dispatch_message"):
                             self._listeners(message.topic.value, message.payload)
+        except MqttCodeError as err:
+            self._diagnostics.increment(f"connect_failure:{err.rc}")
+            if start_future and not start_future.done():
+                _LOGGER.debug("MQTT error starting session: %s", err)
+                start_future.set_exception(err)
+            else:
+                _LOGGER.debug("MQTT error: %s", err)
+            if err.rc == MqttReasonCode.RC_ERROR_UNAUTHORIZED and self._unauthorized_hook:
+                _LOGGER.info("MQTT unauthorized/rate-limit error received, setting backoff to maximum")
+                self._unauthorized_hook()
+                self._backoff = MAX_BACKOFF_INTERVAL
+            raise
         except MqttError as err:
+            self._diagnostics.increment("connect_failure:unknown")
             if start_future and not start_future.done():
                 _LOGGER.info("MQTT error starting session: %s", err)
                 start_future.set_exception(err)
             else:
                 _LOGGER.info("MQTT error: %s", err)
             raise
-        except Exception as err:
+        except (RuntimeError, Exception) as err:
+            self._diagnostics.increment("connect_failure:uncaught")
             # This error is thrown when the MQTT loop is cancelled
             # and the generator is not stopped.
             if "generator didn't stop" in str(err) or "generator didn't yield" in str(err):

roborock/mqtt/session.py

@@ -10,6 +10,8 @@
 
 DEFAULT_TIMEOUT = 30.0
 
+SessionUnauthorizedHook = Callable[[], None]
+
 
 @dataclass
 class MqttParams:
@@ -41,6 +43,14 @@ class MqttParams:
     shared MQTT session diagnostics are included in the overall diagnostics.
     """
 
+    unauthorized_hook: SessionUnauthorizedHook | None = None
+    """Optional hook invoked when an unauthorized error is received.
+    
+    This may be invoked by the background reconnect logic when an
+    unauthorized error is received from the broker. The caller may use
+    this hook to refresh credentials or take other actions as needed.
+    """
+
 
 class MqttSession(ABC):
     """An MQTT session for sending and receiving messages."""

tests/mqtt/test_roborock_session.py

@@ -4,6 +4,7 @@
 import copy
 import datetime
 from collections.abc import Callable, Generator
+from typing import Any
 from unittest.mock import AsyncMock, Mock, patch
 
 import aiomqtt
@@ -31,17 +32,62 @@ def auto_fast_backoff(fast_backoff_fixture: None) -> None:
     """Automatically use the fast backoff fixture."""
 
 
-@pytest.fixture
-def mock_mqtt_client() -> Generator[AsyncMock, None, None]:
-    """Fixture to create a mock MQTT client with patched aiomqtt.Client."""
+class FakeAsyncIterator:
+    """Fake async iterator that waits for messages to arrive, but they never do.
+
+    This is used for testing exceptions in other client functions.
+    """
+
+    def __init__(self) -> None:
+        self.loop = True
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self) -> None:
+        """Iterator that does not generate any messages."""
+        while self.loop:
+            await asyncio.sleep(0.01)
+
+
+@pytest.fixture(name="message_iterator")
+def message_iterator_fixture() -> FakeAsyncIterator:
+    """Fixture to provide a side effect for creating the MQTT client."""
+    return FakeAsyncIterator()
+
+
+@pytest.fixture(name="mock_client")
+def mock_client_fixture(message_iterator: FakeAsyncIterator) -> AsyncMock:
+    """Fixture to provide a side effect for creating the MQTT client."""
     mock_client = AsyncMock()
-    mock_client.messages = FakeAsyncIterator()
+    mock_client.messages = message_iterator
+    return mock_client
+
+
+@pytest.fixture(name="create_client_side_effect")
+def create_client_side_effect_fixture() -> Exception | None:
+    """Fixture to provide a side effect for creating the MQTT client."""
+    return None
+
 
+@pytest.fixture(name="mock_aenter_client")
+def mock_aenter_client_fixture(mock_client: AsyncMock, create_client_side_effect: Exception | None) -> AsyncMock:
+    """Fixture to provide a side effect for creating the MQTT client."""
     mock_aenter = AsyncMock()
     mock_aenter.return_value = mock_client
+    mock_aenter.side_effect = create_client_side_effect
+    return mock_aenter
+
+
+@pytest.fixture
+def mock_mqtt_client(
+    mock_client: AsyncMock,
+    mock_aenter_client: AsyncMock,
+) -> Generator[AsyncMock, None, None]:
+    """Fixture to create a mock MQTT client with patched aiomqtt.Client."""
 
     mock_shim = Mock()
-    mock_shim.return_value.__aenter__ = mock_aenter
+    mock_shim.return_value.__aenter__ = mock_aenter_client
     mock_shim.return_value.__aexit__ = AsyncMock()
 
     with patch("roborock.mqtt.roborock_session.aiomqtt.Client", mock_shim):
@@ -114,21 +160,6 @@ async def test_publish_command(push_response: Callable[[bytes], None]) -> None:
     assert not session.connected
 
 
-class FakeAsyncIterator:
-    """Fake async iterator that waits for messages to arrive, but they never do.
-
-    This is used for testing exceptions in other client functions.
-    """
-
-    def __aiter__(self):
-        return self
-
-    async def __anext__(self) -> None:
-        """Iterator that does not generate any messages."""
-        while True:
-            await asyncio.sleep(1)
-
-
 async def test_publish_failure(mock_mqtt_client: AsyncMock) -> None:
     """Test an MQTT error is received when publishing a message."""
 
@@ -432,3 +463,64 @@ async def test_diagnostics_data(push_response: Callable[[bytes], None]) -> None:
     assert data.get("subscribe_count") == 2
     assert data.get("dispatch_message_count") == 3
     assert data.get("close") == 1
+
+
+@pytest.mark.parametrize(
+    ("create_client_side_effect"),
+    [
+        # Unauthorized
+        aiomqtt.MqttCodeError(rc=135),
+    ],
+)
+async def test_session_unauthorized_hook(mock_mqtt_client: AsyncMock, push_response: Callable[[bytes], None]) -> None:
+    """Test the MQTT session."""
+
+    unauthorized = asyncio.Event()
+
+    params = copy.deepcopy(FAKE_PARAMS)
+    params.unauthorized_hook = unauthorized.set
+
+    with pytest.raises(MqttSessionUnauthorized):
+        await create_mqtt_session(params)
+
+    assert unauthorized.is_set()
+
+
+async def test_session_unauthorized_after_start(
+    mock_aenter_client: AsyncMock,
+    message_iterator: FakeAsyncIterator,
+    mock_mqtt_client: AsyncMock,
+    push_response: Callable[[bytes], None],
+) -> None:
+    """Test the MQTT session."""
+
+    # Configure a hook that is notified of unauthorized errors
+    unauthorized = asyncio.Event()
+    params = copy.deepcopy(FAKE_PARAMS)
+    params.unauthorized_hook = unauthorized.set
+
+    # The client will succeed on first connection attempt, then fail with
+    # unauthorized messages on all future attempts.
+    request_count = 0
+
+    def succeed_then_fail_unauthorized() -> Any:
+        nonlocal request_count
+        request_count += 1
+        if request_count == 1:
+            return mock_mqtt_client
+        raise aiomqtt.MqttCodeError(rc=135)
+
+    mock_aenter_client.side_effect = succeed_then_fail_unauthorized
+    # Don't produce messages, just exit and restart to reconnect
+    message_iterator.loop = False
+
+    push_response(mqtt_packet.gen_connack(rc=0, flags=2))
+
+    session = await create_mqtt_session(params)
+    assert session.connected
+
+    try:
+        async with asyncio.timeout(10):
+            assert await unauthorized.wait()
+    finally:
+        await session.close()