android-midscene-automation 0.2.4 → 1.1.0

package/artemis-main/artemis/tools/history/__init__.py

@@ -0,0 +1,398 @@
+# 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.
+
+"""Shared history search, step replay, and screenshot tools.
+
+Each tool's args_schema defines both its LangChain interface for Pro and its
+ToolDeclaration for Flash. The MCP trace inspector calls the underlying
+functions with an OfflineHistoryReader.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from langchain_core.tools import BaseTool
+from pydantic import BaseModel, Field
+
+from artemis.core.tool_failure import ToolFailure
+from artemis.core.tool_declaration import ToolDeclaration
+from artemis.data_engine.trace import trace_langchain_tool
+from artemis.tools.base import ArtemisTool
+from artemis.tools.history.replay import (
+    DEFAULT_MAX_STEPS as DEFAULT_REPLAY_MAX_STEPS,
+)
+from artemis.tools.history.replay import (
+    DEFAULT_MAX_TOKENS as DEFAULT_REPLAY_MAX_TOKENS,
+)
+from artemis.tools.history.replay import replay_steps_text
+from artemis.tools.history.screenshot import ScreenshotResult, load_step_screenshot
+from artemis.tools.history.search import search_history_text
+from artemis.tools.tool_wrapper import ToolWrapper
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+__all__ = [
+    "HISTORY_TOOLS",
+    "HISTORY_TOOL_DECLARATIONS",
+    "HISTORY_TOOL_NAMES",
+    "HISTORY_TOOL_WRAPPERS",
+    "GET_STEP_SCREENSHOT_TOOL",
+    "REPLAY_STEPS_TOOL",
+    "SEARCH_HISTORY_TOOL",
+    "GetStepScreenshotArgs",
+    "GetStepScreenshotTool",
+    "ReplayStepsArgs",
+    "ReplayStepsTool",
+    "ScreenshotResult",
+    "SearchHistoryArgs",
+    "SearchHistoryTool",
+    "get_history_tools",
+    "get_replay_steps_tool",
+    "get_search_history_tool",
+    "get_step_screenshot",
+    "get_step_screenshot_tool",
+    "history_available",
+    "history_tool_by_name",
+    "history_tool_declarations",
+    "load_step_screenshot",
+    "replay_steps",
+    "replay_steps_text",
+    "search_history",
+    "search_history_available",
+    "search_history_text",
+]
+
+
+# --- Configuration -----------------------------------------------------------------------
+
+
+def _memory_config(warn: bool = True) -> Any:
+    try:
+        from artemis.config import load_agent_config
+
+        return load_agent_config().memory
+    except Exception as e:
+        if warn:
+            logger.debug(f"memory config unavailable, using defaults: {e}")
+        return None
+
+
+def _recall_config(warn: bool = True) -> Any:
+    return getattr(_memory_config(warn), "recall", None)
+
+
+def _replay_config(warn: bool = True) -> Any:
+    return getattr(_memory_config(warn), "replay", None)
+
+
+def history_available(ctx: Any) -> bool:
+    """History tools need a stored history to read (a DataEngine session)."""
+    return ctx is not None and getattr(ctx, "data_engine", None) is not None
+
+
+def search_history_available(ctx: Any) -> bool:
+    """``search_history`` is additionally gated by ``agent.memory.recall.enabled``."""
+    if not history_available(ctx):
+        return False
+    return bool(getattr(_recall_config(warn=False), "enabled", True))
+
+
+# --- search_history ----------------------------------------------------------------------
+
+
+class SearchHistoryArgs(BaseModel):
+    """Arguments schema for the search_history tool."""
+
+    query: str = Field(
+        default="",
+        description=(
+            "Keywords to search for across the whole execution history"
+            " (screen descriptions, actions, results, reasoning, tool calls,"
+            " notes, on-screen text, package/activity names). Case-insensitive;"
+            " every whitespace-separated term is matched independently. May be"
+            " empty when step_range is given."
+        ),
+    )
+    step_range: list[int] | None = Field(
+        default=None,
+        description=(
+            "Optional [start, end] step-number range (inclusive) to restrict the"
+            " search — e.g. the range shown in a compressed-history block. Also"
+            " returns the full per-step action ledger of the range."
+        ),
+    )
+    max_results: int = Field(
+        default=5,
+        ge=1,
+        description="Maximum results to return (server-side cap still applies).",
+    )
+
+
+SEARCH_HISTORY_DESCRIPTION = (
+    "Deterministic keyword / step-range lookup over the full stored execution"
+    " history: step summaries, exact actions and results, reasoning, tool"
+    " calls, notes, on-screen text (OCR/UI tree, incl. package/activity names)"
+    " and compressed-history ledgers. Every hit carries its step number; use"
+    " `replay_steps` for the full record of a step and `get_step_screenshot`"
+    " for its image."
+)
+
+
+class SearchHistoryTool(ArtemisTool):
+    """Deterministic search over the stored execution history."""
+
+    def __init__(self):
+        super().__init__(
+            name="search_history",
+            description=SEARCH_HISTORY_DESCRIPTION,
+            args_schema=SearchHistoryArgs,
+            category="memory",
+        )
+
+    def is_available(self, ctx: Any = None) -> bool:
+        return search_history_available(ctx)
+
+    # pylint: disable=too-many-arguments
+    async def execute(
+        self,
+        driver: Any = None,  # pylint: disable=unused-argument
+        ctx: Any = None,
+        query: str = "",
+        step_range: list[int] | None = None,
+        max_results: int = 5,
+        **kwargs: Any,
+    ) -> str:
+        reader = getattr(ctx, "data_engine", None) if ctx else None
+        if reader is None:
+            return ToolFailure("search_history unavailable: no active execution history.")
+        if not query and not step_range:
+            return ToolFailure(
+                "search_history needs a query and/or a step_range — e.g."
+                ' search_history(query="login timeout") or'
+                ' search_history(query="", step_range=[1, 40]).'
+            )
+        try:
+            return search_history_text(
+                reader,
+                query=query or "",
+                step_range=step_range,
+                max_results=int(max_results or 5),
+                recall_config=_recall_config(),
+            )
+        except Exception as e:
+            logger.error(f"search_history failed: {e}")
+            return ToolFailure(f"search_history failed: {e}")
+
+
+# --- replay_steps ------------------------------------------------------------------------
+
+
+class ReplayStepsArgs(BaseModel):
+    """Arguments schema for the replay_steps tool."""
+
+    start_step: int = Field(..., description="First step to replay (1-based step number).")
+    end_step: int | None = Field(
+        default=None,
+        description=(
+            "Last step of the inclusive range to replay; omit to replay the single"
+            " start_step. Reversed bounds are swapped."
+        ),
+    )
+
+
+REPLAY_STEPS_DESCRIPTION = (
+    "Replays recorded steps exactly as the executing agent's own context showed"
+    " them: what the screen showed, its reasoning, every tool call (name,"
+    " arguments, result), the planned action, any pre-execution interception"
+    " and the execution result. Coordinates normalized [x, y]. Up to"
+    f" {DEFAULT_REPLAY_MAX_STEPS} steps per call. Screenshots are not included"
+    " — fetch them with `get_step_screenshot`."
+)
+
+
+class ReplayStepsTool(ArtemisTool):
+    """Full replay of a range of recorded steps."""
+
+    def __init__(self):
+        super().__init__(
+            name="replay_steps",
+            description=REPLAY_STEPS_DESCRIPTION,
+            args_schema=ReplayStepsArgs,
+            category="memory",
+        )
+
+    def is_available(self, ctx: Any = None) -> bool:
+        return history_available(ctx)
+
+    async def execute(
+        self,
+        driver: Any = None,  # pylint: disable=unused-argument
+        ctx: Any = None,
+        start_step: Any = None,
+        end_step: Any = None,
+        **kwargs: Any,
+    ) -> str:
+        reader = getattr(ctx, "data_engine", None) if ctx else None
+        if reader is None:
+            return ToolFailure("Error: no execution history available.")
+        cfg = _replay_config()
+        try:
+            return replay_steps_text(
+                reader,
+                start_step,
+                end_step,
+                max_steps=int(getattr(cfg, "max_steps", None) or DEFAULT_REPLAY_MAX_STEPS),
+                max_tokens=int(getattr(cfg, "max_tokens", None) or DEFAULT_REPLAY_MAX_TOKENS),
+            )
+        except Exception as e:
+            logger.error(f"replay_steps failed: {e}")
+            return ToolFailure(f"replay_steps failed: {e}")
+
+
+# --- get_step_screenshot -----------------------------------------------------------------
+
+
+class GetStepScreenshotArgs(BaseModel):
+    """Arguments schema for the get_step_screenshot tool."""
+
+    step_number: int = Field(..., description="The step whose screenshot to attach (1-based).")
+    which: Literal["pre", "post", "overlay"] = Field(
+        default="pre",
+        description=(
+            "'pre' = screen observed at the start of the step; 'post' = screen"
+            " after the action (when recorded); 'overlay' = the pre screenshot"
+            " with the step's action drawn on it."
+        ),
+    )
+
+
+GET_STEP_SCREENSHOT_DESCRIPTION = (
+    "Attaches one recorded screenshot of a step: 'pre' (screen observed at the"
+    " start of the step), 'post' (screen after the action, when recorded) or"
+    " 'overlay' (the pre screenshot with the step's action drawn on it — tap"
+    " point / swipe path — to check the action landed on the intended element)."
+)
+
+
+class GetStepScreenshotTool(ArtemisTool):
+    """One recorded step screenshot, attached to the conversation."""
+
+    def __init__(self):
+        super().__init__(
+            name="get_step_screenshot",
+            description=GET_STEP_SCREENSHOT_DESCRIPTION,
+            args_schema=GetStepScreenshotArgs,
+            category="memory",
+        )
+
+    def is_available(self, ctx: Any = None) -> bool:
+        return history_available(ctx)
+
+    async def execute(
+        self,
+        driver: Any = None,  # pylint: disable=unused-argument
+        ctx: Any = None,
+        step_number: Any = None,
+        which: str = "pre",
+        **kwargs: Any,
+    ) -> list[dict[str, Any]] | str:
+        reader = getattr(ctx, "data_engine", None) if ctx else None
+        if reader is None:
+            return ToolFailure("Error: no execution history available.")
+        try:
+            return load_step_screenshot(reader, step_number, which).to_content_blocks()
+        except Exception as e:
+            logger.error(f"get_step_screenshot failed: {e}")
+            return ToolFailure(f"get_step_screenshot failed: {e}")
+
+
+# --- Instances, declarations, exports ------------------------------------------------------
+
+search_history = SearchHistoryTool()
+replay_steps = ReplayStepsTool()
+get_step_screenshot = GetStepScreenshotTool()
+
+HISTORY_TOOLS: tuple[ArtemisTool, ...] = (search_history, replay_steps, get_step_screenshot)
+HISTORY_TOOL_NAMES: frozenset[str] = frozenset(t.name for t in HISTORY_TOOLS)
+
+#: Flash-profile declarations, derived from the very same args schemas.
+SEARCH_HISTORY_TOOL: ToolDeclaration = search_history.to_tool_declaration()
+REPLAY_STEPS_TOOL: ToolDeclaration = replay_steps.to_tool_declaration()
+GET_STEP_SCREENSHOT_TOOL: ToolDeclaration = get_step_screenshot.to_tool_declaration()
+HISTORY_TOOL_DECLARATIONS: tuple[ToolDeclaration, ...] = (
+    SEARCH_HISTORY_TOOL,
+    REPLAY_STEPS_TOOL,
+    GET_STEP_SCREENSHOT_TOOL,
+)
+
+
+def history_tool_by_name(name: str) -> ArtemisTool | None:
+    """The shared tool instance behind a history tool name (Flash dispatch)."""
+    raw = name.split(":")[-1] if ":" in name else name
+    return next((t for t in HISTORY_TOOLS if t.name == raw), None)
+
+
+def history_tool_declarations(ctx: Any) -> list[ToolDeclaration]:
+    """Declarations of the history tools available for ``ctx`` (Flash runner)."""
+    return [tool.to_tool_declaration() for tool in HISTORY_TOOLS if tool.is_available(ctx)]
+
+
+def get_search_history_tool(ctx: Any) -> BaseTool:
+    """Exports search_history as a traced LangChain BaseTool."""
+    return trace_langchain_tool(search_history.to_langchain_tool(ctx), ctx)
+
+
+def get_replay_steps_tool(ctx: Any) -> BaseTool:
+    """Exports replay_steps as a traced LangChain BaseTool."""
+    return trace_langchain_tool(replay_steps.to_langchain_tool(ctx), ctx)
+
+
+def get_step_screenshot_tool(ctx: Any) -> BaseTool:
+    """Exports get_step_screenshot as a traced LangChain BaseTool."""
+    return trace_langchain_tool(get_step_screenshot.to_langchain_tool(ctx), ctx)
+
+
+def get_history_tools(ctx: Any) -> list[BaseTool]:
+    """All three history tools as LangChain tools (read-only agents mount them as a set)."""
+    return [get_search_history_tool(ctx), get_replay_steps_tool(ctx), get_step_screenshot_tool(ctx)]
+
+
+search_history_wrapper = ToolWrapper(
+    tool_fn_getter=get_search_history_tool,
+    on_success_fn=lambda *a, **k: "Searched history",
+    on_failure_fn=lambda err: f"search_history failed: {err}",
+    is_available_fn=search_history_available,
+)
+
+replay_steps_wrapper = ToolWrapper(
+    tool_fn_getter=get_replay_steps_tool,
+    on_success_fn=lambda *a, **k: "Replayed steps",
+    on_failure_fn=lambda err: f"replay_steps failed: {err}",
+    is_available_fn=history_available,
+)
+
+get_step_screenshot_wrapper = ToolWrapper(
+    tool_fn_getter=get_step_screenshot_tool,
+    on_success_fn=lambda *a, **k: "Fetched step screenshot",
+    on_failure_fn=lambda err: f"get_step_screenshot failed: {err}",
+    is_available_fn=history_available,
+)
+
+HISTORY_TOOL_WRAPPERS: tuple[ToolWrapper, ...] = (
+    search_history_wrapper,
+    replay_steps_wrapper,
+    get_step_screenshot_wrapper,
+)

package/artemis-main/artemis/tools/history/replay.py

@@ -0,0 +1,120 @@
+# 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.
+
+"""replay_steps: full replay of a step range, exactly as the executing agent saw it.
+
+Each step is rendered with :func:`~artemis.utils.task_tree.render_step_replay`
+(``T+mm:ss`` step header, ``[Screen]`` description, reasoning, every
+non-internal tool call with name / arguments / result, the executed action,
+safety-net interception and the execution result). Two loose bounds keep one tool result finite (config
+``agent.memory.replay``): at most ``max_steps`` steps per call, and a
+``max_tokens`` budget applied by dropping *whole trailing steps* — a replayed
+step is never cut mid-sentence. Screenshots are deliberately not part of the
+replay; ``get_step_screenshot`` fetches them one at a time.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from artemis.core.tool_failure import ToolFailure
+from artemis.utils.task_tree import render_step_replay
+
+DEFAULT_MAX_STEPS = 5
+DEFAULT_MAX_TOKENS = 12000
+
+
+def _estimate_tokens(text: str) -> int:
+    return len(text) // 4
+
+
+def replay_steps_text(
+    reader: Any,
+    start_step: Any,
+    end_step: Any = None,
+    *,
+    max_steps: int = DEFAULT_MAX_STEPS,
+    max_tokens: int = DEFAULT_MAX_TOKENS,
+) -> str:
+    """Replays steps ``start_step``..``end_step`` (inclusive, 1-based) from ``reader``."""
+    if reader is None:
+        return ToolFailure("Error: no execution history available.")
+    try:
+        start = int(start_step)
+        end = start if end_step is None else int(end_step)
+    except (TypeError, ValueError):
+        return ToolFailure(
+            f"Error: start_step and end_step must be integers, got {start_step!r} and {end_step!r}."
+        )
+    if start > end:
+        start, end = end, start
+
+    max_steps = max(1, int(max_steps or DEFAULT_MAX_STEPS))
+    notes: list[str] = []
+    if end - start + 1 > max_steps:
+        end = start + max_steps - 1
+        notes.append(
+            f"(Only {max_steps} steps are replayed per call: showing Steps {start}–{end}."
+            f" Call again from Step {end + 1} for the rest.)"
+        )
+
+    try:
+        steps: list[dict] = reader.get_agent_friendly_steps_in_range(start, end) or []
+    except Exception as e:
+        return ToolFailure(f"Error loading steps {start}–{end}: {e}")
+    if not steps:
+        if start == end:
+            return ToolFailure(f"Error: step {start} not found.")
+        return ToolFailure(f"Error: no recorded steps in range {start}–{end}.")
+
+    budget_tokens = max(1, int(max_tokens or DEFAULT_MAX_TOKENS))
+    # Step headers carry the session clock (``T+mm:ss``) every agent prompt
+    # and search_history use; without a session start they fall back to the
+    # stored relative time.
+    session_start = getattr(reader, "session_start_time", None)
+    rendered: list[str] = []
+    shown: list[dict] = []
+    used = 0
+    dropped = 0
+    for step in steps:
+        if dropped:
+            dropped += 1
+            continue
+        try:
+            text = render_step_replay(step, session_start=session_start)
+        except Exception as e:
+            return ToolFailure(f"Error rendering steps {start}–{end}: {e}")
+        cost = _estimate_tokens(text) + 1
+        # The first step is always replayed; afterwards the budget drops whole
+        # trailing steps rather than truncating a step in the middle.
+        if rendered and used + cost > budget_tokens:
+            dropped = 1
+            continue
+        rendered.append(text)
+        shown.append(step)
+        used += cost
+
+    if dropped:
+        last_shown = shown[-1].get("step_number", start)
+        first_shown = shown[0].get("step_number", start)
+        notes.append(
+            f"(Token budget reached: Steps {first_shown}–{last_shown} replayed;"
+            f" {dropped} more step(s) in the range not shown."
+            f" Call again from Step {last_shown + 1}.)"
+        )
+
+    text = "\n\n".join(rendered)
+    if notes:
+        text += "\n\n" + "\n".join(notes)
+    return text

package/artemis-main/artemis/tools/history/screenshot.py

@@ -0,0 +1,139 @@
+# 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.
+
+"""get_step_screenshot: one recorded screenshot of one step, on demand.
+
+``pre`` is the screen observed at the start of the step, ``post`` the screen
+after its action (when recorded), ``overlay`` the pre screenshot with the
+step's action drawn on it (tap point / swipe path) so a reviewer can check
+that the action landed on the intended element. The overlay is drawn from the
+step's *raw* stored action (physical pixels), never from the normalized twin.
+"""
+
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+from pathlib import Path
+import sqlite3
+from typing import Any
+
+from artemis.core.tool_failure import ToolFailure
+from artemis.utils.visualization import overlay_action_on_screenshot
+from third_party.mobile_use.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+WHICH_CHOICES = ("pre", "post", "overlay")
+
+
+@dataclass
+class ScreenshotResult:
+    """One screenshot lookup: a caption plus the JPEG bytes when a file exists."""
+
+    step_number: int
+    which: str
+    description: str
+    image_bytes: bytes | None = None
+    #: True when ``which == "overlay"`` and the action was actually drawn.
+    overlay_drawn: bool = False
+    #: True when the lookup could not serve the request (bad arguments, no
+    #: history, unknown step, unreadable file). A step that simply has no
+    #: screenshot of the requested kind is an answer, not an error.
+    is_error: bool = False
+
+    def to_content_blocks(self) -> list[dict[str, Any]] | str:
+        """The tool's return value: content blocks with the image, the
+        plain-text answer when there is no image to attach, or a
+        :class:`ToolFailure` (status ``error`` downstream) when the lookup
+        failed."""
+        if self.is_error:
+            return ToolFailure(self.description)
+        if not self.image_bytes:
+            return self.description
+        encoded = base64.b64encode(self.image_bytes).decode("utf-8")
+        return [
+            {"type": "text", "text": self.description},
+            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded}"}},
+        ]
+
+
+def _failure(step_number: int, which: str, description: str) -> ScreenshotResult:
+    return ScreenshotResult(step_number, which, description, is_error=True)
+
+
+def load_step_screenshot(reader: Any, step_number: Any, which: str = "pre") -> ScreenshotResult:
+    """Loads (and for ``overlay`` annotates) one step screenshot from ``reader``."""
+    which = str(which or "pre").lower()
+    try:
+        number = int(step_number)
+    except (TypeError, ValueError):
+        return _failure(-1, which, f"Error: step_number must be an integer, got {step_number!r}.")
+    if which not in WHICH_CHOICES:
+        return _failure(
+            number, which, f"Error: 'which' must be 'pre', 'post' or 'overlay', got '{which}'."
+        )
+    if reader is None:
+        return _failure(number, which, "Error: no execution history available.")
+
+    source = "pre" if which == "overlay" else which
+    try:
+        path = reader.get_step_image_path(number, source)
+    except (sqlite3.Error, OSError, ValueError) as e:
+        return _failure(number, which, f"Error looking up step {number}: {e}")
+    if path is None:
+        try:
+            record = reader.get_step_record(number)
+        except (sqlite3.Error, ValueError):
+            record = None
+        if record is None:
+            return _failure(number, which, f"Error: step {number} not found.")
+        return ScreenshotResult(
+            number, which, f"No {source}-action screenshot recorded for step {number}."
+        )
+    try:
+        image_bytes = Path(path).read_bytes()
+    except OSError as e:
+        return _failure(number, which, f"Error reading screenshot: {e}")
+
+    if which != "overlay":
+        return ScreenshotResult(
+            number, which, f"Screenshot of step {number} ({which}-action) is attached.", image_bytes
+        )
+
+    try:
+        record = reader.get_step_record(number)
+    except (sqlite3.Error, ValueError):
+        record = None
+    action = getattr(record, "action_taken", None) if record is not None else None
+    try:
+        annotated = overlay_action_on_screenshot(image_bytes, action)
+    except (TypeError, ValueError, OSError) as e:
+        logger.debug(f"Action overlay for step {number} failed: {e}")
+        annotated = None
+    if annotated is None:
+        return ScreenshotResult(
+            number,
+            which,
+            f"Screenshot of step {number} (pre-action; the step's action could not be"
+            " drawn, plain screenshot) is attached.",
+            image_bytes,
+        )
+    return ScreenshotResult(
+        number,
+        which,
+        f"Screenshot of step {number} (pre-action with the step's action drawn on it) is attached.",
+        annotated,
+        overlay_drawn=True,
+    )