411 lines
16 KiB
Python
411 lines
16 KiB
Python
# 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 hashlib
|
|
import re
|
|
import subprocess
|
|
import threading
|
|
from typing import Iterable
|
|
|
|
from adbutils import AdbClient
|
|
from filelock import FileLock, Timeout
|
|
|
|
from artemis.config import settings
|
|
from artemis.runtime.adb_endpoint import adb_command, current_adb_endpoint
|
|
from artemis.service_registry import registry_directory
|
|
from artemis.runtime.awake_lease import ScreenAwakeLease
|
|
from artemis.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 belongs to the long-lived host. A queued task worker
|
|
# must not wait on another process's optional legacy wake-lock cleanup.
|
|
if os.environ.get("ARTEMIS_TASK_WORKER") != "1":
|
|
try:
|
|
ScreenAwakeLease(device_id).cleanup_unowned_references()
|
|
except OSError as exc:
|
|
logger.warning(f"Could not prepare optional legacy awake cleanup on {device_id}: {exc}")
|
|
|
|
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 []
|
|
|
|
# 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 and (not target or serial == target)]
|
|
|
|
|
|
def _owner_lock(device_id: str) -> FileLock:
|
|
# Shared by all checkouts and processes using this user's ADB endpoint.
|
|
scope = f"{current_adb_endpoint().identity}/{device_id}"
|
|
key = hashlib.sha256(scope.encode()).hexdigest()
|
|
directory = registry_directory().parent / "awake"
|
|
directory.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
return FileLock(directory / f"{key}.lock", thread_local=False)
|
|
|
|
|
|
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
|
|
self._owners: dict[str, FileLock] = {}
|
|
self._ownership_warnings: set[str] = set()
|
|
|
|
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
|
|
try:
|
|
owner = _owner_lock(device_id)
|
|
owner.acquire(timeout=0)
|
|
except Timeout:
|
|
# Do not cache a follower as owner: discovery can take over if
|
|
# the process holding the OS lock exits or releases the device.
|
|
return None
|
|
except OSError as exc:
|
|
# Awake is optional. Never bypass ownership or let an ACL / I/O
|
|
# failure abort device connection, resume, or the host monitor.
|
|
if device_id not in self._ownership_warnings:
|
|
logger.warning(
|
|
f"Could not acquire awake ownership on {device_id}; skipping optional "
|
|
f"stay-awake policy: {exc}"
|
|
)
|
|
self._ownership_warnings.add(device_id)
|
|
return None
|
|
self._ownership_warnings.discard(device_id)
|
|
try:
|
|
strategy = _configure_usb_stay_awake(device_id)
|
|
except BaseException:
|
|
owner.release()
|
|
raise
|
|
if strategy is not None:
|
|
self._owners[device_id] = owner
|
|
self._strategies[device_id] = strategy
|
|
if strategy == AWAKE_STRATEGY_HEARTBEAT:
|
|
self._start_heartbeat(device_id)
|
|
else:
|
|
owner.release()
|
|
return strategy
|
|
|
|
def _start_heartbeat(self, device_id: str) -> None:
|
|
if device_id in self._heartbeat_threads:
|
|
return
|
|
stop_event = threading.Event()
|
|
owner = self._owners.get(device_id)
|
|
|
|
def heartbeat() -> None:
|
|
try:
|
|
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",
|
|
)
|
|
finally:
|
|
# Release only after the last in-flight command has returned.
|
|
if owner is not None:
|
|
owner.release()
|
|
|
|
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:
|
|
owner = self._owners.get(device_id)
|
|
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)
|
|
with self._lock:
|
|
if thread is None or not thread.is_alive():
|
|
if owner is not None:
|
|
owner.release()
|
|
if self._owners.get(device_id) is owner:
|
|
self._owners.pop(device_id, None)
|
|
|
|
def release_device(self, device_id: str) -> None:
|
|
"""Release this process's awake ownership when a task gives up its device."""
|
|
self._stop_device(device_id)
|
|
|
|
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)
|
|
for device_id, owner in list(self._owners.items()):
|
|
# Live heartbeat finalizers retain ownership until their ADB call
|
|
# exits; non-heartbeat policies have no pending command to drain.
|
|
if not any(
|
|
t.is_alive() and t.name == f"artemis-awake-heartbeat-{device_id}"
|
|
for t in heartbeat_threads
|
|
):
|
|
owner.release()
|
|
self._owners.pop(device_id, None)
|
|
|
|
@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()
|