fix(state): wire up wait_for_stable_ui dead config field

mobilerun/agent/droid/droid_agent.py

@@ -607,6 +607,7 @@ async def start_handler(
                 stealth=False,
                 vision_enabled=vision_enabled,
                 vision_resize_policy=vision_resize_policy,
+                wait_for_stable_ui=self.config.agent.wait_for_stable_ui,
             )
         elif is_ios:
             self.state_provider = IOSStateProvider(
@@ -626,6 +627,7 @@ async def start_handler(
                 stealth=stealth_enabled,
                 vision_enabled=vision_enabled,
                 vision_resize_policy=vision_resize_policy,
+                wait_for_stable_ui=self.config.agent.wait_for_stable_ui,
             )
 
         # ── 3. Build tool registry ────────────────────────────────────

mobilerun/tools/ui/provider.py

@@ -38,13 +38,54 @@
 # With the schedule above, this fires after ~11s (1+2+3+5).
 _RECOVERY_AFTER_ATTEMPT = 5
 
+# Max number of re-fetches used to confirm the a11y tree has stopped
+# changing. Bounds the extra latency ``wait_for_stable_ui`` can add per
+# get_state() call to roughly max_stability_checks * wait_for_stable_ui.
+_MAX_STABILITY_CHECKS = 3
+
+
+async def _wait_for_stable_ui_tree(
+    fetch: Callable[[], Awaitable[Dict[str, Any]]],
+    initial: Dict[str, Any],
+    wait_for_stable_ui: float,
+    max_checks: int,
+) -> Dict[str, Any]:
+    """Re-fetch state until the a11y tree stops changing between fetches.
+
+    A well-formed dump can still be stale: e.g. a WebView's Chromium
+    accessibility bridge lagging behind a client-side route change reports
+    the previous screen with no error and all required keys present, so
+    ``fetch_state_with_retry``'s error/missing-key checks accept it
+    immediately. Polling until two consecutive fetches agree catches that
+    case instead of handing the agent a screen it already navigated away
+    from.
+    """
+    current = initial
+    for _ in range(max_checks):
+        await asyncio.sleep(wait_for_stable_ui)
+        try:
+            next_data = await fetch()
+        except DeviceDisconnectedError:
+            raise
+        except Exception as e:
+            logger.debug(f"UI-stability re-fetch failed, using last state: {e}")
+            break
+        if "error" in next_data or "a11y_tree" not in next_data:
+            break
+        if next_data["a11y_tree"] == current["a11y_tree"]:
+            return next_data
+        current = next_data
+    return current
+
 
 async def fetch_state_with_retry(
     fetch: Callable[[], Awaitable[Dict[str, Any]]],
     recovery: Optional[Callable[[], Awaitable[None]]] = None,
     max_retries: int = _MAX_RETRIES,
     retry_delays: Optional[List[float]] = None,
     recovery_after: int = _RECOVERY_AFTER_ATTEMPT,
+    wait_for_stable_ui: float = 0.0,
+    max_stability_checks: int = _MAX_STABILITY_CHECKS,
 ) -> Dict[str, Any]:
     """Fetch raw device state with retries, backoff, and mid-retry recovery.
 
@@ -56,6 +97,12 @@ async def fetch_state_with_retry(
         retry_delays: Per-attempt delays. If shorter than max_retries - 1,
             the last value is reused for remaining delays.
         recovery_after: Trigger *recovery* after this many failures.
+        wait_for_stable_ui: Seconds to wait between re-fetches when
+            confirming the a11y tree has stopped changing. 0 disables the
+            stability check (the first well-formed dump is trusted as-is).
+        max_stability_checks: Cap on re-fetches while waiting for the tree
+            to stabilize, so a screen that never settles (e.g. an animation
+            loop) doesn't hang get_state() indefinitely.
 
     Returns:
         The raw state dict (guaranteed to contain ``a11y_tree``,
@@ -90,6 +137,11 @@ async def fetch_state_with_retry(
             if missing:
                 raise Exception(f"Missing data in state: {', '.join(missing)}")
 
+            if wait_for_stable_ui > 0:
+                combined_data = await _wait_for_stable_ui_tree(
+                    fetch, combined_data, wait_for_stable_ui, max_stability_checks
+                )
+
             return combined_data
 
         except DeviceDisconnectedError:
@@ -190,7 +242,10 @@ class AndroidStateProvider(StateProvider):
 
     Includes retry logic with exponential backoff and mid-retry recovery
     (accessibility service restart) for robustness against intermittent
-    Portal/a11y failures.
+    Portal/a11y failures, plus an optional post-fetch stability check
+    (``wait_for_stable_ui``) that re-polls until the a11y tree stops
+    changing, guarding against well-formed but stale dumps (e.g. a WebView
+    a11y bridge lagging behind a client-side route change).
     """
 
     supported = {"element_index", "convert_point"}
@@ -205,11 +260,13 @@ def __init__(
         ui_cls: "type[UIState] | None" = None,
         vision_enabled: bool = False,
         vision_resize_policy: Any = None,
+        wait_for_stable_ui: float = 0.0,
     ) -> None:
         super().__init__(driver)
         self.tree_filter = tree_filter
         self.tree_formatter = tree_formatter
         self.use_normalized = use_normalized
+        self.wait_for_stable_ui = wait_for_stable_ui
         self._ui_cls = ui_cls or (StealthUIState if stealth else UIState)
         self.vision_enabled = vision_enabled
         # Resolves the exact screenshot dims the active vision model grounds on
@@ -291,6 +348,7 @@ async def get_state(self) -> UIState:
         combined_data = await fetch_state_with_retry(
             fetch=self.driver.get_ui_tree,
             recovery=self._recover_portal,
+            wait_for_stable_ui=self.wait_for_stable_ui,
         )
 
         device_context = combined_data["device_context"]