Fix temporary chat selector

browser_utils/initialization/core.py

@@ -560,33 +560,134 @@ async def signal_camoufox_shutdown() -> None:  # pragma: no cover
         )
 
 
-async def enable_temporary_chat_mode(page: AsyncPage) -> None:  # pragma: no cover
-    """
-    Check and enable "Temporary chat" mode in the AI Studio interface.
-    This is an independent UI operation and should be called after the page is fully stable.
+async def _is_temporary_chat_active(page: AsyncPage) -> bool:
+    """Best-effort check for whether Temporary chat is currently enabled.
+
+    Supports both the legacy toggle-class signal and the newer UI variants:
+    - menu item checkmark (`data-test-incognito-checkmark`)
+    - header indicator (`ms-incognito-mode-indicator`)
     """
+
     try:
-        incognito_button_locator = page.locator(
-            'button[aria-label="Temporary chat toggle"]'
+        # Newer UI: header indicator appears only when Temporary chat is active.
+        indicator_locator = page.locator(
+            "ms-incognito-mode-indicator [data-test-id='main-text'], "
+            "ms-incognito-mode-indicator .main-text"
         )
+        indicator_count = await indicator_locator.count()
+        for i in range(indicator_count):
+            try:
+                indicator_item = indicator_locator.nth(i)
+                if not await indicator_item.is_visible(timeout=500):
+                    continue
+                indicator_text = (await indicator_item.inner_text()).strip().lower()
+                if indicator_text == "temporary chat":
+                    return True
+            except Exception:
+                continue
+
+        # Menu button variants (legacy and gray-release UI).
+        toggle_locator = page.locator(
+            "button[data-test-incognito-toggle], "
+            'button[aria-label="Temporary chat toggle"], '
+            'button[aria-label="Toggle temporary chat"]'
+        )
+        if await toggle_locator.count() > 0:
+            # New UI: active state is represented by a checkmark inside the menu item.
+            try:
+                checkmark_locator = page.locator(
+                    "button[data-test-incognito-toggle] [data-test-incognito-checkmark], "
+                    'button[aria-label="Temporary chat toggle"] [data-test-incognito-checkmark], '
+                    'button[aria-label="Toggle temporary chat"] [data-test-incognito-checkmark]'
+                )
+                if (
+                    await checkmark_locator.count() > 0
+                ):
+                    return True
+            except Exception:
+                pass
+
+            # Legacy/alternative signals.
+            for attr_name in ("aria-pressed", "aria-checked"):
+                try:
+                    attr_value = await toggle_locator.get_attribute(attr_name)
+                    if attr_value and attr_value.lower() == "true":
+                        return True
+                except Exception:
+                    pass
+
+            try:
+                button_classes = await toggle_locator.get_attribute("class") or ""
+                if "ms-button-active" in button_classes:
+                    return True
+            except Exception:
+                pass
+
+        return False
+    except Exception:
+        return False
+
 
-        await incognito_button_locator.wait_for(state="visible", timeout=10000)
+async def enable_temporary_chat_mode(page: AsyncPage) -> bool:  # pragma: no cover
+    """
+    Check and enable "Temporary chat" mode in the AI Studio interface.
+    Supports both direct UI visibility and collapsed menu visibility.
+    """
+    incognito_selector = (
+        "button[data-test-incognito-toggle], "
+        'button[aria-label="Temporary chat toggle"], '
+        'button[aria-label="Toggle temporary chat"]'
+    )
+    menu_trigger_selector = 'button[aria-label="View more actions"]'
 
-        button_classes = await incognito_button_locator.get_attribute("class")
+    incognito_locator = page.locator(incognito_selector)
+    menu_trigger = page.locator(menu_trigger_selector)
 
-        if button_classes and "ms-button-active" in button_classes:
+    async def _close_menu_if_needed() -> None:
+        """Best-effort menu cleanup that tolerates UI variants without a menu trigger."""
+        try:
+            if await menu_trigger.count() == 0:
+                return
+            if await menu_trigger.is_visible():
+                if await menu_trigger.get_attribute("aria-expanded") == "true":
+                    logger.debug("[UI] Closing menu to restore UI state")
+                    await page.keyboard.press("Escape")
+        except Exception:
+            pass
+
+    try:
+        # Fast path
+        logger.debug("[UI] Searching for temporary chat button (Fast path)")
+        try:
+            await incognito_locator.wait_for(state="visible", timeout=3000)
+        except Exception:
+            # Fallback
+            logger.debug("[UI] Button not visible, attempting to open menu")
+            if await menu_trigger.is_visible():
+                await menu_trigger.click()
+                # Wait for the menu item to appear
+                await incognito_locator.wait_for(state="visible", timeout=5000)
+            else:
+                logger.warning("[UI] Neither button nor menu trigger found")
+                return False
+
+        if await _is_temporary_chat_active(page):
             logger.debug("[UI] Temporary chat mode already active")
         else:
-            await incognito_button_locator.click(timeout=5000, force=True)
+            logger.debug("[UI] Enabling temporary chat mode")
+            await incognito_locator.click(timeout=5000, force=True)
             await asyncio.sleep(1)
 
-            updated_classes = await incognito_button_locator.get_attribute("class")
-            if updated_classes and "ms-button-active" in updated_classes:
-                logger.debug("[UI] Temporary chat mode enabled")
-            else:
-                logger.warning("[UI] Failed to enable temporary chat mode")
+        enabled = await _is_temporary_chat_active(page)
+
+        # Recovery: Close menu if still expanded
+        await _close_menu_if_needed()
+        return enabled
 
     except asyncio.CancelledError:
         raise
     except Exception as e:
         logger.warning(f"[UI] Error in temporary chat mode: {e}")
+        # Final safety attempt to clear any stuck UI
+        await _close_menu_if_needed()
+        return await _is_temporary_chat_active(page)

browser_utils/models/switcher.py

@@ -160,41 +160,20 @@ async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) ->
             if page_display_match:
                 try:
                     logger.debug("[Model] Re-enabling temporary chat mode...")
-                    incognito_button_locator = page.locator(
-                        'button[aria-label="Temporary chat toggle"]'
-                    )
-
-                    await incognito_button_locator.wait_for(
-                        state="visible", timeout=5000
-                    )
+                    from browser_utils.initialization import enable_temporary_chat_mode
 
-                    button_classes = await incognito_button_locator.get_attribute(
-                        "class"
-                    )
-
-                    if button_classes and "ms-button-active" in button_classes:
-                        logger.debug("[Model] Temporary chat mode already active")
+                    enabled = await enable_temporary_chat_mode(page)
+                    if enabled:
+                        logger.debug("[Model] Temporary chat mode enabled or already active")
                     else:
-                        logger.debug("[Model] Clicking to open temporary chat mode...")
-                        await incognito_button_locator.click(timeout=3000)
-                        await asyncio.sleep(0.5)
-
-                        updated_classes = await incognito_button_locator.get_attribute(
-                            "class"
+                        logger.warning(
+                            "[Model] Temporary chat mode state verification failed"
                         )
-                        if updated_classes and "ms-button-active" in updated_classes:
-                            logger.debug("[Model] Temporary chat mode enabled")
-                        else:
-                            logger.warning(
-                                "Temporary chat mode state verification failed after click, may not have opened successfully."
-                            )
 
                 except asyncio.CancelledError:
                     raise
                 except Exception as e:
-                    logger.warning(
-                        f"Failed to re-enable temporary chat mode after model switching: {e}"
-                    )
+                    logger.warning(f"Failed to re-enable temporary chat mode after model switching: {e}")
 
                 # Invalidate function calling cache on model switch
                 try: