fix: harden Q10 map rendering state

Author: tubededentifrice

roborock/cli.py

@@ -603,6 +603,7 @@ async def maps(ctx, device_id: str):
 async def _await_q10_map_push(
     properties: Q10PropertiesApi,
     predicate: Callable[[], bool],
+    generation: Callable[[], int],
     *,
     timeout: float = _Q10_MAP_PUSH_TIMEOUT,
     allow_cached_on_timeout: bool = False,
@@ -617,9 +618,10 @@ async def _await_q10_map_push(
     """
     loop = asyncio.get_running_loop()
     updated: asyncio.Future[None] = loop.create_future()
+    initial_generation = generation()
 
     def on_update() -> None:
-        if predicate() and not updated.done():
+        if generation() != initial_generation and predicate() and not updated.done():
             updated.set_result(None)
 
     unsub = properties.map.add_update_listener(on_update)
@@ -648,6 +650,7 @@ async def map_image(ctx, device_id: str, output_file: str):
         await _await_q10_map_push(
             properties,
             lambda: properties.map.image_content is not None,
+            lambda: properties.map.map_generation,
             allow_cached_on_timeout=True,
         )
         image_content = properties.map.image_content
@@ -706,7 +709,11 @@ async def q10_position(ctx, device_id: str, include_path: bool):
         click.echo("Feature not supported by device")
         return
     properties = device.b01_q10_properties
-    got_trace = await _await_q10_map_push(properties, lambda: bool(properties.map.path))
+    got_trace = await _await_q10_map_push(
+        properties,
+        lambda: bool(properties.map.path),
+        lambda: properties.map.trace_generation,
+    )
     if not got_trace:
         click.echo("No live trace available (the robot only reports position while cleaning).")
         return
@@ -871,6 +878,7 @@ async def rooms(ctx, device_id: str):
         await _await_q10_map_push(
             properties,
             lambda: properties.map.image_content is not None,
+            lambda: properties.map.map_generation,
             allow_cached_on_timeout=True,
         )
         click.echo(dump_json({room.id: room.name for room in properties.map.rooms}))

roborock/devices/traits/b01/q10/__init__.py

@@ -40,6 +40,7 @@
 ]
 
 _LOGGER = logging.getLogger(__name__)
+_MAP_LIST_REQUEST_TIMEOUT = 30.0
 
 
 def _map_id_from_list_response(response: Any) -> int | str | None:
@@ -134,7 +135,9 @@ def __init__(self, channel: B01Q10Channel) -> None:
             self._map_dps,
         ]
         self._subscribe_task: asyncio.Task[None] | None = None
-        self._map_list_requested = False
+        self._map_request_lock = asyncio.Lock()
+        self._map_list_request_token: object | None = None
+        self._map_list_requested_at: float | None = None
 
     async def start(self) -> None:
         """Start any necessary subscriptions for the trait."""
@@ -166,15 +169,27 @@ async def request_map(self) -> None:
         matching ``get`` command; the resulting protocol-301 map packet is
         routed to :attr:`map`.
         """
-        self._map_list_requested = True
-        try:
-            await self.command.send(
-                B01_Q10_DP.COMMON,
-                {str(B01_Q10_DP.MULTI_MAP.code): {"op": "list"}},
-            )
-        except RoborockException:
-            self._map_list_requested = False
-            raise
+        async with self._map_request_lock:
+            now = asyncio.get_running_loop().time()
+            if (
+                self._map_list_request_token is not None
+                and self._map_list_requested_at is not None
+                and now - self._map_list_requested_at < _MAP_LIST_REQUEST_TIMEOUT
+            ):
+                return
+            token = object()
+            self._map_list_request_token = token
+            self._map_list_requested_at = now
+            try:
+                await self.command.send(
+                    B01_Q10_DP.COMMON,
+                    {str(B01_Q10_DP.MULTI_MAP.code): {"op": "list"}},
+                )
+            except BaseException:
+                if self._map_list_request_token is token:
+                    self._map_list_request_token = None
+                    self._map_list_requested_at = None
+                raise
 
     async def _subscribe_loop(self) -> None:
         """Persistent loop dispatching decoded messages to the read-model traits."""
@@ -197,18 +212,25 @@ async def _handle_message(self, message: Q10Message) -> None:
             _LOGGER.debug("Received Q10 status update: %s", message.dps)
             # Notify all read-model traits about the new message; each trait
             # only updates the fields that it is responsible for.
-            for trait in self._updatable_traits:
-                trait.update_from_dps(message.dps)
+            self.map.begin_source_update()
+            try:
+                for trait in self._updatable_traits:
+                    trait.update_from_dps(message.dps)
+            finally:
+                self.map.end_source_update()
             await self._request_map_from_list_response(message.dps)
 
     async def _request_map_from_list_response(self, decoded_dps: dict[B01_Q10_DP, Any]) -> None:
         """Request map content after receiving our pending map-list response."""
         response = decoded_dps.get(B01_Q10_DP.MULTI_MAP)
-        if not self._map_list_requested or not isinstance(response, dict) or response.get("op") != "list":
+        if self._map_list_request_token is None or not isinstance(response, dict) or response.get("op") != "list":
             return
 
-        self._map_list_requested = False
+        token = self._map_list_request_token
         if (map_id := _map_id_from_list_response(response)) is None:
+            if self._map_list_request_token is token:
+                self._map_list_request_token = None
+                self._map_list_requested_at = None
             _LOGGER.debug("Q10 map list response did not contain a usable map ID")
             return
 
@@ -225,6 +247,10 @@ async def _request_map_from_list_response(self, decoded_dps: dict[B01_Q10_DP, An
         except RoborockException as ex:
             # A failed follow-up must not kill the persistent subscribe loop.
             _LOGGER.debug("Failed to request Q10 map content: %s", ex)
+        finally:
+            if self._map_list_request_token is token:
+                self._map_list_request_token = None
+                self._map_list_requested_at = None
 
 
 def create(channel: B01Q10Channel) -> Q10PropertiesApi:

roborock/devices/traits/b01/q10/map.py

@@ -37,6 +37,7 @@
 from .status import StatusTrait
 
 _LOGGER = logging.getLogger(__name__)
+_DOCKED_STATES = {YXDeviceState.CHARGING, YXDeviceState.EMPTYING_THE_BIN}
 
 
 @dataclass
@@ -91,10 +92,14 @@ def __init__(
         self._config = map_parser_config or B01Q10MapParserConfig()
         self._map_dps = map_dps
         self._status = status or StatusTrait()
-        self._robot_at_dock = self._status.status == YXDeviceState.CHARGING
+        self._robot_at_dock = self._status.status in _DOCKED_STATES
         self._map_packet: Q10MapPacket | None = None
         self._trace_packet: Q10TracePacket | None = None
         self._image_content: bytes | None = None
+        self._map_generation = 0
+        self._trace_generation = 0
+        self._source_update_depth = 0
+        self._source_update_pending = False
         self._map_dps.add_update_listener(self._map_dps_updated)
         self._status.add_update_listener(self._status_updated)
 
@@ -123,35 +128,71 @@ def robot_heading(self) -> int | None:
         """Current heading for orienting a robot marker on a caller-rendered map."""
         return self._trace_packet.heading if self._trace_packet else None
 
+    @property
+    def map_generation(self) -> int:
+        """Number of map packets received by this trait."""
+        return self._map_generation
+
+    @property
+    def trace_generation(self) -> int:
+        """Number of trace packets received by this trait."""
+        return self._trace_generation
+
     def update_from_map_packet(self, packet: Q10MapPacket) -> None:
         """Store a map-protocol update and render the latest sources."""
         self._map_packet = packet
+        self._map_generation += 1
         self._render()
         self._notify_update()
 
     def update_from_trace_packet(self, packet: Q10TracePacket) -> None:
         """Store a trace-protocol update and render the latest sources."""
-        self._trace_packet = packet
+        # A late packet from the completed clean cannot move a robot that the
+        # status stream already confirmed is docked.
+        self._trace_packet = None if self._robot_at_dock else packet
+        self._trace_generation += 1
         self._render()
         self._notify_update()
 
-    def _map_dps_updated(self) -> None:
-        """Render after the low-level map DPS source changes."""
-        if self._map_packet is None:
+    def begin_source_update(self) -> None:
+        """Defer dependent rendering while one DPS message is applied."""
+        self._source_update_depth += 1
+
+    def end_source_update(self) -> None:
+        """Render once after all traits consumed the same DPS message."""
+        self._source_update_depth -= 1
+        if self._source_update_depth == 0 and self._source_update_pending:
+            self._source_update_pending = False
+            self._render_and_notify()
+
+    def _source_updated(self, *, notify_without_map: bool = False) -> None:
+        """Render now, or defer until the current DPS update is complete."""
+        if self._map_packet is None and not notify_without_map:
             return
+        if self._source_update_depth:
+            self._source_update_pending = True
+            return
+        self._render_and_notify()
+
+    def _render_and_notify(self) -> None:
+        """Recompose the current map and publish one update."""
         self._render()
         self._notify_update()
 
+    def _map_dps_updated(self) -> None:
+        """Render after the low-level map DPS source changes."""
+        self._source_updated()
+
     def _status_updated(self) -> None:
         """Render only when the status changes whether the robot is docked."""
-        robot_at_dock = self._status.status == YXDeviceState.CHARGING
+        robot_at_dock = self._status.status in _DOCKED_STATES
         if robot_at_dock == self._robot_at_dock:
             return
         self._robot_at_dock = robot_at_dock
-        if self._map_packet is None:
-            return
-        self._render()
-        self._notify_update()
+        trace_cleared = robot_at_dock and self._trace_packet is not None
+        if robot_at_dock:
+            self._trace_packet = None
+        self._source_updated(notify_without_map=trace_cleared)
 
     def _render(self) -> None:
         """Render the required map with the latest optional trace and overlays."""

roborock/map/map_parser.py

@@ -2,6 +2,7 @@
 
 import io
 import logging
+import threading
 from dataclasses import dataclass, field
 
 from vacuum_map_parser_base.config.color import Color, ColorsPalette, SupportedColor
@@ -105,10 +106,18 @@ def parse(self, map_bytes: bytes) -> ParsedMapData | None:
 class _AdjacencyAwareRoborockImageParser(RoborockImageParser):
     """Apply the shared adjacency color policy to V1 room cells."""
 
-    def __init__(self, palette: ColorsPalette, image_config: ImageConfig) -> None:
+    def __init__(
+        self,
+        palette: ColorsPalette,
+        image_config: ImageConfig,
+        *,
+        recolor_rooms: bool = True,
+    ) -> None:
         super().__init__(palette, image_config)
         self._room_palette = palette
         self._base_room_colors = palette.cached_room_colors.copy()
+        self._recolor_rooms = recolor_rooms
+        self._palette_lock = threading.Lock()
 
     def parse(
         self,
@@ -119,19 +128,22 @@ def parse(
         removed_map: set[int] | None = None,
     ):
         """Assign non-conflicting room colors before the V1 image pass."""
-        self._room_palette.cached_room_colors.clear()
-        self._room_palette.cached_room_colors.update(self._base_room_colors)
+        with self._palette_lock:
+            self._room_palette.cached_room_colors.clear()
+            self._room_palette.cached_room_colors.update(self._base_room_colors)
+
+            if self._recolor_rooms:
 
-        def room_id(value: int) -> int | None:
-            if value in (self.MAP_OUTSIDE, self.MAP_WALL, self.MAP_INSIDE, self.MAP_SCAN):
-                return None
-            return self._get_room_number(value) if value & 0x07 == 0x07 else None
+                def room_id(value: int) -> int | None:
+                    if value in (self.MAP_OUTSIDE, self.MAP_WALL, self.MAP_INSIDE, self.MAP_SCAN):
+                        return None
+                    return self._get_room_number(value) if value & 0x07 == 0x07 else None
 
-        room_colors = adjacency_aware_room_colors(raw_data, width, self._room_palette, room_id)
-        for number, color in room_colors.items():
-            self._room_palette.cached_room_colors[number] = color
-            self._room_palette.cached_room_colors[str(number)] = color
-        return super().parse(raw_data, width, height, carpet_map, removed_map)
+                room_colors = adjacency_aware_room_colors(raw_data, width, self._room_palette, room_id)
+                for number, color in room_colors.items():
+                    self._room_palette.cached_room_colors[number] = color
+                    self._room_palette.cached_room_colors[str(number)] = color
+            return super().parse(raw_data, width, height, carpet_map, removed_map)
 
 
 def _create_map_data_parser(config: MapParserConfig) -> RoborockMapDataParser:
@@ -144,7 +156,11 @@ def _create_map_data_parser(config: MapParserConfig) -> RoborockMapDataParser:
         image_config,
         [],
     )
-    parser._image_parser = _AdjacencyAwareRoborockImageParser(palette, image_config)
+    parser._image_parser = _AdjacencyAwareRoborockImageParser(
+        palette,
+        image_config,
+        recolor_rooms=config.show_rooms,
+    )
     return parser
 
 
@@ -184,7 +200,7 @@ def _create_rendering_components(
         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)}
+        room_colors = {str(room_id): (0, 0, 0, 0) for room_id in map(int, ColorsPalette.ROOM_COLORS)}
 
     return (
         ColorsPalette(color_dicts, room_colors),