File: v1_channel.py

package info (click to toggle)
python-roborock 2.49.1-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 1,360 kB
  • sloc: python: 11,539; makefile: 17
file content (279 lines) | stat: -rw-r--r-- 12,095 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
"""V1 Channel for Roborock devices.

This module provides a unified channel interface for V1 protocol devices,
handling both MQTT and local connections with automatic fallback.
"""
import asyncio
import datetime
import logging
from collections.abc import Callable
from typing import TypeVar

from roborock.containers import HomeDataDevice, NetworkInfo, RoborockBase, UserData
from roborock.exceptions import RoborockException
from roborock.mqtt.session import MqttParams, MqttSession
from roborock.protocols.v1_protocol import (
    SecurityData,
    create_security_data,
)
from roborock.roborock_message import RoborockMessage
from roborock.roborock_typing import RoborockCommand

from .cache import Cache
from .channel import Channel
from .local_channel import LocalChannel, LocalSession, create_local_session
from .mqtt_channel import MqttChannel
from .v1_rpc_channel import (
    PickFirstAvailable,
    V1RpcChannel,
    create_local_rpc_channel,
    create_mqtt_rpc_channel,
)

_LOGGER = logging.getLogger(__name__)

__all__ = [
    "V1Channel",
]

_T = TypeVar("_T", bound=RoborockBase)

# Exponential backoff parameters for reconnecting to local
MIN_RECONNECT_INTERVAL = datetime.timedelta(minutes=1)
MAX_RECONNECT_INTERVAL = datetime.timedelta(minutes=10)
RECONNECT_MULTIPLIER = 1.5
# After this many hours, the network info is refreshed
NETWORK_INFO_REFRESH_INTERVAL = datetime.timedelta(hours=12)
# Interval to check that the local connection is healthy
LOCAL_CONNECTION_CHECK_INTERVAL = datetime.timedelta(seconds=15)


class V1Channel(Channel):
    """Unified V1 protocol channel with automatic MQTT/local connection handling.

    This channel abstracts away the complexity of choosing between MQTT and local
    connections, and provides high-level V1 protocol methods. It automatically
    handles connection setup, fallback logic, and protocol encoding/decoding.
    """

    def __init__(
        self,
        device_uid: str,
        security_data: SecurityData,
        mqtt_channel: MqttChannel,
        local_session: LocalSession,
        cache: Cache,
    ) -> None:
        """Initialize the V1Channel.

        Args:
            mqtt_channel: MQTT channel for cloud communication
            local_session: Factory that creates LocalChannels for a hostname.
        """
        self._device_uid = device_uid
        self._mqtt_channel = mqtt_channel
        self._mqtt_rpc_channel = create_mqtt_rpc_channel(mqtt_channel, security_data)
        self._local_session = local_session
        self._local_channel: LocalChannel | None = None
        self._local_rpc_channel: V1RpcChannel | None = None
        # Prefer local, fallback to MQTT
        self._combined_rpc_channel = PickFirstAvailable(
            [lambda: self._local_rpc_channel, lambda: self._mqtt_rpc_channel]
        )
        self._mqtt_unsub: Callable[[], None] | None = None
        self._local_unsub: Callable[[], None] | None = None
        self._callback: Callable[[RoborockMessage], None] | None = None
        self._cache = cache
        self._reconnect_task: asyncio.Task[None] | None = None
        self._last_network_info_refresh: datetime.datetime | None = None

    @property
    def is_connected(self) -> bool:
        """Return whether any connection is available."""
        return self.is_mqtt_connected or self.is_local_connected

    @property
    def is_local_connected(self) -> bool:
        """Return whether local connection is available."""
        return self._local_channel is not None and self._local_channel.is_connected

    @property
    def is_mqtt_connected(self) -> bool:
        """Return whether MQTT connection is available."""
        return self._mqtt_unsub is not None and self._mqtt_channel.is_connected

    @property
    def rpc_channel(self) -> V1RpcChannel:
        """Return the combined RPC channel prefers local with a fallback to MQTT."""
        return self._combined_rpc_channel

    @property
    def mqtt_rpc_channel(self) -> V1RpcChannel:
        """Return the MQTT RPC channel."""
        return self._mqtt_rpc_channel

    async def subscribe(self, callback: Callable[[RoborockMessage], None]) -> Callable[[], None]:
        """Subscribe to all messages from the device.

        This will establish MQTT connection first, and also attempt to set up
        local connection if possible. Any failures to subscribe to MQTT will raise
        a RoborockException. A local connection failure will not raise an exception,
        since the local connection is optional.
        """
        if self._callback is not None:
            raise ValueError("Only one subscription allowed at a time")

        # Make an initial, optimistic attempt to connect to local with the
        # cache. The cache information will be refreshed by the background task.
        try:
            await self._local_connect(use_cache=True)
        except RoborockException as err:
            _LOGGER.warning("Could not establish local connection for device %s: %s", self._device_uid, err)

        # Start a background task to manage the local connection health. This
        # happens independent of whether we were able to connect locally now.
        _LOGGER.info("self._reconnect_task=%s", self._reconnect_task)
        if self._reconnect_task is None:
            loop = asyncio.get_running_loop()
            self._reconnect_task = loop.create_task(self._background_reconnect())

        if not self.is_local_connected:
            # We were not able to connect locally, so fallback to MQTT and at least
            # establish that connection explicitly. If this fails then raise an
            # error and let the caller know we failed to subscribe.
            self._mqtt_unsub = await self._mqtt_channel.subscribe(self._on_mqtt_message)
            _LOGGER.debug("V1Channel connected to device %s via MQTT", self._device_uid)

        def unsub() -> None:
            """Unsubscribe from all messages."""
            if self._reconnect_task:
                self._reconnect_task.cancel()
                self._reconnect_task = None
            if self._mqtt_unsub:
                self._mqtt_unsub()
                self._mqtt_unsub = None
            if self._local_unsub:
                self._local_unsub()
                self._local_unsub = None
            _LOGGER.debug("Unsubscribed from device %s", self._device_uid)

        self._callback = callback
        return unsub

    async def _get_networking_info(self, *, use_cache: bool = True) -> NetworkInfo:
        """Retrieve networking information for the device.

        This is a cloud only command used to get the local device's IP address.
        """
        cache_data = await self._cache.get()
        if use_cache and cache_data.network_info and (network_info := cache_data.network_info.get(self._device_uid)):
            _LOGGER.debug("Using cached network info for device %s", self._device_uid)
            return network_info
        try:
            network_info = await self._mqtt_rpc_channel.send_command(
                RoborockCommand.GET_NETWORK_INFO, response_type=NetworkInfo
            )
        except RoborockException as e:
            raise RoborockException(f"Network info failed for device {self._device_uid}") from e
        _LOGGER.debug("Network info for device %s: %s", self._device_uid, network_info)
        self._last_network_info_refresh = datetime.datetime.now(datetime.timezone.utc)
        cache_data.network_info[self._device_uid] = network_info
        await self._cache.set(cache_data)
        return network_info

    async def _local_connect(self, *, use_cache: bool = True) -> None:
        """Set up local connection if possible."""
        _LOGGER.debug(
            "Attempting to connect to local channel for device %s (use_cache=%s)", self._device_uid, use_cache
        )
        networking_info = await self._get_networking_info(use_cache=use_cache)
        host = networking_info.ip
        _LOGGER.debug("Connecting to local channel at %s", host)
        # Create a new local channel and connect
        local_channel = self._local_session(host)
        try:
            await local_channel.connect()
        except RoborockException as e:
            raise RoborockException(f"Error connecting to local device {self._device_uid}: {e}") from e
        # Wire up the new channel
        self._local_channel = local_channel
        self._local_rpc_channel = create_local_rpc_channel(self._local_channel)
        self._local_unsub = await self._local_channel.subscribe(self._on_local_message)
        _LOGGER.info("Successfully connected to local device %s", self._device_uid)

    async def _background_reconnect(self) -> None:
        """Task to run in the background to manage the local connection."""
        _LOGGER.debug("Starting background task to manage local connection for %s", self._device_uid)
        reconnect_backoff = MIN_RECONNECT_INTERVAL
        local_connect_failures = 0

        while True:
            try:
                if self.is_local_connected:
                    await asyncio.sleep(LOCAL_CONNECTION_CHECK_INTERVAL.total_seconds())
                    continue

                # Not connected, so wait with backoff before trying to connect.
                # The first time through, we don't sleep, we just try to connect.
                local_connect_failures += 1
                if local_connect_failures > 1:
                    await asyncio.sleep(reconnect_backoff.total_seconds())
                    reconnect_backoff = min(reconnect_backoff * RECONNECT_MULTIPLIER, MAX_RECONNECT_INTERVAL)

                use_cache = self._should_use_cache(local_connect_failures)
                await self._local_connect(use_cache=use_cache)
                # Reset backoff and failures on success
                reconnect_backoff = MIN_RECONNECT_INTERVAL
                local_connect_failures = 0

            except asyncio.CancelledError:
                _LOGGER.debug("Background reconnect task cancelled")
                if self._local_channel:
                    self._local_channel.close()
                return
            except RoborockException as err:
                _LOGGER.debug("Background reconnect failed: %s", err)
            except Exception:
                _LOGGER.exception("Unhandled exception in background reconnect task")

    def _should_use_cache(self, local_connect_failures: int) -> bool:
        """Determine whether to use cached network info on retries.

        On the first retry we'll avoid the cache to handle the case where
        the network ip may have recently changed. Otherwise, use the cache
        if available then expire at some point.
        """
        if local_connect_failures == 1:
            return False
        elif self._last_network_info_refresh and (
            datetime.datetime.now(datetime.timezone.utc) - self._last_network_info_refresh
            > NETWORK_INFO_REFRESH_INTERVAL
        ):
            return False
        return True

    def _on_mqtt_message(self, message: RoborockMessage) -> None:
        """Handle incoming MQTT messages."""
        _LOGGER.debug("V1Channel received MQTT message from device %s: %s", self._device_uid, message)
        if self._callback:
            self._callback(message)

    def _on_local_message(self, message: RoborockMessage) -> None:
        """Handle incoming local messages."""
        _LOGGER.debug("V1Channel received local message from device %s: %s", self._device_uid, message)
        if self._callback:
            self._callback(message)


def create_v1_channel(
    user_data: UserData,
    mqtt_params: MqttParams,
    mqtt_session: MqttSession,
    device: HomeDataDevice,
    cache: Cache,
) -> V1Channel:
    """Create a V1Channel for the given device."""
    security_data = create_security_data(user_data.rriot)
    mqtt_channel = MqttChannel(mqtt_session, device.duid, device.local_key, user_data.rriot, mqtt_params)
    local_session = create_local_session(device.local_key)
    return V1Channel(device.duid, security_data, mqtt_channel, local_session=local_session, cache=cache)