roborock-cli 0.1.1__py3-none-any.whl

roborock_cli/_vendor/roborock/map/map_parser.py

@@ -0,0 +1,123 @@
+"""Module for parsing v1 Roborock map content."""
+
+import io
+import logging
+from dataclasses import dataclass, field
+
+from vacuum_map_parser_base.config.color import ColorsPalette, SupportedColor
+from vacuum_map_parser_base.config.drawable import Drawable
+from vacuum_map_parser_base.config.image_config import ImageConfig
+from vacuum_map_parser_base.config.size import Size, Sizes
+from vacuum_map_parser_base.map_data import MapData
+from vacuum_map_parser_roborock.map_data_parser import RoborockMapDataParser
+
+from roborock_cli._vendor.roborock.exceptions import RoborockException
+
+_LOGGER = logging.getLogger(__name__)
+
+DEFAULT_DRAWABLES = {
+    Drawable.CHARGER: True,
+    Drawable.CLEANED_AREA: False,
+    Drawable.GOTO_PATH: False,
+    Drawable.IGNORED_OBSTACLES: False,
+    Drawable.IGNORED_OBSTACLES_WITH_PHOTO: False,
+    Drawable.MOP_PATH: False,
+    Drawable.NO_CARPET_AREAS: False,
+    Drawable.NO_GO_AREAS: False,
+    Drawable.NO_MOPPING_AREAS: False,
+    Drawable.OBSTACLES: False,
+    Drawable.OBSTACLES_WITH_PHOTO: False,
+    Drawable.PATH: True,
+    Drawable.PREDICTED_PATH: False,
+    Drawable.VACUUM_POSITION: True,
+    Drawable.VIRTUAL_WALLS: False,
+    Drawable.ZONES: False,
+}
+DEFAULT_MAP_SCALE = 4
+MAP_FILE_FORMAT = "PNG"
+
+
+def _default_drawable_factory() -> list[Drawable]:
+    return [drawable for drawable, default_value in DEFAULT_DRAWABLES.items() if default_value]
+
+
+@dataclass
+class MapParserConfig:
+    """Configuration for the Roborock map parser."""
+
+    drawables: list[Drawable] = field(default_factory=_default_drawable_factory)
+    """List of drawables to include in the map rendering."""
+
+    show_background: bool = True
+    """Whether to show the background of the map."""
+
+    show_walls: bool = True
+    """Whether to show the walls of the map."""
+
+    show_rooms: bool = True
+    """Whether to show the rooms of the map."""
+
+    map_scale: int = DEFAULT_MAP_SCALE
+    """Scale factor for the map."""
+
+
+@dataclass
+class ParsedMapData:
+    """Roborock Map Data.
+
+    This class holds the parsed map data and the rendered image.
+    """
+
+    image_content: bytes | None
+    """The rendered image of the map in PNG format."""
+
+    map_data: MapData | None
+    """The parsed map data which contains metadata for points on the map."""
+
+
+class MapParser:
+    """Roborock Map Parser.
+
+    This class is used to parse the map data from the device and render it into an image.
+    """
+
+    def __init__(self, config: MapParserConfig) -> None:
+        """Initialize the MapParser."""
+        self._map_parser = _create_map_data_parser(config)
+
+    def parse(self, map_bytes: bytes) -> ParsedMapData | None:
+        """Parse map_bytes and return MapData and the image."""
+        try:
+            parsed_map = self._map_parser.parse(map_bytes)
+        except (IndexError, ValueError) as err:
+            raise RoborockException("Failed to parse map data") from err
+        if parsed_map.image is None:
+            raise RoborockException("Failed to render map image")
+        img_byte_arr = io.BytesIO()
+        parsed_map.image.data.save(img_byte_arr, format=MAP_FILE_FORMAT)
+        return ParsedMapData(image_content=img_byte_arr.getvalue(), map_data=parsed_map)
+
+
+def _create_map_data_parser(config: MapParserConfig) -> RoborockMapDataParser:
+    """Create a RoborockMapDataParser based on the config entry."""
+    color_dicts = {}
+    room_colors = {}
+
+    if not config.show_background:
+        color_dicts[SupportedColor.MAP_OUTSIDE] = (0, 0, 0, 0)
+
+    if not config.show_walls:
+        color_dicts[SupportedColor.GREY_WALL] = (0, 0, 0, 0)
+        color_dicts[SupportedColor.MAP_WALL] = (0, 0, 0, 0)
+        color_dicts[SupportedColor.MAP_WALL_V2] = (0, 0, 0, 0)
+
+    if not config.show_rooms:
+        room_colors = {str(x): (0, 0, 0, 0) for x in range(1, 32)}
+
+    return RoborockMapDataParser(
+        ColorsPalette(color_dicts, room_colors),
+        Sizes({k: v * config.map_scale for k, v in Sizes.SIZES.items() if k != Size.MOP_PATH_WIDTH}),
+        config.drawables,
+        ImageConfig(scale=config.map_scale),
+        [],
+    )

roborock_cli/_vendor/roborock/mqtt/__init__.py

@@ -0,0 +1,10 @@
+"""This module contains the low level MQTT client for the Roborock vacuum cleaner.
+
+This is not meant to be used directly, but rather as a base for the higher level
+modules.
+"""
+
+# This module is part of the Roborock Python library, which provides a way to
+# interact with Roborock devices using MQTT. It is not intended to be used directly,
+# but rather as a base for higher level modules.
+__all__: list[str] = []

roborock_cli/_vendor/roborock/mqtt/health_manager.py

@@ -0,0 +1,60 @@
+"""A health manager for monitoring MQTT connections to Roborock devices.
+
+We observe a problem where sometimes the MQTT connection appears to be alive but
+no messages are being received. To mitigate this, we track consecutive timeouts
+and restart the connection if too many timeouts occur in succession.
+"""
+
+import datetime
+import logging
+from collections.abc import Awaitable, Callable
+
+_LOGGER = logging.getLogger(__name__)
+
+# Number of consecutive timeouts before considering the connection unhealthy.
+TIMEOUT_THRESHOLD = 3
+
+# We won't restart the session more often than this interval.
+RESTART_COOLDOWN = datetime.timedelta(minutes=30)
+
+
+class HealthManager:
+    """Manager for monitoring the health of MQTT connections.
+
+    This tracks communication timeouts and can trigger restarts of the MQTT
+    session if too many timeouts occur in succession.
+    """
+
+    def __init__(self, restart: Callable[[], Awaitable[None]]) -> None:
+        """Initialize the health manager.
+
+        Args:
+            restart: A callable to restart the MQTT session.
+        """
+        self._consecutive_timeouts = 0
+        self._restart = restart
+        self._last_restart: datetime.datetime | None = None
+
+    async def on_success(self) -> None:
+        """Record a successful communication event."""
+        self._consecutive_timeouts = 0
+
+    async def on_timeout(self) -> None:
+        """Record a timeout event.
+
+        This may trigger a restart of the MQTT session if too many timeouts
+        have occurred in succession.
+        """
+        self._consecutive_timeouts += 1
+        if self._consecutive_timeouts >= TIMEOUT_THRESHOLD:
+            now = datetime.datetime.now(datetime.UTC)
+            since_last = (now - self._last_restart) if self._last_restart else None
+            if since_last is None or since_last >= RESTART_COOLDOWN:
+                _LOGGER.debug(
+                    "Restarting MQTT session after %d consecutive timeouts (duration since last restart %s)",
+                    self._consecutive_timeouts,
+                    since_last,
+                )
+                await self._restart()
+                self._last_restart = now
+                self._consecutive_timeouts = 0

roborock_cli/_vendor/roborock/mqtt/roborock_session.py

@@ -0,0 +1,463 @@
+"""An MQTT session for sending and receiving messages.
+
+See create_mqtt_session for a factory function to create an MQTT session.
+
+This is a thin wrapper around the async MQTT client that handles dispatching messages
+from a topic to a callback function, since the async MQTT client does not
+support this out of the box. It also handles the authentication process and
+receiving messages from the vacuum cleaner.
+"""
+
+import asyncio
+import datetime
+import logging
+import ssl
+from collections.abc import Callable
+from contextlib import asynccontextmanager
+
+import aiomqtt
+import certifi
+from aiomqtt import MqttCodeError, MqttError, TLSParameters
+
+from roborock_cli._vendor.roborock.callbacks import CallbackMap
+from roborock_cli._vendor.roborock.diagnostics import Diagnostics, redact_topic_name
+
+from .health_manager import HealthManager
+from .session import MqttParams, MqttSession, MqttSessionException, MqttSessionUnauthorized
+
+_LOGGER = logging.getLogger(__name__)
+_MQTT_LOGGER = logging.getLogger(f"{__name__}.aiomqtt")
+
+CLIENT_KEEPALIVE = datetime.timedelta(seconds=45)
+TOPIC_KEEPALIVE = datetime.timedelta(seconds=60)
+
+# Exponential backoff parameters
+MIN_BACKOFF_INTERVAL = datetime.timedelta(seconds=10)
+MAX_BACKOFF_INTERVAL = datetime.timedelta(hours=6)
+BACKOFF_MULTIPLIER = 1.5
+
+
+class MqttReasonCode:
+    """MQTT Reason Codes used by Roborock devices.
+
+    This is a subset of paho.mqtt.reasoncodes.ReasonCode where we would like
+    different error handling behavior.
+    """
+
+    RC_ERROR_UNAUTHORIZED = 135
+
+
+class RoborockMqttSession(MqttSession):
+    """An MQTT session for sending and receiving messages.
+
+    You can start a session invoking the start() method which will connect to
+    the MQTT broker. A caller may subscribe to a topic, and the session keeps
+    track of which callbacks to invoke for each topic.
+
+    The client is run as a background task that will run until shutdown. Once
+    connected, the client will wait for messages to be received in a loop. If
+    the connection is lost, the client will be re-created and reconnected. There
+    is backoff to avoid spamming the broker with connection attempts.
+
+    Reconnect attempts are deferred while there are no active subscriptions,
+    which avoids unnecessary reconnect churn for idle sessions. Reconnects
+    resume as soon as a subscription is added again. The client automatically
+    re-establishes any existing subscriptions when the connection returns.
+    """
+
+    def __init__(
+        self,
+        params: MqttParams,
+        topic_idle_timeout: datetime.timedelta = TOPIC_KEEPALIVE,
+    ):
+        self._params = params
+        self._reconnect_task: asyncio.Task[None] | None = None
+        self._healthy = False
+        self._stop = False
+        self._backoff = MIN_BACKOFF_INTERVAL
+        self._client: aiomqtt.Client | None = None
+        self._client_subscribed_topics: set[str] = set()
+        self._client_lock = asyncio.Lock()
+        self._listeners: CallbackMap[str, bytes] = CallbackMap(_LOGGER)
+        self._connection_task: asyncio.Task[None] | None = None
+        self._topic_idle_timeout = topic_idle_timeout
+        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:
+        """True if the session is connected to the broker."""
+        return self._healthy
+
+    @property
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._health_manager
+
+    async def start(self) -> None:
+        """Start the MQTT session.
+
+        This has special behavior for the first connection attempt where any
+        failures are raised immediately. This is to allow the caller to
+        handle the failure and retry if desired itself. Once connected,
+        the session will retry connecting in the background.
+        """
+        self._diagnostics.increment("start_attempt")
+        start_future: asyncio.Future[None] = asyncio.Future()
+        loop = asyncio.get_event_loop()
+        self._reconnect_task = loop.create_task(self._run_reconnect_loop(start_future))
+        try:
+            await start_future
+        except MqttCodeError as err:
+            self._diagnostics.increment(f"start_failure:{err.rc}")
+            if err.rc == MqttReasonCode.RC_ERROR_UNAUTHORIZED:
+                raise MqttSessionUnauthorized(f"Authorization error starting MQTT session: {err}") from err
+            raise MqttSessionException(f"Error starting MQTT session: {err}") from err
+        except MqttError as err:
+            self._diagnostics.increment("start_failure:unknown")
+            raise MqttSessionException(f"Error starting MQTT session: {err}") from err
+        except Exception as err:
+            self._diagnostics.increment("start_failure:uncaught")
+            raise MqttSessionException(f"Unexpected error starting session: {err}") from err
+        else:
+            self._diagnostics.increment("start_success")
+            _LOGGER.debug("MQTT session started successfully")
+
+    async def close(self) -> None:
+        """Cancels the MQTT loop and shutdown the client library."""
+        self._diagnostics.increment("close")
+        self._stop = True
+        tasks = [task for task in [self._connection_task, self._reconnect_task, *self._idle_timers.values()] if task]
+        self._connection_task = None
+        self._reconnect_task = None
+        self._idle_timers.clear()
+
+        for task in tasks:
+            task.cancel()
+        try:
+            await asyncio.gather(*tasks, return_exceptions=True)
+        except asyncio.CancelledError:
+            pass
+
+        self._healthy = False
+
+    async def restart(self) -> None:
+        """Force the session to disconnect and reconnect.
+
+        The active connection task will be cancelled and restarted in the background, retried by
+        the reconnect loop. This is a no-op if there is no active connection.
+        """
+        _LOGGER.info("Forcing MQTT session restart")
+        self._diagnostics.increment("restart")
+        if self._connection_task:
+            self._connection_task.cancel()
+        else:
+            _LOGGER.debug("No message loop task to cancel")
+
+    async def _run_reconnect_loop(self, start_future: asyncio.Future[None] | None) -> None:
+        """Run the MQTT loop."""
+        _LOGGER.info("Starting MQTT session")
+        self._diagnostics.increment("start_loop")
+        while True:
+            try:
+                self._connection_task = asyncio.create_task(self._run_connection(start_future))
+                await self._connection_task
+            except asyncio.CancelledError:
+                _LOGGER.debug("MQTT connection task cancelled")
+            except Exception:
+                # Exceptions are logged and handled in _run_connection.
+                # There is a special case for exceptions on startup where we return
+                # immediately. Otherwise, we let the reconnect loop retry with
+                # backoff when the reconnect loop is active.
+                if start_future and start_future.done() and start_future.exception():
+                    return
+
+            self._healthy = False
+            start_future = None
+            if self._stop:
+                _LOGGER.debug("MQTT session closed, stopping retry loop")
+                return
+            if not self._client_subscribed_topics and not self._listeners.keys():
+                _LOGGER.debug("MQTT session disconnected with no active subscriptions, deferring reconnect")
+                self._diagnostics.increment("reconnect_deferred")
+                while not self._stop and not self._client_subscribed_topics and not self._listeners.keys():
+                    await asyncio.sleep(0.1)
+                if self._stop:
+                    _LOGGER.debug("MQTT session closed while waiting for active subscriptions")
+                    return
+                self._backoff = MIN_BACKOFF_INTERVAL
+                continue
+            _LOGGER.info("MQTT session disconnected, retrying in %s seconds", self._backoff.total_seconds())
+            self._diagnostics.increment("reconnect_wait")
+            await asyncio.sleep(self._backoff.total_seconds())
+            self._backoff = min(self._backoff * BACKOFF_MULTIPLIER, MAX_BACKOFF_INTERVAL)
+
+    async def _run_connection(self, start_future: asyncio.Future[None] | None) -> None:
+        """Connect to the MQTT broker and listen for messages.
+
+        This is the primary connection loop for the MQTT session that is
+        long running and processes incoming messages. If the connection
+        is lost, this method will exit.
+        """
+        try:
+            with self._diagnostics.timer("connection"):
+                async with self._mqtt_client(self._params) as client:
+                    self._backoff = MIN_BACKOFF_INTERVAL
+                    self._healthy = True
+                    _LOGGER.info("MQTT Session connected.")
+                    if start_future and not start_future.done():
+                        start_future.set_result(None)
+
+                    _LOGGER.debug("Processing MQTT messages")
+                    async for message in client.messages:
+                        _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:
+            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):
+                _LOGGER.debug("MQTT loop was cancelled")
+                return
+            if start_future and not start_future.done():
+                _LOGGER.error("Uncaught error starting MQTT session: %s", err)
+                start_future.set_exception(err)
+            else:
+                _LOGGER.exception("Uncaught error during MQTT session: %s", err)
+            raise
+
+    @asynccontextmanager
+    async def _mqtt_client(self, params: MqttParams) -> aiomqtt.Client:
+        """Connect to the MQTT broker and listen for messages."""
+        _LOGGER.debug("Connecting to %s:%s for %s", params.host, params.port, params.username)
+        tls_params = None
+        if params.tls:
+            tls_params = TLSParameters(
+                ca_certs=certifi.where(),
+                cert_reqs=ssl.CERT_REQUIRED if params.verify_tls else ssl.CERT_NONE,
+            )
+        try:
+            async with aiomqtt.Client(
+                hostname=params.host,
+                port=params.port,
+                username=params.username,
+                password=params.password,
+                keepalive=int(CLIENT_KEEPALIVE.total_seconds()),
+                protocol=aiomqtt.ProtocolVersion.V5,
+                tls_params=tls_params,
+                timeout=params.timeout,
+                logger=_MQTT_LOGGER,
+            ) as client:
+                _LOGGER.debug("Connected to MQTT broker")
+                # Re-establish any existing subscriptions
+                async with self._client_lock:
+                    self._client = client
+                    for topic in self._client_subscribed_topics:
+                        self._diagnostics.increment("resubscribe")
+                        _LOGGER.debug("Re-establishing subscription to topic %s", redact_topic_name(topic))
+                        # TODO: If this fails it will break the whole connection. Make
+                        # this retry again in the background with backoff.
+                        await client.subscribe(topic)
+
+                yield client
+        finally:
+            async with self._client_lock:
+                self._client = None
+
+    async def subscribe(self, topic: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
+        """Subscribe to messages on the specified topic and invoke the callback for new messages.
+
+        The callback will be called with the message payload as a bytes object. The callback
+        should not block since it runs in the async loop. It should not raise any exceptions.
+
+        The returned callable unsubscribes from the topic when called, but will delay actual
+        unsubscription for the idle timeout period. If a new subscription comes in during the
+        timeout, the timer is cancelled and the subscription is reused.
+        """
+        _LOGGER.debug("Subscribing to topic %s", redact_topic_name(topic))
+
+        # If there is an idle timer for this topic, cancel it (reuse subscription)
+        if idle_timer := self._idle_timers.pop(topic, None):
+            self._diagnostics.increment("unsubscribe_idle_cancel")
+            idle_timer.cancel()
+            _LOGGER.debug("Cancelled idle timer for topic %s (reused subscription)", redact_topic_name(topic))
+
+        unsub = self._listeners.add_callback(topic, callback)
+
+        async with self._client_lock:
+            if topic not in self._client_subscribed_topics:
+                self._client_subscribed_topics.add(topic)
+                if self._client:
+                    _LOGGER.debug("Establishing subscription to topic %s", topic)
+                    try:
+                        with self._diagnostics.timer("subscribe"):
+                            await self._client.subscribe(topic)
+                    except MqttError as err:
+                        # Clean up the callback if subscription fails
+                        unsub()
+                        self._client_subscribed_topics.discard(topic)
+                        raise MqttSessionException(f"Error subscribing to topic: {err}") from err
+                else:
+                    self._diagnostics.increment("subscribe_pending")
+                    _LOGGER.debug("Client not connected, will establish subscription later")
+
+        def schedule_unsubscribe() -> None:
+            async def idle_unsubscribe():
+                try:
+                    await asyncio.sleep(self._topic_idle_timeout.total_seconds())
+                    # Only unsubscribe if there are no callbacks left for this topic
+                    if not self._listeners.get_callbacks(topic):
+                        async with self._client_lock:
+                            # Check again if we have listeners, in case a subscribe happened
+                            # while we were waiting for the lock or after we popped the timer.
+                            if self._listeners.get_callbacks(topic):
+                                _LOGGER.debug("Skipping unsubscribe for %s, new listeners added", topic)
+                                return
+
+                            self._idle_timers.pop(topic, None)
+                            self._client_subscribed_topics.discard(topic)
+
+                            if self._client:
+                                _LOGGER.debug("Idle timeout expired, unsubscribing from topic %s", topic)
+                                try:
+                                    await self._client.unsubscribe(topic)
+                                except MqttError as err:
+                                    _LOGGER.warning("Error unsubscribing from topic %s: %s", topic, err)
+                except asyncio.CancelledError:
+                    _LOGGER.debug("Idle unsubscribe for topic %s cancelled", topic)
+
+            # Start the idle timer task
+            task = asyncio.create_task(idle_unsubscribe())
+            self._idle_timers[topic] = task
+
+        def delayed_unsub():
+            self._diagnostics.increment("unsubscribe")
+            unsub()  # Remove the callback from CallbackMap
+            # If no more callbacks for this topic, start idle timer
+            if not self._listeners.get_callbacks(topic):
+                self._diagnostics.increment("unsubscribe_idle_start")
+                schedule_unsubscribe()
+            else:
+                _LOGGER.debug("Unsubscribing topic %s, still have active callbacks", topic)
+
+        return delayed_unsub
+
+    async def publish(self, topic: str, message: bytes) -> None:
+        """Publish a message on the topic."""
+        _LOGGER.debug("Sending message to topic %s: %s", topic, message)
+        client: aiomqtt.Client
+        async with self._client_lock:
+            if self._client is None:
+                raise MqttSessionException("Could not publish message, MQTT client not connected")
+            client = self._client
+        try:
+            with self._diagnostics.timer("publish"):
+                await client.publish(topic, message)
+        except MqttError as err:
+            raise MqttSessionException(f"Error publishing message: {err}") from err
+
+
+class LazyMqttSession(MqttSession):
+    """An MQTT session that is started on first attempt to subscribe.
+
+    This is a wrapper around an existing MqttSession that will only start
+    the underlying session when the first attempt to subscribe or publish
+    is made.
+    """
+
+    def __init__(self, session: RoborockMqttSession, diagnostics: Diagnostics) -> None:
+        """Initialize the lazy session with an existing session."""
+        self._lock = asyncio.Lock()
+        self._started = False
+        self._session = session
+        self._diagnostics = diagnostics
+
+    @property
+    def connected(self) -> bool:
+        """True if the session is connected to the broker."""
+        return self._session.connected
+
+    @property
+    def health_manager(self) -> HealthManager:
+        """Return the health manager for the session."""
+        return self._session.health_manager
+
+    async def _maybe_start(self) -> None:
+        """Start the MQTT session if not already started."""
+        async with self._lock:
+            if not self._started:
+                self._diagnostics.increment("start")
+                await self._session.start()
+                self._started = True
+
+    async def subscribe(self, device_id: str, callback: Callable[[bytes], None]) -> Callable[[], None]:
+        """Invoke the callback when messages are received on the topic.
+
+        The returned callable unsubscribes from the topic when called.
+        """
+        await self._maybe_start()
+        return await self._session.subscribe(device_id, callback)
+
+    async def publish(self, topic: str, message: bytes) -> None:
+        """Publish a message on the specified topic.
+
+        This will raise an exception if the message could not be sent.
+        """
+        await self._maybe_start()
+        return await self._session.publish(topic, message)
+
+    async def close(self) -> None:
+        """Cancels the mqtt loop.
+
+        This will close the underlying session and will not allow it to be
+        restarted again.
+        """
+        await self._session.close()
+
+    async def restart(self) -> None:
+        """Force the session to disconnect and reconnect."""
+        await self._session.restart()
+
+
+async def create_mqtt_session(params: MqttParams) -> MqttSession:
+    """Create an MQTT session.
+
+    This function is a factory for creating an MQTT session. This will
+    raise an exception if initial attempt to connect fails. Once connected,
+    the session will retry connecting on failure in the background.
+    """
+    session = RoborockMqttSession(params)
+    await session.start()
+    return session
+
+
+async def create_lazy_mqtt_session(params: MqttParams) -> MqttSession:
+    """Create a lazy MQTT session.
+
+    This function is a factory for creating an MQTT session that will
+    only connect when the first attempt to subscribe or publish is made.
+    """
+    return LazyMqttSession(RoborockMqttSession(params), diagnostics=params.diagnostics.subkey("lazy_mqtt"))