refactor: Centralize UI selectors for CDK overlay, thinking mode, and budget toggles.

Author: NikkeTryHard

browser_utils/initialization/core.py

@@ -23,6 +23,7 @@
 from config import (
     AI_STUDIO_URL_PATTERN,
     INPUT_SELECTOR,
+    MODEL_NAME_SELECTOR,
     USER_INPUT_END_MARKER_SERVER,
     USER_INPUT_START_MARKER_SERVER,
 )
@@ -322,7 +323,7 @@ async def initialize_page_logic(  # pragma: no cover
             )
             logger.info("-> 核心输入区域可见。")
 
-            model_name_locator = found_page.locator('[data-test-id="model-name"]')
+            model_name_locator = found_page.locator(MODEL_NAME_SELECTOR)
             try:
                 model_name_on_page = await model_name_locator.first.inner_text(
                     timeout=5000

browser_utils/models/startup.py

@@ -9,7 +9,7 @@
 from playwright.async_api import Page as AsyncPage
 from playwright.async_api import expect as expect_async
 
-from config import INPUT_SELECTOR
+from config import INPUT_SELECTOR, MODEL_NAME_SELECTOR
 
 from .ui_state import _verify_and_apply_ui_state, _verify_ui_state_settings
 
@@ -166,7 +166,7 @@ async def _set_model_from_page_display(page: AsyncPage, set_storage: bool = Fals
 
     try:
         logger.info("   尝试从页面显示元素读取当前模型名称...")
-        model_name_locator = page.locator('[data-test-id="model-name"]')
+        model_name_locator = page.locator(MODEL_NAME_SELECTOR)
         displayed_model_name_from_page_raw = await model_name_locator.first.inner_text(
             timeout=7000
         )

browser_utils/models/switcher.py

@@ -11,7 +11,7 @@
 from playwright.async_api import Page as AsyncPage
 from playwright.async_api import expect as expect_async
 
-from config import AI_STUDIO_URL_PATTERN, INPUT_SELECTOR
+from config import AI_STUDIO_URL_PATTERN, INPUT_SELECTOR, MODEL_NAME_SELECTOR
 
 from .ui_state import _verify_and_apply_ui_state
 
@@ -129,7 +129,7 @@ async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) ->
                         break
 
             try:
-                model_name_locator = page.locator('[data-test-id="model-name"]')
+                model_name_locator = page.locator(MODEL_NAME_SELECTOR)
                 actual_displayed_model_id_on_page_raw = (
                     await model_name_locator.first.inner_text(timeout=5000)
                 )
@@ -209,7 +209,7 @@ async def switch_ai_studio_model(page: AsyncPage, model_id: str, req_id: str) ->
         current_displayed_name_for_revert_stripped = "无法读取"
 
         try:
-            model_name_locator_revert = page.locator('[data-test-id="model-name"]')
+            model_name_locator_revert = page.locator(MODEL_NAME_SELECTOR)
             current_displayed_name_for_revert_raw = (
                 await model_name_locator_revert.first.inner_text(timeout=5000)
             )

browser_utils/operations_modules/interactions.py

@@ -16,6 +16,7 @@
 
 from browser_utils.operations_modules.errors import save_error_snapshot
 from config import (
+    CHAT_TURN_SELECTOR,
     CLICK_TIMEOUT_MS,
     DEBUG_LOGS_ENABLED,
     INITIAL_WAIT_MS_BEFORE_POLLING,
@@ -79,7 +80,7 @@ async def get_response_via_edit_button(
 ) -> Optional[str]:
     """通过编辑按钮获取响应"""
     logger.info(" (Helper) 尝试通过编辑按钮获取响应...")
-    last_message_container = page.locator("ms-chat-turn").last
+    last_message_container = page.locator(CHAT_TURN_SELECTOR).last
     edit_button = last_message_container.get_by_label("Edit")
     finish_edit_button = last_message_container.get_by_label("Stop editing")
     autosize_textarea_locator = last_message_container.locator("ms-autosize-textarea")
@@ -239,7 +240,7 @@ async def get_response_via_copy_button(
 ) -> Optional[str]:
     """通过复制按钮获取响应"""
     logger.info(" (Helper) 尝试通过复制按钮获取响应...")
-    last_message_container = page.locator("ms-chat-turn").last
+    last_message_container = page.locator(CHAT_TURN_SELECTOR).last
     more_options_button = last_message_container.get_by_label("Open options")
     copy_markdown_button = page.get_by_role("menuitem", name="Copy markdown")
 

browser_utils/page_controller_modules/input.py

@@ -8,6 +8,7 @@
 
 from browser_utils.operations import save_error_snapshot
 from config import (
+    CDK_OVERLAY_CONTAINER_SELECTOR,
     PROMPT_TEXTAREA_SELECTOR,
     RESPONSE_CONTAINER_SELECTOR,
     SUBMIT_BUTTON_SELECTOR,
@@ -188,7 +189,7 @@ async def _open_upload_menu_and_choose_file(self, files_list: List[str]) -> bool
             trigger = self.page.locator(UPLOAD_BUTTON_SELECTOR).first
             await expect_async(trigger).to_be_visible(timeout=3000)
             await trigger.click()
-            menu_container = self.page.locator("div.cdk-overlay-container")
+            menu_container = self.page.locator(CDK_OVERLAY_CONTAINER_SELECTOR)
             # 等待菜单显示
             try:
                 await expect_async(
@@ -261,7 +262,7 @@ async def _open_upload_menu_and_choose_file(self, files_list: List[str]) -> bool
     async def _handle_post_upload_dialog(self):
         """处理上传后可能出现的授权/版权确认对话框，优先点击同意类按钮，不主动关闭重要对话框。"""
         try:
-            overlay_container = self.page.locator("div.cdk-overlay-container")
+            overlay_container = self.page.locator(CDK_OVERLAY_CONTAINER_SELECTOR)
             if await overlay_container.count() == 0:
                 return
 

browser_utils/page_controller_modules/thinking.py

@@ -15,9 +15,13 @@
     ENABLE_THINKING_MODE_TOGGLE_SELECTOR,
     SET_THINKING_BUDGET_TOGGLE_SELECTOR,
     THINKING_BUDGET_INPUT_SELECTOR,
+    THINKING_BUDGET_TOGGLE_OLD_ROOT_SELECTOR,
+    THINKING_BUDGET_TOGGLE_PARENT_SELECTOR,
     THINKING_LEVEL_OPTION_HIGH_SELECTOR,
     THINKING_LEVEL_OPTION_LOW_SELECTOR,
     THINKING_LEVEL_SELECT_SELECTOR,
+    THINKING_MODE_TOGGLE_OLD_ROOT_SELECTOR,
+    THINKING_MODE_TOGGLE_PARENT_SELECTOR,
 )
 from models import ClientDisconnectedError
 
@@ -469,14 +473,14 @@ async def _control_thinking_mode_toggle(
                     try:
                         # 新版UI: 尝试直接点击带 aria-label 的开关父容器
                         alt_toggle = self.page.locator(
-                            'mat-slide-toggle:has(button[aria-label="Toggle thinking mode"])'
+                            THINKING_MODE_TOGGLE_PARENT_SELECTOR
                         )
                         if await alt_toggle.count() > 0:
                             await alt_toggle.click(timeout=CLICK_TIMEOUT_MS)
                         else:
                             # 旧版UI回退: data-test-toggle
                             root = self.page.locator(
-                                'mat-slide-toggle[data-test-toggle="enable-thinking"]'
+                                THINKING_MODE_TOGGLE_OLD_ROOT_SELECTOR
                             )
                             label = root.locator("label.mdc-label")
                             await expect_async(label).to_be_visible(timeout=2000)
@@ -563,14 +567,14 @@ async def _control_thinking_budget_toggle(
                     try:
                         # 新版UI: 尝试直接点击带 aria-label 的开关父容器
                         alt_toggle = self.page.locator(
-                            'mat-slide-toggle:has(button[aria-label="Toggle thinking budget between auto and manual"])'
+                            THINKING_BUDGET_TOGGLE_PARENT_SELECTOR
                         )
                         if await alt_toggle.count() > 0:
                             await alt_toggle.click(timeout=CLICK_TIMEOUT_MS)
                         else:
                             # 旧版UI回退: data-test-toggle
                             root = self.page.locator(
-                                'mat-slide-toggle[data-test-toggle="manual-budget"]'
+                                THINKING_BUDGET_TOGGLE_OLD_ROOT_SELECTOR
                             )
                             label = root.locator("label.mdc-label")
                             await expect_async(label).to_be_visible(timeout=2000)

config/__init__.py

@@ -83,6 +83,15 @@
     "NO_PROXY_ENV",
     "ENABLE_SCRIPT_INJECTION",
     "USERSCRIPT_PATH",
+    # 新增页面元素选择器
+    "MODEL_NAME_SELECTOR",
+    "CDK_OVERLAY_CONTAINER_SELECTOR",
+    "CHAT_TURN_SELECTOR",
+    # 思考模式回退选择器
+    "THINKING_MODE_TOGGLE_PARENT_SELECTOR",
+    "THINKING_MODE_TOGGLE_OLD_ROOT_SELECTOR",
+    "THINKING_BUDGET_TOGGLE_PARENT_SELECTOR",
+    "THINKING_BUDGET_TOGGLE_OLD_ROOT_SELECTOR",
     # 工具函数
     "get_environment_variable",
     "get_boolean_env",

config/selectors.py

@@ -118,3 +118,28 @@
 GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR = (
     'div[data-test-id="searchAsAToolTooltip"] mat-slide-toggle button'
 )
+
+# --- 页面元素选择器 ---
+# 模型名称显示元素
+MODEL_NAME_SELECTOR = '[data-test-id="model-name"]'
+# CDK Overlay 容器（用于菜单、对话框等）
+CDK_OVERLAY_CONTAINER_SELECTOR = "div.cdk-overlay-container"
+# 聊天轮次容器
+CHAT_TURN_SELECTOR = "ms-chat-turn"
+
+# --- 思考模式回退选择器 ---
+# 这些选择器用于 thinking.py 中的回退逻辑
+# 主思考开关父容器（新版UI）
+THINKING_MODE_TOGGLE_PARENT_SELECTOR = (
+    'mat-slide-toggle:has(button[aria-label="Toggle thinking mode"])'
+)
+# 主思考开关旧版根元素
+THINKING_MODE_TOGGLE_OLD_ROOT_SELECTOR = (
+    'mat-slide-toggle[data-test-toggle="enable-thinking"]'
+)
+# 思考预算开关父容器（新版UI）
+THINKING_BUDGET_TOGGLE_PARENT_SELECTOR = 'mat-slide-toggle:has(button[aria-label="Toggle thinking budget between auto and manual"])'
+# 思考预算开关旧版根元素
+THINKING_BUDGET_TOGGLE_OLD_ROOT_SELECTOR = (
+    'mat-slide-toggle[data-test-toggle="manual-budget"]'
+)