perf: 优化延迟、优化冗余代码

Author: Mag1cFall

src/browser/operations.py

@@ -4,12 +4,44 @@
 import os
 import re
 import logging
-from typing import Optional, Any, List, Dict, Callable, Set
+import functools
+from typing import Optional, Any, List, Dict, Callable, Set, TypeVar
 from playwright.async_api import Page as AsyncPage, Locator, Error as PlaywrightAsyncError
 from config import *
 from models import ClientDisconnectedError, ElementClickError
 logger = logging.getLogger('AIStudioProxyServer')
 
+T = TypeVar('T')
+
+def check_disconnect(check_client_disconnected: Callable, stage: str, req_id: str):
+    if check_client_disconnected and check_client_disconnected(stage):
+        raise ClientDisconnectedError(f'[{req_id}] Client disconnected at stage: {stage}')
+
+async def retry_async(
+    func: Callable,
+    max_retries: int = 3,
+    delay: float = 0.15,
+    req_id: str = 'unknown',
+    operation_name: str = 'operation',
+    check_client_disconnected: Callable = None
+) -> Any:
+    last_error = None
+    for attempt in range(1, max_retries + 1):
+        try:
+            if check_client_disconnected:
+                check_disconnect(check_client_disconnected, f'{operation_name} - 尝试 {attempt}', req_id)
+            return await func()
+        except ClientDisconnectedError:
+            raise
+        except Exception as e:
+            last_error = e
+            if attempt < max_retries:
+                logger.warning(f'[{req_id}] {operation_name} 失败 (尝试 {attempt}): {e}')
+                await asyncio.sleep(delay)
+    logger.error(f'[{req_id}] {operation_name} 最终失败，已重试 {max_retries} 次')
+    raise last_error if last_error else Exception(f'{operation_name} failed')
+
+
 async def get_model_name_from_page_parallel(page: AsyncPage, selectors: list, timeout: int = 2000, req_id: str = 'unknown', expected_model_name: str = None) -> Optional[str]:
     if not selectors:
         return None

src/browser/page_controller.py

@@ -795,9 +795,9 @@ async def _upload_images_via_file_input(self, images: List[Dict[str, str]], chec
                 await file_input.set_input_files(tf.name)
                 uploaded_count += 1
 
-                await asyncio.sleep(0.8)
+                await asyncio.sleep(0.2)
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.1)
 
             asyncio.create_task(self._cleanup_temp_files(temp_files))
 

src/media/imagen_controller.py

@@ -59,15 +59,15 @@ async def set_number_of_images(self, count: int, check_client_disconnected: Call
                     self.logger.warning(f'[{self.req_id}] 未找到数量输入框')
                     return
                 await input_locator.fill(str(count))
-                await asyncio.sleep(0.2)
+                await asyncio.sleep(0.1)
                 self.logger.info(f'[{self.req_id}] ✅ 图片数量已设置: {count}')
                 return
             except Exception as e:
                 if isinstance(e, ClientDisconnectedError):
                     raise
                 self.logger.warning(f'[{self.req_id}] 设置数量失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
 
     async def set_aspect_ratio(self, aspect_ratio: str, check_client_disconnected: Callable):
         self.logger.info(f'[{self.req_id}] 设置宽高比: {aspect_ratio}')
@@ -87,7 +87,7 @@ async def set_aspect_ratio(self, aspect_ratio: str, check_client_disconnected: C
                     raise
                 self.logger.warning(f'[{self.req_id}] 设置宽高比失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
 
     async def set_negative_prompt(self, negative_prompt: str, check_client_disconnected: Callable):
         if not negative_prompt:
@@ -101,35 +101,35 @@ async def set_negative_prompt(self, negative_prompt: str, check_client_disconnec
                     self.logger.warning(f'[{self.req_id}] 未找到负面提示词输入框')
                     return
                 await textarea.fill(negative_prompt)
-                await asyncio.sleep(0.2)
+                await asyncio.sleep(0.1)
                 self.logger.info(f'[{self.req_id}] ✅ 负面提示词已设置')
                 return
             except Exception as e:
                 if isinstance(e, ClientDisconnectedError):
                     raise
                 self.logger.warning(f'[{self.req_id}] 设置负面提示词失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
 
     async def fill_prompt(self, prompt: str, check_client_disconnected: Callable):
         self.logger.info(f'[{self.req_id}] 填充提示词 ({len(prompt)} chars)')
         max_retries = 3
         for attempt in range(1, max_retries + 1):
             try:
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
                 text_input = self.page.locator(IMAGEN_PROMPT_INPUT_SELECTOR)
                 await safe_click(text_input, '输入框', self.req_id)
                 await text_input.fill(prompt)
-                await asyncio.sleep(0.2)
+                await asyncio.sleep(0.1)
                 self.logger.info(f'[{self.req_id}] ✅ 提示词已填充')
                 return
             except Exception as e:
                 if isinstance(e, ClientDisconnectedError):
                     raise
                 self.logger.warning(f'[{self.req_id}] 填充提示词失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
         raise Exception('填充提示词失败')
 
     async def run_generation(self, check_client_disconnected: Callable):
@@ -138,7 +138,7 @@ async def run_generation(self, check_client_disconnected: Callable):
         for attempt in range(1, max_retries + 1):
             try:
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
                 run_btn = self.page.locator(IMAGEN_RUN_BUTTON_SELECTOR)
                 await expect_async(run_btn).to_be_visible(timeout=5000)
                 await expect_async(run_btn).to_be_enabled(timeout=5000)
@@ -154,7 +154,7 @@ async def run_generation(self, check_client_disconnected: Callable):
                     raise
                 self.logger.warning(f'[{self.req_id}] 点击 Run 失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
         raise Exception('点击 Run 按钮失败')
 
     async def wait_for_images(self, expected_count: int, check_client_disconnected: Callable, timeout_seconds: int = 180) -> List[GeneratedImage]:
@@ -191,4 +191,4 @@ async def wait_for_images(self, expected_count: int, check_client_disconnected:
             except Exception as e:
                 self.logger.warning(f'[{self.req_id}] 检查图片时出错: {e}')
 
-            await asyncio.sleep(2)
+            await asyncio.sleep(1)

src/media/nano_controller.py

@@ -61,7 +61,7 @@ async def set_aspect_ratio(self, aspect_ratio: str, check_client_disconnected: C
                     return
                 if not await safe_click(dropdown, '宽高比下拉框', self.req_id):
                     continue
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
                 option = self.page.locator(f'mat-option:has-text("{aspect_ratio}")')
                 if await option.count() > 0:
                     if await safe_click(option.first, f'宽高比选项 {aspect_ratio}', self.req_id):
@@ -76,7 +76,7 @@ async def set_aspect_ratio(self, aspect_ratio: str, check_client_disconnected: C
                     raise
                 self.logger.warning(f'[{self.req_id}] 设置宽高比失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
         self.logger.warning(f'[{self.req_id}] 宽高比设置失败，使用默认值')
 
     async def upload_image(self, image_bytes: bytes, mime_type: str, check_client_disconnected: Callable):
@@ -92,7 +92,7 @@ async def upload_image(self, image_bytes: bytes, mime_type: str, check_client_di
                     if attempt < max_retries:
                         continue
                     return
-                await asyncio.sleep(0.5)
+                await asyncio.sleep(0.25)
                 await self._check_disconnect(check_client_disconnected, '插入菜单展开后')
 
                 ext = 'png' if 'png' in mime_type else 'jpg'
@@ -102,29 +102,29 @@ async def upload_image(self, image_bytes: bytes, mime_type: str, check_client_di
                     'mimeType': mime_type,
                     'buffer': image_bytes
                 })
-                await asyncio.sleep(1.5)
+                await asyncio.sleep(0.8)
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.5)
+                await asyncio.sleep(0.25)
                 self.logger.info(f'[{self.req_id}] ✅ 图片已上传')
                 return
             except Exception as e:
                 if isinstance(e, ClientDisconnectedError):
                     raise
                 self.logger.warning(f'[{self.req_id}] 上传图片失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.5)
+                await asyncio.sleep(0.25)
 
     async def fill_prompt(self, prompt: str, check_client_disconnected: Callable):
         self.logger.info(f'[{self.req_id}] 填充提示词 ({len(prompt)} chars)')
         max_retries = 3
         for attempt in range(1, max_retries + 1):
             try:
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
                 text_input = self.page.locator(PROMPT_TEXTAREA_SELECTOR)
                 await safe_click(text_input, '输入框', self.req_id)
                 await text_input.fill(prompt)
-                await asyncio.sleep(0.2)
+                await asyncio.sleep(0.1)
                 actual = await text_input.input_value()
                 if prompt in actual or actual in prompt:
                     self.logger.info(f'[{self.req_id}] ✅ 提示词已填充')
@@ -135,7 +135,7 @@ async def fill_prompt(self, prompt: str, check_client_disconnected: Callable):
                     raise
                 self.logger.warning(f'[{self.req_id}] 填充提示词失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
         raise Exception('填充提示词失败')
 
     async def run_generation(self, check_client_disconnected: Callable):
@@ -144,7 +144,7 @@ async def run_generation(self, check_client_disconnected: Callable):
         for attempt in range(1, max_retries + 1):
             try:
                 await self.page.keyboard.press('Escape')
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
                 run_btn = self.page.locator(SUBMIT_BUTTON_SELECTOR)
                 await expect_async(run_btn).to_be_visible(timeout=5000)
                 await expect_async(run_btn).to_be_enabled(timeout=5000)
@@ -160,7 +160,7 @@ async def run_generation(self, check_client_disconnected: Callable):
                     raise
                 self.logger.warning(f'[{self.req_id}] 点击 Run 失败 (尝试 {attempt}): {e}')
             if attempt < max_retries:
-                await asyncio.sleep(0.3)
+                await asyncio.sleep(0.15)
         raise Exception('点击 Run 按钮失败')
 
     async def wait_for_content(self, check_client_disconnected: Callable, timeout_seconds: int = 120) -> GeneratedContent:
@@ -229,7 +229,7 @@ async def wait_for_content(self, check_client_disconnected: Callable, timeout_se
             except Exception as e:
                 self.logger.warning(f'[{self.req_id}] 检查内容时出错: {e}')
 
-            await asyncio.sleep(0.5)
+            await asyncio.sleep(0.25)
 
     async def _extract_images_via_download(self, count: int) -> List[GeneratedImage]:
         import tempfile
@@ -244,9 +244,9 @@ async def _extract_images_via_download(self, count: int) -> List[GeneratedImage]
                 img = chunk.locator('img')
                 if await img.count() > 0:
                     await img.first.hover()
-                    await asyncio.sleep(0.5)
+                    await asyncio.sleep(0.25)
                     await chunk.evaluate('el => el.dispatchEvent(new MouseEvent("mouseenter", {bubbles: true}))')
-                    await asyncio.sleep(0.3)
+                    await asyncio.sleep(0.15)
 
                 download_btn = chunk.locator('button[aria-label="Download"]')
                 if await download_btn.count() == 0: