fix(a01): restart stale MQTT session on command timeouts

Author: piitaya

roborock/devices/rpc/a01_channel.py

@@ -98,8 +98,10 @@ def find_response(response_message: RoborockMessage) -> None:
         try:
             await asyncio.wait_for(finished.wait(), timeout=_TIMEOUT)
         except TimeoutError as ex:
+            await mqtt_channel.health_manager.on_timeout()
             raise RoborockException(f"Command timed out after {_TIMEOUT}s") from ex
     finally:
         unsub()
 
+    await mqtt_channel.health_manager.on_success()
     return result  # type: ignore[return-value]

roborock/devices/traits/a01/__init__.py

@@ -8,12 +8,14 @@
 A01 devices expose a single API object that handles all device interactions. This API is
 available on the device instance (typically via `device.a01_properties`).
 
-The API provides two main methods:
+The API provides these methods:
 1.  **query_values(protocols)**: Fetches current state for specific data points.
     You must pass a list of protocol enums (e.g. `RoborockDyadDataProtocol` or
     `RoborockZeoProtocol`) to request specific data.
 2.  **set_value(protocol, value)**: Sends a command to the device to change a setting
     or perform an action.
+3.  **add_listener(callback)**: Subscribes to state the device pushes on its own (for
+    example when its state changes), invoking the callback with decoded values.
 
 Note that these APIs fetch data directly from the device upon request and do not
 cache state internally.
@@ -51,7 +53,9 @@
 from roborock.devices.rpc.a01_channel import send_decoded_command
 from roborock.devices.traits import Trait
 from roborock.devices.transport.mqtt_channel import MqttChannel
-from roborock.roborock_message import RoborockDyadDataProtocol, RoborockZeoProtocol
+from roborock.exceptions import RoborockException
+from roborock.protocols.a01_protocol import decode_rpc_response
+from roborock.roborock_message import RoborockDyadDataProtocol, RoborockMessage, RoborockZeoProtocol
 
 __init__ = [
     "DyadApi",
@@ -134,6 +138,12 @@ def convert_zeo_value(protocol_value: RoborockZeoProtocol, value: Any) -> Any:
     return None
 
 
+# RoborockDyadDataProtocol._missing_ maps any unknown code to its first member
+# instead of raising, so incoming data points must be checked against this set
+# before being converted to a protocol.
+_DYAD_PROTOCOL_VALUES = frozenset(protocol.value for protocol in RoborockDyadDataProtocol)
+
+
 class DyadApi(Trait):
     """API for interacting with Dyad devices."""
 
@@ -155,6 +165,32 @@ async def set_value(self, protocol: RoborockDyadDataProtocol, value: Any) -> dic
         params = {protocol: value}
         return await send_decoded_command(self._channel, params)
 
+    async def add_listener(
+        self, callback: Callable[[dict[RoborockDyadDataProtocol, Any]], None]
+    ) -> Callable[[], None]:
+        """Listen for state the device pushes on its own.
+
+        The callback is invoked with decoded values whenever the device sends a
+        message, including unsolicited pushes when its state changes. Only known
+        protocols are delivered. Returns a callable to remove the listener.
+        """
+
+        def on_message(message: RoborockMessage) -> None:
+            try:
+                datapoints = decode_rpc_response(message)
+            except RoborockException:
+                return
+            values: dict[RoborockDyadDataProtocol, Any] = {}
+            for code, value in datapoints.items():
+                if code not in _DYAD_PROTOCOL_VALUES:
+                    continue
+                protocol = RoborockDyadDataProtocol(code)
+                values[protocol] = convert_dyad_value(protocol, value)
+            if values:
+                callback(values)
+
+        return await self._channel.subscribe(on_message)
+
 
 class ZeoApi(Trait):
     """API for interacting with Zeo devices."""

tests/devices/rpc/test_a01_channel.py

@@ -1,10 +1,12 @@
 """Tests for the a01_channel."""
 
 from typing import Any
+from unittest.mock import AsyncMock
 
 import pytest
 
 from roborock.devices.rpc.a01_channel import send_decoded_command
+from roborock.exceptions import RoborockException
 from roborock.protocols.a01_protocol import encode_mqtt_payload
 from roborock.roborock_message import (
     RoborockDyadDataProtocol,
@@ -51,3 +53,37 @@ async def test_id_query(mock_mqtt_channel: FakeChannel):
     }
     mock_mqtt_channel.publish.assert_awaited_once()
     mock_mqtt_channel.subscribe.assert_awaited_once()
+
+
+async def test_query_marks_session_healthy(mock_mqtt_channel: FakeChannel):
+    """A completed query reports success to the health manager."""
+    mock_mqtt_channel.health_manager.on_success = AsyncMock()  # type: ignore[method-assign]
+
+    params: dict[RoborockDyadDataProtocol, Any] = {
+        RoborockDyadDataProtocol.ID_QUERY: [RoborockDyadDataProtocol.POWER]
+    }
+    encoded = encode_mqtt_payload({RoborockDyadDataProtocol.POWER: 75}, value_encoder=lambda x: x)
+    mock_mqtt_channel.response_queue.append(
+        RoborockMessage(
+            protocol=RoborockMessageProtocol.RPC_RESPONSE, payload=encoded.payload, version=encoded.version
+        )
+    )
+
+    await send_decoded_command(mock_mqtt_channel, params)  # type: ignore[call-overload]
+
+    mock_mqtt_channel.health_manager.on_success.assert_awaited_once()
+
+
+async def test_query_timeout_reports_to_health_manager(mock_mqtt_channel: FakeChannel, monkeypatch):
+    """A timed-out query reports to the health manager so a stale session can be restarted."""
+    monkeypatch.setattr("roborock.devices.rpc.a01_channel._TIMEOUT", 0.01)
+    mock_mqtt_channel.health_manager.on_timeout = AsyncMock()  # type: ignore[method-assign]
+
+    params: dict[RoborockDyadDataProtocol, Any] = {
+        RoborockDyadDataProtocol.ID_QUERY: [RoborockDyadDataProtocol.POWER]
+    }
+
+    with pytest.raises(RoborockException, match="timed out"):
+        await send_decoded_command(mock_mqtt_channel, params)  # type: ignore[call-overload]
+
+    mock_mqtt_channel.health_manager.on_timeout.assert_awaited_once()

tests/devices/traits/a01/test_init.py

@@ -136,6 +136,27 @@ async def test_dyad_invalid_response_value(
     assert result == expected_result
 
 
+async def test_dyad_add_listener(dyad_api: DyadApi, fake_channel: FakeChannel):
+    """add_listener delivers decoded values for pushed messages and skips unknown codes."""
+    received: list[dict[RoborockDyadDataProtocol, Any]] = []
+    unsub = await dyad_api.add_listener(received.append)
+
+    # 999 is not a known protocol: it must be skipped, not mapped to the first enum member.
+    fake_channel.notify_subscribers(build_a01_message({206: 3, 209: 80, 216: 0, 999: 1}))
+
+    assert received == [
+        {
+            RoborockDyadDataProtocol.SUCTION: "l3",
+            RoborockDyadDataProtocol.POWER: 80,
+            RoborockDyadDataProtocol.ERROR: "none",
+        }
+    ]
+
+    unsub()
+    fake_channel.notify_subscribers(build_a01_message({206: 1}))
+    assert len(received) == 1  # no callback fires after unsubscribing
+
+
 async def test_zeo_api_query_values(zeo_api: ZeoApi, fake_channel: FakeChannel):
     """Test that ZeoApi currently returns raw values without conversion."""
     fake_channel.response_queue.append(