feat: unify history tools and improve execution reliability

This commit is contained in:
somew1nd committed 2026-09-03 16:05:47 -07:00
1 parent d23a4f6472
commit 2fea9865ee
116 files changed
+8291 -2503

No files matched your search

+191
View File
@@ -438,3 +438,194 @@ def test_record_step_without_app_data_stamps_nothing(tmp_path):
step = engine.storage.get_steps(engine.current_session_id)[0]
assert "foreground_app" not in (step.extra_metadata or {})
# --- Agent-friendly single-step / range lookups (step replay backing) ----------------
def _replay_engine(tmp_path):
mock_ctx = MagicMock(spec=ArtemisContext)
mock_execution_setup = MagicMock()
mock_execution_setup.traces_path = str(tmp_path)
mock_ctx.execution_setup = mock_execution_setup
mock_ctx.device = None
engine = DataEngine(mock_ctx)
engine.start_session("replay lookup session")
return engine
def _join_pending(engine):
for t in list(engine._pending_threads):
t.join()
def test_friendly_step_range_expands_tool_traces_and_matches_full_listing(tmp_path):
engine = _replay_engine(tmp_path)
engine.record_step(
summary="Asked the explorer about the toggle.",
action_taken={"action": "click", "target": [540, 1200], "target_text": "Save"},
last_execution_result={"status": "success"},
)
engine.record_trace(
type="tool",
name="ask_explorer",
payload={"args": {"question": "toggle?"}, "result": "The toggle is ON"},
step_id=engine.last_recorded_step_id,
)
engine.record_step(summary="Second step.", action_taken={"action": "swipe"})
engine.record_step(summary="Third step.", action_taken={"action": "press_key", "key": "BACK"})
_join_pending(engine)
rows = engine.get_agent_friendly_steps_in_range(1, 2)
assert [r["step_number"] for r in rows] == [1, 2]
calls = rows[0]["tool_calls"]
assert calls and calls[0]["name"] == "ask_explorer"
assert calls[0]["args"] == {"question": "toggle?"}
assert calls[0]["result"] == "The toggle is ON"
assert any(e["type"] == "tool_call" for e in rows[0]["interleaved_events"])
# The single-step path renders exactly what the whole-session listing does
# (same trace expansion, same coordinate normalization).
full = engine.get_agent_friendly_steps()
assert rows[0]["action_taken"] == full[0]["action_taken"]
assert rows[0]["interleaved_events"] == full[0]["interleaved_events"]
# Reversed bounds are tolerated; a single step defaults end to start.
assert [r["step_number"] for r in engine.get_agent_friendly_steps_in_range(3, 2)] == [2, 3]
assert [r["step_number"] for r in engine.get_agent_friendly_steps_in_range(3)] == [3]
one = engine.get_agent_friendly_step(3)
assert one is not None and one["summary"] == "Third step."
assert engine.get_agent_friendly_step(99) is None
def test_friendly_step_lookups_without_session_are_empty(tmp_path):
mock_ctx = MagicMock(spec=ArtemisContext)
mock_execution_setup = MagicMock()
mock_execution_setup.traces_path = str(tmp_path)
mock_ctx.execution_setup = mock_execution_setup
mock_ctx.device = None
engine = DataEngine(mock_ctx)
assert engine.get_agent_friendly_steps_in_range(1, 5) == []
assert engine.get_agent_friendly_step(1) is None
def test_tool_result_images_are_described_as_step_screenshots(tmp_path):
"""An image a tool result embedded (stored as an <ImageRef: sha256=…>
placeholder) is rewritten into the description of the step screenshot it
is, never dropped silently; unknown images get the generic label."""
import hashlib
import io
from PIL import Image
from artemis.data_engine.engine import GENERIC_IMAGE_LABEL
buf = io.BytesIO()
Image.new("RGB", (32, 64), "red").save(buf, format="JPEG")
pre_bytes = buf.getvalue()
pre_hash = hashlib.sha256(pre_bytes).hexdigest()
engine = _replay_engine(tmp_path)
engine.record_step(
pre_screenshot_bytes=pre_bytes,
summary="Opened the alarms list.",
action_taken={"action": "click"},
ocr_result=[{"text": "Alarm 07:30 needleocr", "bounds": [0, 0, 10, 10]}],
)
engine.record_step(summary="Recalled the earlier screen.", action_taken={"action": "click"})
engine.record_trace(
type="tool",
name="get_step_screenshot",
payload={
"args": {"step_number": 1, "which": "pre"},
"result": [
{"type": "text", "text": "Screenshot of step 1 (pre-action) is attached."},
{
"type": "image_url",
"image_url": {"url": f"<ImageRef: sha256={pre_hash} length=99999>"},
},
{
"type": "image_url",
"image_url": {"url": "<ImageRef: sha256=" + "0" * 64 + " length=5>"},
},
],
},
step_id=engine.last_recorded_step_id,
)
_join_pending(engine)
step2 = engine.get_agent_friendly_step(2)
result = step2["tool_calls"][0]["result"]
assert all(block["type"] == "text" for block in result)
assert result[1]["text"] == "[screenshot: pre-action of Step 1]"
assert result[1]["image_name"] == pre_hash
assert result[2]["text"] == GENERIC_IMAGE_LABEL
# Whole-session listing describes the same way.
assert engine.get_agent_friendly_steps()[1]["tool_calls"][0]["result"] == result
# --- Coordinate space of agent-friendly steps -------------------------
def test_flash_normalized_record_survives_agent_friendly_steps_unchanged(tmp_path):
"""Flash records the model's 0–1000 target with an explicit space marker;
the agent-friendly view must not normalize it a second time."""
from artemis.utils.coordinates import COORDINATE_SPACE_KEY, COORDINATE_SPACE_NORMALIZED
engine = _replay_engine(tmp_path)
engine.record_step(
action_taken={
"action": "click",
"coordinates": [320, 399],
"normalized_coordinates": [320, 399],
COORDINATE_SPACE_KEY: COORDINATE_SPACE_NORMALIZED,
"args": {"target": [320, 399]},
},
last_execution_result={"status": "success", "result": "click executed"},
)
_join_pending(engine)
row = engine.get_agent_friendly_steps()[0]
assert row["action_taken"]["coordinates"] == [320, 399]
assert row["action_taken"][COORDINATE_SPACE_KEY] == COORDINATE_SPACE_NORMALIZED
assert engine.get_agent_friendly_step(1)["action_taken"]["coordinates"] == [320, 399]
def test_pro_pixel_record_is_normalized_exactly_once(tmp_path):
"""Pro records physical pixels (with the frame size in extra_metadata);
the friendly view converts them once, and the result is stamped so a
downstream pass (MCP inspector, replay) leaves it alone."""
from artemis.utils.coordinates import (
COORDINATE_SPACE_KEY,
COORDINATE_SPACE_NORMALIZED,
COORDINATE_SPACE_PIXEL,
normalize_any_structure,
)
engine = _replay_engine(tmp_path)
action = {
"action": "tap",
"coordinates": [540, 1200],
"normalized_coordinates": [500, 500],
COORDINATE_SPACE_KEY: COORDINATE_SPACE_PIXEL,
"target_text": "Battery",
}
engine.record_step(
action_taken=[action],
last_execution_result={
"status": "success",
"execution": [dict(action, attempts=["Success"])],
},
extra_metadata={"width": 1080, "height": 2400},
)
_join_pending(engine)
row = engine.get_agent_friendly_steps()[0]
friendly = row["action_taken"][0]
assert friendly["coordinates"] == [500, 500]
assert friendly[COORDINATE_SPACE_KEY] == COORDINATE_SPACE_NORMALIZED
assert row["last_execution_result"]["execution"][0]["coordinates"] == [500, 500]
again = normalize_any_structure(row["action_taken"], 1080, 2400)
assert again == row["action_taken"]
@@ -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.
"""OfflineHistoryReader: the MCP inspector's read-only twin of the live engine.
The offline reader must produce exactly the records the live ``DataEngine``
produces (same ``friendly_step``), open the database read-only, and satisfy
the ``HistoryReader`` protocol the history tools are written against.
"""
import io
import sqlite3
from unittest.mock import MagicMock
import pytest
from PIL import Image
from artemis.context import ArtemisContext
from artemis.data_engine.engine import DataEngine
from artemis.data_engine.history_reader import HistoryReader, OfflineHistoryReader
def _jpeg(color: str) -> bytes:
buf = io.BytesIO()
Image.new("RGB", (64, 128), color).save(buf, format="JPEG")
return buf.getvalue()
def _engine(tmp_path):
ctx = MagicMock(spec=ArtemisContext)
ctx.execution_setup = MagicMock()
ctx.execution_setup.traces_path = str(tmp_path)
ctx.device = None
engine = DataEngine(ctx)
engine.start_session("offline reader test")
return engine
def _flush(engine):
for t in list(engine._pending_threads):
t.join()
@pytest.fixture
def recorded(tmp_path):
engine = _engine(tmp_path)
engine.record_step(
pre_screenshot_bytes=_jpeg("red"),
post_screenshot_bytes=_jpeg("blue"),
summary="Tapped Save after asking the explorer.",
action_taken={"action": "click", "target": [540, 1200], "target_text": "Save"},
last_execution_result={"status": "success"},
operator_raw_thinking="Confirm the toggle first.",
extra_metadata={"width": 1080, "height": 2400},
)
engine.record_trace(
type="tool",
name="ask_explorer",
payload={"args": {"question": "toggle?"}, "result": "The toggle is ON."},
step_id=engine.last_recorded_step_id,
)
engine.record_step(
summary="Went back.",
action_taken={"action": "press_key", "key": "back"},
last_execution_result={"status": "success"},
)
engine.record_history_chunk(
start_step_number=1,
end_step_number=2,
version=1,
status="ready",
band2=" - Steps 1–2: saved the toggle",
band3="- Step 1 (T+00:00): click -> executed",
rendered_text="[Chunk 1 | Steps 1–2]",
)
_flush(engine)
return engine
def _reader(engine) -> OfflineHistoryReader:
return OfflineHistoryReader(engine.db_path, engine.global_base_dir, engine.current_session_id)
def test_offline_reader_matches_live_friendly_steps(recorded):
reader = _reader(recorded)
assert isinstance(reader, HistoryReader)
assert reader.session_start_time == recorded.session_start_time
assert reader.base_dir == recorded.base_dir
live = recorded.get_agent_friendly_steps()
offline = reader.get_agent_friendly_steps()
assert offline == live
assert offline[0]["tool_calls"][0]["name"] == "ask_explorer"
# Coordinates normalized exactly like the live engine.
assert offline[0]["action_taken"]["target"] == live[0]["action_taken"]["target"]
assert offline[0]["relative_time"] == live[0]["relative_time"]
assert reader.get_agent_friendly_steps_in_range(2, 1) == live
assert reader.get_agent_friendly_step(2) == live[1]
assert reader.get_agent_friendly_step(9) is None
def test_offline_reader_step_records_images_and_chunks(recorded):
reader = _reader(recorded)
record = reader.get_step_record(1)
assert record is not None and record.step_number == 1
assert reader.get_step_record(7) is None
pre = reader.get_step_image_path(1, "pre")
post = reader.get_step_image_path(1, "post")
assert pre is not None and pre.exists()
assert post is not None and post.exists()
assert pre == recorded.get_image_path(record.pre_image_name)
assert reader.get_step_image_path(2, "pre") is None
chunks = reader.get_history_chunks()
assert len(chunks) == 1
assert chunks[0].start_step_number == 1
def test_offline_reader_is_read_only(recorded):
reader = _reader(recorded)
with reader.storage._get_connection() as conn:
with pytest.raises(sqlite3.OperationalError):
conn.execute("DELETE FROM steps")
# No directories or tables are created for a missing database.
with pytest.raises(FileNotFoundError):
OfflineHistoryReader(recorded.global_base_dir / "nope.db", recorded.global_base_dir, "x")