# 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. import asyncio from datetime import datetime import logging import os import shutil import sys import time from typing import Any import uuid try: from admin_console.core.state import state from admin_console.database.repositories.session_repository import session_repo from admin_console.services import worker_process_io from admin_console.services.media_service import media_service except ImportError: from apps.admin_console.core.state import state from apps.admin_console.database.repositories.session_repository import session_repo from apps.admin_console.services import worker_process_io from apps.admin_console.services.media_service import media_service from artemis.config import ( PAUSE_FILE, TEST_DATA_DIR, TEST_OUTPUTS_DIR, WORKSPACE_ROOT, ) from artemis.runtime import ( AdbEndpoint, AdbTarget, DeviceExecutionLock, current_adb_endpoint, process_supervisor, trace_store, ) logger = logging.getLogger(__name__) class TaskQueueService: """Service managing FIFO task execution, background worker, subprocess lifecycle, and startup tasks. """ # Strong references to in-flight _execute_task_item tasks (asyncio itself only # keeps weak references to running tasks). _run_tasks: set[asyncio.Task] = set() @staticmethod def _task_target(task_item: dict[str, Any]) -> AdbTarget: endpoint_data = task_item.get("adb_endpoint") endpoint = ( AdbEndpoint.from_mapping(endpoint_data) if isinstance(endpoint_data, dict) else current_adb_endpoint() ) serial = task_item.get("device_serial") return AdbTarget(endpoint=endpoint, serial=str(serial) if serial else None) @classmethod def _broadcast_event(cls, event_type: str, data: Any): """Broadcasts an event safely to all registered subscribers.""" for cb in list(state.ipc_subscribers): try: cb(event_type, data) except Exception: # One broken subscriber must not block the others, but a # silent drop hides it entirely. logger.warning( "Event subscriber %r failed for event %s", cb, event_type, exc_info=True, ) @classmethod def _broadcast_startup_progress( cls, session_id: str | None, stage: str, message: str ) -> None: if not session_id: return data = { "session_id": str(session_id), "stage": stage, "message": message, "timestamp": time.time(), } state.record_startup_progress(data) cls._broadcast_event("startup_progress", data) @classmethod def _get_next_pending_task(cls) -> dict[str, Any] | None: """Finds and returns the first pending task from queue_items.""" for item in state.queue_items: if isinstance(item, dict) and item.get("status") == "pending": return item return None @classmethod def _remove_task(cls, session_id: str | None): """Removes a task from queue_items by session_id.""" if not session_id: return removed_items = [ t for t in state.queue_items if isinstance(t, dict) and t.get("session_id") == session_id ] for item in removed_items: DeviceExecutionLock.cancel_reservation(item.get("queue_ticket")) state.queue_items = [ t for t in state.queue_items if not (isinstance(t, dict) and t.get("session_id") == session_id) ] @staticmethod def _resolve_terminal_status( current_status: str | None, returncode: int, was_stopped_manually: bool, ) -> tuple[str, bool]: """Resolve final status while preserving an authoritative task result. The worker exit code is only a fallback for sessions that have not reached a terminal state. ``success`` is the DataEngine alias for the UI-facing ``completed`` status and is normalized here. Returns: A tuple of ``(resolved_status, should_persist)``. """ if was_stopped_manually: return "cancelled", True normalized = current_status.lower().strip() if isinstance(current_status, str) else None if normalized == "success": return "completed", True if normalized in {"completed", "failed", "cancelled"}: return normalized, False return ("completed" if returncode == 0 else "failed"), True # Worker subprocess I/O plumbing lives in worker_process_io; the historical # private names stay bound here so callers and tests keep working unchanged. _subprocess_creation_kwargs = staticmethod(worker_process_io.subprocess_creation_kwargs) _forward_worker_output = staticmethod(worker_process_io.forward_worker_output) _finish_output_forwarder = staticmethod(worker_process_io.finish_output_forwarder) _wait_for_worker_process = staticmethod(worker_process_io.wait_for_worker_process) @staticmethod async def _terminate_worker_process(proc: asyncio.subprocess.Process | None) -> None: """Stop a worker without changing the established POSIX behavior.""" if proc is None or proc.returncode is not None: return if sys.platform == "win32": await process_supervisor.stop_process(proc) return try: proc.kill() except ProcessLookupError: pass @classmethod def ensure_worker_running(cls): """Guarantees that the background queue worker task is active and running.""" try: loop = asyncio.get_running_loop() if state.worker_task is None or state.worker_task.done(): state.worker_task = loop.create_task(cls.queue_worker()) try: state.wake_event.set() except AttributeError: # Partially initialized state (unit tests): the worker # loop polls anyway. pass except RuntimeError: pass @classmethod def _concurrency_mode(cls) -> str: """Resolve the scheduler concurrency mode. Mirrors DeviceExecutionLock's contract (ARTEMIS_CONCURRENCY_MODE / ARTEMIS_MAX_CONCURRENT_TASKS): - "global": strict serial -- one task at a time across all devices. - "per_device": one task per device -- distinct devices run concurrently. Defaults to per-device concurrency, matching the lock layer's default. """ mode = os.environ.get("ARTEMIS_CONCURRENCY_MODE", "").strip().lower() if mode in ("global", "serial", "1"): return "global" if mode in ("per_device", "device", "parallel", "0"): return "per_device" try: if int(os.environ.get("ARTEMIS_MAX_CONCURRENT_TASKS", 0) or 0) == 1: return "global" except (TypeError, ValueError): pass return "per_device" @classmethod async def queue_worker(cls): """Persistent dispatcher scheduling pending tasks onto devices. Scans the queue and launches each eligible task as an independent coroutine. Admission is governed by _concurrency_mode(): "global" admits one task at a time overall, "per_device" admits one task per device so distinct devices execute concurrently. Per-device FIFO ordering and cross-process mutual exclusion remain enforced by DeviceExecutionLock inside each worker process. """ print( "[QueueWorker] Dispatcher initialized " f"(concurrency mode: {cls._concurrency_mode()})." ) try: DeviceExecutionLock.cleanup_stale_locks() except Exception as exc: print(f"[QueueWorker] Initial stale lock cleanup notice: {exc}") try: while True: try: cls._dispatch_pending_tasks() except Exception as exc: print(f"[QueueWorker] Dispatch error: {exc}") try: await asyncio.wait_for(state.wake_event.wait(), timeout=0.3) state.wake_event.clear() except TimeoutError: pass except asyncio.CancelledError: print("[QueueWorker] Dispatcher received cancellation signal.") for run in list(state.active_runs.values()): run_proc = run.get("process") if run_proc is not None and run_proc.returncode is None: await cls._terminate_worker_process(run_proc) raise @classmethod def _dispatch_pending_tasks(cls) -> None: """Launch every pending task admissible under the current concurrency mode.""" mode = cls._concurrency_mode() state.prune_finished_runs() if mode == "global" and state.is_running: return # Dispatched-but-not-yet-spawned runs are only visible as queue items in # "running" state, so admission must count those too -- active_runs alone # lags behind by the subprocess startup latency. in_flight = [ i for i in state.queue_items if isinstance(i, dict) and i.get("status") == "running" ] busy_devices = state.busy_device_ids | { cls._task_target(i).lock_key for i in in_flight if i.get("device_serial") } dispatched_any = False loop = asyncio.get_running_loop() for item in list(state.queue_items): if not (isinstance(item, dict) and item.get("status") == "pending"): continue sess_id = item.get("session_id") if sess_id and str(sess_id) in state.cancelled_session_ids: cls._remove_task(sess_id) continue device = item.get("device_serial") target = cls._task_target(item) if mode == "per_device": # A task without a resolved device may bind to any serial, so it # only launches on an otherwise idle scheduler; the device lock # then allocates freely without contending against active runs. if device is None and (state.active_runs or in_flight or dispatched_any): continue if device is not None and target.lock_key in busy_devices: continue item["status"] = "running" dispatched_any = True if device is not None: busy_devices.add(target.lock_key) run_task = loop.create_task(cls._execute_task_item(item)) # Hold a strong reference: asyncio keeps only weak refs to running # tasks, and a collected run would strand its queue item forever. cls._run_tasks.add(run_task) run_task.add_done_callback(cls._run_tasks.discard) if mode == "global": break @classmethod def _begin_task_run( cls, task_item: dict[str, Any], run_key: str, sess_id: Any, goal: str, profile: str, ) -> None: """Mark the task as running and announce the launch to subscribers.""" task_item["status"] = "running" task_item["start_time"] = time.time() # 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 cls._broadcast_startup_progress( sess_id, "launching", "Starting the execution process" ) # Broadcast session_started so all connected clients know the task has started cls._broadcast_event( "session_started", { "session_id": sess_id, "initial_goal": goal, "profile": profile, "device_serial": task_item.get("device_serial"), }, ) @classmethod def _build_worker_invocation( cls, task_item: dict[str, Any], run_key: str, sess_id: Any, goal: str, profile: str, target: AdbTarget, ) -> tuple[list[str], dict[str, str]]: """Assemble the worker subprocess command line and environment.""" expected_output = task_item.get("expected_output") enable_outputter = task_item.get("enable_outputter") locked_app = task_item.get("locked_app_package") or task_item.get("locked_app") app_path = task_item.get("app_path") test_name = f"web_{int(time.time())}_{run_key[:8]}" env = os.environ.copy() pythonpath_parts = [ str(WORKSPACE_ROOT), str(WORKSPACE_ROOT / "apps" / "admin_console"), str(WORKSPACE_ROOT / "apps" / "cloud_service"), env.get("PYTHONPATH", ""), ] env["PYTHONPATH"] = os.pathsep.join([p for p in pythonpath_parts if p]) env["PYTHONUTF8"] = "1" env["PYTHONUNBUFFERED"] = "1" if state.ipc_port is not None: env["ARTEMIS_IPC_PORT"] = str(state.ipc_port) if sess_id: env["ARTEMIS_SESSION_ID"] = str(sess_id) env["ARTEMIS_TASK_INGRESS"] = str(task_item.get("ingress", "frontend")) env["ARTEMIS_TASK_WORKER"] = "1" target.endpoint.apply_to_environment(env) env[DeviceExecutionLock.LOCK_SCOPE_ENV] = target.lock_scope queue_ticket = task_item.get("queue_ticket") if queue_ticket: env[DeviceExecutionLock.QUEUE_TICKET_ENV] = str(queue_ticket) cmd = [ sys.executable, "-m", "artemis.main", goal, "--profile", profile, "--test-name", test_name, ] if sess_id: cmd.extend(["--session-id", str(sess_id)]) if expected_output: cmd.extend(["--output-description", str(expected_output)]) if enable_outputter is not None: cmd.append("--enable-outputter" if enable_outputter else "--disable-outputter") if locked_app: cmd.extend(["--locked-app", str(locked_app)]) if app_path: cmd.extend(["--app-path", str(app_path)]) device_serial = task_item.get("device_serial") if device_serial: cmd.extend(["--device-serial", str(device_serial)]) env["ADB_DEVICE_SERIAL"] = str(device_serial) return cmd, env @classmethod def _register_worker_run( cls, task_item: dict[str, Any], run_key: str, sess_id: Any, goal: str, profile: str, target: AdbTarget, proc: asyncio.subprocess.Process, ) -> None: """Record the spawned worker in shared state and hand it the device reservation.""" device_serial = task_item.get("device_serial") state.current_process = proc task_item["pid"] = proc.pid state.active_runs[run_key] = { "process": proc, "device_id": str(device_serial) if device_serial else None, "lock_key": target.lock_key if device_serial else None, "adb_endpoint": target.endpoint.to_dict(), "goal": goal, "profile": profile, } if sess_id: try: st = trace_store.read_status(str(sess_id)) if st: st["pid"] = proc.pid trace_store.write_status(str(sess_id), st) except OSError as exc: # Status probes fall back to DB PID bookkeeping, but note it. print( f"[QueueWorker] Could not record worker pid in status.json for {sess_id}: {exc}" ) cls._broadcast_startup_progress( sess_id, "process_ready", "Execution process started" ) ingress_type = str(task_item.get("ingress", "frontend")) DeviceExecutionLock.transfer_reservation( str(task_item.get("queue_ticket")), proc.pid, description=f"{ingress_type} task: {goal[:120]}", device_id=device_serial or "pending", session_id=str(sess_id) if sess_id else None, ingress=ingress_type, lock_scope=target.lock_scope, ) @classmethod def _start_output_forwarder( cls, sess_id: Any, proc: asyncio.subprocess.Process ) -> asyncio.Task[None] | None: """Start forwarding the worker's output; returns the forwarder task, if any.""" # Forward the worker's combined output and tee it into the trace's # stdout.log so the log paths advertised by the MCP API exist. log_path = None if sess_id: try: log_path = trace_store.get_trace_stdout_log_path(str(sess_id)) except Exception: log_path = None if isinstance(proc.stdout, asyncio.StreamReader): return asyncio.create_task( cls._forward_worker_output(proc.stdout, log_path) ) return None @classmethod async def _terminate_if_cancelled_during_launch( cls, run_key: str, sess_id: Any, proc: asyncio.subprocess.Process ) -> None: """Terminate a worker whose task was cancelled while it was being launched.""" if sess_id and ( str(sess_id) in getattr(state, "cancelled_session_ids", set()) or run_key in state.manually_stopped_run_ids ): print( f"[QueueWorker] Task [{sess_id}] was cancelled during launch. Terminating." ) await cls._terminate_worker_process(proc) @classmethod async def _persist_terminal_session_status( cls, sess_id: Any, returncode: int, manual_stop: bool ) -> str: """Resolve and persist the session's terminal status in the DB and trace store.""" current_status = session_repo.get_session_status(sess_id) new_status, should_persist = cls._resolve_terminal_status( current_status, returncode, manual_stop, ) if should_persist: if session_repo.update_session_status(sess_id, new_status, time.time()): print(f"[QueueWorker] Updated session {sess_id} status to '{new_status}'") else: # The DB row is the fallback MCP pollers reconcile # against when status.json is stale, so a failed DB # write must not pass silently. logger.error( "[QueueWorker] Could not persist terminal DB status " "'%s' for session %s", new_status, sess_id, ) else: print( f"[QueueWorker] Preserved authoritative session {sess_id} " f"status '{new_status}'" ) # A stale status.json would leave MCP pollers seeing "running" # until their next DB reconcile, so retry transient write # failures before giving up. for attempt in range(3): try: if trace_store.read_status(str(sess_id)): canonical_mcp_status = "completed" if new_status in ("completed", "success") else new_status trace_store.update_trace_status( str(sess_id), canonical_mcp_status, ) break except OSError as exc: if attempt < 2: await asyncio.sleep(0.5 * (attempt + 1)) continue logger.error( "[QueueWorker] Could not persist terminal MCP status " "for %s after %d attempts: %s", sess_id, attempt + 1, exc, ) return new_status @classmethod async def _recover_or_fail_recording(cls, sess_id: Any) -> None: """Attempt recovery of a not-yet-ready recording, else mark it failed.""" rec_info = session_repo.get_video_recording_for_session(sess_id) rec_status = (rec_info or {}).get("status") recovered_video_path = None if rec_status != "ready": try: video_rec_map = session_repo.get_video_recordings_map() video_idx = await asyncio.to_thread(media_service.build_video_index) recovered_url = await asyncio.to_thread( media_service.resolve_video_url, {"session_id": sess_id}, video_rec_map, video_idx, ) if recovered_url: session_repo.mark_recording_ready(sess_id, recovered_url) cls._broadcast_event( "recording_ready", {"session_id": sess_id, "video_url": recovered_url}, ) recovered_video_path = recovered_url except Exception as rec_err: print(f"[QueueWorker] Error attempting recording recovery: {rec_err}") if not recovered_video_path and rec_status != "ready": recording_error = "Task worker exited before recording finalization completed" if session_repo.mark_recording_failed_if_pending(sess_id, recording_error): cls._broadcast_event( "recording_failed", {"session_id": sess_id, "error": recording_error}, ) @classmethod def _announce_session_end( cls, task_item: dict[str, Any], sess_id: Any, goal: str, new_status: str, manual_stop: bool, ) -> None: """Broadcast session_ended and dispatch the external completion notification.""" state.active_connections.pop(sess_id, None) cls._broadcast_event( "session_ended", { "session_id": sess_id, "status": new_status, "was_stopped_manually": manual_stop, }, ) conversation_id = task_item.get("conversation_id") if task_item else None if conversation_id or task_item.get("ingress") == "mcp": try: from mcp_server.notifiers import notify notify( conversation_id=conversation_id or "", message=f"Artemis autonomous task '{goal}' finished with status '{new_status}'.\nTrace ID: {sess_id}", title=f"Task {new_status.capitalize()}: {goal[:40]}", event_type=new_status, payload={ "trace_id": sess_id, "session_id": sess_id, "status": new_status, "goal": goal, }, ) except Exception as notif_err: print(f"[QueueWorker] Notification dispatch notice: {notif_err}") @classmethod def _release_run_slot( cls, sess_id: Any, run_key: str, proc: asyncio.subprocess.Process | None ) -> None: """Clean up the finished task and release this run's scheduling slot.""" if sess_id: 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 if not state.active_runs: state.active_session_id = None state.current_goal = None state.current_profile = None # 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 async def _execute_task_item(cls, task_item: dict[str, Any]) -> None: """Run one queued task to completion in its own worker subprocess.""" sess_id = task_item.get("session_id") run_key = str(sess_id) if sess_id else uuid.uuid4().hex goal = task_item.get("goal") profile = task_item.get("profile", "flash") proc: asyncio.subprocess.Process | None = None output_task: asyncio.Task[None] | None = None try: if not isinstance(goal, str) or not goal.strip(): raise ValueError("Queued task must contain a non-empty string goal.") cls._begin_task_run(task_item, run_key, sess_id, goal, profile) target = cls._task_target(task_item) cmd, env = cls._build_worker_invocation( task_item, run_key, sess_id, goal, profile, target ) device_serial = task_item.get("device_serial") print( f"[QueueWorker] Starting task [{sess_id}]: '{goal}' (profile: {profile}, device: {device_serial or 'auto'}, outputter: {bool(task_item.get('expected_output') or task_item.get('enable_outputter'))})" ) proc = await asyncio.create_subprocess_exec( *cmd, cwd=str(WORKSPACE_ROOT), env=env, **cls._subprocess_creation_kwargs(), ) cls._register_worker_run( task_item, run_key, sess_id, goal, profile, target, proc ) output_task = cls._start_output_forwarder(sess_id, proc) await cls._terminate_if_cancelled_during_launch(run_key, sess_id, proc) # 3. Await subprocess completion returncode = await cls._wait_for_worker_process(proc) print(f"[QueueWorker] Task [{sess_id}] exited with returncode {returncode}") manual_stop = run_key in state.manually_stopped_run_ids or bool( sess_id and str(sess_id) in state.cancelled_session_ids ) # 4. Perform fallback database status update and notification if sess_id: new_status = await cls._persist_terminal_session_status( sess_id, returncode, manual_stop ) await cls._recover_or_fail_recording(sess_id) cls._announce_session_end( task_item, sess_id, goal, new_status, manual_stop ) except asyncio.CancelledError: print(f"[QueueWorker] Task [{sess_id}] received cancellation signal.") if proc is not None and proc.returncode is None: await cls._terminate_worker_process(proc) raise except Exception as e: print(f"[QueueWorker] Unexpected error executing task [{sess_id}]: {e}") finally: await cls._finish_output_forwarder(output_task) # 5. Clean up the finished task and release this run's scheduling slot cls._release_run_slot(sess_id, run_key, proc) @classmethod def _find_duplicate_submission( cls, goals: list[str], session_id: str | None, device_serial: str | None, endpoint: AdbEndpoint, now: float, ) -> dict[str, Any] | None: """Return the short-circuit response for a duplicate submission, if any.""" # 1. Deduplication by session_id: if session_id is already running or queued, do not re-enqueue if session_id and len(goals) == 1: target_sid = str(session_id) if state.active_session_id and str(state.active_session_id) == target_sid: return { "status": "started", "tasks": [{"session_id": target_sid, "status": "running"}], "enqueued_count": 0, "total_queued": len(state.queue_tasks), } existing_item = next( ( item for item in state.queue_items if isinstance(item, dict) and str(item.get("session_id")) == target_sid ), None, ) if existing_item: return { "status": existing_item.get("status", "queued"), "tasks": [existing_item], "enqueued_count": 0, "total_queued": len(state.queue_tasks), } # 2. Debounce duplicate rapid submissions (e.g. UI double-click or network retry within 1s) if len(goals) == 1: first_goal = goals[0] recent_duplicate = next( ( item for item in reversed(state.queue_items) if isinstance(item, dict) and item.get("status") == "pending" and item.get("goal") == first_goal and (not device_serial or item.get("device_serial") == device_serial) and item.get("adb_endpoint", {}).get("identity") == endpoint.identity and (now - float(item.get("created_at", 0))) < 1.0 ), None, ) if recent_duplicate: return { "status": "queued", "tasks": [recent_duplicate], "enqueued_count": 0, "total_queued": len(state.queue_tasks), } return None @classmethod async def _reject_unavailable_device( cls, device_serial: str | None ) -> dict[str, Any] | None: """Return the rejection response for an unattached explicit serial, if any.""" # Strict device binding: reject an explicitly requested serial that is not # attached and authorized, instead of silently running on another device. # The shared validator fails open on an indeterminate/empty enumeration so # the task can proceed and fail downstream with a clear no-device error. if device_serial: try: from artemis.runtime import device_pool rejection = await device_pool.validate_explicit_serial_async(device_serial) except Exception: rejection = None if rejection: return { "status": "rejected", "error": rejection, "tasks": [], "enqueued_count": 0, "total_queued": len(state.queue_tasks), } return None @classmethod def _create_queue_item( cls, goal: str, index: int, now: float, endpoint: AdbEndpoint, single_session_id: str | None, profile: str, expected_output: str | None, enable_outputter: bool | None, locked_app_package: str | None, app_path: str | None, device_serial: str | None, ingress: str, conversation_id: str | None, ) -> dict[str, Any]: """Reserve a device slot and build one pending queue item for a goal.""" sess_id = single_session_id if single_session_id else str(uuid.uuid4()) assigned_serial = device_serial if not assigned_serial: try: from artemis.runtime import device_pool assigned_serial = device_pool.select_device() except Exception: assigned_serial = None queue_ticket = DeviceExecutionLock.reserve( description=f"{ingress} task: {goal[:120]}", device_id=assigned_serial or "pending", session_id=sess_id, ingress=ingress, lock_scope=endpoint.identity, ) return { "session_id": sess_id, "goal": goal, "profile": profile or "flash", "expected_output": expected_output, "enable_outputter": enable_outputter, "locked_app_package": locked_app_package, "app_path": app_path, "device_serial": assigned_serial, "adb_endpoint": endpoint.to_dict(), "ingress": ingress, "conversation_id": conversation_id, "status": "pending", "queue_ticket": queue_ticket, "created_at": now + index * 0.001, "start_time": now + index * 0.001, } @classmethod async def enqueue_tasks( cls, goals: list[str], profile: str = "flash", expected_output: str | None = None, enable_outputter: bool | None = None, locked_app_package: str | None = None, app_path: str | None = None, device_serial: str | None = None, ingress: str = "frontend", session_id: str | None = None, conversation_id: str | None = None, ) -> dict[str, Any]: """Enqueues one or more goals and wakes up the background worker.""" cls.ensure_worker_running() enqueued_tasks = [] now = time.time() endpoint = current_adb_endpoint() duplicate_response = cls._find_duplicate_submission( goals, session_id, device_serial, endpoint, now ) if duplicate_response is not None: return duplicate_response rejection_response = await cls._reject_unavailable_device(device_serial) if rejection_response is not None: return rejection_response single_session_id = session_id if (session_id and len(goals) == 1) else None for i, goal in enumerate(goals): task_item = cls._create_queue_item( goal, i, now, endpoint, single_session_id, profile, expected_output, enable_outputter, locked_app_package, app_path, device_serial, ingress, conversation_id, ) state.queue_items.append(task_item) enqueued_tasks.append(task_item) cls._broadcast_startup_progress( task_item["session_id"], "queued", "Task received and queued" ) # Wake worker immediately state.wake_event.set() return { "status": "queued" if state.is_running else "started", "tasks": enqueued_tasks, "enqueued_count": len(goals), "total_queued": len(state.queue_tasks), } @staticmethod def _clear_pause_file() -> None: """Remove a leftover pause marker after a stop request.""" if PAUSE_FILE.exists(): try: PAUSE_FILE.unlink() except OSError: # Best-effort cleanup of the pause marker; a leftover file # only pauses until the next resume request. pass @classmethod def _terminate_all_device_owners(cls) -> None: """Terminate every device lock owner and persist each cancellation.""" active_owners: dict[str, Any] = {} try: active_owners = DeviceExecutionLock.get_active_owners() except OSError as exc: # Unreadable lock dir: fall back to the single-owner probe below. print(f"[stop_tasks] Could not enumerate device owners: {exc}") fallback = DeviceExecutionLock.get_active_owner() if fallback and not active_owners: active_owners["default"] = fallback for dev_owner in list(active_owners.values()): if dev_owner and DeviceExecutionLock.is_active_owner(dev_owner): process_supervisor.terminate_tree_verified( dev_owner.pid, dev_owner.process_created_at, ) DeviceExecutionLock.cleanup_stale_locks(dev_owner.device_id) sid = dev_owner.session_id if sid: session_repo.update_session_status(str(sid), "cancelled", time.time()) try: is_mcp = bool(dev_owner and dev_owner.ingress == "mcp") if is_mcp or trace_store.read_status(str(sid)): trace_store.update_trace_status( str(sid), "cancelled", error="Task stopped from the Artemis frontend.", ) except OSError as exc: print( f"[stop_tasks] Could not persist cancelled MCP status for {sid}: {exc}" ) cls._broadcast_event( "session_ended", { "session_id": sid, "status": "cancelled", "was_stopped_manually": True, }, ) @classmethod def _kill_all_local_runs(cls) -> None: """Mark every locally-managed run manually stopped and kill its process.""" for run_key, run in list(state.active_runs.items()): # 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: try: run_proc.kill() except (ProcessLookupError, OSError): # Process already exited between the check and the kill. pass # active_runs entries are popped by each run's finalizer once the # process exit is observed; clearing them here would free the device # slots before the processes are actually gone. if state.current_process: try: state.current_process.kill() except (ProcessLookupError, OSError): # Process already exited between the check and the kill. pass state.current_process = None @classmethod def _stop_all_tasks(cls) -> bool: """Terminate all active device owners and clear pending queue submissions.""" # 1. Cancel local queue reservations for item in state.queue_items: if isinstance(item, dict) and item.get("status") != "running": DeviceExecutionLock.cancel_reservation(item.get("queue_ticket")) state.clear_queue() # 2. Terminate all active owners across all devices cls._terminate_all_device_owners() cls._kill_all_local_runs() state.active_connections.clear() state.active_session_id = None state.current_goal = None state.current_profile = None cls._clear_pause_file() cls.ensure_worker_running() state.wake_event.set() return True @classmethod def _resolve_stop_owner( cls, active_owners: dict[str, Any], target_sid: str | None, target_device: str | None, ) -> Any: """Resolve the device lock owner targeted by this stop request, if any.""" owner = None if target_sid: for dev_owner in active_owners.values(): if dev_owner.session_id and str(dev_owner.session_id) == target_sid: owner = dev_owner break if owner is None and target_device: clean_target_dev = DeviceExecutionLock._normalize_device_id(target_device) for dev_key, dev_owner in active_owners.items(): if dev_key == clean_target_dev or dev_owner.device_id == target_device: owner = dev_owner break if owner is None: # Covers scoped/legacy lock records that get_active_owners cannot # key by session or device. A stale owner for a different session # is never adopted: with no live owner, a session-targeted stop # falls through to the queue/session-repository fallbacks below. fallback_owner = DeviceExecutionLock.get_active_owner(target_device) if fallback_owner and ( not target_sid or (fallback_owner.session_id and str(fallback_owner.session_id) == target_sid) ): owner = fallback_owner return owner @classmethod def _resolve_local_run( cls, target_sid: str | None, target_device: str | None ) -> tuple[str | None, dict[str, Any] | None, Any]: """Resolve the locally-managed run for this target. Concurrent scheduling keeps one subprocess per run in state.active_runs. Returns ``(local_run_key, local_run, local_proc)``. """ 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_key, local_run = next( ( (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: # Untargeted stop (legacy single-device gesture): with exactly one # live run the target is unambiguous. local_run_key, local_run = next(iter(state.active_runs.items())) local_proc = (local_run or {}).get("process") if local_proc is None and not target_sid and not target_device: # Untargeted legacy fallback only: the last-started process may # stand in when no per-run entry resolved. An explicitly targeted # stop that matched no run must never grab current_process -- it # mirrors the newest run, which can be an unrelated concurrent one. local_proc = 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, ) return local_run_key, local_run, local_proc @classmethod def _find_target_queue_item(cls, target_sid: str | None) -> dict[str, Any] | None: """Find the running (preferred) or pending queue item for this target.""" running_item = next( ( item for item in state.queue_items if isinstance(item, dict) and item.get("status") == "running" and (not target_sid or str(item.get("session_id")) == target_sid) ), None, ) return running_item or next( ( item for item in state.queue_items if isinstance(item, dict) and item.get("status") == "pending" and (not target_sid or str(item.get("session_id")) == target_sid) ), None, ) @classmethod def _resolve_stopped_session_id( cls, owner: Any, target_sid: str | None, local_run_key: str | None, is_local_owner: bool, local_item: dict[str, Any] | None, ) -> Any: """Attribute the stop request to a session id, if one can be resolved.""" # 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. return ( 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") if local_item else None ) @classmethod def _terminate_stop_target( cls, owner: Any, owner_record_exists: bool, target_sid: str | None, local_run: dict[str, Any] | None, local_run_key: str | None, local_proc: Any, local_pid: Any, local_item: dict[str, Any] | None, is_local_owner: bool, ) -> tuple[bool, bool, bool] | None: """Terminate the resolved stop target. Returns ``(stopped, is_local_owner, reservation_cancelled)``, or ``None`` when the caller must refuse the stop outright. """ stopped = False reservation_cancelled = False if owner and DeviceExecutionLock.is_active_owner(owner): stopped = process_supervisor.terminate_tree_verified( owner.pid, owner.process_created_at, ) DeviceExecutionLock.cleanup_stale_locks(owner.device_id) elif owner is None and owner_record_exists and not target_sid: # Never fall back to a frontend PID while another process has an # owner record that is still being published or cannot be parsed. return None elif owner is None and local_proc is not None and ( local_run is not None or not target_sid or str(state.active_session_id) == target_sid ): # The locally-managed worker can be stopped during its short # initialization window before Agent acquires the device lease. if target_sid: state.cancelled_session_ids.add(target_sid) 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) except Exception: stopped = False try: local_proc.kill() stopped = True except (ProcessLookupError, OSError): # Process already exited; terminate_tree above may have got it. pass stopped = True is_local_owner = True elif owner is None and local_item and ( not target_sid or str(local_item.get("session_id")) == target_sid ): # Cancel a frontend submission before its worker has started. This # does not touch pending reservations created by other ingresses. DeviceExecutionLock.cancel_reservation(local_item.get("queue_ticket")) reservation_cancelled = True stopped = True is_local_owner = True elif target_sid: # Fallback 1: check global device queue global_queued = DeviceExecutionLock.get_queued_tasks() for q_item in global_queued: if str(q_item.get("session_id")) == target_sid: ticket = q_item.get("queue_ticket") if ticket: DeviceExecutionLock.cancel_reservation(ticket) stopped = True break # Fallback 2: check session repository for a running session with a live worker PID if not stopped: row = session_repo.get_session_by_id(target_sid) if row and row.get("status") == "running": row_pid = row.get("pid") if row_pid and session_repo.process_is_alive(row_pid): try: process_supervisor.terminate_tree(int(row_pid)) except Exception as exc: # Report it: a surviving worker keeps the device busy. print( f"[stop_tasks] Could not terminate worker pid {row_pid}: {exc}" ) DeviceExecutionLock.cleanup_stale_locks() stopped = True return stopped, is_local_owner, reservation_cancelled @classmethod def _finalize_targeted_stop( cls, stopped_session_id: Any, owner: Any, owner_pid: Any, is_local_owner: bool, local_run_key: str | None, local_proc: Any, ) -> None: """Propagate the stop into shared state, the DB, and event subscribers.""" if is_local_owner: if stopped_session_id: # 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)) 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 for sid, conn_info in list(state.active_connections.items()): if (stopped_session_id and str(sid) == str(stopped_session_id)) or ( owner_pid and conn_info.get("pid") == owner_pid ): state.active_connections.pop(sid, None) if stopped_session_id: session_repo.update_session_status( str(stopped_session_id), "cancelled", time.time() ) try: is_mcp = bool(owner and owner.ingress == "mcp") if is_mcp or trace_store.read_status(str(stopped_session_id)): trace_store.update_trace_status( str(stopped_session_id), "cancelled", error="Task stopped from the Artemis frontend.", ) except Exception as exc: print(f"Failed to update MCP cancellation status: {exc}") cls._broadcast_event( "session_ended", { "session_id": stopped_session_id, "status": "cancelled", "was_stopped_manually": True, }, ) if state.active_session_id and ( not stopped_session_id or str(state.active_session_id) == str(stopped_session_id) ): state.active_session_id = None state.current_goal = None state.current_profile = None @classmethod def _remove_stopped_queue_item( cls, stopped_session_id: Any, reservation_cancelled: bool ) -> None: """Drop the stopped session's queue item and cancel its reservation.""" if not stopped_session_id: return stopped_item = next( ( item for item in state.queue_items if isinstance(item, dict) and str(item.get("session_id")) == str(stopped_session_id) ), None, ) if stopped_item and not reservation_cancelled: DeviceExecutionLock.cancel_reservation(stopped_item.get("queue_ticket")) state.queue_items = [ item for item in state.queue_items if not ( isinstance(item, dict) and str(item.get("session_id")) == str(stopped_session_id) ) ] @classmethod def _stop_targeted_task( cls, target_sid: str | None, target_device: str | None ) -> bool: """Stop a specific task (or default single-device active task).""" active_owners = {} try: active_owners = DeviceExecutionLock.get_active_owners() except OSError as exc: # Unreadable lock dir: the local-process fallbacks below still apply. print(f"[stop_tasks] Could not enumerate device owners: {exc}") owner = cls._resolve_stop_owner(active_owners, target_sid, target_device) owner_record_exists = DeviceExecutionLock.has_owner_record(target_device) owner_pid = owner.pid if owner else None local_run_key, local_run, local_proc = cls._resolve_local_run( target_sid, target_device ) local_pid = getattr(local_proc, "pid", None) is_local_owner = bool(owner_pid and local_pid and owner_pid == local_pid) local_item = cls._find_target_queue_item(target_sid) stopped_session_id = cls._resolve_stopped_session_id( owner, target_sid, local_run_key, is_local_owner, local_item ) outcome = cls._terminate_stop_target( owner, owner_record_exists, target_sid, local_run, local_run_key, local_proc, local_pid, local_item, is_local_owner, ) if outcome is None: return False stopped, is_local_owner, reservation_cancelled = outcome if not stopped and not target_sid: return False cls._finalize_targeted_stop( stopped_session_id, owner, owner_pid, is_local_owner, local_run_key, local_proc, ) cls._remove_stopped_queue_item(stopped_session_id, reservation_cancelled) cls._clear_pause_file() cls.ensure_worker_running() state.wake_event.set() return True @classmethod def stop_tasks( cls, clear_all: bool = False, session_id: str | None = None, device_id: str | None = None, ) -> bool: """Stop the active task controlling a mobile device or all tasks. The active lease is shared by frontend, MCP, CLI, SDK, and other UI processes across all connected devices. - If ``session_id`` or ``device_id`` is specified, only the corresponding running or queued task is cancelled. - If ``clear_all`` is requested, all active device owners are terminated and pending queue submissions are cleared. - If no specific task is specified and ``clear_all`` is False, stops the currently active task in single-device mode for backward compatibility. """ target_sid = str(session_id).strip() if session_id else None target_device = str(device_id).strip() if device_id else None if clear_all: return cls._stop_all_tasks() return cls._stop_targeted_task(target_sid, target_device) @classmethod def resume_task(cls) -> bool: if PAUSE_FILE.exists(): PAUSE_FILE.unlink() return True return False @staticmethod def archive_older_replays_on_launch(): """Archives all existing step replay output folders on server launch.""" if TEST_OUTPUTS_DIR.exists(): archive_dir = TEST_DATA_DIR / "older" for item in TEST_OUTPUTS_DIR.iterdir(): if item.is_dir(): archive_dir.mkdir(parents=True, exist_ok=True) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") target_name = f"{item.name}_{timestamp}" target_path = archive_dir / target_name counter = 1 while target_path.exists(): target_name = f"{item.name}_{timestamp}_{counter}" target_path = archive_dir / target_name counter += 1 print( f"Archiving older replay record on server launch: " f"{item.name} -> {target_path}" ) try: shutil.move(str(item), str(target_path)) except Exception as e: print( f"Warning: Failed to archive {item.name} during server launch: {e}", file=sys.stderr, ) @staticmethod def verify_chunks_exist_on_launch(replay_manager): """Verifies that chunked directories exist for all sessions in the database.""" print("Verifying session chunks on launch...") try: sessions = session_repo.get_all_sessions() session_ids = [s.get("session_id") for s in sessions if s.get("session_id")] print(f"Found {len(session_ids)} sessions in database to verify.") for session_id in session_ids: replay_manager._ensure_session_chunked(str(session_id)) except Exception as e: print( f"Warning: Failed to verify chunks during server launch: {e}", file=sys.stderr, ) task_queue_service = TaskQueueService()