Fix AT-SPI crash and screenshot hangs, add screen capture after sleep

Framework/Built_In_Automation/Desktop/Linux/BuiltInFunctions.py

@@ -11,44 +11,116 @@
 
 from Framework.module_installer import install_missing_modules
 
+# ---------------------------------------------------------------------------
+# Optional system dependencies.
+#
+# pyatspi (AT-SPI2) is a thin Python layer over the GObject-Introspection
+# "Atspi" typelib, so the pip side (python3-pyatspi + pygobject) is only half of
+# it: the system typelib (gir1.2-atspi-2.0 / at-spi2-core) and a running
+# accessibility bus are required too. Having a display is not enough -- AT-SPI
+# is just as absent on a VNC/Xvfb desktop as on a bare server. python-xlib is
+# only used by the XComposite screenshot path, which already falls back to xwd.
+#
+# Importing this module therefore must never abort the process: CommonUtil
+# imports it just to take a desktop screenshot, and screenshot capture
+# (xdotool/xwd/XComposite based) does not need AT-SPI at all. A missing
+# dependency is recorded here and reported by the AT-SPI actions themselves as
+# an ordinary step failure.
+# ---------------------------------------------------------------------------
+
 try:
     import pyatspi
     from pyatspi.action import Action
     from pyatspi.editabletext import EditableText, Text
-except ImportError:
+
+    ATSPI_IMPORT_ERROR: Optional[str] = None
+except Exception:
     install_missing_modules(["python3-pyatspi==1.19.0", "pygobject==3.50.1"])
     try:
         import pyatspi
         from pyatspi.action import Action
         from pyatspi.editabletext import EditableText, Text
-    except ImportError:
-        sys.stderr.write(
-            "Error: system dependency is not installed. Install them by running Installer/setup_linux_inspector.sh.\n"
-        )
-        sys.exit(1)
+
+        ATSPI_IMPORT_ERROR = None
+    except Exception as atspi_error:
+        # pyatspi raises ImportError when the module is missing and ValueError
+        # ("Namespace Atspi not available") when the GI typelib is missing.
+        ATSPI_IMPORT_ERROR = f"{type(atspi_error).__name__}: {atspi_error}"
+
+        # The Accessible protocol below annotates methods with these interfaces,
+        # and annotations in a class body are evaluated at class-creation time,
+        # so the names must exist even when the real interfaces do not.
+        Action = EditableText = Text = Any  # type: ignore[misc,assignment]
+
+        class _AtspiUnavailable:
+            """Stand-in for the `pyatspi` module when AT-SPI is not installed.
+
+            Every attribute access raises, so an AT-SPI-dependent action fails as
+            a normal step failure with an actionable message (actions wrap their
+            body in try/except) instead of taking the node down.
+            """
+
+            def __getattr__(self, name: str):
+                raise RuntimeError(atspi_unavailable_message())
+
+        pyatspi = _AtspiUnavailable()  # type: ignore[assignment]
 
 try:
     from Xlib import X, display as xlib_display
     from Xlib.ext import composite as xlib_composite  # noqa: F401  registers ext methods
     from Xlib.error import XError
-except ImportError:
+
+    XLIB_IMPORT_ERROR: Optional[str] = None
+except Exception:
     install_missing_modules(["python-xlib==0.33"])
     try:
         from Xlib import X, display as xlib_display
         from Xlib.ext import composite as xlib_composite  # noqa: F401
         from Xlib.error import XError
-    except ImportError:
-        sys.stderr.write(
-            "Error: python-xlib is not installed. Install it by running Installer/setup_linux_inspector.sh.\n"
-        )
-        sys.exit(1)
+
+        XLIB_IMPORT_ERROR = None
+    except Exception as xlib_error:
+        XLIB_IMPORT_ERROR = f"{type(xlib_error).__name__}: {xlib_error}"
+        X = xlib_display = None  # type: ignore[assignment]
+
+        class XError(Exception):  # type: ignore[no-redef]
+            """Placeholder so `except XError` stays valid without python-xlib."""
 
 from Framework.Utilities import CommonUtil
 from Framework.Built_In_Automation.Shared_Resources import (
     BuiltInFunctionSharedResources as Shared_Resources,
 )
 from Framework.Utilities.decorators import logger
 
+_SETUP_HINT = "Installer/setup_linux_inspector.sh"
+
+
+def is_atspi_available() -> bool:
+    """True when the AT-SPI2 bindings imported successfully."""
+    return ATSPI_IMPORT_ERROR is None
+
+
+def atspi_unavailable_message() -> str:
+    """Actionable message for callers that need AT-SPI but cannot have it."""
+    return (
+        "Linux desktop automation requires AT-SPI2 (pyatspi), which is not importable on "
+        f"this machine ({ATSPI_IMPORT_ERROR}). Beyond the Python bindings it needs the system "
+        "AT-SPI2 typelib and a running accessibility bus -- on Debian/Ubuntu install "
+        "'at-spi2-core gir1.2-atspi-2.0' (or your distro's equivalent). "
+        f"See {_SETUP_HINT} for the X tools and accessibility settings."
+    )
+
+
+# Warn once, at import, rather than per action: the node keeps running, but
+# whoever reads the console should know why Linux desktop actions will fail.
+if ATSPI_IMPORT_ERROR:
+    sys.stderr.write(f"Warning: {atspi_unavailable_message()}\n")
+if XLIB_IMPORT_ERROR:
+    sys.stderr.write(
+        f"Warning: python-xlib is not available ({XLIB_IMPORT_ERROR}); "
+        f"screenshots fall back to xwd. See {_SETUP_HINT}.\n"
+    )
+
 
 class Collection: ...
 
@@ -328,6 +400,14 @@ def _capture_via_composite(file_path: str, winid: str) -> bool:
     Composite extension absent, deps missing, X errors) so the caller can
     fall back to a different strategy.
     """
+    if xlib_display is None:
+        CommonUtil.ExecLog(
+            MODULE_NAME,
+            f"Composite capture skipped: python-xlib unavailable ({XLIB_IMPORT_ERROR})",
+            4,
+        )
+        return False
+
     try:
         from PIL import Image
     except ImportError as e:
@@ -1013,7 +1093,13 @@ def _get_frame_geometry_for_window(window_id: str) -> dict | None:
     Used to align captured screenshots with the AT-SPI accessibility tree, whose
     coordinates may differ from the X11 window's geometry by the window
     decoration (title bar) offset.
+
+    Returns None when AT-SPI is unavailable, so the caller keeps the full
+    window pixmap instead of losing the screenshot entirely.
     """
+    if not is_atspi_available():
+        return None
+
     geometry = _get_window_geometry(window_id)
     if not geometry:
         return None

Framework/Built_In_Automation/Sequential_Actions/action_declarations/common.py

@@ -1,6 +1,6 @@
 declarations = (
     { "name": "step result",                                 "function": "step_result",                           "screenshot": "none" },
-    { "name": "sleep",                                       "function": "Sleep",                                 "screenshot": "none" },
+    { "name": "sleep",                                       "function": "Sleep",                                 "screenshot": "auto" },
     { "name": "wait",                                        "function": "Wait_For_Element",                      "screenshot": "none" },
     { "name": "wait disable",                                "function": "Wait_For_Element",                      "screenshot": "none" },
     { "name": "save text",                                   "function": "Save_Text",                             "screenshot": "none" },

Framework/Utilities/CommonUtil.py

@@ -825,6 +825,27 @@ def PhysicalAvailableMemory():
 )  # Initialize global variables for TakeScreenShot()
 
 
+AUTO_SCREEN_CAPTURE = "auto"
+
+
+def _resolve_auto_screen_capture(shared_variables):
+    """Pick the capture type for an action that is not tied to a platform.
+
+    Common actions (e.g. "sleep") are shared by every module, so their
+    declaration cannot name a platform up front. "auto" defers that choice to
+    here, where the drivers the test actually has open are visible.
+
+    Only web is resolved today: a sleep is most often used to let a page settle,
+    and the whole point of the capture is to show the page state once the wait
+    is over. Everything else stays "none" so no other module changes behaviour.
+    Add a "mobile" branch here if the same is ever wanted for Appium runs.
+    """
+    for driver_key in ("selenium_driver", "playwright_page"):
+        if shared_variables.get(driver_key) is not None:
+            return "web"
+    return "none"
+
+
 def set_screenshot_vars(shared_variables):
     """ Save screen capture type and selenium/appium driver objects as global variables, so TakeScreenShot() can access them """
     # We can't import Shared Variables due to cyclic imports causing local runs to break, so this is the work around
@@ -836,6 +857,8 @@ def set_screenshot_vars(shared_variables):
     try:
         if "screen_capture" in shared_variables:  # Type of screenshot (desktop/mobile)
             screen_capture_type = shared_variables["screen_capture"]
+        if screen_capture_type == AUTO_SCREEN_CAPTURE:
+            screen_capture_type = _resolve_auto_screen_capture(shared_variables)
         if screen_capture_type == "mobile":  # Appium driver object
             if "device_id" in shared_variables:
                 device_id = shared_variables[
@@ -1041,6 +1064,94 @@ def _get_window_screenshot_bbox():
     return None
 
 
+_linux_capture_screenshot = None  # None = not resolved yet, False = resolution failed
+
+
+def _get_linux_capture_screenshot():
+    """Resolve the Linux desktop screenshot function once per process.
+
+    The Linux module pulls in optional system dependencies (AT-SPI, python-xlib)
+    that are absent on most nodes, and a module that fails to import is never
+    cached in sys.modules — so importing it per screenshot re-ran its dependency
+    install and re-raised per screenshot. Catch BaseException (it used to call
+    sys.exit(), and SystemExit slips straight past `except Exception`) and cache
+    the outcome, so a missing desktop stack costs one warning, not a traceback
+    on every capture.
+    """
+    global _linux_capture_screenshot
+
+    if _linux_capture_screenshot is None:
+        try:
+            from Framework.Built_In_Automation.Desktop.Linux.BuiltInFunctions import (
+                capture_screenshot as linux_capture_screenshot,
+            )
+
+            _linux_capture_screenshot = linux_capture_screenshot
+        except BaseException as e:
+            ExecLog(
+                MODULE_NAME,
+                "Linux desktop screenshot support is unavailable: %s: %s" % (type(e).__name__, e),
+                3,
+            )
+            _linux_capture_screenshot = False
+
+    return _linux_capture_screenshot or None
+
+
+SCREENSHOT_CAPTURE_TIMEOUT_SECONDS = 60
+
+
+def _log_capture_timeout(sModuleInfo, function_name, Method):
+    ExecLog(
+        sModuleInfo,
+        "Screenshot for Action: %s Method: %s did not finish within %s seconds and was "
+        "abandoned -- the browser/device is not responding. Continuing the run."
+        % (function_name, Method, SCREENSHOT_CAPTURE_TIMEOUT_SECONDS),
+        2,
+    )
+
+
+def _capture_with_timeout(capture, sModuleInfo, function_name, Method):
+    """Run one blocking screen capture, waiting at most SCREENSHOT_CAPTURE_TIMEOUT_SECONDS.
+
+    Selenium and Appium screenshots are synchronous HTTP calls with no read
+    timeout of their own -- `RemoteConnection._timeout` defaults to
+    `socket._GLOBAL_DEFAULT_TIMEOUT`, so urllib3 waits forever -- and
+    `TakeScreenShot` runs outside `_run_action_with_timeout`. A wedged browser
+    therefore blocked the capture (and the test case's final join on it)
+    indefinitely, with no timeout able to recover it.
+
+    The capture runs on a *daemon* thread and is simply abandoned on timeout --
+    the same trade-off `_ActionTimeoutWorker` makes for a hung action -- so an
+    abandoned capture can never hold up node shutdown. Any exception is
+    re-raised on the calling thread so the existing handlers in
+    `Thread_ScreenShot` still see it. Playwright pages are not passed here:
+    they must stay on their own thread, and use Playwright's own timeout.
+
+    Returns True when the capture finished in time.
+    """
+    done = threading.Event()
+    failure = {}
+
+    def runner():
+        try:
+            capture()
+        except BaseException:
+            failure["exc_info"] = sys.exc_info()
+        finally:
+            done.set()
+
+    threading.Thread(target=runner, name="zeuz_screenshot_capture", daemon=True).start()
+
+    if not done.wait(SCREENSHOT_CAPTURE_TIMEOUT_SECONDS):
+        _log_capture_timeout(sModuleInfo, function_name, Method)
+        return False
+
+    if "exc_info" in failure:
+        raise failure["exc_info"][1].with_traceback(failure["exc_info"][2])
+    return True
+
+
 def Thread_ScreenShot(function_name, image_folder, Method, Driver, image_name, skip_delay=False):
     """ Capture screen of mobile or desktop.
 
@@ -1087,39 +1198,48 @@ def Thread_ScreenShot(function_name, image_folder, Method, Driver, image_name, s
         if should_delay_before_capture and not skip_delay and not _wait_for_debug_screenshot_delay(sModuleInfo, function_name, Method):
             return
 
+        # Every capture below is bounded: a hung browser/device must not park the run.
+        def _desktop_grab():
+            image = ImageGrab_Mac_Win.grab(_get_window_screenshot_bbox())
+            image.save(ImageName, format="PNG")  # Save to disk
+
         # Capture screenshot of desktop
         if Method == "desktop":
             if sys.platform in ("linux", "linux2"):
-                # Import Linux screenshot function for AT-SPI desktop automation
-                try:
-                    if sys.platform in ("linux", "linux2"):
-                        from Framework.Built_In_Automation.Desktop.Linux.BuiltInFunctions import capture_screenshot as linux_capture_screenshot
-                except Exception:
-                    linux_capture_screenshot = None
-                if linux_capture_screenshot:
-                    linux_capture_screenshot(ImageName)
-                else:
+                linux_capture_screenshot = _get_linux_capture_screenshot()
+                if linux_capture_screenshot is None:
                     ExecLog(
                         sModuleInfo,
                         "Linux screenshot module not available",
                         3,
                     )
                     return
+                if not _capture_with_timeout(
+                    lambda: linux_capture_screenshot(ImageName), sModuleInfo, function_name, Method
+                ):
+                    return
             elif sys.platform == "win32" or sys.platform == "darwin":
-                bbox = _get_window_screenshot_bbox()
-                image = ImageGrab_Mac_Win.grab(bbox)
-                image.save(ImageName, format="PNG")  # Save to disk
+                if not _capture_with_timeout(_desktop_grab, sModuleInfo, function_name, Method):
+                    return
 
         # Capture screenshot of web browser
         elif Method == "web":
             if type(Driver).__module__.startswith("playwright."):
-                Driver.screenshot(path=ImageName)
-            else:
-                Driver.get_screenshot_as_file(ImageName)  # Must be .png, otherwise an exception occurs
+                # Playwright objects must stay on their own thread; bound it with its own timeout.
+                Driver.screenshot(path=ImageName, timeout=SCREENSHOT_CAPTURE_TIMEOUT_SECONDS * 1000)
+            elif not _capture_with_timeout(
+                lambda: Driver.get_screenshot_as_file(ImageName),  # Must be .png, otherwise an exception occurs
+                sModuleInfo, function_name, Method,
+            ):
+                return
 
         # Capture screenshot of mobile
         elif Method == "mobile":
-            Driver.save_screenshot(ImageName)  # Must be .png, otherwise an exception occurs
+            # Must be .png, otherwise an exception occurs
+            if not _capture_with_timeout(
+                lambda: Driver.save_screenshot(ImageName), sModuleInfo, function_name, Method
+            ):
+                return
         else:
             ExecLog(
                 sModuleInfo,

Framework/deploy_handler/long_poll_handler.py

@@ -6,6 +6,7 @@
 import random
 import json
 import httpx
+import requests
 from colorama import Fore
 from pathlib import Path
 from urllib.parse import urlparse

server/linux.py

@@ -49,8 +49,8 @@ def inspect(app_name: str | None = None, window_id: str | None = None):
     window of the app (an app may have multiple windows; see /linux/apps).
     """
     from Framework.Built_In_Automation.Desktop.Linux import BuiltInFunctions
-    if BuiltInFunctions is None:
-        return InspectorResponse(status="error", error="Linux automation module not available")
+    if not BuiltInFunctions.is_atspi_available():
+        return InspectorResponse(status="error", error=BuiltInFunctions.atspi_unavailable_message())
 
     try:
         # Determine app name
@@ -94,7 +94,7 @@ def inspect(app_name: str | None = None, window_id: str | None = None):
 def get_apps():
     """Return available Linux applications visible to AT-SPI."""
     from Framework.Built_In_Automation.Desktop.Linux import BuiltInFunctions
-    if BuiltInFunctions is None:
+    if not BuiltInFunctions.is_atspi_available():
         return []
 
     try:
@@ -117,7 +117,7 @@ def get_apps():
 async def upload_linux_ui_dump():
     """Continuously upload Linux UI dump if changed."""
     from Framework.Built_In_Automation.Desktop.Linux import BuiltInFunctions
-    if BuiltInFunctions is None:
+    if not BuiltInFunctions.is_atspi_available():
         return
 
     prev_xml_hash = ""