Files
Artemis-Eapil/tests/unit/agents/test_flash_runner_ledger.py
T
somew1nd 67248c887e feat: improve task execution, verification and packaging
Add execution incidents, action bursts and shared Flash history.

Expose verification and Explorer settings with streamed check results.

Support graceful cancellation, recording recovery and bundled UI resources.

Add CI quality checks and clean up comments and documentation.
2026-09-02 21:45:51 -07:00

334 lines
14 KiB
Python

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""FlashRunner over the session transcript ledger (Pro-identical history).
Covers: the ``video_analyzer`` availability gate (same as the Pro graph) and
its prompt segment, the unbounded-by-default turn limit, the Pro-shaped
observation tail (``# CURRENT OBSERVATION [T+mm:ss]`` header, screenshot,
``--- Visible UI Elements ---`` list), and the end-to-end loop: committed
turns carry text-only tool messages plus an ``--- Action Execution Result
(T+mm:ss) ---`` message, every recorded step of a multi-action turn is
registered on the committed turn, the final-turn tool restriction only
applies to bounded loops, and a no-tool-call turn is nudged, not terminated.
"""
import base64
import re
from unittest.mock import AsyncMock, Mock, patch
import pytest
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
from artemis.agents.flash.runner import FlashRunner
from artemis.agents.validator.tool_declarations import ToolExecutionResult
from artemis.context import ArtemisContext
from artemis.graph.state import State
from artemis.memory.transcript import (
EXECUTION_RESULT_MARKER,
PRO_UI_LIST_MARKER,
TranscriptLedger,
)
OBSERVATION_HEADER_RE = re.compile(r"^# CURRENT OBSERVATION \[T\+\d{2,}:\d{2}\]$")
RESULT_RE = re.compile(
rf"^{re.escape(EXECUTION_RESULT_MARKER)} \(T\+\d{{2,}}:\d{{2}}\) ---\nStatus: success$"
)
@pytest.fixture
def mock_context():
ctx = Mock(spec=ArtemisContext)
ctx.llm_config = Mock()
mock_llm_cfg = Mock()
mock_llm_cfg.model = "gemini-2.5-flash"
mock_llm_cfg.temperature = 0.1
ctx.llm_config.get_agent.return_value = mock_llm_cfg
ctx.device = Mock()
ctx.device.device_width = 1080
ctx.device.device_height = 2400
ctx.data_engine = None
ctx.adb_client = None
ctx.driver = Mock()
return ctx
# ---------------------------------------------------------------------------
# video_analyzer availability (same gate as the Pro graph)
# ---------------------------------------------------------------------------
def test_video_analyzer_bound_only_when_recording_tools_enabled(mock_context):
with patch("artemis.controllers.unified_controller.get_driver"):
mock_context.execution_setup = Mock()
mock_context.execution_setup.video_recording_tools_enabled = True
names = [t.name for t in FlashRunner(mock_context, goal="g")._get_tools()]
assert "video_analyzer" in names
assert names[-1] == "report_task_status"
mock_context.execution_setup.video_recording_tools_enabled = False
names = [t.name for t in FlashRunner(mock_context, goal="g")._get_tools()]
assert "video_analyzer" not in names
# A Mock attribute (truthy but not True) must not enable the tool.
mock_context.execution_setup = Mock()
names = [t.name for t in FlashRunner(mock_context, goal="g")._get_tools()]
assert "video_analyzer" not in names
def test_recall_history_bound_only_with_a_data_engine_session(mock_context):
with patch("artemis.controllers.unified_controller.get_driver"):
# No DataEngine: nothing to search, the tool stays out (as in Pro).
names = [t.name for t in FlashRunner(mock_context, goal="g")._get_tools()]
assert "recall_history" not in names
mock_context.data_engine = Mock()
with patch("artemis.tools.history_recall._recall_config", return_value=None):
runner = FlashRunner(mock_context, goal="g")
tools = runner._get_tools()
names = [t.name for t in tools]
assert "recall_history" in names
declaration = next(t for t in tools if t.name == "recall_history")
assert set(declaration.parameters["properties"]) == {
"query",
"step_range",
"include_details",
"include_images",
"max_results",
}
assert declaration.parameters["required"] == ["query"]
assert "`recall_history`" in runner._render_system_prompt(tools)
def test_video_analyzer_prompt_segment_follows_availability(mock_context):
with patch("artemis.controllers.unified_controller.get_driver"):
runner = FlashRunner(mock_context, goal="g")
mock_context.execution_setup = Mock()
mock_context.execution_setup.video_recording_tools_enabled = True
with_video = runner._render_system_prompt(runner._get_tools())
assert "`video_analyzer`" in with_video
mock_context.execution_setup.video_recording_tools_enabled = False
without_video = runner._render_system_prompt(runner._get_tools())
assert "video_analyzer" not in without_video
# Session-clock / history teaching is unconditional.
assert "T+mm:ss" in without_video
assert "Action Execution Result" in without_video
# ---------------------------------------------------------------------------
# Turn limit: unbounded by default, explicit caps honoured
# ---------------------------------------------------------------------------
def test_turn_limit_semantics(mock_context):
with patch("artemis.controllers.unified_controller.get_driver"):
assert FlashRunner(mock_context, goal="g", max_turns=0).turn_limit is None
assert FlashRunner(mock_context, goal="g", max_turns=7).turn_limit == 7
with patch("artemis.agents.flash.runner.load_agent_config", side_effect=RuntimeError):
assert FlashRunner(mock_context, goal="g").turn_limit is None
# ---------------------------------------------------------------------------
# Observation tail: Pro shape with the session-relative header
# ---------------------------------------------------------------------------
def test_observation_tail_has_pro_shape(mock_context):
with patch("artemis.controllers.unified_controller.get_driver"):
runner = FlashRunner(mock_context, goal="Open Settings")
ledger = TranscriptLedger()
tail = runner._build_tail(ledger, 1, b"IMG", "[1] Settings")
texts = [b["text"] for b in tail.content if b["type"] == "text"]
assert texts[0] == "Your objective is: Open Settings"
assert OBSERVATION_HEADER_RE.match(texts[1])
assert texts[2] == "--- Current Screenshot ---"
assert texts[3] == f"{PRO_UI_LIST_MARKER}\n[1] Settings"
assert sum(1 for b in tail.content if b["type"] == "image_url") == 1
later = runner._build_tail(
ledger, 5, None, None, injected="[INJECTED] stop", notices=["nudge"], is_final=True
)
texts = [b["text"] for b in later.content]
assert OBSERVATION_HEADER_RE.match(texts[0])
assert "Your objective is" not in "".join(texts)
assert texts[1:] == [
"nudge",
"[INJECTED] stop",
"[WARNING] This is your final turn; only 'report_task_status' is available.",
]
# ---------------------------------------------------------------------------
# End-to-end loop over the ledger (mocked model / executor)
# ---------------------------------------------------------------------------
def _exec_result(tc_id, name, post_bytes, ui="[1] next"):
return ToolExecutionResult(
tool_call_id=tc_id,
tool_name=name,
status="success",
text_summary=f"{name} executed",
screenshot_bytes=post_bytes,
ui_elements_text=ui,
)
def _make_runner(mock_context, responses, max_turns=None):
runner = FlashRunner(mock_context, goal="Open Settings", max_turns=max_turns)
runner.summarizer = None
runner._init_llm = Mock(return_value=Mock())
# Snapshot each invocation: the runner appends to the same message list
# after the call, so the recorded call args would otherwise drift.
runner.calls = []
pending = list(responses)
async def _invoke(llm, tools, messages):
runner.calls.append({"tools": list(tools), "messages": list(messages)})
return pending.pop(0)
runner._invoke_model = AsyncMock(side_effect=_invoke)
runner.executor.execute = AsyncMock(
side_effect=lambda name, args, tc_id, state: _exec_result(
tc_id, name, f"IMG-{tc_id}".encode()
)
)
return runner
def _report(status="completed", tc_id="tc-report", **extra):
return AIMessage(
content="Done.",
tool_calls=[
{"name": "report_task_status", "args": {"status": status, **extra}, "id": tc_id}
],
)
_INITIAL_OBSERVATION = ("shot0.png", b"IMG0", "[1] Settings")
@pytest.mark.asyncio
async def test_run_builds_prompt_from_ledger_with_session_offsets(mock_context):
"""Turn 2 must see: the static system prefix, the committed turn 1 (its
tail, the AI message, text-only tool messages and an execution-result
message with a ``T+`` offset) and a fresh observation tail carrying the
last action's post screenshot — the Pro transcript shape, with no cap."""
responses = [
AIMessage(
content="I will tap then wait.",
tool_calls=[
{"name": "click", "args": {"target": [500, 600]}, "id": "tc1"},
{"name": "wait_for_delay", "args": {"seconds": 1}, "id": "tc2"},
],
),
_report(explanation="ok"),
]
with (
patch("artemis.controllers.unified_controller.get_driver"),
patch(
"artemis.agents.flash.runner.capture_screenshot_and_parse_ui",
AsyncMock(return_value=_INITIAL_OBSERVATION),
),
):
runner = _make_runner(mock_context, responses, max_turns=0)
assert runner.turn_limit is None
result = await runner.run(State(initial_goal="Open Settings"))
assert result == {"status": "completed", "explanation": "ok"}
assert runner._invoke_model.await_count == 2
turn2_messages = runner.calls[1]["messages"]
assert isinstance(turn2_messages[0], SystemMessage)
assert "T+mm:ss" in turn2_messages[0].content
tail1 = turn2_messages[1]
assert isinstance(tail1, HumanMessage)
tail1_texts = [b["text"] for b in tail1.content if b["type"] == "text"]
assert tail1_texts[0] == "Your objective is: Open Settings"
assert OBSERVATION_HEADER_RE.match(tail1_texts[1])
assert isinstance(turn2_messages[2], AIMessage)
tool_msgs = [m for m in turn2_messages if isinstance(m, ToolMessage)]
assert [m.content for m in tool_msgs] == ["click executed", "wait_for_delay executed"]
result_msg = turn2_messages[5]
assert isinstance(result_msg, HumanMessage)
assert RESULT_RE.match(result_msg.content[0]["text"])
tail2 = turn2_messages[6]
tail2_texts = [b["text"] for b in tail2.content if b["type"] == "text"]
assert OBSERVATION_HEADER_RE.match(tail2_texts[0])
assert "Your objective is" not in "".join(tail2_texts)
images = [b for b in tail2.content if b["type"] == "image_url"]
assert len(images) == 1
assert images[0]["image_url"]["url"].endswith(base64.b64encode(b"IMG-tc2").decode())
assert tail2_texts[1] == "--- Current Screenshot ---"
assert tail2_texts[2] == f"{PRO_UI_LIST_MARKER}\n[1] next"
assert len(turn2_messages) == 7
# Both actions of the multi-action turn were registered on the committed turn.
ledger = mock_context.transcript_ledger
assert ledger.unchunked_turns()[0]["step_keys"] == ["tc1", "tc2"]
@pytest.mark.asyncio
async def test_run_final_turn_restricts_tools_and_fails_without_report(mock_context):
responses = [AIMessage(content="I give up.", tool_calls=[])]
with (
patch("artemis.controllers.unified_controller.get_driver"),
patch(
"artemis.agents.flash.runner.capture_screenshot_and_parse_ui",
AsyncMock(return_value=_INITIAL_OBSERVATION),
),
):
runner = _make_runner(mock_context, responses, max_turns=1)
result = await runner.run(State(initial_goal="Open Settings"))
assert result == {"status": "failed", "explanation": "I give up."}
tools = runner.calls[0]["tools"]
assert [t.name for t in tools] == ["report_task_status"]
tail = runner.calls[0]["messages"][-1]
assert any("final turn" in b.get("text", "") for b in tail.content)
@pytest.mark.asyncio
async def test_run_no_tool_call_turn_is_nudged_not_terminated(mock_context):
responses = [AIMessage(content="thinking only", tool_calls=[]), _report()]
with (
patch("artemis.controllers.unified_controller.get_driver"),
patch(
"artemis.agents.flash.runner.capture_screenshot_and_parse_ui",
AsyncMock(return_value=_INITIAL_OBSERVATION),
),
):
runner = _make_runner(mock_context, responses, max_turns=0)
result = await runner.run(State(initial_goal="Open Settings"))
assert result == {"status": "completed"}
turn2_messages = runner.calls[1]["messages"]
# Unbounded loop: no final-turn restriction was ever applied.
tools = runner.calls[1]["tools"]
assert "click" in [t.name for t in tools]
tail2_texts = [b["text"] for b in turn2_messages[-1].content if b["type"] == "text"]
assert any("did not call any tools" in t for t in tail2_texts)
# A helper-only / empty turn commits without an execution-result message.
assert not any(
isinstance(m, HumanMessage)
and any(EXECUTION_RESULT_MARKER in b.get("text", "") for b in m.content)
for m in turn2_messages
)