Skip to content

Commit 1f08b0c

Browse files
committed
fix: 支持新旧前端选择器并发匹配,修复ms-prompt-input-wrapper兼容性
1 parent 364dbc8 commit 1f08b0c

File tree

10 files changed

+211
-43
lines changed

10 files changed

+211
-43
lines changed

src/browser/initialization.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from typing import Optional, Any, Dict, Tuple
77
from playwright.async_api import Page as AsyncPage, Browser as AsyncBrowser, BrowserContext as AsyncBrowserContext, Error as PlaywrightAsyncError, expect as expect_async
88
from config import *
9+
from config.selectors import PROMPT_TEXTAREA_SELECTORS
910
from models import ClientDisconnectedError
1011
logger = logging.getLogger('AIStudioProxyServer')
1112

@@ -305,10 +306,19 @@ async def _initialize_page_logic(browser: AsyncBrowser):
305306
logger.info(f'✅ 确认位于 AI Studio: {current_url}')
306307
await found_page.bring_to_front()
307308
try:
308-
input_wrapper_locator = found_page.locator('ms-prompt-box')
309-
await expect_async(input_wrapper_locator).to_be_visible(timeout=35000)
310-
await expect_async(found_page.locator(INPUT_SELECTOR)).to_be_visible(timeout=10000)
311-
logger.info('✅ 核心输入区域可见。')
309+
from browser.selector_utils import wait_for_any_selector
310+
wrapper_selectors = ['ms-prompt-input-wrapper', 'ms-prompt-box']
311+
wrapper_locator, wrapper_matched = await wait_for_any_selector(found_page, wrapper_selectors, timeout=35000)
312+
if wrapper_locator:
313+
logger.info(f'✅ 输入框wrapper可见 (匹配: {wrapper_matched})')
314+
else:
315+
logger.warning('⚠️ 未找到任何wrapper,尝试直接查找输入框')
316+
input_locator, matched = await wait_for_any_selector(found_page, PROMPT_TEXTAREA_SELECTORS, timeout=10000)
317+
if input_locator:
318+
logger.info(f'✅ 核心输入区域可见 (匹配: {matched})')
319+
else:
320+
await expect_async(found_page.locator(INPUT_SELECTOR)).to_be_visible(timeout=10000)
321+
logger.info('✅ 核心输入区域可见 (默认选择器)')
312322
try:
313323
from config.selectors import MODEL_SELECTORS_LIST
314324
from browser.operations import get_model_name_from_page_parallel

src/browser/page_controller.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,11 @@
55
import re
66
import os
77
from playwright.async_api import Page as AsyncPage, expect as expect_async, TimeoutError, Locator
8-
from config import TEMPERATURE_INPUT_SELECTOR, MAX_OUTPUT_TOKENS_SELECTOR, STOP_SEQUENCE_INPUT_SELECTOR, MAT_CHIP_REMOVE_BUTTON_SELECTOR, TOP_P_INPUT_SELECTOR, SUBMIT_BUTTON_SELECTOR, OVERLAY_SELECTOR, PROMPT_TEXTAREA_SELECTOR, RESPONSE_CONTAINER_SELECTOR, RESPONSE_TEXT_SELECTOR, EDIT_MESSAGE_BUTTON_SELECTOR, USE_URL_CONTEXT_SELECTOR, UPLOAD_BUTTON_SELECTOR, INSERT_BUTTON_SELECTOR, THINKING_MODE_TOGGLE_SELECTOR, SET_THINKING_BUDGET_TOGGLE_SELECTOR, THINKING_BUDGET_INPUT_SELECTOR, GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR, ZERO_STATE_SELECTOR, SYSTEM_INSTRUCTIONS_BUTTON_SELECTOR, SYSTEM_INSTRUCTIONS_TEXTAREA_SELECTOR, SKIP_PREFERENCE_VOTE_BUTTON_SELECTOR, CLICK_TIMEOUT_MS, WAIT_FOR_ELEMENT_TIMEOUT_MS, CLEAR_CHAT_VERIFY_TIMEOUT_MS, DEFAULT_TEMPERATURE, DEFAULT_MAX_OUTPUT_TOKENS, DEFAULT_STOP_SEQUENCES, DEFAULT_TOP_P, ENABLE_URL_CONTEXT, ENABLE_THINKING_BUDGET, DEFAULT_THINKING_BUDGET, ENABLE_GOOGLE_SEARCH, THINKING_LEVEL_SELECT_SELECTOR, THINKING_LEVEL_OPTION_HIGH_SELECTOR, THINKING_LEVEL_OPTION_LOW_SELECTOR, DEFAULT_THINKING_LEVEL, ADVANCED_SETTINGS_EXPANDER_SELECTOR
8+
from config import TEMPERATURE_INPUT_SELECTOR, MAX_OUTPUT_TOKENS_SELECTOR, STOP_SEQUENCE_INPUT_SELECTOR, MAT_CHIP_REMOVE_BUTTON_SELECTOR, TOP_P_INPUT_SELECTOR, SUBMIT_BUTTON_SELECTOR, SUBMIT_BUTTON_SELECTORS, OVERLAY_SELECTOR, PROMPT_TEXTAREA_SELECTOR, PROMPT_TEXTAREA_SELECTORS, RESPONSE_CONTAINER_SELECTOR, RESPONSE_TEXT_SELECTOR, EDIT_MESSAGE_BUTTON_SELECTOR, USE_URL_CONTEXT_SELECTOR, UPLOAD_BUTTON_SELECTOR, INSERT_BUTTON_SELECTOR, INSERT_BUTTON_SELECTORS, THINKING_MODE_TOGGLE_SELECTOR, SET_THINKING_BUDGET_TOGGLE_SELECTOR, THINKING_BUDGET_INPUT_SELECTOR, GROUNDING_WITH_GOOGLE_SEARCH_TOGGLE_SELECTOR, ZERO_STATE_SELECTOR, SYSTEM_INSTRUCTIONS_BUTTON_SELECTOR, SYSTEM_INSTRUCTIONS_TEXTAREA_SELECTOR, SKIP_PREFERENCE_VOTE_BUTTON_SELECTOR, CLICK_TIMEOUT_MS, WAIT_FOR_ELEMENT_TIMEOUT_MS, CLEAR_CHAT_VERIFY_TIMEOUT_MS, DEFAULT_TEMPERATURE, DEFAULT_MAX_OUTPUT_TOKENS, DEFAULT_STOP_SEQUENCES, DEFAULT_TOP_P, ENABLE_URL_CONTEXT, ENABLE_THINKING_BUDGET, DEFAULT_THINKING_BUDGET, ENABLE_GOOGLE_SEARCH, THINKING_LEVEL_SELECT_SELECTOR, THINKING_LEVEL_OPTION_HIGH_SELECTOR, THINKING_LEVEL_OPTION_LOW_SELECTOR, DEFAULT_THINKING_LEVEL, ADVANCED_SETTINGS_EXPANDER_SELECTOR
99
from models import ClientDisconnectedError, ElementClickError
1010
from .operations import save_error_snapshot, _wait_for_response_completion, _get_final_response_content, click_element
1111
from .thinking_normalizer import parse_reasoning_param, describe_config
12+
from .selector_utils import wait_for_any_selector, get_first_visible_locator
1213

1314
class PageController:
1415

@@ -704,8 +705,7 @@ async def _verify_chat_cleared(self, check_client_disconnected: Callable):
704705
async def _robust_click_insert_assets(self, check_client_disconnected: Callable) -> bool:
705706
self.logger.info(f"[{self.req_id}] 开始寻找并点击媒体添加按钮...")
706707

707-
trigger_selectors = [
708-
INSERT_BUTTON_SELECTOR,
708+
trigger_selectors = INSERT_BUTTON_SELECTORS + [
709709
'button[aria-label*="Insert"]',
710710
'button[iconname="add_circle"]',
711711
'button[iconname="note_add"]'
@@ -896,9 +896,18 @@ async def _paste_images_via_event(self, images: List[Dict[str, str]], target_loc
896896

897897
async def submit_prompt(self, prompt: str, image_list: List, check_client_disconnected: Callable):
898898
self.logger.info(f'[{self.req_id}] 📤 提交提示 ({len(prompt)} chars)...')
899-
prompt_textarea_locator = self.page.locator(PROMPT_TEXTAREA_SELECTOR)
899+
prompt_textarea_locator, matched_selector = await get_first_visible_locator(self.page, PROMPT_TEXTAREA_SELECTORS, timeout=5000)
900+
if not prompt_textarea_locator:
901+
self.logger.warning(f'[{self.req_id}] 未找到输入框,尝试默认选择器')
902+
prompt_textarea_locator = self.page.locator(PROMPT_TEXTAREA_SELECTOR)
903+
else:
904+
self.logger.info(f'[{self.req_id}] 找到输入框 (匹配: {matched_selector})')
900905
autosize_wrapper_locator = self.page.locator('ms-prompt-box .text-wrapper')
901-
submit_button_locator = self.page.locator(SUBMIT_BUTTON_SELECTOR)
906+
submit_button_locator, submit_matched = await get_first_visible_locator(self.page, SUBMIT_BUTTON_SELECTORS, timeout=3000)
907+
if not submit_button_locator:
908+
submit_button_locator = self.page.locator(SUBMIT_BUTTON_SELECTOR)
909+
else:
910+
self.logger.info(f'[{self.req_id}] 找到提交按钮 (匹配: {submit_matched})')
902911
try:
903912
await expect_async(prompt_textarea_locator).to_be_visible(timeout=5000)
904913
await self._check_disconnect(check_client_disconnected, '输入框可见后')

src/browser/selector_utils.py

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
import asyncio
2+
from typing import Optional, List, Tuple
3+
from playwright.async_api import Page as AsyncPage, Locator
4+
5+
6+
async def wait_for_any_selector(
7+
page: AsyncPage,
8+
selectors: List[str],
9+
timeout: int = 5000,
10+
state: str = 'visible'
11+
) -> Tuple[Optional[Locator], Optional[str]]:
12+
async def check_one(selector: str) -> Tuple[bool, str]:
13+
try:
14+
locator = page.locator(selector)
15+
await locator.wait_for(state=state, timeout=timeout)
16+
return (True, selector)
17+
except:
18+
return (False, selector)
19+
20+
tasks = [asyncio.create_task(check_one(sel)) for sel in selectors]
21+
22+
try:
23+
done, pending = await asyncio.wait(
24+
tasks,
25+
return_when=asyncio.FIRST_COMPLETED,
26+
timeout=timeout / 1000 + 1
27+
)
28+
29+
for task in done:
30+
success, selector = task.result()
31+
if success:
32+
for p in pending:
33+
p.cancel()
34+
return (page.locator(selector), selector)
35+
36+
for p in pending:
37+
p.cancel()
38+
return (None, None)
39+
40+
except asyncio.TimeoutError:
41+
for t in tasks:
42+
if not t.done():
43+
t.cancel()
44+
return (None, None)
45+
46+
47+
async def get_first_visible_locator(
48+
page: AsyncPage,
49+
selectors: List[str],
50+
timeout: int = 3000
51+
) -> Tuple[Optional[Locator], Optional[str]]:
52+
for selector in selectors:
53+
try:
54+
locator = page.locator(selector)
55+
if await locator.count() > 0:
56+
first = locator.first
57+
if await first.is_visible():
58+
return (first, selector)
59+
except:
60+
continue
61+
62+
return await wait_for_any_selector(page, selectors, timeout)
63+
64+
65+
async def click_first_available(
66+
page: AsyncPage,
67+
selectors: List[str],
68+
timeout: int = 5000
69+
) -> Tuple[bool, Optional[str]]:
70+
locator, selector = await get_first_visible_locator(page, selectors, timeout)
71+
if locator:
72+
try:
73+
await locator.click(timeout=1000)
74+
return (True, selector)
75+
except:
76+
try:
77+
await locator.evaluate('el => el.click()')
78+
return (True, selector)
79+
except:
80+
pass
81+
return (False, None)

src/config/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,4 @@
22
from .timeouts import *
33
from .selectors import *
44
from .settings import *
5-
__all__ = ['MODEL_NAME', 'CHAT_COMPLETION_ID_PREFIX', 'DEFAULT_FALLBACK_MODEL_ID', 'DEFAULT_TEMPERATURE', 'DEFAULT_MAX_OUTPUT_TOKENS', 'DEFAULT_TOP_P', 'DEFAULT_STOP_SEQUENCES', 'AI_STUDIO_URL_PATTERN', 'MODELS_ENDPOINT_URL_CONTAINS', 'USER_INPUT_START_MARKER_SERVER', 'USER_INPUT_END_MARKER_SERVER', 'EXCLUDED_MODELS_FILENAME', 'STREAM_TIMEOUT_LOG_STATE', 'RESPONSE_COMPLETION_TIMEOUT', 'INITIAL_WAIT_MS_BEFORE_POLLING', 'POLLING_INTERVAL', 'POLLING_INTERVAL_STREAM', 'SILENCE_TIMEOUT_MS', 'POST_SPINNER_CHECK_DELAY_MS', 'FINAL_STATE_CHECK_TIMEOUT_MS', 'POST_COMPLETION_BUFFER', 'CLEAR_CHAT_VERIFY_TIMEOUT_MS', 'CLEAR_CHAT_VERIFY_INTERVAL_MS', 'CLICK_TIMEOUT_MS', 'CLIPBOARD_READ_TIMEOUT_MS', 'WAIT_FOR_ELEMENT_TIMEOUT_MS', 'PSEUDO_STREAM_DELAY', 'PROMPT_TEXTAREA_SELECTOR', 'INPUT_SELECTOR', 'INPUT_SELECTOR2', 'SUBMIT_BUTTON_SELECTOR', 'INSERT_BUTTON_SELECTOR', 'RESPONSE_CONTAINER_SELECTOR', 'RESPONSE_TEXT_SELECTOR', 'LOADING_SPINNER_SELECTOR', 'OVERLAY_SELECTOR', 'ERROR_TOAST_SELECTOR', 'EDIT_MESSAGE_BUTTON_SELECTOR', 'MESSAGE_TEXTAREA_SELECTOR', 'FINISH_EDIT_BUTTON_SELECTOR', 'MORE_OPTIONS_BUTTON_SELECTOR', 'COPY_MARKDOWN_BUTTON_SELECTOR', 'COPY_MARKDOWN_BUTTON_SELECTOR_ALT', 'MAX_OUTPUT_TOKENS_SELECTOR', 'STOP_SEQUENCE_INPUT_SELECTOR', 'MAT_CHIP_REMOVE_BUTTON_SELECTOR', 'TOP_P_INPUT_SELECTOR', 'TEMPERATURE_INPUT_SELECTOR', 'USE_URL_CONTEXT_SELECTOR', 'UPLOAD_BUTTON_SELECTOR', 'DEBUG_LOGS_ENABLED', 'TRACE_LOGS_ENABLED', 'AUTO_SAVE_AUTH', 'AUTH_SAVE_TIMEOUT', 'AUTO_CONFIRM_LOGIN', 'AUTH_PROFILES_DIR', 'ACTIVE_AUTH_DIR', 'SAVED_AUTH_DIR', 'LOG_DIR', 'APP_LOG_FILE_PATH', 'NO_PROXY_ENV', 'ENABLE_SCRIPT_INJECTION', 'USERSCRIPT_PATH', 'get_environment_variable', 'get_boolean_env', 'get_int_env']
5+
__all__ = ['MODEL_NAME', 'CHAT_COMPLETION_ID_PREFIX', 'DEFAULT_FALLBACK_MODEL_ID', 'DEFAULT_TEMPERATURE', 'DEFAULT_MAX_OUTPUT_TOKENS', 'DEFAULT_TOP_P', 'DEFAULT_STOP_SEQUENCES', 'AI_STUDIO_URL_PATTERN', 'MODELS_ENDPOINT_URL_CONTAINS', 'USER_INPUT_START_MARKER_SERVER', 'USER_INPUT_END_MARKER_SERVER', 'EXCLUDED_MODELS_FILENAME', 'STREAM_TIMEOUT_LOG_STATE', 'RESPONSE_COMPLETION_TIMEOUT', 'INITIAL_WAIT_MS_BEFORE_POLLING', 'POLLING_INTERVAL', 'POLLING_INTERVAL_STREAM', 'SILENCE_TIMEOUT_MS', 'POST_SPINNER_CHECK_DELAY_MS', 'FINAL_STATE_CHECK_TIMEOUT_MS', 'POST_COMPLETION_BUFFER', 'CLEAR_CHAT_VERIFY_TIMEOUT_MS', 'CLEAR_CHAT_VERIFY_INTERVAL_MS', 'CLICK_TIMEOUT_MS', 'CLIPBOARD_READ_TIMEOUT_MS', 'WAIT_FOR_ELEMENT_TIMEOUT_MS', 'PSEUDO_STREAM_DELAY', 'PROMPT_TEXTAREA_SELECTOR', 'PROMPT_TEXTAREA_SELECTORS', 'INPUT_SELECTOR', 'INPUT_SELECTOR2', 'SUBMIT_BUTTON_SELECTOR', 'SUBMIT_BUTTON_SELECTORS', 'INSERT_BUTTON_SELECTOR', 'INSERT_BUTTON_SELECTORS', 'UPLOAD_BUTTON_SELECTOR', 'UPLOAD_BUTTON_SELECTORS', 'HIDDEN_FILE_INPUT_SELECTOR', 'HIDDEN_FILE_INPUT_SELECTORS', 'RESPONSE_CONTAINER_SELECTOR', 'RESPONSE_TEXT_SELECTOR', 'LOADING_SPINNER_SELECTOR', 'LOADING_SPINNER_SELECTORS', 'OVERLAY_SELECTOR', 'ERROR_TOAST_SELECTOR', 'EDIT_MESSAGE_BUTTON_SELECTOR', 'MESSAGE_TEXTAREA_SELECTOR', 'FINISH_EDIT_BUTTON_SELECTOR', 'MORE_OPTIONS_BUTTON_SELECTOR', 'COPY_MARKDOWN_BUTTON_SELECTOR', 'COPY_MARKDOWN_BUTTON_SELECTOR_ALT', 'MAX_OUTPUT_TOKENS_SELECTOR', 'STOP_SEQUENCE_INPUT_SELECTOR', 'MAT_CHIP_REMOVE_BUTTON_SELECTOR', 'TOP_P_INPUT_SELECTOR', 'TEMPERATURE_INPUT_SELECTOR', 'USE_URL_CONTEXT_SELECTOR', 'DEBUG_LOGS_ENABLED', 'TRACE_LOGS_ENABLED', 'AUTO_SAVE_AUTH', 'AUTH_SAVE_TIMEOUT', 'AUTO_CONFIRM_LOGIN', 'AUTH_PROFILES_DIR', 'ACTIVE_AUTH_DIR', 'SAVED_AUTH_DIR', 'LOG_DIR', 'APP_LOG_FILE_PATH', 'NO_PROXY_ENV', 'ENABLE_SCRIPT_INJECTION', 'USERSCRIPT_PATH', 'get_environment_variable', 'get_boolean_env', 'get_int_env']

src/config/imagen_selectors.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,20 @@
77
IMAGEN_TOOLBAR_SELECTOR = 'ms-toolbar'
88
IMAGEN_TOOLBAR_TITLE_SELECTOR = 'ms-toolbar h1.mode-title'
99

10-
IMAGEN_PROMPT_INPUT_SELECTOR = 'ms-prompt-box textarea[aria-label="Enter a prompt to generate an image"]'
11-
IMAGEN_RUN_BUTTON_SELECTOR = 'ms-prompt-box ms-run-button button[aria-label="Run"]'
10+
IMAGEN_PROMPT_INPUT_SELECTORS = [
11+
'ms-prompt-input-wrapper textarea[aria-label="Enter a prompt to generate an image"]',
12+
'textarea[aria-label="Enter a prompt to generate an image"]',
13+
'ms-prompt-box textarea[aria-label="Enter a prompt to generate an image"]',
14+
]
15+
IMAGEN_PROMPT_INPUT_SELECTOR = IMAGEN_PROMPT_INPUT_SELECTORS[0]
16+
17+
IMAGEN_RUN_BUTTON_SELECTORS = [
18+
'ms-run-button button[aria-label="Run"]',
19+
'button[aria-label="Run"].run-button',
20+
'ms-prompt-input-wrapper ms-run-button button[aria-label="Run"]',
21+
'ms-prompt-box ms-run-button button[aria-label="Run"]',
22+
]
23+
IMAGEN_RUN_BUTTON_SELECTOR = IMAGEN_RUN_BUTTON_SELECTORS[0]
1224

1325
IMAGEN_GALLERY_CONTAINER_SELECTOR = 'ms-image-generation-gallery'
1426
IMAGEN_GALLERY_ITEM_SELECTOR = 'ms-image-generation-gallery-image'

src/config/selectors.py

Lines changed: 40 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,49 @@
1-
PROMPT_TEXTAREA_SELECTOR = 'ms-prompt-box textarea'
1+
PROMPT_TEXTAREA_SELECTORS = [
2+
'ms-prompt-input-wrapper ms-autosize-textarea textarea',
3+
'ms-prompt-box textarea',
4+
]
5+
PROMPT_TEXTAREA_SELECTOR = PROMPT_TEXTAREA_SELECTORS[0]
26
INPUT_SELECTOR = PROMPT_TEXTAREA_SELECTOR
37
INPUT_SELECTOR2 = PROMPT_TEXTAREA_SELECTOR
4-
SUBMIT_BUTTON_SELECTOR = 'ms-prompt-box ms-run-button button'
5-
INSERT_BUTTON_SELECTOR = 'button[data-test-id="add-media-button"]'
6-
UPLOAD_BUTTON_SELECTOR = 'button[role="menuitem"]:has-text("Upload a file")'
7-
HIDDEN_FILE_INPUT_SELECTOR = 'input[type="file"][data-test-upload-file-input]'
8+
9+
SUBMIT_BUTTON_SELECTORS = [
10+
'button[aria-label="Run"].run-button',
11+
'ms-run-button button[aria-label="Run"]',
12+
'ms-prompt-box ms-run-button button',
13+
]
14+
SUBMIT_BUTTON_SELECTOR = SUBMIT_BUTTON_SELECTORS[0]
15+
16+
INSERT_BUTTON_SELECTORS = [
17+
'button[aria-label="Insert assets such as images, videos, files, or audio"]',
18+
'button[data-test-add-chunk-menu-button]',
19+
'button[data-test-id="add-media-button"]',
20+
]
21+
INSERT_BUTTON_SELECTOR = INSERT_BUTTON_SELECTORS[0]
22+
23+
UPLOAD_BUTTON_SELECTORS = [
24+
'button[role="menuitem"]:has-text("Upload a file")',
25+
'button[aria-label="Upload File"]',
26+
]
27+
UPLOAD_BUTTON_SELECTOR = UPLOAD_BUTTON_SELECTORS[0]
28+
29+
HIDDEN_FILE_INPUT_SELECTORS = [
30+
'input[type="file"][data-test-upload-file-input]',
31+
'input.file-input[type="file"]',
32+
]
33+
HIDDEN_FILE_INPUT_SELECTOR = HIDDEN_FILE_INPUT_SELECTORS[0]
34+
835
UPLOADED_MEDIA_ITEM_SELECTOR = 'ms-prompt-box .multi-media-row ms-media-chip'
936
SKIP_PREFERENCE_VOTE_BUTTON_SELECTOR = 'button[data-test-id="skip-button"][aria-label="Skip preference vote"]'
1037
RESPONSE_CONTAINER_SELECTOR = 'ms-chat-turn .chat-turn-container.model'
1138
RESPONSE_TEXT_SELECTOR = 'ms-cmark-node.cmark-node'
12-
LOADING_SPINNER_SELECTOR = 'ms-prompt-box ms-run-button button svg .stoppable-spinner'
39+
40+
LOADING_SPINNER_SELECTORS = [
41+
'button[aria-label="Run"].run-button svg .stoppable-spinner',
42+
'ms-run-button button svg .stoppable-spinner',
43+
'ms-prompt-box ms-run-button button svg .stoppable-spinner',
44+
]
45+
LOADING_SPINNER_SELECTOR = LOADING_SPINNER_SELECTORS[0]
46+
1347
OVERLAY_SELECTOR = '.mat-mdc-dialog-inner-container'
1448
ZERO_STATE_SELECTOR = 'ms-zero-state'
1549
ERROR_TOAST_SELECTOR = 'div.toast.warning, div.toast.error'

src/config/tts_selectors.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,11 @@
1010
TTS_FOOTER_AUDIO_PLAYER_WRAPPER_SELECTOR = '.speech-prompt-footer-actions-left'
1111
TTS_AUDIO_PLAYER_SELECTOR = '.speech-prompt-footer audio[controls]'
1212
TTS_RUN_BUTTON_WRAPPER_SELECTOR = '.speech-prompt-footer .button-wrapper'
13-
TTS_RUN_BUTTON_SELECTOR = 'ms-run-button button[aria-label="Run"]'
13+
TTS_RUN_BUTTON_SELECTORS = [
14+
'ms-run-button button[aria-label="Run"]',
15+
'.speech-prompt-footer button[aria-label="Run"].run-button',
16+
]
17+
TTS_RUN_BUTTON_SELECTOR = TTS_RUN_BUTTON_SELECTORS[0]
1418

1519
TTS_SINGLE_SPEAKER_BUILDER_SELECTOR = '.single-speaker-prompt-builder-wrapper'
1620
TTS_SINGLE_SPEAKER_STYLE_INPUT_SELECTOR = 'ms-autosize-textarea.style-instructions-textarea textarea'

src/media/imagen_controller.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,15 @@
44
from playwright.async_api import Page as AsyncPage, Locator, expect as expect_async
55
from config.imagen_selectors import (
66
IMAGEN_PAGE_URL_TEMPLATE, IMAGEN_SUPPORTED_MODELS, IMAGEN_ROOT_SELECTOR,
7-
IMAGEN_PROMPT_INPUT_SELECTOR, IMAGEN_RUN_BUTTON_SELECTOR,
7+
IMAGEN_PROMPT_INPUT_SELECTOR, IMAGEN_PROMPT_INPUT_SELECTORS,
8+
IMAGEN_RUN_BUTTON_SELECTOR, IMAGEN_RUN_BUTTON_SELECTORS,
89
IMAGEN_GALLERY_CONTAINER_SELECTOR, IMAGEN_GALLERY_ITEM_SELECTOR,
910
IMAGEN_GENERATED_IMAGE_SELECTOR, IMAGEN_SETTINGS_PANEL_SELECTOR,
1011
IMAGEN_SETTINGS_NUM_RESULTS_INPUT_SELECTOR, IMAGEN_SETTINGS_ASPECT_RATIO_BUTTON_SELECTOR,
1112
IMAGEN_SETTINGS_NEGATIVE_PROMPT_SELECTOR
1213
)
1314
from browser.operations import safe_click
15+
from browser.selector_utils import wait_for_any_selector, get_first_visible_locator
1416
from .models import ImageGenerationConfig, GeneratedImage
1517
from models import ClientDisconnectedError
1618

@@ -118,7 +120,10 @@ async def fill_prompt(self, prompt: str, check_client_disconnected: Callable):
118120
try:
119121
await self.page.keyboard.press('Escape')
120122
await asyncio.sleep(0.15)
121-
text_input = self.page.locator(IMAGEN_PROMPT_INPUT_SELECTOR)
123+
text_input, matched = await get_first_visible_locator(self.page, IMAGEN_PROMPT_INPUT_SELECTORS)
124+
if not text_input:
125+
raise Exception('未找到输入框')
126+
self.logger.info(f'[{self.req_id}] 找到输入框 (匹配: {matched})')
122127
await safe_click(text_input, '输入框', self.req_id)
123128
await text_input.fill(prompt)
124129
await asyncio.sleep(0.1)
@@ -139,8 +144,10 @@ async def run_generation(self, check_client_disconnected: Callable):
139144
try:
140145
await self.page.keyboard.press('Escape')
141146
await asyncio.sleep(0.15)
142-
run_btn = self.page.locator(IMAGEN_RUN_BUTTON_SELECTOR)
143-
await expect_async(run_btn).to_be_visible(timeout=10000)
147+
run_btn, matched = await wait_for_any_selector(self.page, IMAGEN_RUN_BUTTON_SELECTORS, timeout=10000)
148+
if not run_btn:
149+
raise Exception('未找到Run按钮')
150+
self.logger.info(f'[{self.req_id}] 找到Run按钮 (匹配: {matched})')
144151
await expect_async(run_btn).to_be_enabled(timeout=10000)
145152
if not await safe_click(run_btn, 'Run 按钮', self.req_id):
146153
if attempt < max_retries:

0 commit comments

Comments
 (0)