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.
334 lines
14 KiB
Python
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
|
|
)
|