fix: bound task failures and release stalled workers

artemis/drivers/android/adb_driver.py

@@ -19,6 +19,7 @@ import base64
 from io import BytesIO
 from pathlib import Path
 from typing import Any, Literal
+from PIL import Image
 
 from adbutils import AdbClient, AdbDevice
 from artemis.clients.ui_automator_client import (
@@ -26,6 +27,7 @@ from artemis.clients.ui_automator_client import (
     _parse_hierarchy_xml_to_elements,
 )
 from artemis.config.agent import load_agent_config
+from artemis.config.settings import settings
 from artemis.config.paths import get_temp_dir
 from artemis.drivers.base import BaseDeviceDriver, KeyCode, ScreenData, SwipeDirection
 from artemis.toolchain import find_ffmpeg, find_scrcpy
@@ -91,6 +93,7 @@ class AndroidAdbDriver(BaseDeviceDriver):
         self._recording_process: asyncio.subprocess.Process | None = None
         self._recording_output_path: Path | None = None
         self._screenshot_config = load_agent_config().screenshot
+        self._screen_capture_task: asyncio.Task[ScreenData] | None = None
 
     @property
     def device_id(self) -> str:
@@ -122,6 +125,24 @@ class AndroidAdbDriver(BaseDeviceDriver):
             await asyncio.to_thread(self._ui_adb_client.disconnect)
 
     async def get_screen_data(self, skip_settling: bool = False) -> ScreenData:
+        """Bound observation latency and share unfinished capture work across retries."""
+        task = self._screen_capture_task
+        if task is None or task.done():
+            task = asyncio.create_task(self._get_screen_data(skip_settling))
+            # A timed-out caller does not cancel a blocking SDK thread or start
+            # another concurrent read on the same device. Retrieve late errors.
+            task.add_done_callback(lambda done: None if done.cancelled() else done.exception())
+            self._screen_capture_task = task
+        try:
+            return await asyncio.wait_for(
+                asyncio.shield(task), timeout=settings.ARTEMIS_OBSERVATION_TIMEOUT_SECONDS
+            )
+        except TimeoutError as exc:
+            raise TimeoutError(
+                f"设备屏幕观测超时（{settings.ARTEMIS_OBSERVATION_TIMEOUT_SECONDS:g} 秒）"
+            ) from exc
+
+    async def _get_screen_data(self, skip_settling: bool = False) -> ScreenData:
         """Captures screenshot and XML hierarchy concurrently."""
         if not skip_settling:
             await asyncio.sleep(0.3)
@@ -131,11 +152,14 @@ class AndroidAdbDriver(BaseDeviceDriver):
         screenshot_base64: str = ""
         ui_hierarchy_xml: str | None = None
         ui_elements: list[dict[str, Any]] = []
+        capture_errors: list[str] = []
 
         # Check if ui_adb_client has direct screen data support
         if self._ui_adb_client and hasattr(self._ui_adb_client, "get_screen_data"):
             try:
-                raw_ui_data = self._ui_adb_client.get_screen_data()
+                raw_ui_data = await asyncio.to_thread(self._ui_adb_client.get_screen_data)
+                if asyncio.iscoroutine(raw_ui_data):
+                    raw_ui_data = await raw_ui_data
                 if (
                     hasattr(raw_ui_data, "base64")
                     and isinstance(raw_ui_data.base64, str)
@@ -158,11 +182,18 @@ class AndroidAdbDriver(BaseDeviceDriver):
                     elif isinstance(raw_elements, list):
                         ui_elements = raw_elements
             except Exception as e:
+                capture_errors.append(str(e))
                 logger.debug(f"Direct ui_adb_client.get_screen_data error: {e}")
 
         if not screenshot_base64 or not isinstance(screenshot_base64, str):
             try:
-                pil_img = await asyncio.to_thread(self.device.screenshot)
+                raw_image = await asyncio.to_thread(
+                    self.device.shell,
+                    ["screencap", "-p"],
+                    encoding=None,
+                    timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                )
+                pil_img = Image.open(BytesIO(raw_image))
                 if hasattr(pil_img, "save"):
                     buf = BytesIO()
                     pil_img.save(buf, format="PNG")
@@ -170,18 +201,24 @@ class AndroidAdbDriver(BaseDeviceDriver):
                     self._width, self._height = pil_img.size
                     screenshot_base64 = base64.b64encode(screenshot_bytes).decode("utf-8")
             except Exception as e:
+                capture_errors.append(str(e))
                 logger.warning(f"Device screenshot capture failed on {self._device_id}: {e}")
 
         if not screenshot_base64 or not isinstance(screenshot_base64, str):
-            # Fallback for headless testing environments
-            screenshot_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
-            screenshot_bytes = base64.b64decode(screenshot_base64)
+            raise RuntimeError(
+                "设备截图失败：ADB 与屏幕辅助服务均未返回有效截图；" + "; ".join(capture_errors)
+            )
+
+        with Image.open(BytesIO(screenshot_bytes)) as image:
+            if min(image.size) <= 1:
+                raise RuntimeError("设备截图无效：收到占位图，禁止基于旧画面继续操作")
+            image.verify()
 
         # 2. Capture UI Hierarchy via UIAutomator
         if self._ui_adb_client and not ui_elements:
             try:
                 if hasattr(self._ui_adb_client, "get_hierarchy"):
-                    res = self._ui_adb_client.get_hierarchy()
+                    res = await asyncio.to_thread(self._ui_adb_client.get_hierarchy)
                     hierarchy = await res if asyncio.iscoroutine(res) else res
                     if isinstance(hierarchy, str):
                         ui_hierarchy_xml = hierarchy
@@ -189,7 +226,7 @@ class AndroidAdbDriver(BaseDeviceDriver):
                     else:
                         ui_elements = hierarchy
                 elif hasattr(self._ui_adb_client, "get_ui_elements"):
-                    res = self._ui_adb_client.get_ui_elements()
+                    res = await asyncio.to_thread(self._ui_adb_client.get_ui_elements)
                     ui_elements = await res if asyncio.iscoroutine(res) else res
             except Exception as e:
                 logger.debug(f"UI hierarchy extraction failed: {e}")
@@ -208,16 +245,8 @@ class AndroidAdbDriver(BaseDeviceDriver):
             screen_width=w,
             screen_height=h,
         )
-        clean_b64 = (
-            screenshot_base64
-            if isinstance(screenshot_base64, str) and screenshot_base64
-            else "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
-        )
-        clean_bytes = (
-            screenshot_bytes
-            if isinstance(screenshot_bytes, bytes) and screenshot_bytes
-            else base64.b64decode(clean_b64)
-        )
+        clean_b64 = screenshot_base64
+        clean_bytes = screenshot_bytes
         if self._screenshot_config.enabled:
             limited_bytes = limit_screenshot_long_edge(
                 clean_bytes, self._screenshot_config.max_long_edge
@@ -255,7 +284,9 @@ class AndroidAdbDriver(BaseDeviceDriver):
                     cmd = f" && sleep {delay_ms / 1000.0:.3f} && ".join(taps)
 
             logger.info(f"[ADB] {cmd}")
-            await asyncio.to_thread(self.device.shell, cmd)
+            await asyncio.to_thread(
+                self.device.shell, cmd, timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS
+            )
             return True
         except Exception as e:
             logger.error(f"Tap failed at ({x}, {y}): {e}")
@@ -275,7 +306,9 @@ class AndroidAdbDriver(BaseDeviceDriver):
         try:
             cmd = f"input swipe {start_x} {start_y} {end_x} {end_y} {duration_ms}"
             logger.info(f"[ADB] {cmd}")
-            await asyncio.to_thread(self.device.shell, cmd)
+            await asyncio.to_thread(
+                self.device.shell, cmd, timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS
+            )
             return True
         except Exception as e:
             logger.error(f"Swipe failed from ({start_x},{start_y}) to ({end_x},{end_y}): {e}")
@@ -324,7 +357,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
                     "input keyevent 67 && "
                     "input keyevent 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67 67"
                 )
-                await asyncio.to_thread(self.device.shell, clear_cmd)
+                await asyncio.to_thread(
+                    self.device.shell,
+                    clear_cmd,
+                    timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                )
             else:
                 # Append mode: move cursor to the very end of existing text
                 await asyncio.to_thread(
@@ -346,7 +383,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
                         set_clip_ok = True
 
                     if set_clip_ok:
-                        await asyncio.to_thread(self.device.shell, "input keyevent 279")
+                        await asyncio.to_thread(
+                            self.device.shell,
+                            "input keyevent 279",
+                            timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                        )
                         return True
                 except Exception as e:
                     logger.debug(f"Clipboard paste fallback to ADB input: {e}")
@@ -359,7 +400,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
                 if "adbkeyboard" in str(default_ime).lower():
                     b64_text = base64.b64encode(norm_text.encode("utf-8")).decode("utf-8")
                     broadcast_cmd = f"am broadcast -a ADB_INPUT_B64 --es msg '{b64_text}'"
-                    await asyncio.to_thread(self.device.shell, broadcast_cmd)
+                    await asyncio.to_thread(
+                        self.device.shell,
+                        broadcast_cmd,
+                        timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                    )
                     return True
             except Exception as e:
                 # ADBKeyboard probe/broadcast failed; fall through to native input.
@@ -370,10 +415,18 @@ class AndroidAdbDriver(BaseDeviceDriver):
             for i, line in enumerate(lines):
                 if i > 0:
                     # Send Enter key between lines
-                    await asyncio.to_thread(self.device.shell, "input keyevent 66")
+                    await asyncio.to_thread(
+                        self.device.shell,
+                        "input keyevent 66",
+                        timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                    )
                 if line:
                     escaped = _escape_for_adb_text(line)
-                    await asyncio.to_thread(self.device.shell, f"input text {escaped}")
+                    await asyncio.to_thread(
+                        self.device.shell,
+                        f"input text {escaped}",
+                        timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                    )
             return True
         except Exception as e:
             logger.error(f"Input text failed for '{text}': {e}")
@@ -386,7 +439,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
                 key_name = str(key).lower().replace("keycode.", "")
                 keycode_val = ANDROID_KEYCODE_MAP.get(key_name, key)
 
-            await asyncio.to_thread(self.device.shell, f"input keyevent {keycode_val}")
+            await asyncio.to_thread(
+                self.device.shell,
+                f"input keyevent {keycode_val}",
+                timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+            )
             return True
         except Exception as e:
             logger.error(f"Press key failed for '{key}': {e}")
@@ -395,7 +452,9 @@ class AndroidAdbDriver(BaseDeviceDriver):
     async def launch_app(self, package_name: str) -> bool:
         try:
             cmd = f"monkey -p {package_name} -c android.intent.category.LAUNCHER 1"
-            await asyncio.to_thread(self.device.shell, cmd)
+            await asyncio.to_thread(
+                self.device.shell, cmd, timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS
+            )
             return True
         except Exception as e:
             logger.error(f"Launch app failed for '{package_name}': {e}")
@@ -403,7 +462,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
 
     async def stop_app(self, package_name: str) -> bool:
         try:
-            await asyncio.to_thread(self.device.shell, f"am force-stop {package_name}")
+            await asyncio.to_thread(
+                self.device.shell,
+                f"am force-stop {package_name}",
+                timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+            )
             return True
         except Exception as e:
             logger.error(f"Stop app failed for '{package_name}': {e}")
@@ -439,7 +502,11 @@ class AndroidAdbDriver(BaseDeviceDriver):
     async def execute_shell(self, command: str, timeout_seconds: float = 15.0) -> str:
         try:
             return await asyncio.wait_for(
-                asyncio.to_thread(self.device.shell, command),
+                asyncio.to_thread(
+                    self.device.shell,
+                    command,
+                    timeout=settings.ARTEMIS_DEVICE_COMMAND_TIMEOUT_SECONDS,
+                ),
                 timeout=timeout_seconds,
             )
         except Exception as e: