diff --git a/apps/admin_console/core/state.py b/apps/admin_console/core/state.py index 951b1de..f88e9f0 100644 --- a/apps/admin_console/core/state.py +++ b/apps/admin_console/core/state.py @@ -17,29 +17,7 @@ from collections.abc import Callable from typing import Any from artemis.config import PAUSE_FILE - - -def _pid_exists(pid: int) -> bool: - """Best-effort liveness probe; returns True when liveness cannot be determined. - - psutil is preferred because on Windows ``os.kill(pid, 0)`` does not report - dead pids reliably. - """ - try: - import psutil - - return psutil.pid_exists(pid) - except Exception: - pass - try: - import os - - os.kill(pid, 0) - return True - except PermissionError: - return True - except (ProcessLookupError, OSError): - return False +from artemis.runtime.process_probe import pid_is_alive class ServerState: @@ -59,7 +37,12 @@ class ServerState: self.current_profile: str | None = None self.active_connections: dict[str, dict[str, Any]] = {} self.active_session_id: str | None = None - self.was_stopped_manually: bool = False + # Run keys (session id, or the synthetic run key of session-less runs) + # that were stopped manually. Tracked per run so a manual stop of one + # device's task never pollutes the terminal-status resolution of a + # concurrently running task on another device. Entries are discarded by + # each run's finalizer. + self.manually_stopped_run_ids: set[str] = set() self.cancelled_session_ids: set[str] = set() self.startup_progress: dict[str, list[dict[str, Any]]] = {} @@ -143,7 +126,7 @@ class ServerState: self.active_runs.pop(sid, None) continue pid = getattr(proc, "pid", None) - if pid and not _pid_exists(pid): + if pid and not pid_is_alive(pid): self.active_runs.pop(sid, None) @property @@ -168,15 +151,8 @@ class ServerState: self.current_process = None else: pid = getattr(self.current_process, "pid", None) - if pid: - try: - import os - - os.kill(pid, 0) - has_proc = True - except (ProcessLookupError, PermissionError, OSError): - has_proc = False - self.current_process = None + if pid and pid_is_alive(pid): + has_proc = True else: has_proc = False self.current_process = None @@ -189,12 +165,9 @@ class ServerState: for sid, conn in list(self.active_connections.items()): c_pid = conn.get("pid") if c_pid: - try: - import os - - os.kill(c_pid, 0) + if pid_is_alive(c_pid): has_live_connection = True - except (ProcessLookupError, PermissionError, OSError): + else: self.active_connections.pop(sid, None) if not has_proc and not has_running_item and not has_live_connection: diff --git a/apps/admin_console/database/repositories/session_repository.py b/apps/admin_console/database/repositories/session_repository.py index 4cfe55d..a0ae1d1 100644 --- a/apps/admin_console/database/repositories/session_repository.py +++ b/apps/admin_console/database/repositories/session_repository.py @@ -41,14 +41,19 @@ class SessionRepository: @staticmethod def process_is_alive(pid: Any) -> bool: - """Return whether a session's worker PID is still alive.""" - try: - import psutil + """Return whether a session's worker PID is still alive. - process = psutil.Process(int(pid)) - return process.is_running() and process.status() != psutil.STATUS_ZOMBIE - except Exception: + Delegates to the repository-wide probe: indeterminate liveness + (e.g. AccessDenied) counts as alive so a running session is never + reaped by mistake. + """ + try: + pid_int = int(pid) + except (TypeError, ValueError): return False + from artemis.runtime.process_probe import pid_is_alive + + return pid_is_alive(pid_int) def get_all_sessions(self) -> list[dict[str, Any]]: with db_session(self.db_path) as conn: diff --git a/apps/admin_console/services/task_queue_service.py b/apps/admin_console/services/task_queue_service.py index 0737e3e..8f5892d 100644 --- a/apps/admin_console/services/task_queue_service.py +++ b/apps/admin_console/services/task_queue_service.py @@ -415,10 +415,12 @@ class TaskQueueService: task_item["status"] = "running" task_item["start_time"] = time.time() - # A fresh launch clears a stale stop request from a previous task, - # matching the historical serial worker's per-task reset. Per-session - # stops are tracked in cancelled_session_ids and are unaffected. - state.was_stopped_manually = False + # A fresh launch clears a stale stop request left over for this run + # key from a previous task. Manual stops are tracked per run in + # manually_stopped_run_ids so stopping one device's task never + # affects concurrent runs. Per-session stops are tracked in + # cancelled_session_ids and are unaffected. + state.manually_stopped_run_ids.discard(run_key) state.current_goal = goal state.current_profile = profile state.active_session_id = sess_id @@ -546,7 +548,7 @@ class TaskQueueService: if sess_id and ( str(sess_id) in getattr(state, "cancelled_session_ids", set()) - or state.was_stopped_manually + or run_key in state.manually_stopped_run_ids ): print( f"[QueueWorker] Task [{sess_id}] was cancelled during launch. Terminating." @@ -557,7 +559,7 @@ class TaskQueueService: returncode = await cls._wait_for_worker_process(proc) print(f"[QueueWorker] Task [{sess_id}] exited with returncode {returncode}") - manual_stop = state.was_stopped_manually or bool( + manual_stop = run_key in state.manually_stopped_run_ids or bool( sess_id and str(sess_id) in state.cancelled_session_ids ) @@ -663,6 +665,7 @@ class TaskQueueService: cls._remove_task(sess_id) state.cancelled_session_ids.discard(str(sess_id)) state.cancelled_session_ids.discard(run_key) + state.manually_stopped_run_ids.discard(run_key) state.active_runs.pop(run_key, None) if proc is not None and state.current_process is proc: state.current_process = None @@ -670,7 +673,9 @@ class TaskQueueService: state.active_session_id = None state.current_goal = None state.current_profile = None - state.was_stopped_manually = False + # No runs left: any not-yet-consumed manual-stop markers are + # stale and must not leak into future runs. + state.manually_stopped_run_ids.clear() state.wake_event.set() @classmethod @@ -892,11 +897,13 @@ class TaskQueueService: }, ) - if state.active_runs or state.current_process: - state.was_stopped_manually = True for run_key, run in list(state.active_runs.items()): - # Track the stop per run so each finalizer resolves its terminal - # status as cancelled even if the global flag is reset meanwhile. + # Mark each run individually: the global "stopped manually" + # boolean was shared process-wide and polluted the terminal + # status of unrelated concurrent runs. + state.manually_stopped_run_ids.add(str(run_key)) + # Also track it per session so each finalizer resolves its + # terminal status as cancelled. state.cancelled_session_ids.add(str(run_key)) run_proc = run.get("process") if run_proc is not None and run_proc.returncode is None: @@ -969,20 +976,35 @@ class TaskQueueService: # Resolve the locally-managed run for this target (concurrent scheduling # keeps one subprocess per run in state.active_runs). local_run = None + local_run_key: str | None = None if target_sid: local_run = state.active_runs.get(target_sid) + if local_run is not None: + local_run_key = target_sid elif target_device: - local_run = next( + local_run_key, local_run = next( ( - run - for run in state.active_runs.values() + (key, run) + for key, run in state.active_runs.items() if run.get("device_id") and str(run["device_id"]) == target_device ), + (None, None), + ) + elif len(state.active_runs) == 1: + local_run_key, local_run = next(iter(state.active_runs.items())) + local_proc = (local_run or {}).get("process") or state.current_process + if local_run is None and local_proc is not None: + # Legacy path: current_process without a resolved run entry. Find + # the run that owns this process so a manual stop can be attributed + # to it instead of to the whole scheduler. + local_run_key = next( + ( + str(key) + for key, run in state.active_runs.items() + if run.get("process") is local_proc + ), None, ) - elif len(state.active_runs) == 1: - local_run = next(iter(state.active_runs.values())) - local_proc = (local_run or {}).get("process") or state.current_process local_pid = getattr(local_proc, "pid", None) is_local_owner = bool(owner_pid and local_pid and owner_pid == local_pid) @@ -1006,11 +1028,16 @@ class TaskQueueService: ), None, ) + # active_session_id mirrors the most recently launched run, so with + # concurrent runs the resolved local run key (== its session id) must + # take precedence to avoid attributing the stop to an unrelated run. stopped_session_id = ( owner.session_id if owner and owner.session_id else target_sid if target_sid + else local_run_key + if local_run_key is not None else state.active_session_id if is_local_owner else local_item.get("session_id") @@ -1038,8 +1065,8 @@ class TaskQueueService: # initialization window before Agent acquires the device lease. if target_sid: state.cancelled_session_ids.add(target_sid) - else: - state.was_stopped_manually = True + elif local_run_key is not None: + state.manually_stopped_run_ids.add(str(local_run_key)) if local_pid: try: stopped = process_supervisor.terminate_tree(local_pid) @@ -1092,8 +1119,8 @@ class TaskQueueService: # Per-session cancellation: the run's own finalizer resolves the # terminal status without affecting other concurrent runs. state.cancelled_session_ids.add(str(stopped_session_id)) - else: - state.was_stopped_manually = True + elif local_run_key is not None: + state.manually_stopped_run_ids.add(str(local_run_key)) if local_proc is not None and state.current_process is local_proc: state.current_process = None diff --git a/artemis/runtime/__init__.py b/artemis/runtime/__init__.py index bb0d537..8199fcb 100644 --- a/artemis/runtime/__init__.py +++ b/artemis/runtime/__init__.py @@ -28,6 +28,7 @@ from artemis.runtime.adb_endpoint import ( current_adb_endpoint, ) from artemis.runtime.device_pool import DevicePool, DeviceStatus, device_pool +from artemis.runtime.process_probe import pid_is_alive from artemis.runtime.daemon_client import ( ensure_daemon_running, get_daemon_session, @@ -78,6 +79,7 @@ __all__ = [ "get_server_status", "is_daemon_running", "is_port_in_use", + "pid_is_alive", "ProcessSupervisor", "process_supervisor", "read_server_info", diff --git a/artemis/runtime/awake_lease.py b/artemis/runtime/awake_lease.py index f7312e1..371804c 100644 --- a/artemis/runtime/awake_lease.py +++ b/artemis/runtime/awake_lease.py @@ -28,6 +28,7 @@ import time import uuid from artemis.config.paths import get_temp_dir +from artemis.runtime.process_probe import pid_is_alive from artemis.utils.logger import get_logger logger = get_logger(__name__) @@ -77,16 +78,11 @@ class ScreenAwakeLease: @staticmethod def _owner_is_alive(payload: dict) -> bool: try: - import psutil - pid = int(payload["pid"]) - process = psutil.Process(pid) - if not process.is_running(): - return False - created_at = float(payload.get("process_created_at", 0.0)) - return created_at <= 0 or abs(process.create_time() - created_at) < 1.0 - except Exception: + 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: diff --git a/artemis/runtime/device_lock.py b/artemis/runtime/device_lock.py index 94e89bb..38429a0 100644 --- a/artemis/runtime/device_lock.py +++ b/artemis/runtime/device_lock.py @@ -27,6 +27,7 @@ from typing import Any import uuid from artemis.config.paths import get_temp_dir +from artemis.runtime.process_probe import pid_is_alive class DeviceBusyError(RuntimeError): @@ -202,22 +203,22 @@ class DeviceExecutionLock: def _owner_is_alive(owner: DeviceLockOwner) -> bool: if owner.pid <= 0: return False - try: - import psutil + if owner.process_created_at <= 0: + # Legacy owner payloads without a create time cannot use PID-reuse + # protection; require an Artemis-looking process before trusting + # the record (a recycled PID must not keep a lock alive forever). + try: + import psutil - process = psutil.Process(owner.pid) - if not process.is_running(): - return False - if owner.process_created_at <= 0: - try: - name = process.name().lower() - cmdline = " ".join(process.cmdline()).lower() - return any(k in name or k in cmdline for k in ("python", "artemis", "pytest")) - except Exception: + process = psutil.Process(owner.pid) + if not process.is_running(): return False - return abs(process.create_time() - owner.process_created_at) < 1.0 - except Exception: - return False + name = process.name().lower() + cmdline = " ".join(process.cmdline()).lower() + return any(k in name or k in cmdline for k in ("python", "artemis", "pytest")) + except Exception: + return False + return pid_is_alive(owner.pid, owner.process_created_at) @staticmethod def _read_owner(path: Path) -> DeviceLockOwner | None: @@ -268,6 +269,108 @@ class DeviceExecutionLock: payload["lock_scope"] = str(lock_scope) return payload + @classmethod + def _ticket_targets_device( + cls, + ticket_owner: DeviceLockOwner, + *, + target_lock_id: str, + target_scope: str | None, + include_pending: bool, + ) -> bool: + """Return whether a queue ticket belongs to the given device's queue. + + A concrete ticket matches only its own device. A ticket whose device id + is still ``pending``/``any`` is a submission placeholder that no worker + has claimed yet; with ``include_pending=True`` it matches any device + with a compatible lock scope (used by listings, and by + ``_build_device_queue`` which then narrows the match down to a single + parked device), while ``include_pending=False`` excludes it. + """ + ticket_dev = cls._normalize_device_id(ticket_owner.device_id) + ticket_lock_id = cls._normalize_lock_id( + ticket_owner.device_id, + ticket_owner.lock_scope, + ) + if ticket_lock_id == target_lock_id: + return True + if not include_pending: + return False + same_scope = ( + not ticket_owner.lock_scope + or not target_scope + or ticket_owner.lock_scope == target_scope + ) + return ticket_dev in ("pending", "any") and same_scope + + @classmethod + def _build_device_queue( + cls, + wait_files: list[Path], + *, + target_lock_id: str, + target_scope: str | None, + ) -> list[Path]: + """Return the FIFO wait-ticket queue that gates ``target_lock_id``. + + Concrete tickets belong to exactly their own device's queue. An + unclaimed ``pending``/``any`` ticket is *parked* on at most one device: + pending tickets, taken in FIFO order, occupy the idle scope-compatible + candidate devices in sorted lock-id order. Historically such a ticket + was counted into every device's queue, so a single unclaimed + head-of-queue ticket blocked all devices at once. Parking preserves + submission order on the one device the ticket is presumed to claim + while leaving every other idle device schedulable. + + The pending submission itself cannot starve: its ticket file keeps its + original timestamp, so once a worker claims a concrete device the + ticket sorts ahead of every younger ticket for that device. + """ + entries: list[tuple[Path, str | None]] = [] + known_lock_ids: set[str] = set() + for wait_path in wait_files: + owner = cls._read_owner(wait_path) + if owner is None: + continue + if cls._normalize_device_id(owner.device_id) in ("pending", "any"): + if cls._ticket_targets_device( + owner, + target_lock_id=target_lock_id, + target_scope=target_scope, + include_pending=True, + ): + entries.append((wait_path, None)) + continue + same_scope = ( + not owner.lock_scope + or not target_scope + or owner.lock_scope == target_scope + ) + if not same_scope: + continue + lock_id = cls._normalize_lock_id(owner.device_id, owner.lock_scope) + known_lock_ids.add(lock_id) + entries.append((wait_path, lock_id)) + + lock_dir = get_temp_dir("device-locks") + idle_candidates = sorted( + lock_id + for lock_id in known_lock_ids + if not (lock_dir / f"artemis-device-{lock_id}.lock").exists() + ) + parked: dict[Path, str] = {} + next_candidate = 0 + for wait_path, lock_id in entries: + if lock_id is None and next_candidate < len(idle_candidates): + parked[wait_path] = idle_candidates[next_candidate] + next_candidate += 1 + + return [ + wait_path + for wait_path, lock_id in entries + if lock_id == target_lock_id or parked.get(wait_path) == target_lock_id + ] + @classmethod def reserve( cls, @@ -547,27 +650,15 @@ class DeviceExecutionLock: ) is_eligible = is_reentrant or bool(all_wait_files and all_wait_files[0] == self._queue_path) else: - # In multi-device parallel mode, per-device FIFO applies - target_dev = self.clean_device_id - device_queue: list[Path] = [] - for wait_path in all_wait_files: - ticket_owner = self._read_owner(wait_path) - if ticket_owner is None: - continue - ticket_dev = self._normalize_device_id(ticket_owner.device_id) - ticket_lock_id = self._normalize_lock_id( - ticket_owner.device_id, - ticket_owner.lock_scope, - ) - same_scope = ( - not ticket_owner.lock_scope - or not self.lock_scope - or ticket_owner.lock_scope == self.lock_scope - ) - if ticket_lock_id == target_dev or ( - ticket_dev in ("pending", "any") and same_scope - ): - device_queue.append(wait_path) + # In multi-device parallel mode, per-device FIFO applies. + # Unclaimed pending/any tickets are parked on at most one + # idle device instead of gating the head of every device's + # queue (see _build_device_queue). + device_queue = self._build_device_queue( + all_wait_files, + target_lock_id=self.clean_device_id, + target_scope=self.lock_scope, + ) is_eligible = bool(device_queue and device_queue[0] == self._queue_path) if is_eligible and self._try_acquire_owner_lock(): @@ -729,15 +820,11 @@ class DeviceExecutionLock: if owner.token in active_tokens: continue - target_dev = cls._normalize_device_id(owner.device_id) - target_lock_id = cls._normalize_lock_id(owner.device_id, owner.lock_scope) - same_scope = ( - not target_scope - or not owner.lock_scope - or owner.lock_scope == target_scope - ) - if clean_target is not None and target_lock_id != clean_target and not ( - target_dev in ("pending", "any") and same_scope + if clean_target is not None and not cls._ticket_targets_device( + owner, + target_lock_id=clean_target, + target_scope=target_scope, + include_pending=True, ): continue diff --git a/artemis/runtime/process_probe.py b/artemis/runtime/process_probe.py new file mode 100644 index 0000000..f10fd81 --- /dev/null +++ b/artemis/runtime/process_probe.py @@ -0,0 +1,98 @@ +# 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. + +"""Single source of truth for PID liveness probing across the Artemis codebase. + +Semantics (decided by the maintainers): + +- psutil only. No bare ``os.kill(pid, 0)`` fallback: on Windows, CPython + implements ``os.kill(pid, 0)`` for non-CTRL signals via + ``OpenProcess``/``TerminateProcess`` handle probing, which is both unreliable + for dead pids and semantically surprising. +- ``psutil.NoSuchProcess`` -> dead. Zombie processes -> dead. +- ``psutil.AccessDenied`` / ``PermissionError`` -> **alive**. When liveness + cannot be determined the probe defaults to "alive": it is always safer to + skip reaping a resource than to reap one that is still owned. +- ``created_at`` (a ``psutil``-style process create time) enables PID-reuse + protection: when provided and the probed process's create time differs by + one second or more, the original process is dead and its PID was recycled. +- Any other unexpected error -> alive (default-alive), with a warning log. +""" + +from __future__ import annotations + +import logging + +logger = logging.getLogger(__name__) + + +def pid_is_alive(pid: int, created_at: float | None = None) -> bool: + """Return whether the process identified by ``pid`` is still alive. + + Args: + pid: The process id to probe. Non-positive or unparsable pids are dead. + created_at: Optional expected process create time (as reported by + ``psutil.Process.create_time()``). When given and > 0, a create-time + mismatch of >= 1 second means the PID was recycled -> dead. + + Returns: + True when the process is alive or liveness cannot be determined + (default-alive); False when the process is definitively gone, a + zombie, or its PID has been recycled by another process. + """ + try: + pid = int(pid) + except (TypeError, ValueError): + return False + if pid <= 0: + return False + + try: + import psutil + except Exception as exc: # pragma: no cover - psutil is a hard dependency + logger.warning(f"psutil unavailable while probing pid {pid}: {exc}; assuming alive") + return True + + try: + process = psutil.Process(pid) + if not process.is_running(): + return False + + try: + if process.status() == psutil.STATUS_ZOMBIE: + return False + except psutil.NoSuchProcess: + return False + except (psutil.AccessDenied, PermissionError): + # Cannot inspect the status; keep the default-alive verdict. + pass + + if created_at is not None and created_at > 0: + try: + if abs(process.create_time() - float(created_at)) >= 1.0: + # Same PID, different birth time: the PID was recycled. + return False + except psutil.NoSuchProcess: + return False + except (psutil.AccessDenied, PermissionError): + pass + + return True + except psutil.NoSuchProcess: + return False + except (psutil.AccessDenied, PermissionError): + return True + except Exception as exc: + logger.warning(f"Could not determine liveness of pid {pid}: {exc}; assuming alive") + return True diff --git a/artemis/runtime/server_lifecycle.py b/artemis/runtime/server_lifecycle.py index e57c970..2ad40db 100644 --- a/artemis/runtime/server_lifecycle.py +++ b/artemis/runtime/server_lifecycle.py @@ -168,31 +168,9 @@ def _reconcile_orphaned_sessions() -> int: def _any_pid_alive(pids: list[int]) -> bool: """Return whether any target process is still running (zombies excluded).""" - if not pids: - return False - try: - import psutil + from artemis.runtime.process_probe import pid_is_alive - for pid in pids: - try: - process = psutil.Process(pid) - if process.is_running() and process.status() != psutil.STATUS_ZOMBIE: - return True - except psutil.NoSuchProcess: - continue - except psutil.AccessDenied: - return True - return False - except Exception: - return any(_pid_exists_portably(pid) for pid in pids) - - -def _pid_exists_portably(pid: int) -> bool: - try: - os.kill(pid, 0) - return True - except OSError: - return False + return any(pid_is_alive(pid) for pid in pids) def find_server_pids(port: int = 8000) -> list[int]: @@ -210,17 +188,10 @@ def find_server_pids(port: int = 8000) -> list[int]: if info and info.get("port") == port: saved_pid = info.get("pid") if saved_pid and isinstance(saved_pid, int): - try: - import psutil + from artemis.runtime.process_probe import pid_is_alive - if psutil.pid_exists(saved_pid): - discovered.add(saved_pid) - except Exception: - try: - os.kill(saved_pid, 0) - discovered.add(saved_pid) - except OSError: - pass + if pid_is_alive(saved_pid): + discovered.add(saved_pid) # 2. Check port listeners with platform-native tools # 2a. lsof (macOS & Linux) diff --git a/mcp_server/tools/task_manager.py b/mcp_server/tools/task_manager.py index 6f80fe6..d5f83b7 100644 --- a/mcp_server/tools/task_manager.py +++ b/mcp_server/tools/task_manager.py @@ -45,16 +45,12 @@ def _pid_alive(pid: int | None) -> bool: if not pid: return False try: - import psutil + pid_int = int(pid) + except (TypeError, ValueError): + return False + from artemis.runtime.process_probe import pid_is_alive - proc = psutil.Process(int(pid)) - return proc.is_running() and proc.status() != psutil.STATUS_ZOMBIE - except Exception: - try: - os.kill(int(pid), 0) - return True - except OSError: - return False + return pid_is_alive(pid_int) def _session_tracked_by_lock(trace_id: str) -> bool: diff --git a/mcp_server/utils/trace_store.py b/mcp_server/utils/trace_store.py index 379cd2f..3938983 100644 --- a/mcp_server/utils/trace_store.py +++ b/mcp_server/utils/trace_store.py @@ -12,12 +12,30 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Trace storage and lifecycle management for MCP background automation tasks.""" +"""Trace storage and lifecycle management for MCP background automation tasks. +status.json is written and updated concurrently by real independent processes +(the admin console daemon, MCP server tools, and spawned worker subprocesses), +so this module provides: + +- Atomic writes (temp file + fsync + ``os.replace``) so readers never observe + a torn/truncated file. +- A PID-aware ``status.json.lock`` file (O_CREAT | O_EXCL with bounded retry + and stale-lock detection) serializing read-modify-write updates. +- Three-state reads: missing file (silent ``None``), corrupt JSON (WARNING + + the file is quarantined as ``status.json.corrupt``), other IO errors + (WARNING). +""" + +from contextlib import contextmanager import json +import logging import os import time from typing import Any +import uuid + +logger = logging.getLogger(__name__) try: from artemis.config.paths import get_traces_dir @@ -27,6 +45,153 @@ except Exception: PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) TRACES_DIR = os.path.join(PROJECT_ROOT, "traces") +_LOCK_TIMEOUT_SECONDS = 5.0 +_LOCK_POLL_SECONDS = 0.05 +# A lock younger than this is assumed to be held by a live writer and is not +# probed for staleness. Probing opens the lock file, and on Windows an open +# reader handle blocks the owner's unlink (no FILE_SHARE_DELETE), so eager +# probing can leak the lock forever under contention. +_LOCK_STALE_PROBE_MIN_AGE_SECONDS = 0.5 +_LOCK_STALE_PROBE_INTERVAL_SECONDS = 1.0 +_MALFORMED_LOCK_GRACE_SECONDS = 10.0 +_REPLACE_RETRIES = 5 +_REPLACE_RETRY_DELAY_SECONDS = 0.02 +_UNLINK_RETRIES = 10 +_UNLINK_RETRY_DELAY_SECONDS = 0.01 + + +def _own_create_time() -> float: + try: + import psutil + + return float(psutil.Process(os.getpid()).create_time()) + except Exception: + return 0.0 + + +def _replace_with_retry(src: str, dst: str) -> None: + """``os.replace`` with retries for transient Windows sharing violations.""" + last_error: OSError | None = None + for attempt in range(_REPLACE_RETRIES): + try: + os.replace(src, dst) + return + except OSError as exc: + last_error = exc + if attempt < _REPLACE_RETRIES - 1: + time.sleep(_REPLACE_RETRY_DELAY_SECONDS) + raise last_error # type: ignore[misc] + + +def _atomic_write_json(path: str, data: dict[str, Any]) -> None: + """Write JSON atomically: temp file in the same directory + fsync + replace.""" + os.makedirs(os.path.dirname(path), exist_ok=True) + temp_path = f"{path}.{uuid.uuid4().hex}.tmp" + try: + with open(temp_path, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2, ensure_ascii=False) + f.flush() + os.fsync(f.fileno()) + _replace_with_retry(temp_path, path) + finally: + # A leftover temp file (e.g. after a crash between write and replace) + # never affects readers; remove it best-effort on the success path. + try: + os.unlink(temp_path) + except OSError: + pass + + +def _unlink_with_retry(path: str) -> bool: + """Best-effort unlink retrying transient Windows sharing violations.""" + for attempt in range(_UNLINK_RETRIES): + try: + os.unlink(path) + return True + except FileNotFoundError: + return True + except OSError: + if attempt < _UNLINK_RETRIES - 1: + time.sleep(_UNLINK_RETRY_DELAY_SECONDS) + return False + + +def _lock_is_stale(lock_path: str) -> bool: + """A status lock is stale when its owner process is gone or it is malformed and old.""" + try: + if time.time() - os.path.getmtime(lock_path) < _LOCK_STALE_PROBE_MIN_AGE_SECONDS: + # Young lock: assume a live writer holds it. Do not open the file + # (an open reader handle would block the owner's unlink on Windows). + return False + except OSError: + return False + try: + with open(lock_path, encoding="utf-8") as f: + payload = json.load(f) + pid = int(payload["pid"]) + created_at = float(payload.get("process_created_at", 0.0) or 0.0) + except FileNotFoundError: + return False + except Exception: + try: + return time.time() - os.path.getmtime(lock_path) >= _MALFORMED_LOCK_GRACE_SECONDS + except OSError: + return False + from artemis.runtime.process_probe import pid_is_alive + + return not pid_is_alive(pid, created_at if created_at > 0 else None) + + +@contextmanager +def _status_lock(path: str): + """Bounded cross-process mutex for read-modify-write updates of ``path``. + + Yields True when the lock was acquired. On timeout the caller proceeds + without mutual exclusion (with a WARNING) rather than deadlocking: a + last-writer-wins update is still strictly better than blocking MCP tools + forever on a leaked lock. + """ + lock_path = f"{path}.lock" + os.makedirs(os.path.dirname(lock_path), exist_ok=True) + payload = json.dumps({"pid": os.getpid(), "process_created_at": _own_create_time()}) + deadline = time.monotonic() + _LOCK_TIMEOUT_SECONDS + acquired = False + last_stale_probe = 0.0 + while not acquired: + try: + fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + now = time.monotonic() + if now - last_stale_probe >= _LOCK_STALE_PROBE_INTERVAL_SECONDS: + last_stale_probe = now + if _lock_is_stale(lock_path): + _unlink_with_retry(lock_path) + continue + if now >= deadline: + logger.warning( + f"Timed out acquiring trace status lock {lock_path}; " + "proceeding without mutual exclusion" + ) + break + time.sleep(_LOCK_POLL_SECONDS) + except OSError as exc: + logger.warning(f"Could not create trace status lock {lock_path}: {exc}") + break + else: + try: + os.write(fd, payload.encode("utf-8")) + finally: + os.close(fd) + acquired = True + try: + yield acquired + finally: + if acquired and not _unlink_with_retry(lock_path): + logger.warning( + f"Could not remove trace status lock {lock_path}; it will be " + "reaped as stale once its owner exits" + ) + def get_trace_dir(trace_id: str) -> str: """Returns the absolute path to the directory for a given trace_id.""" @@ -82,23 +247,41 @@ def init_trace( def read_status(trace_id: str) -> dict[str, Any] | None: - """Reads and returns the status.json for a given trace_id, or None if it doesn't exist.""" + """Read the status.json for a trace_id, distinguishing three failure states. + + - Missing file: normal (a trace that was never started) -> silent ``None``. + - Corrupt JSON: WARNING with the path; the damaged file is renamed to + ``status.json.corrupt`` as evidence, then ``None`` is returned. + - Other IO errors: WARNING, ``None``. + """ path = get_status_path(trace_id) if not os.path.exists(path): return None try: with open(path, encoding="utf-8") as f: return json.load(f) - except Exception: + except FileNotFoundError: + return None + except (json.JSONDecodeError, UnicodeDecodeError, ValueError) as exc: + logger.warning( + f"Corrupt status.json for trace {trace_id} at {path}: {exc}; " + f"quarantining as {path}.corrupt" + ) + try: + _replace_with_retry(path, f"{path}.corrupt") + except OSError as move_exc: + logger.warning( + f"Could not quarantine corrupt status file {path}: {move_exc}" + ) + return None + except OSError as exc: + logger.warning(f"Could not read status.json for trace {trace_id} at {path}: {exc}") return None def write_status(trace_id: str, data: dict[str, Any]) -> None: - """Writes the given status dictionary to the trace's status.json.""" - path = get_status_path(trace_id) - os.makedirs(os.path.dirname(path), exist_ok=True) - with open(path, "w", encoding="utf-8") as f: - json.dump(data, f, indent=2, ensure_ascii=False) + """Atomically writes the given status dictionary to the trace's status.json.""" + _atomic_write_json(get_status_path(trace_id), data) def update_trace_status( @@ -108,37 +291,56 @@ def update_trace_status( result: Any | None = None, device_serial: str | None = None, ) -> dict[str, Any] | None: - """Updates specific fields of the status.json for a given trace_id.""" - data = read_status(trace_id) - if not data: - return None + """Updates specific fields of the status.json for a given trace_id. - # "success" is a legacy alias for the canonical "completed" terminal - # status and must never be persisted into status.json. - if status == "success": - status = "completed" + The read-modify-write cycle holds a cross-process lock so concurrent + updates from the daemon, MCP tools, and worker processes do not lose each + other's fields. + """ + path = get_status_path(trace_id) + with _status_lock(path): + data = read_status(trace_id) + if not data: + if os.path.exists(path) or os.path.exists(f"{path}.corrupt"): + logger.warning( + f"Dropping status update ({status!r}) for trace {trace_id}: " + "existing status.json is unreadable" + ) + return None - data["status"] = status - if status in ("completed", "failed", "cancelled"): - data["end_time"] = time.time() + # "success" is a legacy alias for the canonical "completed" terminal + # status and must never be persisted into status.json. + if status == "success": + status = "completed" - if error is not None: - data["error"] = error - if result is not None: - data["result"] = result - if device_serial is not None: - data["device_serial"] = device_serial + data["status"] = status + if status in ("completed", "failed", "cancelled"): + data["end_time"] = time.time() - write_status(trace_id, data) - return data + if error is not None: + data["error"] = error + if result is not None: + data["result"] = result + if device_serial is not None: + data["device_serial"] = device_serial + + write_status(trace_id, data) + return data def update_trace_device_serial(trace_id: str, device_serial: str) -> dict[str, Any] | None: """Updates the device_serial field of the status.json for a given trace_id.""" - data = read_status(trace_id) - if not data: - return None + path = get_status_path(trace_id) + with _status_lock(path): + data = read_status(trace_id) + if not data: + if os.path.exists(path) or os.path.exists(f"{path}.corrupt"): + logger.warning( + f"Dropping device_serial update for trace {trace_id}: " + "existing status.json is unreadable" + ) + return None - data["device_serial"] = device_serial - write_status(trace_id, data) - return data + data["device_serial"] = device_serial + write_status(trace_id, data) + return data diff --git a/tests/unit/admin_console/test_task_queue_service.py b/tests/unit/admin_console/test_task_queue_service.py index f1466ff..79bddbe 100644 --- a/tests/unit/admin_console/test_task_queue_service.py +++ b/tests/unit/admin_console/test_task_queue_service.py @@ -52,7 +52,7 @@ def clean_state(tmp_path, monkeypatch): state.current_profile = None state.active_session_id = None state.active_connections.clear() - state.was_stopped_manually = False + state.manually_stopped_run_ids.clear() state.cancelled_session_ids.clear() if state.worker_task and not state.worker_task.done(): state.worker_task.cancel() @@ -66,7 +66,7 @@ def clean_state(tmp_path, monkeypatch): state.current_process = None state.current_goal = None state.current_profile = None - state.was_stopped_manually = False + state.manually_stopped_run_ids.clear() if hasattr(state, "recent_submissions"): state.recent_submissions.clear() if state.worker_task and not state.worker_task.done(): @@ -448,7 +448,7 @@ def test_stop_tasks_terminates_external_global_owner_and_preserves_local_waiter( local_waiter.kill.assert_not_called() assert state.current_process is local_waiter assert state.queue_items[0]["session_id"] == "frontend-waiter" - assert state.was_stopped_manually is False + assert not state.manually_stopped_run_ids assert "frontend-waiter" not in state.cancelled_session_ids assert "mcp-session" not in state.active_connections update_status.assert_called_once() @@ -921,6 +921,111 @@ async def test_enqueue_tasks_deduplicates_by_session_id(): assert res2["tasks"][0]["session_id"] == "sid-dedup-1" +@pytest.mark.asyncio +async def test_manual_stop_of_one_run_does_not_pollute_concurrent_run(): + """Stopping run A must not flip run B's terminal status or its payload. + + Regression test for the process-global ``was_stopped_manually`` flag that + used to mark *every* in-flight run as manually stopped. Two concurrent runs + on different devices: device A's task is stopped manually, device B's task + then finishes normally and must still be reported as completed. + """ + ended_payloads: dict[str, dict] = {} + + def capture(event_type, data): + if event_type == "session_ended": + ended_payloads[str(data.get("session_id"))] = dict(data) + + class FakeProc: + def __init__(self, pid): + self.pid = pid + self.returncode = None + self.stdout = None + self._done = asyncio.Event() + + async def wait(self): + await self._done.wait() + return self.returncode + + def finish(self, returncode): + self.returncode = returncode + self._done.set() + + def kill(self): + self.finish(-9) + + procs: dict[str, FakeProc] = {} + + async def fake_subprocess_exec(*args, **kwargs): + goal = args[3] + # Use our own pid so the reaped-process watchdog keeps waiting. + proc = FakeProc(os.getpid()) + procs[goal] = proc + return proc + + item_a = { + "session_id": "run-a", + "goal": "Goal A", + "status": "pending", + "device_serial": "dev-a", + } + item_b = { + "session_id": "run-b", + "goal": "Goal B", + "status": "pending", + "device_serial": "dev-b", + } + state.queue_items = [item_a, item_b] + state.ipc_subscribers.append(capture) + + try: + with ( + patch("asyncio.create_subprocess_exec", side_effect=fake_subprocess_exec), + patch("apps.admin_console.services.task_queue_service.session_repo") as mock_repo, + patch("apps.admin_console.services.task_queue_service.media_service"), + patch( + "apps.admin_console.services.task_queue_service.process_supervisor.terminate_tree", + return_value=True, + ), + ): + mock_repo.get_session_status.return_value = "running" + mock_repo.get_video_recording_for_session.return_value = {"status": "ready"} + + task_a = asyncio.create_task(TaskQueueService._execute_task_item(item_a)) + task_b = asyncio.create_task(TaskQueueService._execute_task_item(item_b)) + for _ in range(40): + if "run-a" in state.active_runs and "run-b" in state.active_runs: + break + await asyncio.sleep(0.02) + assert {"run-a", "run-b"} <= set(state.active_runs) + # B launched last, so the legacy single-task mirror points at B. + assert state.active_session_id == "run-b" + + # Manually stop device A's task without naming its session id. + assert task_queue_service.stop_tasks(clear_all=False, device_id="dev-a") is True + await asyncio.wait_for(task_a, timeout=5.0) + + # B finishes normally afterwards. + procs["Goal B"].finish(0) + await asyncio.wait_for(task_b, timeout=5.0) + + assert ended_payloads["run-a"]["status"] == "cancelled" + assert ended_payloads["run-a"]["was_stopped_manually"] is True + assert ended_payloads["run-b"]["status"] == "completed" + assert ended_payloads["run-b"]["was_stopped_manually"] is False + + persisted = { + call.args[0]: call.args[1] + for call in mock_repo.update_session_status.call_args_list + } + assert persisted.get("run-a") == "cancelled" + assert persisted.get("run-b") == "completed" + assert "run-b" not in state.cancelled_session_ids + assert "run-b" not in state.manually_stopped_run_ids + finally: + state.ipc_subscribers.remove(capture) + + @pytest.mark.asyncio async def test_enqueue_tasks_debounces_rapid_identical_submissions(): state.queue_items = [] diff --git a/tests/unit/mcp/test_trace_store.py b/tests/unit/mcp/test_trace_store.py index cb99db8..eaff9ab 100644 --- a/tests/unit/mcp/test_trace_store.py +++ b/tests/unit/mcp/test_trace_store.py @@ -14,9 +14,12 @@ """Unit tests for MCP trace store.""" +import glob +import logging import os import shutil import tempfile +import threading import uuid import pytest @@ -91,6 +94,125 @@ def test_read_nonexistent_status(temp_trace_env): assert trace_store.read_status(non_existent) is None +def test_leftover_atomic_temp_file_does_not_affect_reads(temp_trace_env): + """A crash between temp-write and replace leaves a .tmp file that readers ignore.""" + trace_id = str(uuid.uuid4()) + trace_store.init_trace(trace_id, "Crash sim", "Flash") + + status_path = trace_store.get_status_path(trace_id) + # Simulate a writer that crashed after writing its temp file. + with open(f"{status_path}.{uuid.uuid4().hex}.tmp", "w", encoding="utf-8") as f: + f.write('{"status": "half-writ') + + data = trace_store.read_status(trace_id) + assert data is not None + assert data["status"] == "running" + assert data["task_desc"] == "Crash sim" + + +def test_failed_atomic_replace_preserves_original(temp_trace_env): + """If the atomic replace fails, the previous status.json stays intact.""" + from unittest.mock import patch + + trace_id = str(uuid.uuid4()) + trace_store.init_trace(trace_id, "Original", "Flash") + + def broken_replace(src, dst): + raise OSError("simulated replace failure") + + with patch.object(trace_store, "_replace_with_retry", broken_replace): + with pytest.raises(OSError): + trace_store.write_status(trace_id, {"status": "clobbered"}) + + data = trace_store.read_status(trace_id) + assert data is not None + assert data["task_desc"] == "Original" + # The failed writer's temp file must not linger. + assert glob.glob(os.path.join(trace_store.get_trace_dir(trace_id), "*.tmp")) == [] + + +def test_corrupt_status_is_quarantined_and_read_returns_none(temp_trace_env, caplog): + trace_id = str(uuid.uuid4()) + trace_store.init_trace(trace_id, "Corrupt me", "Flash") + + status_path = trace_store.get_status_path(trace_id) + with open(status_path, "w", encoding="utf-8") as f: + f.write('{"status": "running", "trace_id"') # torn write + + with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"): + assert trace_store.read_status(trace_id) is None + + assert any("Corrupt status.json" in rec.message for rec in caplog.records) + assert os.path.exists(f"{status_path}.corrupt") + assert not os.path.exists(status_path) + + +def test_missing_status_stays_silent(temp_trace_env, caplog): + with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"): + assert trace_store.read_status(str(uuid.uuid4())) is None + assert caplog.records == [] + + +def test_update_on_corrupt_status_logs_and_drops(temp_trace_env, caplog): + trace_id = str(uuid.uuid4()) + trace_store.init_trace(trace_id, "Corrupt update", "Flash") + + status_path = trace_store.get_status_path(trace_id) + with open(status_path, "w", encoding="utf-8") as f: + f.write("not json at all") + + with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"): + assert trace_store.update_trace_status(trace_id, "completed") is None + + assert any("Dropping status update" in rec.message for rec in caplog.records) + # Updating a trace that never existed stays a silent no-op. + caplog.clear() + with caplog.at_level(logging.WARNING, logger="mcp_server.utils.trace_store"): + assert trace_store.update_trace_status(str(uuid.uuid4()), "completed") is None + assert not any("Dropping status update" in rec.message for rec in caplog.records) + + +def test_concurrent_read_modify_write_does_not_lose_updates(temp_trace_env): + """Interleaved RMW updates from two writers must not lose each other's fields.""" + trace_id = str(uuid.uuid4()) + trace_store.init_trace(trace_id, "Concurrency", "Flash") + iterations = 25 + failures: list[BaseException] = [] + + def set_errors(): + try: + for i in range(iterations): + assert trace_store.update_trace_status( + trace_id, "running", error=f"err-{i}" + ) is not None + except BaseException as exc: # noqa: BLE001 - surfaced in the main thread + failures.append(exc) + + def set_serials(): + try: + for i in range(iterations): + assert trace_store.update_trace_device_serial( + trace_id, f"serial-{i}" + ) is not None + except BaseException as exc: # noqa: BLE001 - surfaced in the main thread + failures.append(exc) + + threads = [threading.Thread(target=set_errors), threading.Thread(target=set_serials)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30.0) + + assert failures == [] + final = trace_store.read_status(trace_id) + assert final is not None + # Every serialized RMW preserves the other writer's latest field, so both + # final values must survive regardless of which thread finished last. + assert final["error"] == f"err-{iterations - 1}" + assert final["device_serial"] == f"serial-{iterations - 1}" + assert final["status"] == "running" + + def test_trace_paths(temp_trace_env): trace_id = "test-paths-id" trace_dir = trace_store.get_trace_dir(trace_id) diff --git a/tests/unit/runtime/test_device_lock.py b/tests/unit/runtime/test_device_lock.py index 64f68fc..6d4acf9 100644 --- a/tests/unit/runtime/test_device_lock.py +++ b/tests/unit/runtime/test_device_lock.py @@ -185,6 +185,78 @@ def test_submission_reservations_preserve_order_before_workers_start(): thread.join(timeout=2.0) +def test_pending_reservation_at_queue_head_does_not_block_other_devices(): + """An unclaimed pending submission must not gate every device's queue head. + + Regression: a reservation whose device is still "pending" used to be + counted into every device's FIFO queue, so one unclaimed head-of-queue + ticket blocked all devices at once. It must now park on at most one idle + candidate device (the first in sorted order) and leave the others + schedulable. + """ + pending_ticket = DeviceExecutionLock.reserve("unclaimed submission") + # device-alpha is a known idle device (a queued submission targets it), so + # the pending ticket parks there; device-beta must stay schedulable. + alpha_ticket = DeviceExecutionLock.reserve( + "task for device alpha", device_id="device-alpha" + ) + try: + beta = DeviceExecutionLock("device-beta", "task for device beta") + # Must acquire immediately even though the older pending ticket is at + # the head of the global queue directory. + beta.acquire(blocking=False) + beta.release() + finally: + DeviceExecutionLock.cancel_reservation(pending_ticket) + DeviceExecutionLock.cancel_reservation(alpha_ticket) + + +def test_pending_reservation_is_served_with_original_fifo_position_after_claim(): + """A pending reservation is not starved: once its worker claims a device it + keeps its original submission-order position in that device's queue.""" + holder = DeviceExecutionLock("emulator-5554", "current holder") + holder.acquire() + + # Submitted while the device is busy; the reservation stays "pending". + first_ticket = DeviceExecutionLock.reserve("first submitted task") + + claimed = DeviceExecutionLock("emulator-5554", "claimed pending task", first_ticket) + youngest = DeviceExecutionLock("emulator-5554", "youngest task") + order = [] + claimed_acquired = threading.Event() + youngest_acquired = threading.Event() + + def run_claimed(): + claimed.acquire() + order.append("claimed") + claimed_acquired.set() + + def run_youngest(): + youngest.acquire() + order.append("youngest") + youngest_acquired.set() + + t1 = threading.Thread(target=run_claimed) + t2 = threading.Thread(target=run_youngest) + t1.start() + t2.start() + time.sleep(0.2) + assert not claimed_acquired.is_set() + assert not youngest_acquired.is_set() + + holder.release() + # The claimed pending ticket kept its original (oldest) FIFO timestamp and + # must win the device before the younger concrete ticket. + assert claimed_acquired.wait(timeout=2.0) + assert not youngest_acquired.is_set() + claimed.release() + assert youngest_acquired.wait(timeout=2.0) + youngest.release() + t1.join(timeout=2.0) + t2.join(timeout=2.0) + assert order == ["claimed", "youngest"] + + def test_cancelled_reservation_cannot_execute(): ticket = DeviceExecutionLock.reserve("cancelled task") assert DeviceExecutionLock.cancel_reservation(ticket) diff --git a/tests/unit/runtime/test_process_probe.py b/tests/unit/runtime/test_process_probe.py new file mode 100644 index 0000000..5a05476 --- /dev/null +++ b/tests/unit/runtime/test_process_probe.py @@ -0,0 +1,65 @@ +# 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. + +"""Unit tests for the unified PID liveness probe.""" + +import os +from unittest.mock import MagicMock, patch + +import psutil + +from artemis.runtime.process_probe import pid_is_alive + + +def test_own_pid_is_alive(): + assert pid_is_alive(os.getpid()) is True + + +def test_own_pid_with_matching_create_time_is_alive(): + created_at = psutil.Process(os.getpid()).create_time() + assert pid_is_alive(os.getpid(), created_at) is True + + +def test_recycled_pid_with_mismatched_create_time_is_dead(): + assert pid_is_alive(os.getpid(), 1.0) is False + + +def test_invalid_pids_are_dead(): + assert pid_is_alive(0) is False + assert pid_is_alive(-5) is False + assert pid_is_alive(None) is False # type: ignore[arg-type] + assert pid_is_alive("not-a-pid") is False # type: ignore[arg-type] + + +def test_no_such_process_is_dead(): + with patch("psutil.Process", side_effect=psutil.NoSuchProcess(424242)): + assert pid_is_alive(424242) is False + + +def test_zombie_process_is_dead(): + proc = MagicMock() + proc.is_running.return_value = True + proc.status.return_value = psutil.STATUS_ZOMBIE + with patch("psutil.Process", return_value=proc): + assert pid_is_alive(1234) is False + + +def test_access_denied_defaults_to_alive(): + with patch("psutil.Process", side_effect=psutil.AccessDenied(1234)): + assert pid_is_alive(1234) is True + + +def test_unexpected_error_defaults_to_alive(): + with patch("psutil.Process", side_effect=RuntimeError("boom")): + assert pid_is_alive(1234) is True