android-midscene-automation 0.2.4 → 1.1.1

package/artemis-main/artemis/runtime/awake_lease.py

@@ -0,0 +1,283 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Cross-process ownership for Android's shell screen-bright wake lock."""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from datetime import UTC, datetime
+import hashlib
+import json
+import os
+from pathlib import Path
+import re
+import subprocess
+import time
+import uuid
+
+from artemis.config.paths import get_temp_dir
+from artemis.runtime.adb_endpoint import adb_command
+from artemis.runtime.process_probe import pid_is_alive
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+_WAKE_LOCK_TYPE = "SCREEN_BRIGHT_WAKE_LOCK"
+_WAKE_LOCK_RE = re.compile(
+    rf"{_WAKE_LOCK_TYPE}:.*\bheld=(true|false)\b.*\brefCount=(\d+)\b",
+    re.IGNORECASE,
+)
+
+
+class ScreenAwakeLease:
+    """Share exactly one Android shell wake-lock reference across clients.
+
+    ``cmd power set-wakelock`` exposes one ref-counted shell wake lock for a
+    display/type pair. Calling it independently from every process leaks
+    references whenever a process crashes. PID-aware lease files make the
+    Android reference a process-shared resource instead: the first live client
+    normalizes stale references and acquires one, later clients only register,
+    and the last client drains the shared reference on disconnect.
+    """
+
+    _MUTEX_TIMEOUT_SECONDS = 10.0
+    _MUTEX_POLL_SECONDS = 0.05
+    _MALFORMED_GRACE_SECONDS = 5.0
+    _MAX_REFERENCE_DRAIN = 256
+
+    def __init__(self, device_id: str):
+        self.device_id = device_id
+        self.token = uuid.uuid4().hex
+        device_hash = hashlib.sha256(device_id.encode("utf-8")).hexdigest()[:16]
+        root = get_temp_dir("awake-leases")
+        self.lease_dir = root / f"{device_hash}.leases"
+        self.mutex_path = root / f"{device_hash}.mutex"
+        self.lease_path = self.lease_dir / f"{self.token}.lease"
+        self._acquired = False
+
+    @staticmethod
+    def _process_created_at(pid: int | None = None) -> float:
+        try:
+            import psutil
+
+            return float(psutil.Process(pid or os.getpid()).create_time())
+        except Exception:
+            return 0.0
+
+    @staticmethod
+    def _owner_is_alive(payload: dict) -> bool:
+        try:
+            pid = int(payload["pid"])
+            created_at = float(payload.get("process_created_at", 0.0) or 0.0)
+        except (KeyError, TypeError, ValueError):
+            return False
+        return pid_is_alive(pid, created_at if created_at > 0 else None)
+
+    @staticmethod
+    def _read_payload(path: Path) -> dict | None:
+        try:
+            value = json.loads(path.read_text(encoding="utf-8"))
+            return value if isinstance(value, dict) else None
+        except (OSError, ValueError, TypeError, json.JSONDecodeError):
+            return None
+
+    @classmethod
+    def _path_is_stale(cls, path: Path) -> bool:
+        payload = cls._read_payload(path)
+        if payload is not None:
+            return not cls._owner_is_alive(payload)
+        try:
+            return time.time() - path.stat().st_mtime >= cls._MALFORMED_GRACE_SECONDS
+        except OSError:
+            return True
+
+    def _payload(self) -> dict:
+        return {
+            "pid": os.getpid(),
+            "process_created_at": self._process_created_at(),
+            "token": self.token,
+            "device_id": self.device_id,
+            "created_at": datetime.now(UTC).isoformat(),
+        }
+
+    @contextmanager
+    def _mutex(self):
+        self.mutex_path.parent.mkdir(parents=True, exist_ok=True)
+        started_at = time.monotonic()
+        acquired = False
+        while not acquired:
+            try:
+                fd = os.open(self.mutex_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
+            except FileExistsError:
+                if self._path_is_stale(self.mutex_path):
+                    try:
+                        self.mutex_path.unlink(missing_ok=True)
+                    except OSError:
+                        pass
+                    continue
+                if time.monotonic() - started_at >= self._MUTEX_TIMEOUT_SECONDS:
+                    raise TimeoutError("Timed out acquiring the Artemis awake-lease mutex.")
+                time.sleep(self._MUTEX_POLL_SECONDS)
+            else:
+                try:
+                    os.write(fd, json.dumps(self._payload()).encode("utf-8"))
+                finally:
+                    os.close(fd)
+                acquired = True
+        try:
+            yield
+        finally:
+            payload = self._read_payload(self.mutex_path)
+            if payload is not None and payload.get("token") == self.token:
+                try:
+                    self.mutex_path.unlink(missing_ok=True)
+                except OSError:
+                    pass
+
+    def _run(self, args: list[str], description: str) -> subprocess.CompletedProcess[str] | None:
+        try:
+            result = subprocess.run(
+                adb_command(["-s", self.device_id, *args]),
+                stdin=subprocess.DEVNULL,
+                capture_output=True,
+                text=True,
+                timeout=10,
+                check=False,
+            )
+            if result.returncode != 0:
+                detail = (result.stderr or result.stdout).strip()
+                logger.warning(
+                    f"Could not {description} on {self.device_id}: {detail or 'ADB command failed'}"
+                )
+            return result
+        except Exception as exc:
+            logger.warning(f"Could not {description} on {self.device_id}: {exc}")
+            return None
+
+    def _wake_lock_state(self) -> tuple[bool, int] | None:
+        result = self._run(
+            ["shell", "cmd", "power", "set-wakelock", "list"],
+            "inspect the shared screen-bright wake lock",
+        )
+        if result is None or result.returncode != 0:
+            return None
+        match = _WAKE_LOCK_RE.search(result.stdout)
+        if match is None:
+            return (False, 0)
+        return (match.group(1).lower() == "true", int(match.group(2)))
+
+    def _set_wake_lock(self, action: str) -> bool:
+        result = self._run(
+            [
+                "shell",
+                "cmd",
+                "power",
+                "set-wakelock",
+                action,
+                "-d",
+                "0",
+                _WAKE_LOCK_TYPE,
+            ],
+            f"{action} the shared screen-bright wake lock",
+        )
+        return result is not None and result.returncode == 0
+
+    def _drain_references(self) -> bool:
+        """Release legacy or crashed-process references until none remain."""
+        for _ in range(self._MAX_REFERENCE_DRAIN):
+            state = self._wake_lock_state()
+            if state is None:
+                return False
+            held, count = state
+            if not held or count <= 0:
+                return True
+            if not self._set_wake_lock("release"):
+                return False
+        logger.warning(
+            f"Refused to drain more than {self._MAX_REFERENCE_DRAIN} shared wake-lock "
+            f"references on {self.device_id}"
+        )
+        return False
+
+    def _live_leases(self) -> list[Path]:
+        self.lease_dir.mkdir(parents=True, exist_ok=True)
+        live: list[Path] = []
+        for path in self.lease_dir.glob("*.lease"):
+            if self._path_is_stale(path):
+                try:
+                    path.unlink(missing_ok=True)
+                except OSError:
+                    pass
+            else:
+                live.append(path)
+        return live
+
+    def acquire(self) -> bool:
+        if self._acquired:
+            return True
+        try:
+            with self._mutex():
+                live = self._live_leases()
+                state = self._wake_lock_state()
+                if state is None:
+                    return False
+
+                if not live:
+                    # No registered owner means every existing reference is
+                    # legacy or belongs to a crashed process.
+                    if not self._drain_references():
+                        return False
+                    state = (False, 0)
+
+                held, count = state
+                if not held or count <= 0:
+                    if not self._set_wake_lock("acquire"):
+                        return False
+                    verified = self._wake_lock_state()
+                    if verified is None or not verified[0] or verified[1] <= 0:
+                        return False
+
+                self.lease_path.write_text(
+                    json.dumps(self._payload(), ensure_ascii=False), encoding="utf-8"
+                )
+                self._acquired = True
+                return True
+        except Exception as exc:
+            logger.warning(f"Could not acquire shared awake lease on {self.device_id}: {exc}")
+            return False
+
+    def release(self) -> None:
+        if not self._acquired:
+            return
+        try:
+            with self._mutex():
+                self.lease_path.unlink(missing_ok=True)
+                if not self._live_leases():
+                    self._drain_references()
+        except Exception as exc:
+            logger.warning(f"Could not release shared awake lease on {self.device_id}: {exc}")
+        finally:
+            self._acquired = False
+
+    def cleanup_unowned_references(self) -> bool:
+        """Drain legacy shell wake locks only when no live old host owns them."""
+        try:
+            with self._mutex():
+                if self._live_leases():
+                    return False
+                return self._drain_references()
+        except Exception as exc:
+            logger.warning(f"Could not clean unowned screen wake locks on {self.device_id}: {exc}")
+            return False

package/artemis-main/artemis/runtime/awake_service.py

@@ -0,0 +1,343 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""USB-scoped screen-awake policy shared by every local Artemis entrypoint."""
+
+from __future__ import annotations
+
+import os
+import re
+import subprocess
+import threading
+from typing import Iterable
+
+from adbutils import AdbClient
+
+from artemis.config import settings
+from artemis.runtime.adb_endpoint import adb_command
+from artemis.runtime.awake_lease import ScreenAwakeLease
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+AWAKE_STRATEGY_USB = "usb_stay_on"
+AWAKE_STRATEGY_HEARTBEAT = "host_heartbeat"
+AWAKE_HEARTBEAT_INTERVAL_SECONDS = 5.0
+DEVICE_DISCOVERY_INTERVAL_SECONDS = 2.0
+
+
+def _awake_enabled() -> bool:
+    enabled = os.environ.get("ARTEMIS_KEEP_DEVICE_AWAKE", "true").strip().lower()
+    return enabled not in {"0", "false", "no", "off"}
+
+
+def _run_awake_adb_command(
+    device_id: str, args: list[str], description: str
+) -> subprocess.CompletedProcess[str] | None:
+    """Run one non-fatal ADB command used by the USB awake policy."""
+    try:
+        result = subprocess.run(
+            adb_command(["-s", device_id, *args]),
+            stdin=subprocess.DEVNULL,
+            capture_output=True,
+            text=True,
+            timeout=10,
+            check=False,
+        )
+        if result.returncode != 0:
+            detail = (result.stderr or result.stdout).strip()
+            logger.warning(
+                f"Could not {description} on {device_id}: {detail or 'ADB command failed'}"
+            )
+        return result
+    except Exception as exc:
+        logger.warning(f"Could not {description} on {device_id}: {exc}")
+        return None
+
+
+def sanitize_device_state(device_id: str) -> None:
+    """Sanitize Android device state by cleaning up orphan screenrecord or modal dialogs."""
+    if not device_id:
+        return
+    # 1. Kill any dangling screenrecord processes left from interrupted tasks
+    _run_awake_adb_command(
+        device_id,
+        ["shell", "pkill", "-f", "screenrecord"],
+        "kill orphan screenrecord processes",
+    )
+    # 2. Close system modal dialogs (ANR / crash popups)
+    _run_awake_adb_command(
+        device_id,
+        ["shell", "am", "broadcast", "-a", "android.intent.action.CLOSE_SYSTEM_DIALOGS"],
+        "dismiss system crash/ANR dialogs",
+    )
+
+
+def _configure_usb_stay_awake(device_id: str) -> str | None:
+    """Enable USB stay-awake, falling back to an effective host heartbeat.
+
+    The policy is intentionally scoped to USB power rather than an Artemis
+    process or task. Android applies the normal screen timeout as soon as USB
+    power is removed, even though no ADB command can be sent after unplugging.
+    AC, wireless, and dock charging retain the user's normal behavior.
+    """
+    if not _awake_enabled() or os.environ.get("ARTEMIS_CLOUD_MODE") == "1":
+        return None
+
+    # Migration cleanup: previous Artemis versions used a persistent shell
+    # wake lock. Drain it only after every live legacy owner has exited.
+    ScreenAwakeLease(device_id).cleanup_unowned_references()
+
+    commands = (
+        (
+            "enable USB-scoped stay-awake policy",
+            ["shell", "svc", "power", "stayon", "usb"],
+        ),
+        ("wake display", ["shell", "input", "keyevent", "KEYCODE_WAKEUP"]),
+        ("dismiss non-secure keyguard", ["shell", "wm", "dismiss-keyguard"]),
+    )
+    for description, args in commands:
+        _run_awake_adb_command(device_id, args, description)
+
+    configured = _run_awake_adb_command(
+        device_id,
+        ["shell", "settings", "get", "global", "stay_on_while_plugged_in"],
+        "verify USB-scoped stay-awake policy",
+    )
+    power_state = _run_awake_adb_command(
+        device_id,
+        ["shell", "dumpsys", "power"],
+        "verify active USB stay-awake state",
+    )
+    try:
+        value = int(configured.stdout.strip()) if configured is not None else -1
+    except (TypeError, ValueError):
+        value = -1
+    stay_on_active = (
+        configured is not None
+        and configured.returncode == 0
+        and value == 2
+        and power_state is not None
+        and power_state.returncode == 0
+        and re.search(r"\bmStayOn=true\b", power_state.stdout) is not None
+    )
+    if stay_on_active:
+        logger.info(f"USB-scoped stay-awake policy verified active on {device_id}")
+        return AWAKE_STRATEGY_USB
+
+    logger.warning(
+        f"USB stay-awake is not active on {device_id} (setting={value}); "
+        f"using a {AWAKE_HEARTBEAT_INTERVAL_SECONDS:g}-second host heartbeat"
+    )
+    primed = _run_awake_adb_command(
+        device_id,
+        ["shell", "input", "keyevent", "KEYCODE_UNKNOWN"],
+        "prime the host stay-awake heartbeat",
+    )
+    if primed is None or primed.returncode != 0:
+        return None
+    return AWAKE_STRATEGY_HEARTBEAT
+
+
+def _discover_connected_device_ids() -> list[str]:
+    """Return the explicitly targeted device, or the connected devices Artemis manages."""
+    if not _awake_enabled() or os.environ.get("ARTEMIS_CLOUD_MODE") == "1":
+        return []
+
+    target = os.environ.get("ARTEMIS_DEVICE_ID") or os.environ.get("ADB_DEVICE_SERIAL")
+    host = os.environ.get("ADB_HOST") or settings.ADB_HOST or "127.0.0.1"
+    port_text = os.environ.get("ADB_PORT")
+    try:
+        port = int(port_text) if port_text else int(settings.ADB_PORT or 5037)
+    except (TypeError, ValueError):
+        port = 5037
+
+    try:
+        devices = AdbClient(host=host, port=port).device_list()
+        serials = [device.serial for device in devices if device.serial]
+    except Exception as exc:
+        logger.debug(f"Could not discover a device for the USB awake policy: {exc}")
+        return []
+
+    if target:
+        return [target] if target in serials else []
+    # Keep awake only the devices Artemis actually manages: serials claimed in
+    # the device pool by an active execution lock or a live queue reservation.
+    # Other devices sharing the same ADB server belong to unrelated tools or
+    # users, and their power policy must not be touched.
+    try:
+        from artemis.runtime.device_pool import device_pool
+
+        claimed = device_pool.get_claimed_serials()
+    except Exception as exc:
+        logger.debug(f"Could not query claimed devices for the awake policy: {exc}")
+        return []
+    return [serial for serial in serials if serial in claimed]
+
+
+class ScreenAwakeService:
+    """Apply USB stay-awake or a connection-scoped host heartbeat."""
+
+    def __init__(self):
+        self._lock = threading.RLock()
+        self._strategies: dict[str, str] = {}
+        self._heartbeat_stops: dict[str, threading.Event] = {}
+        self._heartbeat_threads: dict[str, threading.Thread] = {}
+        self._monitor_stop: threading.Event | None = None
+        self._monitor_thread: threading.Thread | None = None
+        self._shutdown_requested = False
+
+    def start(self, device_ids: Iterable[str] | None = None) -> dict[str, str | None]:
+        if self._shutdown_requested:
+            return {}
+        selected = list(device_ids) if device_ids is not None else _discover_connected_device_ids()
+        results = {}
+        for device_id in dict.fromkeys(selected):
+            if self._shutdown_requested:
+                break
+            if device_id:
+                results[device_id] = self.ensure_device(device_id)
+        if (
+            _awake_enabled()
+            and os.environ.get("ARTEMIS_CLOUD_MODE") != "1"
+            and not self._shutdown_requested
+        ):
+            self._start_device_monitor()
+        return results
+
+    def ensure_device(self, device_id: str) -> str | None:
+        """Configure one device once per process; Android owns unplug behavior."""
+        if (
+            not device_id
+            or not _awake_enabled()
+            or os.environ.get("ARTEMIS_CLOUD_MODE") == "1"
+            or self._shutdown_requested
+        ):
+            return None
+        with self._lock:
+            existing = self._strategies.get(device_id)
+            if existing is not None:
+                return existing
+            strategy = _configure_usb_stay_awake(device_id)
+            if strategy is not None:
+                self._strategies[device_id] = strategy
+                if strategy == AWAKE_STRATEGY_HEARTBEAT:
+                    self._start_heartbeat(device_id)
+            return strategy
+
+    def _start_heartbeat(self, device_id: str) -> None:
+        if device_id in self._heartbeat_threads:
+            return
+        stop_event = threading.Event()
+
+        def heartbeat() -> None:
+            while not stop_event.wait(AWAKE_HEARTBEAT_INTERVAL_SECONDS):
+                _run_awake_adb_command(
+                    device_id,
+                    ["shell", "input", "keyevent", "KEYCODE_UNKNOWN"],
+                    "send the host stay-awake heartbeat",
+                )
+
+        thread = threading.Thread(
+            target=heartbeat,
+            name=f"artemis-awake-heartbeat-{device_id}",
+            daemon=True,
+        )
+        self._heartbeat_stops[device_id] = stop_event
+        self._heartbeat_threads[device_id] = thread
+        thread.start()
+
+    def _start_device_monitor(self) -> None:
+        with self._lock:
+            if self._monitor_thread is not None and self._monitor_thread.is_alive():
+                return
+            stop_event = threading.Event()
+
+            def monitor() -> None:
+                while not stop_event.wait(DEVICE_DISCOVERY_INTERVAL_SECONDS):
+                    connected = set(_discover_connected_device_ids())
+                    if stop_event.is_set():
+                        break
+                    self._reconcile_connected_devices(connected)
+
+            thread = threading.Thread(
+                target=monitor,
+                name="artemis-awake-device-monitor",
+                daemon=True,
+            )
+            self._monitor_stop = stop_event
+            self._monitor_thread = thread
+            thread.start()
+
+    def _reconcile_connected_devices(self, connected: set[str]) -> None:
+        """Stop unplugged heartbeats and enroll newly attached devices."""
+        with self._lock:
+            missing = set(self._strategies) - connected
+        for device_id in missing:
+            self._stop_device(device_id)
+        for device_id in connected:
+            self.ensure_device(device_id)
+
+    def _stop_device(self, device_id: str) -> None:
+        with self._lock:
+            stop_event = self._heartbeat_stops.pop(device_id, None)
+            thread = self._heartbeat_threads.pop(device_id, None)
+            self._strategies.pop(device_id, None)
+            if stop_event is not None:
+                stop_event.set()
+        if thread is not None and thread is not threading.current_thread():
+            thread.join(timeout=2.0)
+
+    def shutdown(self) -> None:
+        """Stop host work; unplug behavior remains owned by Android's USB policy."""
+        self._shutdown_requested = True
+        with self._lock:
+            monitor_stop = self._monitor_stop
+            monitor_thread = self._monitor_thread
+            heartbeat_stops = list(self._heartbeat_stops.values())
+            heartbeat_threads = list(self._heartbeat_threads.values())
+            if monitor_stop is not None:
+                monitor_stop.set()
+            for stop_event in heartbeat_stops:
+                stop_event.set()
+            self._monitor_stop = None
+            self._monitor_thread = None
+            self._heartbeat_stops.clear()
+            self._heartbeat_threads.clear()
+            self._strategies.clear()
+
+        if monitor_thread is not None:
+            monitor_thread.join(timeout=2.0)
+        for thread in heartbeat_threads:
+            thread.join(timeout=2.0)
+
+    @property
+    def device_ids(self) -> tuple[str, ...]:
+        with self._lock:
+            return tuple(self._strategies)
+
+
+screen_awake_service = ScreenAwakeService()
+
+
+def start_awake_service(device_ids: Iterable[str] | None = None) -> dict[str, str | None]:
+    return screen_awake_service.start(device_ids)
+
+
+def ensure_device_awake(device_id: str) -> str | None:
+    return screen_awake_service.ensure_device(device_id)
+
+
+def shutdown_awake_service() -> None:
+    screen_awake_service.shutdown()