From f68d4288e07218c271015655c7a4c77202324b07 Mon Sep 17 00:00:00 2001 From: somew1nd Date: Tue, 1 Sep 2026 10:49:51 -0700 Subject: [PATCH 1/2] refactor(core): remove dead parallel agent architecture --- artemis/__init__.py | 3 - artemis/agents/__init__.py | 24 +- artemis/agents/base.py | 146 --- artemis/agents/operator/loop_detector.py | 42 - artemis/agents/operator/operator_agent.py | 72 -- artemis/agents/planner/planner_agent.py | 53 - artemis/agents/planner/types.py | 85 -- artemis/agents/planner/utils.py | 89 -- artemis/agents/summarizer/summarizer_agent.py | 53 - .../validator/failure_analyzer_agent.py | 69 -- .../agents/validator/recovery_strategies.py | 177 ---- artemis/agents/validator/safety_net.py | 48 - artemis/controllers/android_controller.py | 925 ------------------ artemis/core/__init__.py | 40 +- artemis/core/context.py | 62 -- artemis/core/events.py | 77 -- artemis/core/exceptions.py | 57 -- artemis/core/registry.py | 46 - artemis/core/state.py | 70 -- artemis/engine/__init__.py | 27 - artemis/engine/base_runner.py | 36 - artemis/engine/graph_runner.py | 89 -- artemis/engine/pipeline.py | 95 -- artemis/engine/reactive_runner.py | 112 --- artemis/llm/client_factory.py | 25 - tests/conftest.py | 7 - .../test_reactive_runner_integration.py | 45 - tests/unit/test_agents.py | 146 --- tests/unit/test_prompt_builder.py | 24 +- 29 files changed, 8 insertions(+), 2736 deletions(-) delete mode 100644 artemis/agents/base.py delete mode 100644 artemis/agents/operator/loop_detector.py delete mode 100644 artemis/agents/operator/operator_agent.py delete mode 100644 artemis/agents/planner/planner_agent.py delete mode 100644 artemis/agents/planner/types.py delete mode 100644 artemis/agents/planner/utils.py delete mode 100644 artemis/agents/summarizer/summarizer_agent.py delete mode 100644 artemis/agents/validator/failure_analyzer_agent.py delete mode 100644 artemis/agents/validator/recovery_strategies.py delete mode 100644 artemis/agents/validator/safety_net.py delete mode 100644 artemis/controllers/android_controller.py delete mode 100644 artemis/core/context.py delete mode 100644 artemis/core/events.py delete mode 100644 artemis/core/exceptions.py delete mode 100644 artemis/core/registry.py delete mode 100644 artemis/core/state.py delete mode 100644 artemis/engine/__init__.py delete mode 100644 artemis/engine/base_runner.py delete mode 100644 artemis/engine/graph_runner.py delete mode 100644 artemis/engine/pipeline.py delete mode 100644 artemis/engine/reactive_runner.py delete mode 100644 artemis/llm/client_factory.py delete mode 100644 tests/integration/test_reactive_runner_integration.py delete mode 100644 tests/unit/test_agents.py diff --git a/artemis/__init__.py b/artemis/__init__.py index 54fe90e..97aa8ee 100644 --- a/artemis/__init__.py +++ b/artemis/__init__.py @@ -21,7 +21,6 @@ from artemis.interfaces.sdk.task import StreamEvent, StreamEventType, Task from artemis_client import Capabilities, Device, TaskHandle, TaskResult from artemis.sdk.agent import Agent from artemis.sdk.builders import Builders -from artemis.agents.base import BaseAgent, AgentRegistry from artemis.drivers.base import BaseDeviceDriver from artemis.tools.base import artemis_tool, ToolRegistry @@ -38,8 +37,6 @@ __all__ = [ "Task", "TaskHandle", "TaskResult", - "BaseAgent", - "AgentRegistry", "BaseDeviceDriver", "artemis_tool", "ToolRegistry", diff --git a/artemis/agents/__init__.py b/artemis/agents/__init__.py index bd79617..92c813e 100644 --- a/artemis/agents/__init__.py +++ b/artemis/agents/__init__.py @@ -12,26 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Modular Agent Framework for ARTEMIS.""" +"""Agent implementations for ARTEMIS.""" -from artemis.agents.base import ( - AgentConfig, - AgentRegistry, - AgentResponse, - BaseAgent, -) -from artemis.agents.operator.operator_agent import OperatorAgent -from artemis.agents.planner.planner_agent import PlannerAgent -from artemis.agents.summarizer.summarizer_agent import SummarizerAgent -from artemis.agents.validator.failure_analyzer_agent import FailureAnalyzerAgent - -__all__ = [ - "BaseAgent", - "AgentConfig", - "AgentResponse", - "AgentRegistry", - "PlannerAgent", - "OperatorAgent", - "SummarizerAgent", - "FailureAnalyzerAgent", -] +__all__: list[str] = [] diff --git a/artemis/agents/base.py b/artemis/agents/base.py deleted file mode 100644 index 67ddc8f..0000000 --- a/artemis/agents/base.py +++ /dev/null @@ -1,146 +0,0 @@ -# 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. - -"""Base Agent Protocol & Lifecycle Framework for ARTEMIS. - -Provides a unified interface and lifecycle management for all specialized -multimodal agents (Planner, Operator, FailureAnalyzer, Explorer, Summarizer, etc.). -""" - -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any - -from pydantic import BaseModel, Field - -from artemis.drivers.factory import get_driver -from artemis.utils.logger import get_logger - -if TYPE_CHECKING: - pass - -logger = get_logger(__name__) - - -class AgentConfig(BaseModel): - """Configuration options for an individual agent.""" - - model_name: str | None = Field(default=None, description="LLM/VLM model override") - temperature: float = Field(default=0.0, description="Sampling temperature") - max_tokens: int = Field(default=4096, description="Maximum output tokens") - system_prompt_override: str | None = Field( - default=None, description="Custom system prompt override" - ) - enabled_tools: list[str] = Field( - default_factory=list, description="Tool names enabled for this agent" - ) - - -class AgentResponse(BaseModel): - """Standardized structured response from an agent invocation.""" - - status: str = Field( - default="success", description="Execution status ('success', 'retry', 'fail', 'done')" - ) - thought: str | None = Field(default=None, description="Reasoning and internal monologue") - action_taken: str | None = Field(default=None, description="Action or tool called") - output: Any = Field(default=None, description="Output payload or state updates") - metadata: dict[str, Any] = Field( - default_factory=dict, description="Diagnostic and profiling metadata" - ) - - -class BaseAgent(ABC): - """Abstract Base Class defining the standard lifecycle for all ARTEMIS agents.""" - - def __init__( - self, - name: str, - role: str, - description: str, - config: AgentConfig | None = None, - ctx: Any | None = None, - ): - self.name = name - self.role = role - self.description = description - self.config = config or AgentConfig() - self.ctx: Any | None = ctx - self._is_initialized = False - - @property - def driver(self) -> Any | None: - """Helper to get the active device driver from the context.""" - if self.ctx: - return get_driver(self.ctx) - return None - - async def initialize(self, ctx: Any) -> None: - """Lifecycle hook: Called once when the agent is loaded into the workflow.""" - self.ctx = ctx - self._is_initialized = True - logger.info(f"Agent '{self.name}' ({self.role}) initialized.") - - @abstractmethod - async def process(self, state: Any) -> dict[str, Any]: - """Main execution hook: Receives state, reasons, interacts with tools, and returns state updates.""" - ... - - async def cleanup(self) -> None: - """Lifecycle hook: Called upon session or workflow termination.""" - self._is_initialized = False - logger.debug(f"Agent '{self.name}' cleanup complete.") - - def __repr__(self) -> str: - return f"<{self.__class__.__name__} name={self.name!r} role={self.role!r}>" - - -class AgentRegistry: - """Central registry and factory for ARTEMIS agents.""" - - _registry: dict[str, type[BaseAgent]] = {} - - @classmethod - def register(cls, name: str | None = None): - """Decorator to register an agent class with a unique identifier.""" - - def decorator(subclass: type[BaseAgent]) -> type[BaseAgent]: - reg_name = name or subclass.__name__.lower() - cls._registry[reg_name] = subclass - return subclass - - return decorator - - @classmethod - def get_agent_class(cls, name: str) -> type[BaseAgent] | None: - return cls._registry.get(name.lower()) - - @classmethod - def create_agent( - cls, - name: str, - ctx: Any | None = None, - config: AgentConfig | None = None, - **kwargs: Any, - ) -> BaseAgent: - """Instantiates an agent by registered name.""" - agent_cls = cls.get_agent_class(name) - if not agent_cls: - raise KeyError( - f"Agent '{name}' is not registered. Available: {list(cls._registry.keys())}" - ) - return agent_cls(ctx=ctx, config=config, **kwargs) - - @classmethod - def list_agents(cls) -> list[str]: - return list(cls._registry.keys()) diff --git a/artemis/agents/operator/loop_detector.py b/artemis/agents/operator/loop_detector.py deleted file mode 100644 index 9b345c2..0000000 --- a/artemis/agents/operator/loop_detector.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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. - -"""Loop, deadlock, and action repetition detector for Operator Agent.""" - -from typing import Any -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class LoopDetector: - """Monitors recent actions to detect and break infinite interaction loops.""" - - def __init__(self, repetition_threshold: int = 3): - self.repetition_threshold = repetition_threshold - self._action_signatures: list[str] = [] - - def record_action(self, action_name: str, params: dict[str, Any]) -> None: - sig = f"{action_name}:{sorted(params.items())}" - self._action_signatures.append(sig) - - def is_loop_detected(self) -> bool: - """Returns True if the identical action has been repeated >= threshold times.""" - if len(self._action_signatures) < self.repetition_threshold: - return False - last_sig = self._action_signatures[-1] - return all(s == last_sig for s in self._action_signatures[-self.repetition_threshold :]) - - def reset(self) -> None: - self._action_signatures.clear() diff --git a/artemis/agents/operator/operator_agent.py b/artemis/agents/operator/operator_agent.py deleted file mode 100644 index 4baef59..0000000 --- a/artemis/agents/operator/operator_agent.py +++ /dev/null @@ -1,72 +0,0 @@ -# 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. - -"""Modular Operator Agent implementation adhering to BaseAgent lifecycle.""" - -from typing import Any -from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent -from artemis.agents.operator.loop_detector import LoopDetector -from artemis.agents.operator.prompt_builder import OperatorPromptBuilder -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -@AgentRegistry.register("operator") -class OperatorAgent(BaseAgent): - """Specialized agent responsible for real-time screen analysis and UI interaction decisions.""" - - def __init__( - self, - config: AgentConfig | None = None, - ctx: Any | None = None, - **kwargs: Any, - ): - super().__init__( - name="operator", - role="Mobile UI Interaction Operator", - description="Perceives mobile UI and executes touch/key gestures to progress tasks.", - config=config, - ctx=ctx, - ) - self.loop_detector = LoopDetector(repetition_threshold=3) - - async def process(self, state: Any) -> dict[str, Any]: - """Observes screen state, reasons next action, and dispatches tool call.""" - if not self.driver: - raise RuntimeError("OperatorAgent requires an active device driver in context.") - - # 1. Capture screen data - screen_data = await self.driver.get_screen_data() - - # 2. Check for loop deadlock - if self.loop_detector.is_loop_detected(): - logger.warning( - "Operator detected repetitive interaction loop! Injecting recovery back-press." - ) - await self.driver.press_key("back") - self.loop_detector.reset() - return {"action_taken": "press_key:back", "status": "recovering"} - - # 3. Formulate prompt & messages - goal = getattr(state, "task_goal", "") - human_text = OperatorPromptBuilder.build_human_message(goal=goal) - - return { - "thought": "Observed screen and deciding next interaction.", - "action_taken": "click", - "screen_width": screen_data.width, - "screen_height": screen_data.height, - "status": "success", - } diff --git a/artemis/agents/planner/planner_agent.py b/artemis/agents/planner/planner_agent.py deleted file mode 100644 index 8fb6946..0000000 --- a/artemis/agents/planner/planner_agent.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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. - -"""Planner Agent implementation adhering to BaseAgent lifecycle.""" - -from typing import Any -from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent -from artemis.agents.planner.planner import PlannerNode -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -@AgentRegistry.register("planner") -class PlannerAgent(BaseAgent): - """Specialized agent responsible for high-level goal decomposition and strategic task planning.""" - - def __init__( - self, - config: AgentConfig | None = None, - ctx: Any | None = None, - **kwargs: Any, - ): - super().__init__( - name="planner", - role="Strategic Task Planner", - description="Decomposes user objectives into structured, observable milestone steps.", - config=config, - ctx=ctx, - ) - self._node = PlannerNode(ctx=ctx, **kwargs) if ctx else None - - async def initialize(self, ctx: Any) -> None: - await super().initialize(ctx) - self._node = PlannerNode(ctx=ctx) - - async def process(self, state: Any) -> dict[str, Any]: - if not self._node and self.ctx: - self._node = PlannerNode(ctx=self.ctx) - if not self._node: - raise RuntimeError("PlannerAgent has not been initialized with a ArtemisContext.") - return await self._node(state) diff --git a/artemis/agents/planner/types.py b/artemis/agents/planner/types.py deleted file mode 100644 index 8cc1b5d..0000000 --- a/artemis/agents/planner/types.py +++ /dev/null @@ -1,85 +0,0 @@ -# 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. - -from datetime import datetime -from enum import Enum -from typing import Annotated - -from pydantic import BaseModel - - -class _CyFunctionDetectorMeta(type): - def __instancecheck__(self, instance): - name = type(instance).__name__ - return ( - name - in ( - "cyfunction", - "cython_function_or_method", - "builtin_function_or_method", - ) - or "cyfunction" in name.lower() - ) - - -class CyFunctionDetector(metaclass=_CyFunctionDetectorMeta): - pass - - -class PlannerSubgoalOutput(BaseModel): - model_config = {"ignored_types": (CyFunctionDetector,)} - description: str - - -class PlannerOutput(BaseModel): - model_config = {"ignored_types": (CyFunctionDetector,)} - subgoals: list[PlannerSubgoalOutput] - - -class SubgoalStatus(Enum): - NOT_STARTED = "NOT_STARTED" - PENDING = "PENDING" - SUCCESS = "SUCCESS" - FAILURE = "FAILURE" - - -class Subgoal(BaseModel): - model_config = {"ignored_types": (CyFunctionDetector,)} - id: Annotated[str, "Unique identifier of the subgoal"] - description: Annotated[str, "Description of the subgoal"] - completion_reason: Annotated[ - str | None, "Reason why the subgoal was completed (failure or success)" - ] = None - status: SubgoalStatus - started_at: Annotated[datetime | None, "When the subgoal started"] = None - ended_at: Annotated[datetime | None, "When the subgoal ended"] = None - - def __str__(self): - status_emoji = "❓" - if self.status == SubgoalStatus.SUCCESS: - status_emoji = "✅" - elif self.status == SubgoalStatus.FAILURE: - status_emoji = "❌" - elif self.status == SubgoalStatus.PENDING: - status_emoji = "⏳" - elif self.status == SubgoalStatus.NOT_STARTED: - status_emoji = "(not started yet)" - - output = f"- [ID:{self.id}]: {self.description} : {status_emoji}." - if self.completion_reason: - output += f" Completion reason: {self.completion_reason}" - return output - - def __repr__(self): - return str(self) diff --git a/artemis/agents/planner/utils.py b/artemis/agents/planner/utils.py deleted file mode 100644 index 2d8a0e8..0000000 --- a/artemis/agents/planner/utils.py +++ /dev/null @@ -1,89 +0,0 @@ -# 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. - -try: - from datetime import UTC, datetime -except ImportError: - from datetime import datetime - - UTC = UTC -import random -import string - -from artemis.agents.planner.types import Subgoal, SubgoalStatus - - -def get_current_subgoal(subgoals: list[Subgoal]) -> Subgoal | None: - return next((s for s in subgoals if s.status == SubgoalStatus.PENDING), None) - - -def get_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]: - return [s for s in subgoals if s.id in ids] - - -def get_next_subgoal(subgoals: list[Subgoal]) -> Subgoal | None: - return next((s for s in subgoals if s.status == SubgoalStatus.NOT_STARTED), None) - - -def nothing_started(subgoals: list[Subgoal]) -> bool: - return all(s.status == SubgoalStatus.NOT_STARTED for s in subgoals) - - -def complete_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: - current_subgoal = get_current_subgoal(subgoals) - if not current_subgoal: - return subgoals - current_subgoal.status = SubgoalStatus.SUCCESS - current_subgoal.ended_at = datetime.now(UTC) - return subgoals - - -def complete_subgoals_by_ids(subgoals: list[Subgoal], ids: list[str]) -> list[Subgoal]: - for subgoal in subgoals: - if subgoal.id in ids: - subgoal.status = SubgoalStatus.SUCCESS - subgoal.ended_at = datetime.now(UTC) - return subgoals - - -def fail_current_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: - current_subgoal = get_current_subgoal(subgoals) - if not current_subgoal: - return subgoals - current_subgoal.status = SubgoalStatus.FAILURE - current_subgoal.ended_at = datetime.now(UTC) - return subgoals - - -def all_completed(subgoals: list[Subgoal]) -> bool: - return all(s.status == SubgoalStatus.SUCCESS for s in subgoals) - - -def one_of_them_is_failure(subgoals: list[Subgoal]) -> bool: - return any(s.status == SubgoalStatus.FAILURE for s in subgoals) - - -def start_next_subgoal(subgoals: list[Subgoal]) -> list[Subgoal]: - next_subgoal = get_next_subgoal(subgoals) - if not next_subgoal: - return subgoals - next_subgoal.status = SubgoalStatus.PENDING - next_subgoal.started_at = datetime.now(UTC) - return subgoals - - -def generate_id(length: int = 6) -> str: - """Generates a small and distinct random string ID.""" - chars = string.ascii_lowercase + string.digits - return "".join(random.choice(chars) for _ in range(length)) diff --git a/artemis/agents/summarizer/summarizer_agent.py b/artemis/agents/summarizer/summarizer_agent.py deleted file mode 100644 index 1331d46..0000000 --- a/artemis/agents/summarizer/summarizer_agent.py +++ /dev/null @@ -1,53 +0,0 @@ -# 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. - -"""Summarizer Agent implementation adhering to BaseAgent lifecycle.""" - -from typing import Any -from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent -from artemis.agents.summarizer.summarizer import SummarizerNode -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -@AgentRegistry.register("summarizer") -class SummarizerAgent(BaseAgent): - """Specialized agent responsible for step execution summarization and trajectory logging.""" - - def __init__( - self, - config: AgentConfig | None = None, - ctx: Any | None = None, - **kwargs: Any, - ): - super().__init__( - name="summarizer", - role="Execution History Summarizer", - description="Synthesizes actions and observations into concise milestone summaries.", - config=config, - ctx=ctx, - ) - self._node = SummarizerNode(ctx=ctx) if ctx else None - - async def initialize(self, ctx: Any) -> None: - await super().initialize(ctx) - self._node = SummarizerNode(ctx=ctx) - - async def process(self, state: Any) -> dict[str, Any]: - if not self._node and self.ctx: - self._node = SummarizerNode(ctx=self.ctx) - if not self._node: - raise RuntimeError("SummarizerAgent has not been initialized with a ArtemisContext.") - return await self._node(state) diff --git a/artemis/agents/validator/failure_analyzer_agent.py b/artemis/agents/validator/failure_analyzer_agent.py deleted file mode 100644 index 45ddf91..0000000 --- a/artemis/agents/validator/failure_analyzer_agent.py +++ /dev/null @@ -1,69 +0,0 @@ -# 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. - -"""Modular Failure Analyzer Agent adhering to BaseAgent lifecycle.""" - -from typing import Any -from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent -from artemis.agents.validator.recovery_strategies import RecoveryStrategyRegistry -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -@AgentRegistry.register("failure_analyzer") -class FailureAnalyzerAgent(BaseAgent): - """Specialized agent responsible for diagnosing step failures and triggering autonomous recovery.""" - - def __init__( - self, - config: AgentConfig | None = None, - ctx: Any | None = None, - **kwargs: Any, - ): - super().__init__( - name="failure_analyzer", - role="Autonomous Self-Healing Failure Analyzer", - description="Diagnoses execution exceptions, analyzes screen discrepancies, and runs self-healing sequences.", - config=config, - ctx=ctx, - ) - - async def process(self, state: Any) -> dict[str, Any]: - """Analyzes failure context and executes appropriate recovery strategy.""" - if not self.driver: - raise RuntimeError("FailureAnalyzerAgent requires an active device driver in context.") - - failure_context = getattr(state, "failure_context", {}) or {} - if isinstance(state, dict): - failure_context = state.get("failure_context", {}) - - logger.info(f"FailureAnalyzer diagnosing failure context: {failure_context}") - - recovery_result = await RecoveryStrategyRegistry.attempt_recovery( - driver=self.driver, - failure_context=failure_context, - ) - - if recovery_result and recovery_result.success: - return { - "status": "recovered", - "recovery_action": recovery_result.action_type.value, - "reason": recovery_result.reason, - } - - return { - "status": "unrecoverable", - "reason": "No applicable recovery strategy succeeded.", - } diff --git a/artemis/agents/validator/recovery_strategies.py b/artemis/agents/validator/recovery_strategies.py deleted file mode 100644 index 446c41a..0000000 --- a/artemis/agents/validator/recovery_strategies.py +++ /dev/null @@ -1,177 +0,0 @@ -# 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. - -"""Modular Recovery Strategies for the Validation and Self-Healing Engine.""" - -from abc import ABC, abstractmethod -from enum import Enum -from typing import Any -from artemis.drivers.base import BaseDeviceDriver -from artemis.utils.logger import get_logger -from pydantic import BaseModel, Field - -logger = get_logger(__name__) - - -class RecoveryActionType(str, Enum): - RETRY_DYNAMIC_LOCATOR = "retry_dynamic_locator" - DISMISS_DIALOG = "dismiss_dialog" - NAVIGATE_BACK = "navigate_back" - SCROLL_AND_RETRY = "scroll_and_retry" - RESTART_APP = "restart_app" - FALLBACK_COORDINATES = "fallback_coordinates" - - -class RecoveryResult(BaseModel): - """Result of executing an autonomous self-healing recovery strategy.""" - - success: bool - action_type: RecoveryActionType - reason: str - metadata: dict[str, Any] = Field(default_factory=dict) - - -class BaseRecoveryStrategy(ABC): - """Abstract interface for modular self-healing failure recovery tactics.""" - - @property - @abstractmethod - def name(self) -> str: - """Strategy identifier.""" - ... - - @abstractmethod - async def can_handle(self, failure_context: dict[str, Any]) -> bool: - """Determines if this strategy applies to the observed failure mode.""" - ... - - @abstractmethod - async def execute( - self, - driver: BaseDeviceDriver, - failure_context: dict[str, Any], - ) -> RecoveryResult: - """Executes the self-healing corrective actions.""" - ... - - -class DismissDialogRecoveryStrategy(BaseRecoveryStrategy): - """Strategy to detect and dismiss unexpected modal dialogs, popups, or permission prompts.""" - - @property - def name(self) -> str: - return "dismiss_dialog" - - async def can_handle(self, failure_context: dict[str, Any]) -> bool: - category = failure_context.get("category", "") - error_msg = failure_context.get("error_message", "").lower() - return ( - category == "ui_obstruction" - or "permission" in error_msg - or "dialog" in error_msg - or "popup" in error_msg - ) - - async def execute( - self, - driver: BaseDeviceDriver, - failure_context: dict[str, Any], - ) -> RecoveryResult: - logger.info( - "[Recovery] Attempting to dismiss unexpected dialog or overlay via Back button..." - ) - success = await driver.press_key("back") - return RecoveryResult( - success=success, - action_type=RecoveryActionType.DISMISS_DIALOG, - reason="Sent BACK keyevent to dismiss modal overlay", - ) - - -class NavigateBackRecoveryStrategy(BaseRecoveryStrategy): - """Strategy to navigate back when agent navigates to a wrong screen.""" - - @property - def name(self) -> str: - return "navigate_back" - - async def can_handle(self, failure_context: dict[str, Any]) -> bool: - category = failure_context.get("category", "") - return category in ("wrong_screen", "dead_end", "navigation_error") - - async def execute( - self, - driver: BaseDeviceDriver, - failure_context: dict[str, Any], - ) -> RecoveryResult: - logger.info("[Recovery] Wrong screen detected, rolling back navigation...") - success = await driver.press_key("back") - return RecoveryResult( - success=success, - action_type=RecoveryActionType.NAVIGATE_BACK, - reason="Navigated back from incorrect screen", - ) - - -class ScrollAndSearchRecoveryStrategy(BaseRecoveryStrategy): - """Strategy to scroll viewport when target element is off-screen.""" - - @property - def name(self) -> str: - return "scroll_and_search" - - async def can_handle(self, failure_context: dict[str, Any]) -> bool: - category = failure_context.get("category", "") - return category in ("element_not_found", "off_screen", "target_missing") - - async def execute( - self, - driver: BaseDeviceDriver, - failure_context: dict[str, Any], - ) -> RecoveryResult: - direction = failure_context.get("scroll_direction", "down") - logger.info(f"[Recovery] Target missing, scrolling {direction} to reveal element...") - success = await driver.swipe_direction(direction=direction) - return RecoveryResult( - success=success, - action_type=RecoveryActionType.SCROLL_AND_RETRY, - reason=f"Scrolled {direction} to discover offscreen element", - ) - - -class RecoveryStrategyRegistry: - """Central registry and chain-of-responsibility for self-healing recovery strategies.""" - - _strategies: list[BaseRecoveryStrategy] = [ - DismissDialogRecoveryStrategy(), - NavigateBackRecoveryStrategy(), - ScrollAndSearchRecoveryStrategy(), - ] - - @classmethod - def register(cls, strategy: BaseRecoveryStrategy) -> None: - cls._strategies.append(strategy) - - @classmethod - async def attempt_recovery( - cls, - driver: BaseDeviceDriver, - failure_context: dict[str, Any], - ) -> RecoveryResult | None: - """Finds the first applicable strategy and executes recovery.""" - for strategy in cls._strategies: - if await strategy.can_handle(failure_context): - logger.info(f"Applying recovery strategy: {strategy.name}") - return await strategy.execute(driver=driver, failure_context=failure_context) - return None diff --git a/artemis/agents/validator/safety_net.py b/artemis/agents/validator/safety_net.py deleted file mode 100644 index ff33cc8..0000000 --- a/artemis/agents/validator/safety_net.py +++ /dev/null @@ -1,48 +0,0 @@ -# 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. - -"""Safety net and system dialog auto-recovery.""" - -from typing import Any -from artemis.drivers.base import BaseDeviceDriver -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class SafetyNet: - """Detects and resolves unwanted popups, system ANRs, and permission prompts.""" - - SYSTEM_DIALOG_PATTERNS = [ - "is not responding", - "Wait", - "Close app", - "Allow only while using the app", - "While using the app", - "Don't allow", - ] - - @classmethod - async def check_and_recover(cls, screen_data: Any, driver: BaseDeviceDriver) -> bool: - """Inspects screen elements for crash/permission dialogs and dismisses them.""" - for element in screen_data.ui_elements: - text = element.get("text", "") - if any(pattern.lower() in text.lower() for pattern in cls.SYSTEM_DIALOG_PATTERNS): - logger.info( - f"SafetyNet detected system dialog '{text}', attempting auto-dismiss..." - ) - center = element.get("center", [500, 500]) - await driver.tap(center[0], center[1]) - return True - return False diff --git a/artemis/controllers/android_controller.py b/artemis/controllers/android_controller.py deleted file mode 100644 index 77d3177..0000000 --- a/artemis/controllers/android_controller.py +++ /dev/null @@ -1,925 +0,0 @@ -# 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 -import base64 -from io import BytesIO -from pathlib import Path -import re -import tempfile -import time -from typing import Any -from uuid import uuid4 - -from adbutils import AdbClient, AdbDevice -from PIL import Image - -from artemis.clients.ui_automator_client import ( - UIAutomatorClient, - _parse_hierarchy_xml_to_elements, -) -from artemis.config import get_temp_dir -from artemis.controllers.device_controller import ( - MobileDeviceController, - ScreenDataResponse, -) -from artemis.controllers.types import ( - Bounds, - CoordinatesSelectorRequest, - TapOutput, -) -from artemis.data_engine.models import VideoRecordingRecord -from artemis.utils.logger import get_logger -from artemis.utils.video import ( - DEFAULT_MAX_DURATION_SECONDS, - RecordingSession, - VideoRecordingResult, - build_scrcpy_record_command, - cleanup_video_segments, - concatenate_videos, - get_active_session, - has_active_session, - normalize_recording_to_mp4, - remove_active_session, - set_active_session, - trim_video, -) - -logger = get_logger(__name__) - - -class AndroidDeviceController(MobileDeviceController): - def __init__( - self, - device_id: str, - adb_client: AdbClient, - ui_adb_client: UIAutomatorClient, - device_width: int, - device_height: int, - data_engine_start_time: float | None = None, - data_engine: Any | None = None, - ): - self.device_id = device_id - self.adb_client = adb_client - self.ui_adb_client = ui_adb_client - self.device_width = device_width - self.device_height = device_height - self.data_engine_start_time = data_engine_start_time - self.data_engine = data_engine - self._device: AdbDevice | None = None - self._segment_cache: dict[ - tuple[str, int, float, float], VideoRecordingResult - ] = {} - - @property - def device(self) -> AdbDevice: - if self._device is None: - self._device = self.adb_client.device(serial=self.device_id) - return self._device - - async def tap( - self, - coords: CoordinatesSelectorRequest, - long_press: bool = False, - long_press_duration: int = 1000, - times: int = 1, - delay_ms: int = 100, - ) -> TapOutput: - try: - if long_press: - cmd = ( - f"input swipe {coords.x} {coords.y} {coords.x} {coords.y} {long_press_duration}" - ) - else: - if times <= 1: - cmd = f"input tap {coords.x} {coords.y}" - else: - # Chain multiple taps in a single shell command to eliminate shell spawning overhead - # between rapid consecutive clicks (e.g. 7 taps to enter developer mode) - taps = [f"input tap {coords.x} {coords.y}"] * times - cmd = f" && sleep {delay_ms / 1000.0:.3f} && ".join(taps) - - logger.info(f"Executing ADB shell: {cmd}") - self.device.shell(cmd) - return TapOutput(error=None) - except Exception as e: - return TapOutput(error=f"ADB tap failed: {str(e)}") - - async def swipe( - self, - start: CoordinatesSelectorRequest, - end: CoordinatesSelectorRequest, - duration: int = 400, - ) -> str | None: - try: - cmd = f"input touchscreen swipe {start.x} {start.y} {end.x} {end.y} {duration}" - logger.info(f"Executing ADB command: {cmd}") - self.device.shell(cmd) - return None - except Exception as e: - return f"ADB swipe failed: {str(e)}" - - async def get_screen_data(self) -> ScreenDataResponse: - """Get screen data using the UIAutomator2 client""" - try: - logger.info("Using UIAutomator2 for screen data retrieval") - ui_data = self.ui_adb_client.get_screen_data() - return ScreenDataResponse( - base64=ui_data.base64, - elements=ui_data.elements, - width=ui_data.width, - height=ui_data.height, - platform="android", - ) - except Exception as e: - logger.error(f"Failed to get screen data: {e}") - raise - - async def screenshot(self) -> str: - try: - b64 = self.ui_adb_client.get_screenshot_base64() - if b64 is None: - raise RuntimeError("Failed to capture screenshot") - return b64 - except Exception as e: - logger.error(f"Failed to take screenshot: {e}") - raise - - async def input_text(self, text: str) -> bool: - try: - norm_text = ( - text.replace(r"\r\n", "\n") - .replace(r"\n", "\n") - .replace(r"\r", "\n") - ) - - # 1. Tier 1: Try clipboard injection + KEYCODE_PASTE (Zero IME interference, preserves multiline, works for all charsets) - if self.ui_adb_client: - try: - set_clip_ok = False - if hasattr(self.ui_adb_client, "set_clipboard"): - set_clip_ok = self.ui_adb_client.set_clipboard(norm_text) - elif hasattr(self.ui_adb_client, "_device") and self.ui_adb_client._device: - self.ui_adb_client._device.set_clipboard(norm_text) - set_clip_ok = True - elif hasattr(self.ui_adb_client, "_ensure_connected"): - dev = self.ui_adb_client._ensure_connected() - dev.set_clipboard(norm_text) - set_clip_ok = True - - if set_clip_ok: - self.device.shell("input keyevent 279") - return True - except Exception as e: - logger.debug(f"Clipboard paste fallback: {e}") - - return self._input_text_adb_fallback(norm_text) - except Exception as e: - logger.error(f"Failed to input text: {e}") - return False - - def _input_text_adb_fallback(self, text: str) -> bool: - """Fallback method using ADB shell input text command.""" - try: - norm_text = ( - text.replace(r"\r\n", "\n") - .replace(r"\n", "\n") - .replace(r"\r", "\n") - ) - lines = norm_text.split("\n") - for line_idx, line in enumerate(lines): - if line_idx > 0: - self.device.shell("input keyevent 66") - if not line: - continue - parts = line.split("%s") - for i, part in enumerate(parts): - to_write = "" - for char in part: - if char == " ": - to_write += "%s" - elif char in [ - "&", - "<", - ">", - "|", - ";", - "(", - ")", - "$", - "`", - "\\", - '"', - "'", - ]: - to_write += f"\\{char}" - else: - to_write += char - - if to_write: - self.device.shell(f"input text '{to_write}'") - - if i < len(parts) - 1: - self.device.shell("input keyevent 62") - - return True - except Exception as e: - logger.error(f"Failed to input text via ADB fallback: {e}") - return False - - async def launch_app(self, package_or_bundle_id: str) -> bool: - try: - # Clear foreground obstructions before launching - logger.info( - "Clearing foreground obstructions (status bar, system dialogs," - " keyboard) before launching app..." - ) - try: - # 1. Collapse status bar / notification shade - self.device.shell("cmd statusbar collapse") - except Exception as e: - logger.warning(f"Failed to collapse status bar: {e}") - - try: - # 2. Close system dialogs - self.device.shell("am broadcast -a android.intent.action.CLOSE_SYSTEM_DIALOGS") - except Exception as e: - logger.warning(f"Failed to close system dialogs: {e}") - - # 3. Dismiss keyboard if visible - if self.is_keyboard_visible(): - await self.dismiss_keyboard() - - self.device.app_start(package_or_bundle_id) - return True - except Exception as e: - logger.error(f"Failed to launch app {package_or_bundle_id}: {e}") - return False - - async def terminate_app(self, package_or_bundle_id: str | None) -> bool: - try: - if package_or_bundle_id is None: - current_app = self._get_current_foreground_package() - if current_app: - logger.info(f"Stopping currently running app: {current_app}") - self.device.app_stop(current_app) - else: - logger.warning("No foreground app detected") - return False - else: - self.device.app_stop(package_or_bundle_id) - return True - except Exception as e: - logger.error(f"Failed to terminate app {package_or_bundle_id}: {e}") - return False - - async def open_url(self, url: str) -> bool: - try: - cmd = f"am start -a android.intent.action.VIEW -d {url}" - logger.info(f"Executing ADB shell: {cmd}") - self.device.shell(cmd) - return True - except Exception as e: - logger.error(f"Failed to open URL {url}: {e}") - return False - - def is_keyboard_visible(self) -> bool: - try: - res = self.device.shell("dumpsys input_method") - if isinstance(res, bytes): - res = res.decode("utf-8") - - # Common patterns for open keyboard in dumpsys - patterns = [ - "mInputShown=true", - "mInputViewShowing=true", - "mInputViewShowing=VISIBLE", - ] - for p in patterns: - if p in res: - return True - - # Alternate check using dumpsys window - res_win = self.device.shell( - "dumpsys window | grep -E 'mInputMethodWindow|mShowingInputMethod'" - ) - if isinstance(res_win, bytes): - res_win = res_win.decode("utf-8") - if "mInputMethodWindow" in res_win and "visible=true" in res_win.lower(): - return True - - return False - except Exception as e: - logger.error(f"Failed to check keyboard visibility: {e}") - return False - - async def dismiss_keyboard(self) -> bool: - try: - if not self.is_keyboard_visible(): - return True - - logger.info("Keyboard is visible, attempting to dismiss via Keyevent 111 (Escape)") - self.device.shell("input keyevent 111") - await asyncio.sleep(0.3) - - if not self.is_keyboard_visible(): - logger.info("Keyboard successfully dismissed via Keyevent 111") - return True - - logger.info("Keyboard still visible, trying Keyevent 4 (BACK) as fallback to dismiss") - self.device.shell("input keyevent 4") - await asyncio.sleep(0.3) - - if not self.is_keyboard_visible(): - logger.info("Keyboard successfully dismissed via Keyevent 4") - return True - - logger.warning("Failed to dismiss keyboard after Escape and BACK") - return False - except Exception as e: - logger.error(f"Error in dismiss_keyboard: {e}") - return False - - async def press_back(self) -> bool: - try: - if self.is_keyboard_visible(): - logger.info( - "Intercepted press_back: Keyboard is visible. Attempting to" - " dismiss keyboard first." - ) - return await self.dismiss_keyboard() - - logger.info("Executing normal press_back (keyboard not visible)") - self.device.shell("input keyevent 4") - return True - except Exception as e: - logger.error(f"Failed to press back: {e}") - return False - - async def press_home(self) -> bool: - try: - self.device.shell("input keyevent 3") - return True - except Exception as e: - logger.error(f"Failed to press home: {e}") - return False - - async def press_enter(self) -> bool: - try: - self.device.shell("input keyevent 66") - return True - except Exception as e: - logger.error(f"Failed to press enter: {e}") - return False - - async def press_key(self, keycode: str) -> bool: - try: - # Intercept BACK keycodes to handle keyboard dismissal first - clean_keycode = str(keycode).strip().upper() - if clean_keycode in ("4", "KEYCODE_BACK", "BACK"): - if self.is_keyboard_visible(): - logger.info( - f"Intercepted press_key({keycode}): Keyboard is" - " visible. Dismissing keyboard first." - ) - return await self.dismiss_keyboard() - - self.device.shell(f"input keyevent {keycode}") - return True - except Exception as e: - logger.error(f"Failed to press key {keycode}: {e}") - return False - - async def get_ui_hierarchy(self) -> list[dict]: - - max_attempts = 4 - retry_delay = 0.5 - - for attempt in range(1, max_attempts + 1): - try: - hierarchy_xml = self.ui_adb_client.get_hierarchy() - elements = _parse_hierarchy_xml_to_elements(hierarchy_xml) - - # Check if hierarchy contains actual UI elements besides the root node - if len(elements) > 1: - if attempt > 1: - logger.info( - f"Successfully retrieved non-empty UI hierarchy on attempt {attempt}" - ) - return elements - - logger.warning( - "Retrieved empty or root-only UI hierarchy (length" - f" {len(elements)}) on attempt {attempt}/{max_attempts}." - f" Retrying in {retry_delay}s..." - ) - except Exception as e: - logger.warning( - f"Failed to get UI hierarchy on attempt {attempt}/{max_attempts}: {e}" - ) - if attempt == max_attempts: - logger.error(f"Failed to get UI hierarchy after {max_attempts} attempts: {e}") - return [] - - if attempt < max_attempts: - await asyncio.sleep(retry_delay) - - return [] - - def find_element( - self, - ui_hierarchy: list[dict], - resource_id: str | None = None, - text: str | None = None, - index: int = 0, - ) -> tuple[dict | None, Bounds | None, str | None]: - if not resource_id and not text: - return None, None, "No resource_id or text provided" - - matches = [] - for element in ui_hierarchy: - if resource_id and element.get("resource-id") == resource_id: - matches.append(element) - elif text and (element.get("text") == text or element.get("accessibilityText") == text): - matches.append(element) - - if not matches: - criteria = f"resource_id='{resource_id}'" if resource_id else f"text='{text}'" - return None, None, f"No element found with {criteria}" - - if index >= len(matches): - criteria = f"resource_id='{resource_id}'" if resource_id else f"text='{text}'" - return ( - None, - None, - (f"Index {index} out of range for {criteria} (found {len(matches)} matches)"), - ) - - element = matches[index] - bounds = self._extract_bounds(element) - - return element, bounds, None - - def _get_current_foreground_package(self) -> str | None: - try: - app_info = self.device.current_app() - if app_info and app_info.package: - return app_info.package - return None - except Exception as e: - logger.error( - f"Failed to get current foreground package using adbutils: {e}." - " Falling back to dumpsys parsing." - ) - try: - result = self.device.shell("dumpsys window | grep mCurrentFocus") - - # Convert to string if bytes - if isinstance(result, bytes): - result_str = result.decode("utf-8") - elif isinstance(result, str): - result_str = result - else: - return None - - if result_str and "=" in result_str: - parts = result_str.split("/") - if len(parts) > 0: - package = parts[0].split()[-1] - return package if package else None - except Exception as fallback_err: - logger.error(f"Fallback dumpsys parsing also failed: {fallback_err}") - return None - - def _extract_bounds(self, element: dict) -> Bounds | None: - bounds_str = element.get("bounds") - if not bounds_str or not isinstance(bounds_str, str): - return None - - try: - match = re.match(r"\[(\d+),(\d+)\]\[(\d+),(\d+)\]", bounds_str) - if match: - return Bounds( - x1=int(match.group(1)), - y1=int(match.group(2)), - x2=int(match.group(3)), - y2=int(match.group(4)), - ) - except (ValueError, IndexError): - return None - - return None - - async def erase_text(self, nb_chars: int | None = None) -> bool: - # 1. If clearing the whole field, try fast UIAutomator2 clear first - if nb_chars is None: - logger.info("Attempting fast UIAutomator2 clear_text...") - if self.ui_adb_client.clear_text(): - return True - logger.warning("UIAutomator2 clear_text failed, falling back to optimized ADB shell") - - # 2. Optimized ADB fallback: combine multiple KEYCODE_DEL keyevents into a single command - try: - chars_to_delete = nb_chars if nb_chars is not None else 50 - if chars_to_delete > 0: - # KEYCODE_DEL is 67. Build arguments list "67 67 67 ..." - keycode_args = " ".join(["67"] * chars_to_delete) - cmd = f"input keyevent {keycode_args}" - logger.info(f"Executing combined ADB shell delete: {cmd}") - self.device.shell(cmd) - return True - except Exception as e: - logger.error(f"Failed to erase text via ADB fallback: {e}") - return False - - async def cleanup(self) -> None: - pass - - def get_compressed_b64_screenshot(self, image_base64: str, quality: int = 50) -> str: - if image_base64.startswith("data:image"): - image_base64 = image_base64.split(",")[1] - - image_data = base64.b64decode(image_base64) - image = Image.open(BytesIO(image_data)) - - compressed_io = BytesIO() - image.save(compressed_io, format="JPEG", quality=quality, optimize=True) - - compressed_base64 = base64.b64encode(compressed_io.getvalue()).decode("utf-8") - return compressed_base64 - - async def start_video_recording( - self, - max_duration_seconds: int = DEFAULT_MAX_DURATION_SECONDS, - output_dir: Path | None = None, - ) -> VideoRecordingResult: - """Start screen recording on Android device using scrcpy.""" - self._segment_cache.clear() - if has_active_session(self.device_id): - return VideoRecordingResult( - success=False, - message=(f"Recording already in progress for device {self.device_id}"), - ) - - try: - # Create a temp file for the MKV recording if output_dir not provided - if not output_dir: - output_dir = Path( - tempfile.mkdtemp(prefix="scrcpy_", dir=get_temp_dir("recordings")) - ) - else: - output_dir.mkdir(parents=True, exist_ok=True) - local_video_path = output_dir / "recording.mkv" - - video_id = uuid4() - session = RecordingSession( - video_id=video_id, - device_id=self.device_id, - start_time=time.time(), - data_engine_start_time=self.data_engine_start_time, - local_video_path=local_video_path, - capture_width=self.device_width, - capture_height=self.device_height, - ) - - # Persist to local database if Data Engine is active - if self.data_engine and self.data_engine.storage: - try: - record = VideoRecordingRecord( - video_id=video_id, - session_id=self.data_engine.current_session_id, - device_id=self.device_id, - start_time=session.start_time, - local_video_path=str(local_video_path), - ) - self.data_engine.storage.create_video_recording(record) - except Exception as db_err: - logger.error(f"Failed to persist video recording start to DB: {db_err}") - - # Start scrcpy in background - cmd = build_scrcpy_record_command( - "scrcpy", self.device_id, local_video_path, lock_capture_orientation=False - ) - - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - - session.process = process - set_active_session(self.device_id, session) - - logger.info( - f"Started scrcpy recording on {self.device_id}, saving to {local_video_path}" - ) - - # Wait a bit and check if it's still running. - await asyncio.sleep(1.0) - if process.returncode is not None: - stderr = await process.stderr.read() - logger.error(f"scrcpy failed to start: {stderr.decode()}") - remove_active_session(self.device_id) - return VideoRecordingResult( - success=False, - message=f"scrcpy failed to start: {stderr.decode()}", - ) - - return VideoRecordingResult( - success=True, - message=f"Recording started on {self.device_id}", - ) - - except Exception as e: - logger.error(f"Failed to start scrcpy recording: {e}") - remove_active_session(self.device_id) - return VideoRecordingResult( - success=False, - message=f"Failed to start recording: {e}", - ) - - async def stop_video_recording(self) -> VideoRecordingResult: - """Stop scrcpy recording and return the video file.""" - self._segment_cache.clear() - session = get_active_session(self.device_id) - if not session: - return VideoRecordingResult( - success=False, - message=f"No active recording for device {self.device_id}", - ) - - try: - process = session.process - if process is not None: - process.terminate() - await asyncio.wait_for(process.wait(), timeout=5.0) - - output_path = session.local_video_path - - # Fix: Concatenate previous segments if scrcpy crashed and restarted - if session.android_video_segments: - all_segments = session.android_video_segments + [output_path] - concatenated_mkv = output_path.parent / "full_recording.mkv" - - if await concatenate_videos(all_segments, concatenated_mkv): - cleanup_video_segments(all_segments) - output_path = concatenated_mkv - - # Convert MKV to MP4 - mp4_path = ( - output_path.parent / "recording.mp4" - if output_path.name == "full_recording.mkv" - else output_path.with_suffix(".mp4") - ) - success = await self._convert_mkv_to_mp4(output_path, mp4_path) - - remove_active_session(self.device_id) - - final_video_path = mp4_path if (success and mp4_path.exists()) else output_path - - # Persist update to local database if Data Engine is active - if self.data_engine and self.data_engine.storage: - try: - record = VideoRecordingRecord( - video_id=session.video_id, - session_id=self.data_engine.current_session_id, - device_id=self.device_id, - start_time=session.start_time, - end_time=time.time(), - local_video_path=str(final_video_path), - ) - self.data_engine.storage.update_video_recording(record) - except Exception as db_err: - logger.error(f"Failed to persist video recording stop to DB: {db_err}") - - if success and mp4_path.exists(): - try: - output_path.unlink() - except Exception: - pass - return VideoRecordingResult( - success=True, - message=f"Recording stopped, saved to {mp4_path}", - video_path=mp4_path, - ) - else: - return VideoRecordingResult( - success=True, - message=( - f"Recording stopped, saved as MKV to {output_path} (conversion failed)" - ), - video_path=output_path, - ) - - except Exception as e: - logger.error(f"Failed to stop scrcpy recording: {e}") - remove_active_session(self.device_id) - return VideoRecordingResult( - success=False, - message=f"Failed to stop recording: {e}", - ) - - async def extract_segment_metadata( - self, - start_relative_time: float, - end_relative_time: float | None = None, - ) -> VideoRecordingResult: - """Get a video segment for a specific time range (relative to video start).""" - session = get_active_session(self.device_id) - if not session: - return VideoRecordingResult( - success=False, - message=f"No active recording for device {self.device_id}", - ) - - cache_key = None - if end_relative_time is not None: - cache_key = ( - str(session.video_id), - session.generation, - round(start_relative_time, 1), - round(end_relative_time, 1), - ) - cached_res = self._segment_cache.get(cache_key) - if cached_res and cached_res.success and cached_res.video_path and cached_res.video_path.exists(): - logger.info( - "Reusing generation-scoped trimmed video segment for range" - f" {cache_key[2]}s to {cache_key[3]}s" - ) - return cached_res - - try: - mkv_path = session.local_video_path - if not mkv_path or not mkv_path.exists(): - return VideoRecordingResult( - success=False, - message="Recording file not found", - ) - - # Check if process is still running (best effort for disconnection check) - crashed = False - restart_message = "" - - if session.process and session.process.returncode is not None: - logger.warning(f"scrcpy process for {self.device_id} has terminated unexpectedly") - crashed = True - - # Move current file to segments list (save it before overwriting session state) - session.android_video_segments.append(mkv_path) - - # Try to restart scrcpy - try: - session.android_segment_index += 1 - session.generation = session.android_segment_index - output_dir = mkv_path.parent - new_video_path = output_dir / f"recording_{session.android_segment_index}.mkv" - - cmd = build_scrcpy_record_command( - "scrcpy", self.device_id, new_video_path, lock_capture_orientation=False - ) - - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - ) - - # Wait a bit and check if it's still running. - await asyncio.sleep(1.0) - if process.returncode is None: - session.process = process - session.local_video_path = new_video_path - restart_message = "Successfully restarted scrcpy." - else: - stderr = await process.stderr.read() - restart_message = f"Failed to restart scrcpy: {stderr.decode()}" - except Exception as e: - restart_message = f"Failed to restart scrcpy: {e}" - - current_time = time.time() - video_duration = current_time - session.start_time - - # Calculate offset - offset = 0.0 - if session.data_engine_start_time is not None: - offset = session.start_time - session.data_engine_start_time - if offset < 0: - offset = 0.0 - - video_start_relative_time = start_relative_time - offset - video_end_relative_time = ( - end_relative_time - offset if end_relative_time is not None else None - ) - - # Safety buffer: avoid reading the very last second being written - safe_duration = max(0.0, video_duration - 1.0) - - truncation_warning = None - if video_end_relative_time is not None and video_end_relative_time > safe_duration: - truncation_warning = ( - "Video segment truncated. Requested end time" - f" {end_relative_time:.1f}s exceeded available safe" - f" duration. Truncated to {safe_duration + offset:.1f}s." - ) - - if video_end_relative_time is None or video_end_relative_time > safe_duration: - video_end_relative_time = safe_duration - - if video_start_relative_time < 0: - video_start_relative_time = 0.0 - - if video_start_relative_time >= video_end_relative_time: - return VideoRecordingResult( - success=False, - message=( - "Invalid time range or requested too close to current" - f" time. Available duration: {safe_duration:.1f}s" - ), - ) - - # Trim directly from MKV to MP4 (using mkv_path which points to the file before we potentially changed it for restart) - - trim_output_dir = tempfile.mkdtemp( - prefix="video_trimmed_", - dir=get_temp_dir("trimmed_videos"), - ) - trim_output_path = Path(trim_output_dir) / "segment.mp4" - - # Retry mechanism for ffmpeg stability - success = False - for attempt in range(3): - success = await trim_video( - mkv_path, - video_start_relative_time, - video_end_relative_time, - trim_output_path, - ) - if success and trim_output_path.exists(): - break - logger.warning(f"ffmpeg trim attempt {attempt + 1} failed, retrying in 0.5s...") - await asyncio.sleep(0.5) - - if not success or not trim_output_path.exists(): - return VideoRecordingResult( - success=False, - message="Failed to trim video segment after retries", - ) - - actual_start = video_start_relative_time + offset - actual_end = video_end_relative_time + offset - - message = f"Video segment retrieved for range {actual_start:.1f}s to {actual_end:.1f}s" - if crashed: - message += ( - f". Warning: scrcpy crashed. Data loss may have occurred. {restart_message}" - ) - - file_size_mb = trim_output_path.stat().st_size / (1024 * 1024) - duration = video_end_relative_time - video_start_relative_time - - res = VideoRecordingResult( - success=True, - message=message, - video_path=trim_output_path, - file_size_mb=round(file_size_mb, 2), - duration_seconds=round(duration, 2), - actual_start_relative_time=actual_start, - warning=truncation_warning, - video_id=session.video_id, - generation=session.generation, - sealed_until=session.sealed_until, - source_revision=( - f"{session.video_id}:{session.generation}:{round(actual_end, 3)}" - ), - ) - if cache_key is not None: - self._segment_cache[cache_key] = res - return res - - except Exception as e: - logger.error(f"Failed to get video segment: {e}") - return VideoRecordingResult( - success=False, - message=f"Failed to get video segment: {e}", - ) - - async def _convert_mkv_to_mp4(self, mkv_path: Path, mp4_path: Path) -> bool: - """Normalize MKV into a fixed-size, browser-safe MP4.""" - session = get_active_session(self.device_id) - return await normalize_recording_to_mp4( - mkv_path, - mp4_path, - session.capture_width if session else self.device_width, - session.capture_height if session else self.device_height, - ) diff --git a/artemis/core/__init__.py b/artemis/core/__init__.py index c0dc5fc..bc8705f 100644 --- a/artemis/core/__init__.py +++ b/artemis/core/__init__.py @@ -12,40 +12,10 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Core framework contracts and lifecycle primitives.""" +"""Core framework primitives. -from artemis.core.context import ExecutionContext, DeviceInfoContext -from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord -from artemis.core.events import Event, EventBus, HookType, global_event_bus -from artemis.core.exceptions import ( - ArtemisException, - DeviceDriverException, - DeviceNotFoundException, - AgentExecutionException, - ToolExecutionException, - LLMProviderException, - StateTransitionException, -) -from artemis.core.registry import AgentRegistry, ToolRegistry, DriverRegistry +Live submodules: ``artemis.core.tool_declaration`` and ``artemis.core.diagnostics`` +(imported directly by module path by their consumers). +""" -__all__ = [ - "ExecutionContext", - "DeviceInfoContext", - "ExecutionContextState", - "ExecutionStatus", - "StepRecord", - "Event", - "EventBus", - "HookType", - "global_event_bus", - "ArtemisException", - "DeviceDriverException", - "DeviceNotFoundException", - "AgentExecutionException", - "ToolExecutionException", - "LLMProviderException", - "StateTransitionException", - "AgentRegistry", - "ToolRegistry", - "DriverRegistry", -] +__all__: list[str] = [] diff --git a/artemis/core/context.py b/artemis/core/context.py deleted file mode 100644 index 0815e33..0000000 --- a/artemis/core/context.py +++ /dev/null @@ -1,62 +0,0 @@ -# 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. - -"""Clean Execution Context for ARTEMIS.""" - -from pathlib import Path -from typing import Any -from uuid import uuid4 -from pydantic import BaseModel, Field - - -class DeviceInfoContext(BaseModel): - """Device metadata and configuration.""" - - device_id: str = Field(default="default-device", description="ADB serial or device identifier") - platform: str = Field( - default="android", description="Platform type ('android', 'cloud', 'mock')" - ) - width: int = Field(default=1080, description="Screen width in pixels") - height: int = Field(default=2400, description="Screen height in pixels") - - -class ExecutionContext: - """Central context holding device driver, session identifiers, and telemetry tracer.""" - - def __init__( - self, - task_goal: str, - device_id: str = "default-device", - trace_id: str | None = None, - traces_path: Path | str | None = None, - locked_package: str | None = None, - platform: str = "android", - ): - self.task_goal = task_goal - self.trace_id = trace_id or str(uuid4()) - self.traces_path = Path(traces_path) if traces_path else Path("traces") / self.trace_id - self.locked_package = locked_package - self.device = DeviceInfoContext(device_id=device_id, platform=platform) - self.driver: Any | None = None - self.telemetry: Any | None = None - # Fully-wired ArtemisContext (data_engine, adb clients, execution setup). - # Required by the live (non-mock) paths of GraphRunner/ReactiveRunner. - self.artemis_context: Any | None = None - self.custom_attributes: dict[str, Any] = {} - - def set_attribute(self, key: str, value: Any) -> None: - self.custom_attributes[key] = value - - def get_attribute(self, key: str, default: Any = None) -> Any: - return self.custom_attributes.get(key, default) diff --git a/artemis/core/events.py b/artemis/core/events.py deleted file mode 100644 index dd772b5..0000000 --- a/artemis/core/events.py +++ /dev/null @@ -1,77 +0,0 @@ -# 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. - -"""Event Bus and Lifecycle Hooks System for ARTEMIS.""" - -from collections import defaultdict -from collections.abc import Callable -from enum import Enum -import inspect -import time -from typing import Any -from pydantic import BaseModel, Field -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class HookType(str, Enum): - BEFORE_TASK = "before_task" - AFTER_TASK = "after_task" - BEFORE_STEP = "before_step" - AFTER_STEP = "after_step" - ON_ACTION = "on_action" - ON_ERROR = "on_error" - ON_SCREEN_CHANGE = "on_screen_change" - - -class Event(BaseModel): - """Encapsulates an event payload emitted through the event bus.""" - - event_type: HookType = Field(..., description="Type of the lifecycle event") - timestamp: float = Field(default_factory=time.time, description="Event occurrence timestamp") - trace_id: str = Field(..., description="Active session trace identifier") - payload: dict[str, Any] = Field(default_factory=dict, description="Event data dictionary") - - -class EventBus: - """Asynchronous event bus supporting subscription and broadcast of lifecycle events.""" - - def __init__(self): - self._subscribers: dict[HookType, list[Callable[..., Any]]] = defaultdict(list) - - def subscribe(self, hook: HookType, callback: Callable[..., Any]) -> None: - """Registers a listener callback for a specific lifecycle event.""" - self._subscribers[hook].append(callback) - - def unsubscribe(self, hook: HookType, callback: Callable[..., Any]) -> None: - """Removes a registered callback.""" - if callback in self._subscribers[hook]: - self._subscribers[hook].remove(callback) - - async def emit(self, event: Event) -> None: - """Asynchronously dispatches an event to all subscribed listeners.""" - listeners = self._subscribers.get(event.event_type, []) - for listener in listeners: - try: - if inspect.iscoroutinefunction(listener): - await listener(event) - else: - listener(event) - except Exception as e: - logger.warning(f"Error executing event listener for {event.event_type}: {e}") - - -# Global singleton event bus instance -global_event_bus = EventBus() diff --git a/artemis/core/exceptions.py b/artemis/core/exceptions.py deleted file mode 100644 index 7eb1983..0000000 --- a/artemis/core/exceptions.py +++ /dev/null @@ -1,57 +0,0 @@ -# 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. - -"""Standardized exception hierarchy for ARTEMIS.""" - - -class ArtemisException(Exception): - """Base exception for all ARTEMIS framework errors.""" - - pass - - -class DeviceDriverException(ArtemisException): - """Raised when hardware, emulator, or ADB interactions fail.""" - - pass - - -class DeviceNotFoundException(DeviceDriverException): - """Raised when the specified target device cannot be located or connected.""" - - pass - - -class AgentExecutionException(ArtemisException): - """Raised when an autonomous agent fails during reasoning or execution.""" - - pass - - -class ToolExecutionException(ArtemisException): - """Raised when a tool action fails or receives invalid arguments.""" - - pass - - -class LLMProviderException(ArtemisException): - """Raised when an upstream LLM/VLM API call fails or times out.""" - - pass - - -class StateTransitionException(ArtemisException): - """Raised when graph execution encounters an invalid state transition.""" - - pass diff --git a/artemis/core/registry.py b/artemis/core/registry.py deleted file mode 100644 index 140e877..0000000 --- a/artemis/core/registry.py +++ /dev/null @@ -1,46 +0,0 @@ -# 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. - -"""Global Plugin and Component Registry for Agents, Tools, and Drivers.""" - -from typing import Any, TypeVar - -T = TypeVar("T") - - -class ComponentRegistry: - """Generic typed component registry supporting runtime dynamic plugin loading.""" - - def __init__(self, name: str): - self.name = name - self._items: dict[str, Any] = {} - - def register(self, key: str, item: Any) -> Any: - self._items[key.lower()] = item - return item - - def get(self, key: str) -> Any | None: - return self._items.get(key.lower()) - - def list_all(self) -> dict[str, Any]: - return dict(self._items) - - def contains(self, key: str) -> bool: - return key.lower() in self._items - - -# Global Registries -AgentRegistry = ComponentRegistry("Agents") -ToolRegistry = ComponentRegistry("Tools") -DriverRegistry = ComponentRegistry("Drivers") diff --git a/artemis/core/state.py b/artemis/core/state.py deleted file mode 100644 index 9c3a4e0..0000000 --- a/artemis/core/state.py +++ /dev/null @@ -1,70 +0,0 @@ -# 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. - -"""Core Execution State and Lifecycle Data Structures.""" - -from enum import Enum -import time -from typing import Any -from pydantic import BaseModel, Field - - -class ExecutionStatus(str, Enum): - PENDING = "pending" - RUNNING = "running" - SUCCESS = "success" - FAILED = "failed" - PAUSED = "paused" - INTERRUPTED = "interrupted" - - -class StepRecord(BaseModel): - """Encapsulates a single step's observation, thought, and action.""" - - step_number: int = Field(..., description="1-indexed step turn counter") - thought: str | None = Field(default=None, description="Reasoning text") - action_name: str | None = Field(default=None, description="Triggered action name") - action_params: dict[str, Any] = Field(default_factory=dict, description="Action arguments") - result: str | None = Field(default=None, description="Action execution result output") - screenshot_path: str | None = Field(default=None, description="Saved step screenshot file URI") - timestamp: float = Field(default_factory=time.time, description="Timestamp of step start") - duration_seconds: float = Field(default=0.0, description="Step processing duration in seconds") - - -class ExecutionContextState(BaseModel): - """Full execution state passed across agents and runners.""" - - task_goal: str = Field(..., description="Initial user task objective") - trace_id: str = Field(..., description="Unique task trace identifier") - status: ExecutionStatus = Field( - default=ExecutionStatus.RUNNING, description="Current execution lifecycle status" - ) - current_turn: int = Field(default=0, description="Current step turn count") - max_turns: int = Field(default=30, description="Maximum allowed step turns before timeout") - steps: list[StepRecord] = Field( - default_factory=list, description="Chronological execution step history" - ) - active_plan: list[str] = Field( - default_factory=list, description="Decomposed milestone plan items" - ) - current_step_index: int = Field(default=0, description="Current index within active_plan") - metadata: dict[str, Any] = Field( - default_factory=dict, description="Custom metadata and state parameters" - ) - error_message: str | None = Field( - default=None, description="Failure reason if status is FAILED" - ) - final_output: Any | None = Field( - default=None, description="Final synthesized structured output from Outputter or runner" - ) diff --git a/artemis/engine/__init__.py b/artemis/engine/__init__.py deleted file mode 100644 index 5064667..0000000 --- a/artemis/engine/__init__.py +++ /dev/null @@ -1,27 +0,0 @@ -# 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. - -"""Execution Engine and Pipeline Orchestration layer.""" - -from artemis.engine.base_runner import BaseRunner -from artemis.engine.graph_runner import GraphRunner -from artemis.engine.pipeline import Pipeline -from artemis.engine.reactive_runner import ReactiveRunner - -__all__ = [ - "BaseRunner", - "ReactiveRunner", - "GraphRunner", - "Pipeline", -] diff --git a/artemis/engine/base_runner.py b/artemis/engine/base_runner.py deleted file mode 100644 index 11433a0..0000000 --- a/artemis/engine/base_runner.py +++ /dev/null @@ -1,36 +0,0 @@ -# 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. - -"""Base Execution Runner protocol and lifecycle manager.""" - -from abc import ABC, abstractmethod -from artemis.core.context import ExecutionContext -from artemis.core.state import ExecutionContextState -from artemis.drivers.base import BaseDeviceDriver -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class BaseRunner(ABC): - """Abstract Base Runner for orchestrating mobile agent workflows.""" - - def __init__(self, ctx: ExecutionContext, driver: BaseDeviceDriver): - self.ctx = ctx - self.driver = driver - - @abstractmethod - async def run(self, max_turns: int = 30) -> ExecutionContextState: - """Executes the autonomous automation workflow until completion or limit.""" - ... diff --git a/artemis/engine/graph_runner.py b/artemis/engine/graph_runner.py deleted file mode 100644 index 632bf7f..0000000 --- a/artemis/engine/graph_runner.py +++ /dev/null @@ -1,89 +0,0 @@ -# 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. - -"""Pro Graph Runner orchestrating multi-agent state machines with self-healing.""" - -from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord -from artemis.drivers.mock.mock_driver import MockDeviceDriver -from artemis.engine.base_runner import BaseRunner -from artemis.graph.graph import get_graph -from artemis.graph.state import State -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class GraphRunner(BaseRunner): - """Deep reasoning closed-loop runner coordinating Planner, Operator, Validator, and Summarizer.""" - - async def run(self, max_turns: int = 50) -> ExecutionContextState: - state = ExecutionContextState( - task_goal=self.ctx.task_goal, - trace_id=self.ctx.trace_id, - max_turns=max_turns, - ) - - logger.info( - f"Starting GraphRunner for deep closed-loop goal: '{self.ctx.task_goal}' on device: '{self.driver.device_id}'" - ) - - is_mock = isinstance(self.driver, MockDeviceDriver) - - if not is_mock: - try: - artemis_ctx = self.ctx.artemis_context - if artemis_ctx is None: - # Fail fast with a clear contract violation instead of - # silently building an empty ArtemisContext that drops - # data_engine/llm_config and dies later in perception. - raise RuntimeError( - "GraphRunner requires ExecutionContext.artemis_context" - " to be wired (data_engine, adb clients, execution" - " setup) before running the live graph path." - ) - graph = await get_graph(artemis_ctx) - graph_input = State.initial(self.ctx.task_goal) - async for _ in graph.astream( - input=graph_input, config={"recursion_limit": max_turns} - ): - pass - state.status = ExecutionStatus.SUCCESS - return state - except Exception as e: - logger.error(f"Live GraphRunner execution failed: {e}", exc_info=True) - state.status = ExecutionStatus.FAILED - state.error_message = str(e) - return state - - # Plan-Execute-Validate closed loop on driver - state.active_plan = [ - f"Step 1: Inspect screen for goal '{self.ctx.task_goal}'", - "Step 2: Perform required UI interactions", - "Step 3: Verify target screen state", - ] - - screen_data = await self.driver.get_screen_data() - state.steps.append( - StepRecord( - step_number=1, - thought="Decomposed goal into structured milestones; inspecting initial UI state.", - action_name="get_screen_data", - action_params={}, - result="Screen captured successfully", - duration_seconds=0.05, - ) - ) - state.current_turn = 1 - state.status = ExecutionStatus.SUCCESS - return state diff --git a/artemis/engine/pipeline.py b/artemis/engine/pipeline.py deleted file mode 100644 index 326f23f..0000000 --- a/artemis/engine/pipeline.py +++ /dev/null @@ -1,95 +0,0 @@ -# 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 os -from typing import Any, Literal - -from artemis.config import OutputConfig, load_agent_config -from artemis.core.context import ExecutionContext -from artemis.core.state import ExecutionContextState -from artemis.drivers.base import BaseDeviceDriver -from artemis.drivers.mock.mock_driver import MockDeviceDriver -from artemis.engine.base_runner import BaseRunner -from artemis.engine.graph_runner import GraphRunner -from artemis.engine.reactive_runner import ReactiveRunner -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class Pipeline: - """Orchestrates runner selection and execution flow.""" - - @classmethod - def create_runner( - cls, - profile: Literal["flash", "pro"], - ctx: ExecutionContext, - driver: BaseDeviceDriver, - ) -> BaseRunner: - if profile.lower() == "flash": - return ReactiveRunner(ctx=ctx, driver=driver) - return GraphRunner(ctx=ctx, driver=driver) - - @classmethod - async def execute( - cls, - goal: str, - profile: Literal["flash", "pro"] = "pro", - driver: BaseDeviceDriver | None = None, - device_id: str = "default-device", - output_config: OutputConfig | None = None, - enable_outputter: bool | None = None, - locked_package: str | None = None, - **kwargs: Any, - ) -> ExecutionContextState: - ctx = ExecutionContext(task_goal=goal, device_id=device_id) - if not driver: - driver = MockDeviceDriver(device_id=device_id) - - max_turns = kwargs.get("max_turns", 5) - runner = cls.create_runner(profile=profile, ctx=ctx, driver=driver) - state = await runner.run(max_turns=max_turns) - - # Check Outputter configuration - try: - agent_cfg = load_agent_config() - outputter_cfg = getattr(agent_cfg, "outputter", None) - is_enabled = ( - enable_outputter - if enable_outputter is not None - else (outputter_cfg.enabled if outputter_cfg else True) - ) - force_synth = outputter_cfg.force_synthesis if outputter_cfg else False - - should_synthesize = is_enabled and ( - ( - output_config - and output_config.needs_structured_format(default_enabled=is_enabled) - ) - or force_synth - ) - - if should_synthesize: - logger.info("Executing post-execution Outputter synthesis in Pipeline...") - summary = "\n".join( - f"- Step {s.step_number}: {s.action_name}({s.action_params}) -> {s.result or s.thought or 'completed'}" - for s in state.steps - ) - state.metadata["structured_output"] = summary - state.final_output = summary - except Exception as e: - logger.warning(f"Pipeline outputter post-processing skipped: {e}") - - return state diff --git a/artemis/engine/reactive_runner.py b/artemis/engine/reactive_runner.py deleted file mode 100644 index 9a7a73c..0000000 --- a/artemis/engine/reactive_runner.py +++ /dev/null @@ -1,112 +0,0 @@ -# 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. - -"""Flash Reactive Runner executing rapid Observe-Think-Act cycles.""" - -import time - -from artemis.agents.flash.runner import FlashRunner -from artemis.config import load_agent_config -from artemis.core.state import ExecutionContextState, ExecutionStatus, StepRecord -from artemis.drivers.mock.mock_driver import MockDeviceDriver -from artemis.engine.base_runner import BaseRunner -from artemis.graph.state import State -from artemis.utils.logger import get_logger - -logger = get_logger(__name__) - - -class ReactiveRunner(BaseRunner): - """Fast, single-loop reactive runner for short-to-medium workflows (< 30 steps).""" - - async def run(self, max_turns: int | None = None) -> ExecutionContextState: - if max_turns is None: - try: - cfg = load_agent_config() - max_turns = cfg.flash.max_turns - except Exception: - max_turns = 30 - - state = ExecutionContextState( - task_goal=self.ctx.task_goal, - trace_id=self.ctx.trace_id, - max_turns=max_turns, - ) - - logger.info( - f"Starting ReactiveRunner for goal: '{self.ctx.task_goal}' on device: '{self.driver.device_id}'" - ) - - is_mock = isinstance(self.driver, MockDeviceDriver) - - if not is_mock: - try: - artemis_ctx = self.ctx.artemis_context - if artemis_ctx is None: - # Fail fast with a clear contract violation instead of - # silently building an empty ArtemisContext that drops - # data_engine/llm_config and dies later mid-run. - raise RuntimeError( - "ReactiveRunner requires ExecutionContext.artemis_context" - " to be wired (data_engine, adb clients, execution" - " setup) before running the live Flash path." - ) - flash_runner = FlashRunner( - artemis_ctx, goal=self.ctx.task_goal, max_turns=max_turns - ) - legacy_state = State.initial(self.ctx.task_goal) - result = await flash_runner.run(legacy_state) - - state.current_turn = 1 - if result.get("status") == "completed": - state.status = ExecutionStatus.SUCCESS - else: - state.status = ExecutionStatus.FAILED - state.error_message = result.get("explanation") or "Task failed" - return state - except Exception as e: - logger.error(f"Live FlashRunner execution failed: {e}", exc_info=True) - state.status = ExecutionStatus.FAILED - state.error_message = str(e) - return state - - # Driver-based execution loop (for MockDriver and direct driver automation) - for turn in range(1, max_turns + 1): - state.current_turn = turn - step_start = time.time() - - # 1. Observe screen - screen_data = await self.driver.get_screen_data() - - # 2. Execute action step - action_name = "wait_for_delay" - action_params = {"seconds": 0.1} - await self.driver.wait_for_delay(0.1) - - # 3. Record step - step_record = StepRecord( - step_number=turn, - thought=f"Turn {turn}: Observing screen state and progressing toward goal: '{self.ctx.task_goal}'.", - action_name=action_name, - action_params=action_params, - result="Executed successfully", - duration_seconds=time.time() - step_start, - ) - state.steps.append(step_record) - - if turn >= 1: - state.status = ExecutionStatus.SUCCESS - break - - return state diff --git a/artemis/llm/client_factory.py b/artemis/llm/client_factory.py deleted file mode 100644 index be2cf49..0000000 --- a/artemis/llm/client_factory.py +++ /dev/null @@ -1,25 +0,0 @@ -# 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. - -"""LLM Provider Client Factory.""" - -from langchain_core.language_models.chat_models import BaseChatModel -from artemis.llm.router import ModelEndpoint, default_router - - -def get_chat_model(role: str = "operator", endpoint: ModelEndpoint | None = None) -> BaseChatModel: - """Instantiates a configured chat model for the requested role or endpoint.""" - if endpoint: - return default_router.instantiate_model(endpoint) - return default_router.create_chat_model(role) diff --git a/tests/conftest.py b/tests/conftest.py index e3c2118..39f1c01 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,6 @@ """Global pytest fixtures for ARTEMIS test suite.""" import pytest -from artemis.core.context import ExecutionContext from artemis.drivers.mock.mock_driver import MockDeviceDriver @@ -23,9 +22,3 @@ from artemis.drivers.mock.mock_driver import MockDeviceDriver def mock_driver(): """Provides an isolated MockDeviceDriver instance.""" return MockDeviceDriver(device_id="fixture-mock-device", width=1080, height=2400) - - -@pytest.fixture -def test_context(): - """Provides a fresh ExecutionContext.""" - return ExecutionContext(task_goal="Test sample task objective", device_id="fixture-mock-device") diff --git a/tests/integration/test_reactive_runner_integration.py b/tests/integration/test_reactive_runner_integration.py deleted file mode 100644 index c65c1d0..0000000 --- a/tests/integration/test_reactive_runner_integration.py +++ /dev/null @@ -1,45 +0,0 @@ -# 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. - -"""Integration tests for ReactiveRunner and Pipeline.""" - -import pytest -from artemis.core.context import ExecutionContext -from artemis.core.state import ExecutionStatus -from artemis.drivers.mock.mock_driver import MockDeviceDriver -from artemis.engine.pipeline import Pipeline -from artemis.engine.reactive_runner import ReactiveRunner - - -@pytest.mark.asyncio -async def test_reactive_runner_execution(): - """Verify ReactiveRunner executes workflow cycles on mock driver.""" - ctx = ExecutionContext(task_goal="Turn on Bluetooth in Settings", device_id="test-emulator") - driver = MockDeviceDriver(device_id="test-emulator") - - runner = ReactiveRunner(ctx=ctx, driver=driver) - state = await runner.run(max_turns=5) - - assert state.status == ExecutionStatus.SUCCESS - assert state.current_turn >= 1 - assert len(state.steps) >= 1 - assert state.steps[0].action_name == "wait_for_delay" - - -@pytest.mark.asyncio -async def test_pipeline_dispatch(): - """Verify Pipeline executes with default mock driver.""" - state = await Pipeline.execute(goal="Open YouTube and play video", profile="flash") - assert state.status == ExecutionStatus.SUCCESS - assert state.task_goal == "Open YouTube and play video" diff --git a/tests/unit/test_agents.py b/tests/unit/test_agents.py deleted file mode 100644 index 433961d..0000000 --- a/tests/unit/test_agents.py +++ /dev/null @@ -1,146 +0,0 @@ -# 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 Base Agent Protocol & Lifecycle Registry.""" - -import pytest -from typing import Any -from artemis.agents.base import AgentConfig, AgentRegistry, BaseAgent -from artemis.agents.validator.failure_analyzer_agent import FailureAnalyzerAgent -from artemis.agents.validator.recovery_strategies import ( - RecoveryStrategyRegistry, -) -from artemis.drivers.mock.mock_driver import MockDeviceDriver - - -class DummyTestAgent(BaseAgent): - """Simple test agent implementing BaseAgent protocol.""" - - def __init__(self, config: AgentConfig | None = None, ctx: Any | None = None, **kwargs: Any): - super().__init__( - name="dummy_test_agent", - role="Test Worker", - description="Agent used for testing lifecycle transitions.", - config=config, - ctx=ctx, - ) - - async def process(self, state: Any) -> dict[str, Any]: - return {"status": "success", "processed_state": True} - - -@pytest.mark.asyncio -async def test_base_agent_lifecycle(): - """Verify BaseAgent initialization, execution, and cleanup lifecycle.""" - agent = DummyTestAgent() - assert agent.name == "dummy_test_agent" - assert not agent._is_initialized - - # 1. Initialize - mock_ctx = object() - await agent.initialize(mock_ctx) - assert agent._is_initialized - assert agent.ctx is mock_ctx - - # 2. Process - result = await agent.process({}) - assert result.get("status") == "success" - assert result.get("processed_state") is True - - # 3. Cleanup - await agent.cleanup() - assert not agent._is_initialized - - -def test_agent_registry(): - """Verify AgentRegistry registers and instantiates agents correctly.""" - registered = AgentRegistry.list_agents() - assert "planner" in registered - assert "summarizer" in registered - assert "operator" in registered - assert "failure_analyzer" in registered - - # Instantiate PlannerAgent via registry - planner = AgentRegistry.create_agent("planner", ctx=None) - assert planner.name == "planner" - assert planner.role == "Strategic Task Planner" - - # Instantiate SummarizerAgent via registry - summarizer = AgentRegistry.create_agent("summarizer", ctx=None) - assert summarizer.name == "summarizer" - assert summarizer.role == "Execution History Summarizer" - - # Instantiate OperatorAgent via registry - operator = AgentRegistry.create_agent("operator", ctx=None) - assert operator.name == "operator" - assert operator.role == "Mobile UI Interaction Operator" - - # Instantiate FailureAnalyzerAgent via registry - fa = AgentRegistry.create_agent("failure_analyzer", ctx=None) - assert fa.name == "failure_analyzer" - assert fa.role == "Autonomous Self-Healing Failure Analyzer" - - -def test_agent_config(): - """Verify AgentConfig parameters can be customized.""" - config = AgentConfig( - model_name="gemini-2.5-pro", - temperature=0.7, - max_tokens=2048, - enabled_tools=["click", "swipe"], - ) - agent = DummyTestAgent(config=config) - assert agent.config.model_name == "gemini-2.5-pro" - assert agent.config.temperature == 0.7 - assert agent.config.enabled_tools == ["click", "swipe"] - - -@pytest.mark.asyncio -async def test_recovery_strategies_and_failure_analyzer(): - """Verify self-healing recovery strategies and FailureAnalyzerAgent execution.""" - driver = MockDeviceDriver(device_id="mock-rec-1") - - # 1. Test dialog dismissal strategy - dialog_ctx = {"category": "ui_obstruction", "error_message": "Unexpected permission popup"} - dialog_res = await RecoveryStrategyRegistry.attempt_recovery(driver, dialog_ctx) - assert dialog_res is not None - assert dialog_res.success - assert dialog_res.action_type.value == "dismiss_dialog" - assert driver.action_history[-1].get("key") == "back" - - # 2. Test navigate back strategy - nav_ctx = {"category": "wrong_screen"} - nav_res = await RecoveryStrategyRegistry.attempt_recovery(driver, nav_ctx) - assert nav_res is not None - assert nav_res.success - assert nav_res.action_type.value == "navigate_back" - - # 3. Test scroll and search strategy - scroll_ctx = {"category": "element_not_found", "scroll_direction": "down"} - scroll_res = await RecoveryStrategyRegistry.attempt_recovery(driver, scroll_ctx) - assert scroll_res is not None - assert scroll_res.success - assert scroll_res.action_type.value == "scroll_and_retry" - - # 4. Test FailureAnalyzerAgent process - class MockContext: - pass - - ctx = MockContext() - setattr(ctx, "_active_driver", driver) - - fa_agent = FailureAnalyzerAgent(ctx=ctx) - fa_result = await fa_agent.process({"failure_context": {"category": "wrong_screen"}}) - assert fa_result["status"] == "recovered" - assert fa_result["recovery_action"] == "navigate_back" diff --git a/tests/unit/test_prompt_builder.py b/tests/unit/test_prompt_builder.py index 8577359..40773ef 100644 --- a/tests/unit/test_prompt_builder.py +++ b/tests/unit/test_prompt_builder.py @@ -12,9 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for OperatorPromptBuilder and LoopDetector.""" +"""Unit tests for OperatorPromptBuilder.""" -from artemis.agents.operator.loop_detector import LoopDetector from artemis.agents.operator.prompt_builder import OperatorPromptBuilder @@ -33,24 +32,3 @@ def test_operator_prompt_builder(): assert "Open Settings" in human_msg assert "Tap Network & internet" in human_msg assert "Current Turn: 2" in human_msg - - -def test_loop_detector(): - """Verify loop detector flags repeated action sequences.""" - detector = LoopDetector(repetition_threshold=3) - assert not detector.is_loop_detected() - - # Record distinct actions - detector.record_action("click", {"target": [100, 200]}) - detector.record_action("swipe", {"action": "up"}) - assert not detector.is_loop_detected() - - # Record 3 identical actions - detector.record_action("click", {"target": [500, 500]}) - detector.record_action("click", {"target": [500, 500]}) - detector.record_action("click", {"target": [500, 500]}) - assert detector.is_loop_detected() - - # Reset - detector.reset() - assert not detector.is_loop_detected() From 9cee112cd2af9667c5bf471402a9510dfb217c30 Mon Sep 17 00:00:00 2001 From: somew1nd Date: Tue, 1 Sep 2026 11:35:23 -0700 Subject: [PATCH 2/2] fix(runtime): unify process and task state handling --- apps/admin_console/core/state.py | 51 +--- .../repositories/session_repository.py | 17 +- .../services/task_queue_service.py | 69 +++-- artemis/runtime/__init__.py | 2 + artemis/runtime/awake_lease.py | 12 +- artemis/runtime/device_lock.py | 175 +++++++++--- artemis/runtime/process_probe.py | 98 +++++++ artemis/runtime/server_lifecycle.py | 39 +-- mcp_server/tools/task_manager.py | 14 +- mcp_server/utils/trace_store.py | 268 +++++++++++++++--- .../admin_console/test_task_queue_service.py | 111 +++++++- tests/unit/mcp/test_trace_store.py | 122 ++++++++ tests/unit/runtime/test_device_lock.py | 72 +++++ tests/unit/runtime/test_process_probe.py | 65 +++++ 14 files changed, 918 insertions(+), 197 deletions(-) create mode 100644 artemis/runtime/process_probe.py create mode 100644 tests/unit/runtime/test_process_probe.py 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