test: add comprehensive unit, integration, and e2e test suites

Introduce complete test coverage for the ARTEMIS platform across multiple layers:
Unit tests: Add tests for agent nodes (Planner, Operator, Validator, Diagnoser, FlashRunner, Explorer), tools, utils (OCR/XML fusion, UI hierarchy), and data engine.
Integration tests: Add live tool verification, reactive runner, IME, and standalone workflow tests.
E2E tests: Add end-to-end scenario tests for media playback and failure recovery flows.
Fixtures & mocks: Add mock inputs (audio, video, screenshot) and test utilities.

tests/unit/data_engine/test_end_to_end_thoughts.py

@@ -0,0 +1,167 @@
+# 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 pathlib import Path
+import shutil
+import tempfile
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+from langchain_core.messages import AIMessage, HumanMessage
+from artemis.context import ArtemisContext
+from artemis.data_engine.engine import DataEngine
+from artemis.data_engine.trace import CURRENT_TRACE_ID, DataEngineCallbackHandler
+from artemis.utils.task_tree import build_plan_and_history
+import pytest
+
+
+@pytest.fixture
+def temp_workspace():
+    tmpdir = tempfile.mkdtemp()
+    yield Path(tmpdir)
+    shutil.rmtree(tmpdir)
+
+
+@pytest.mark.asyncio
+async def test_end_to_end_thoughts_non_duplication(temp_workspace):
+    # 1. Setup mock context and DataEngine
+    mock_ctx = MagicMock(spec=ArtemisContext)
+    mock_execution_setup = MagicMock()
+    mock_execution_setup.traces_path = str(temp_workspace / "traces")
+    mock_ctx.execution_setup = mock_execution_setup
+    mock_ctx.device = None
+    mock_ctx.llm_config = MagicMock()
+
+    engine = DataEngine(mock_ctx)
+    mock_ctx.data_engine = engine
+
+    # Start session
+    engine.start_session("Test initial goal")
+
+    # 2. Allocate step ID
+    step_id = engine.allocate_step_id()
+
+    # 3. Simulate Operator Agent run starts
+    operator_trace_id = uuid4()
+    engine.record_trace(
+        type="agent",
+        name="operator",
+        payload={"args": {}},
+        trace_id=operator_trace_id,
+        step_id=step_id,
+        status="running",
+    )
+
+    # Set context variables for parenting
+    token = CURRENT_TRACE_ID.set(operator_trace_id)
+    try:
+        # 4. Simulate LLM Call callback sequence
+        handler = DataEngineCallbackHandler(mock_ctx)
+        llm_run_id = uuid4()
+
+        # Start model
+        handler.on_chat_model_start(
+            serialized={"name": "ChatGoogleGenerativeAI"},
+            messages=[[HumanMessage(content="What should I do?")]],
+            run_id=llm_run_id,
+            parent_run_id=None,
+        )
+
+        # End model (generating thought + action)
+        class MockGeneration:
+            def __init__(self, message):
+                self.message = message
+
+        class MockLLMResult:
+            def __init__(self, generations):
+                self.generations = generations
+
+        response_message = AIMessage(
+            content=[
+                {
+                    "type": "thinking",
+                    "thinking": "I need to tap the search bar to find books.",
+                },
+                {"type": "text", "text": "Tapping search bar."},
+            ]
+        )
+        mock_response = MockLLMResult(generations=[[MockGeneration(response_message)]])
+
+        handler.on_llm_end(mock_response, run_id=llm_run_id)
+    finally:
+        CURRENT_TRACE_ID.reset(token)
+
+    # 5. Finish Operator Agent trace
+    engine.record_trace(
+        type="agent",
+        name="operator",
+        payload={"result": "Success"},
+        trace_id=operator_trace_id,
+        step_id=step_id,
+        status="success",
+    )
+
+    # 6. Record step
+    engine.record_step(
+        pre_screenshot_bytes=b"pre_screenshot",
+        ui_tree=[],
+        ocr_result=[],
+        action_taken=[{"action": "tap", "coordinates": [100, 200]}],
+        operator_raw_thinking="I need to tap the search bar to find books.",
+        operator_native_thinking="I need to tap the search bar to find books.",
+        last_execution_result={"status": "success"},
+        summary="Tapped search bar",
+    )
+
+    # Flush any pending tasks
+    await engine.shutdown()
+
+    # 7. Retrieve agent-friendly steps
+    steps = engine.get_agent_friendly_steps()
+
+    # Assert step was retrieved
+    assert len(steps) == 1
+    step_data = steps[0]
+
+    # Assert thoughts are present in interleaved events (since we associated step_id now!)
+    events = step_data.get("interleaved_events", [])
+    thought_events = [e for e in events if e.get("type") in ("thought", "native_thought")]
+    assert len(thought_events) > 0
+    assert any(
+        e["content"] == "I need to tap the search bar to find books." for e in thought_events
+    )
+
+    # 8. Build history and check for duplications
+    plan = "- [ ] Search for books"
+    history = build_plan_and_history(
+        plan, steps, "default", last_n_detailed=1, for_failure_analyzer=True
+    )
+
+    print("\n--- Generated History Output ---")
+    print(history)
+    print("--------------------------------")
+
+    # Assert thought is printed
+    assert "I need to tap the search bar to find books." in history
+
+    # Assert thought is NOT printed multiple times
+    # Count occurrences of the thought string
+    count = history.count("I need to tap the search bar to find books.")
+    # The thought might appear in native thinking block and raw thinking block if both were stored.
+    # In our test we set BOTH operator_raw_thinking and operator_native_thinking to the same string,
+    # AND the LLM trace generated both "thinking" block and "text" block.
+    # Let's count how many times it was actually rendered in the Operator Decision Loop.
+    # It should appear at most 2 times (once as Native Thought, once as Thought/Raw Thought).
+    # If there was duplicate appending due to fallback failure, it would appear 4+ times.
+    assert count <= 2

tests/unit/data_engine/test_engine.py

@@ -0,0 +1,185 @@
+# 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 unittest.mock import MagicMock
+from artemis.context import ArtemisContext
+from artemis.data_engine.engine import DataEngine
+
+
+def test_get_or_create_image_updates_missing_data(tmp_path):
+    # Setup mock context
+    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
+
+    # Initialize DataEngine
+    engine = DataEngine(mock_ctx)
+
+    image_bytes = b"fake_image_bytes"
+
+    # 1. Save image without OCR/UI tree
+    image_name1 = engine.get_or_create_image(image_bytes)
+
+    # Verify it was saved with None for OCR/UI tree
+    record1 = engine.storage.get_image(image_name1)
+    assert record1 is not None
+    assert record1.ocr_result is None
+    assert record1.ui_tree is None
+
+    # 2. Save same image with OCR/UI tree
+    fake_ocr = [{"text": "hello", "box": [0, 0, 10, 10]}]
+    fake_ui = [{"resource-id": "button1", "text": "click me"}]
+
+    image_name2 = engine.get_or_create_image(image_bytes, ui_tree=fake_ui, ocr_result=fake_ocr)
+
+    assert image_name1 == image_name2
+
+    # Verify it now has OCR/UI tree
+    record2 = engine.storage.get_image(image_name2)
+    assert record2 is not None
+    assert record2.ocr_result == fake_ocr
+    assert record2.ui_tree == fake_ui
+
+
+def test_create_image_duplicate_handling(tmp_path):
+    from artemis.data_engine.storage import StorageManager
+    from artemis.data_engine.models import ImageRecord
+
+    storage = StorageManager(tmp_path / "data_engine.db", tmp_path)
+    image_record = ImageRecord(
+        image_name="test_image_hash",
+        timestamp=123.456,
+        ocr_result=None,
+        ui_tree=None,
+        extra_metadata={},
+    )
+
+    # First insert
+    storage.create_image(image_record)
+
+    # Second insert with same image_name should not raise IntegrityError due to INSERT OR IGNORE
+    storage.create_image(image_record)
+
+    record = storage.get_image("test_image_hash")
+    assert record is not None
+    assert record.image_name == "test_image_hash"
+
+
+def test_video_recording_persistence(tmp_path):
+    from uuid import uuid4
+    from artemis.data_engine.storage import StorageManager
+    from artemis.data_engine.models import VideoRecordingRecord
+
+    storage = StorageManager(tmp_path / "data_engine.db", tmp_path)
+    video_id = uuid4()
+    session_id = uuid4()
+
+    record = VideoRecordingRecord(
+        video_id=video_id,
+        session_id=session_id,
+        device_id="emulator-5554",
+        start_time=100.0,
+        end_time=None,
+        local_video_path="/tmp/recording.mp4",
+    )
+
+    # Test insert
+    storage.create_video_recording(record)
+
+    persisted = storage.get_video_recording(video_id)
+    assert persisted is not None
+    assert persisted.video_id == video_id
+    assert persisted.session_id == session_id
+    assert persisted.device_id == "emulator-5554"
+    assert persisted.start_time == 100.0
+    assert persisted.end_time is None
+    assert persisted.local_video_path == "/tmp/recording.mp4"
+
+    # Test update
+    record.end_time = 150.0
+    record.local_video_path = "/tmp/recording_final.mp4"
+    storage.update_video_recording(record)
+
+    updated = storage.get_video_recording(video_id)
+    assert updated is not None
+    assert updated.end_time == 150.0
+    assert updated.local_video_path == "/tmp/recording_final.mp4"
+
+
+def test_record_step_suppresses_identical_post_screenshot(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("test_goal")
+    same_image_bytes = b"identical_screen_bytes_123"
+
+    step_id = engine.record_step(
+        pre_screenshot_bytes=same_image_bytes,
+        post_screenshot_bytes=same_image_bytes,
+        summary="Test step with identical screenshot",
+    )
+    for t in list(engine._pending_threads):
+        t.join()
+
+    step_record = engine.storage.get_steps(engine.current_session_id)[0]
+    assert step_record is not None
+    assert step_record.pre_image_name is not None
+    # post_image_name must be None when no visual transition occurred
+    assert step_record.post_image_name is None
+
+
+def test_update_step_execution_result_suppresses_identical_post_screenshot(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("test_goal")
+    image_bytes = b"pre_screen_bytes_456"
+
+    step_id = engine.record_step(
+        pre_screenshot_bytes=image_bytes,
+        summary="Initial step",
+    )
+
+    step_record = engine.storage.get_steps(engine.current_session_id)[0]
+    assert step_record.post_image_name is None
+
+    # Update with post_image_name matching pre_image_name
+    engine.storage.update_step_execution_result(
+        step_id,
+        {"status": "success"},
+        post_image_name=step_record.pre_image_name,
+    )
+
+    updated_step = engine.storage.get_steps(engine.current_session_id)[0]
+    assert updated_step.post_image_name is None
+
+    # Update with distinct post_image_name
+    engine.storage.update_step_execution_result(
+        step_id,
+        {"status": "success"},
+        post_image_name="distinct_new_hash_789",
+    )
+
+    updated_step_distinct = engine.storage.get_steps(engine.current_session_id)[0]
+    assert updated_step_distinct.post_image_name == "distinct_new_hash_789"

tests/unit/data_engine/test_latency_optimizations_deep.py

@@ -0,0 +1,566 @@
+# 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
+import hashlib
+import json
+import logging
+from pathlib import Path
+from unittest.mock import AsyncMock, MagicMock, patch
+
+from langchain_core.messages import HumanMessage
+from artemis.context import ArtemisContext
+from artemis.data_engine.engine import DataEngine
+from artemis.graph.state import State
+from artemis.utils.logger import DataEngineHandler
+import pytest
+
+
+@pytest.fixture
+def mock_artemis_ctx(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
+    mock_ctx.task_plan_snapshot = None
+    return mock_ctx
+
+
+@pytest.mark.asyncio
+async def test_data_engine_handler_queue_buffering_and_background_worker(
+    mock_artemis_ctx,
+):
+    """Verify DataEngineHandler.emit buffers log records via in-memory queue without spawning OS threads per log line."""
+    engine = DataEngine(mock_artemis_ctx)
+    session_id = engine.start_session("Test Log Queue")
+
+    handler = DataEngineHandler()
+
+    with patch("sys.modules") as mock_modules:
+        mock_engine_mod = MagicMock()
+        mock_engine_mod._CURRENT_DATA_ENGINE = engine
+        mock_modules.get.return_value = mock_engine_mod
+
+        # Emit 50 rapid log records
+        for i in range(50):
+            record = logging.LogRecord(
+                name=f"test.logger.{i}",
+                level=logging.INFO,
+                pathname="test.py",
+                lineno=10 + i,
+                msg=f"Async queue test log message {i}",
+                args=(),
+                exc_info=None,
+            )
+            handler.emit(record)
+
+    # Wait for the background worker thread of DataEngineHandler and DataEngine tasks to drain
+    await asyncio.sleep(0.4)
+    if engine._pending_tasks:
+        await asyncio.gather(*engine._pending_tasks)
+
+    # Verify that traces were recorded and persisted accurately via direct SQLite query
+    with engine.storage._get_connection() as conn:
+        cursor = conn.execute(
+            "SELECT * FROM traces WHERE session_id = ? AND type = 'log'",
+            (str(session_id),),
+        )
+        log_rows = cursor.fetchall()
+
+    assert len(log_rows) >= 50
+    messages = {json.loads(row["payload"])["message"] for row in log_rows}
+    for i in range(50):
+        assert f"Async queue test log message {i}" in messages
+
+
+@pytest.mark.asyncio
+async def test_update_step_action_and_metadata_background_offloading(mock_artemis_ctx, tmp_path):
+    """Verify that update_step_* methods publish SSE events instantly in memory and offload SQLite & disk I/O to background."""
+    engine = DataEngine(mock_artemis_ctx)
+    session_id = engine.start_session("Test Step Offloading")
+    step_id = engine.allocate_step_id()
+
+    # Create the initial step
+    engine.record_step(summary="Initial step summary")
+    if engine._pending_tasks:
+        await asyncio.gather(*engine._pending_tasks)
+
+    # Mock _publish to verify immediate real-time SSE delivery
+    published_events = []
+    engine._publish = lambda event_type, payload: published_events.append((event_type, payload))
+
+    fake_action = {"action": "tap", "coordinates": [100, 200]}
+    fake_post_jpg = b"fake_post_screenshot_jpg_bytes_12345"
+
+    # Call update_step_action with post screenshot bytes
+    engine.update_step_action(fake_action, post_screenshot_bytes=fake_post_jpg)
+
+    # Verify SSE was published immediately before/during task execution
+    assert len(published_events) == 1
+    assert published_events[0][0] == "step_updated"
+    assert published_events[0][1]["action_taken"] == fake_action
+
+    # Call other update_step_* methods
+    engine.update_step_summary(step_id, "Updated concise summary")
+    engine.update_step_thinking(step_id, "Raw thinking process...")
+    engine.update_step_native_thinking(step_id, "Native thinking process...")
+    engine.update_step_execution_result(
+        step_id, {"status": "success"}, post_image_name="hash_post_123"
+    )
+
+    # Now await all background pending tasks
+    if engine._pending_tasks:
+        await asyncio.gather(*engine._pending_tasks)
+
+    # Verify SQLite record has been correctly updated by background tasks
+    steps = engine.storage.get_steps(session_id)
+    step_record = next((s for s in steps if s.step_id == step_id), None)
+    assert step_record is not None
+    assert step_record.action_taken == fake_action
+    assert step_record.summary == "Updated concise summary"
+    assert step_record.operator_raw_thinking == "Raw thinking process..."
+    assert step_record.operator_native_thinking == "Native thinking process..."
+    assert step_record.last_execution_result == {"status": "success"}
+    assert step_record.post_image_name == "hash_post_123"
+
+    # Verify post.jpg was written to disk inside step directory
+    post_jpg_path = engine.current_step_dir / "post.jpg"
+    assert post_jpg_path.exists()
+    assert post_jpg_path.read_bytes() == fake_post_jpg
+
+
+@pytest.mark.asyncio
+async def test_record_step_background_image_and_step_persistence(
+    mock_artemis_ctx,
+):
+    """Verify record_step immediately calculates SHA-256 hashes and offloads get_or_create_image + create_step to background."""
+    engine = DataEngine(mock_artemis_ctx)
+    session_id = engine.start_session("Test Record Step Image Offloading")
+    step_id = engine.allocate_step_id()
+
+    pre_img_bytes = b"fake_pre_screenshot_data_001"
+    post_img_bytes = b"fake_post_screenshot_data_002"
+    expected_pre_hash = hashlib.sha256(pre_img_bytes).hexdigest()
+    expected_post_hash = hashlib.sha256(post_img_bytes).hexdigest()
+
+    returned_step_id = engine.record_step(
+        pre_screenshot_bytes=pre_img_bytes,
+        post_screenshot_bytes=post_img_bytes,
+        summary="Testing background image persistence",
+    )
+    assert returned_step_id == step_id
+
+    # Await background storage tasks
+    if engine._pending_tasks:
+        await asyncio.gather(*engine._pending_tasks)
+
+    # Verify image records were created in SQLite
+    pre_record = engine.storage.get_image(expected_pre_hash)
+    post_record = engine.storage.get_image(expected_post_hash)
+    assert pre_record is not None
+    assert pre_record.image_name == expected_pre_hash
+    assert post_record is not None
+    assert post_record.image_name == expected_post_hash
+
+    # Verify image files exist on disk
+    pre_path = engine.get_image_path(expected_pre_hash)
+    post_path = engine.get_image_path(expected_post_hash)
+    assert pre_path.exists() and pre_path.read_bytes() == pre_img_bytes
+    assert post_path.exists() and post_path.read_bytes() == post_img_bytes
+
+    # Verify StepRecord has correct hashes and summary
+    steps = engine.storage.get_steps(session_id)
+    step_record = next((s for s in steps if s.step_id == step_id), None)
+    assert step_record is not None
+    assert step_record.pre_image_name == expected_pre_hash
+    assert step_record.post_image_name == expected_post_hash
+    assert step_record.summary == "Testing background image persistence"
+
+
+@pytest.mark.asyncio
+async def test_perception_node_async_offloading(mock_artemis_ctx, tmp_path):
+    """Verify perception_node offloads injected_instruction reading/unlinking and get_or_create_image to background."""
+    from artemis.graph.perception import perception_node
+
+    engine = DataEngine(mock_artemis_ctx)
+    engine.start_session("Test Perception Node Offloading")
+    mock_artemis_ctx.data_engine = engine
+
+    # Create an injected_instruction.json file
+    instruction_file = Path(engine.base_dir) / "injected_instruction.json"
+    instruction_file.write_text(json.dumps({"instruction": "Avoid clicking Ads"}), encoding="utf-8")
+
+    # Prepare mock controller
+    fake_screenshot_bytes = b"fake_screenshot_for_perception_node_123"
+    fake_screenshot_b64 = base64.b64encode(fake_screenshot_bytes).decode("utf-8")
+    expected_hash = hashlib.sha256(fake_screenshot_bytes).hexdigest()
+
+    mock_screen_data = MagicMock()
+    mock_screen_data.base64 = fake_screenshot_b64
+    mock_screen_data.elements = [{"resource-id": "app_bar", "bounds": "[0,0][1080,200]"}]
+    mock_screen_data.width = 1080
+    mock_screen_data.height = 2400
+
+    with (
+        patch("artemis.graph.perception.UnifiedMobileController") as mock_controller_cls,
+        patch("artemis.graph.perception.perform_ocr", new_callable=AsyncMock) as mock_ocr,
+        patch("artemis.graph.perception._detect_status_bar_height", return_value=0),
+        patch("artemis.graph.perception._should_skip_settling", return_value=True),
+    ):
+        mock_controller = mock_controller_cls.return_value
+        mock_controller.get_screen_data = AsyncMock(return_value=mock_screen_data)
+        mock_ocr.return_value = [{"text": "Home", "box": [100, 100, 200, 200]}]
+
+        state = State(
+            messages=[HumanMessage(content="test")],
+            initial_goal="test goal",
+            complete_subgoals_by_ids=[],
+            validator_messages=[],
+            latest_ui_hierarchy=[],
+            latest_screenshot=None,
+            focused_app_info=None,
+            device_date=None,
+            structured_decisions=None,
+        )
+
+        update = await perception_node(state, mock_artemis_ctx)
+
+    # Verify perception_node returned expected update immediately
+    assert update["injected_instruction"] == "Avoid clicking Ads"
+    assert update["latest_screenshot"] == str(engine.get_image_path(expected_hash))
+    assert update["operator_raw_data"]["width"] == 1080
+
+    # Verify injected_instruction.json was unlinked
+    assert not instruction_file.exists()
+
+    # Await background tasks triggered by perception_node
+    await asyncio.sleep(0.1)
+    if engine._pending_tasks:
+        await asyncio.gather(*engine._pending_tasks)
+
+    # Verify image record and file were created in background by perception_node
+    img_record = engine.storage.get_image(expected_hash)
+    assert img_record is not None
+    assert img_record.image_name == expected_hash
+    assert engine.get_image_path(expected_hash).exists()
+
+
+@pytest.mark.asyncio
+async def test_snapshot_copy_and_restore_offloading(mock_artemis_ctx, tmp_path):
+    """Verify operator create_snapshot and graph rollback/restore execute asynchronously without blocking."""
+    from artemis.utils.file import create_snapshot, restore_snapshot
+
+    notes_dir = tmp_path / "notes"
+    notes_dir.mkdir(parents=True, exist_ok=True)
+    (notes_dir / "task_plan.md").write_text("# Plan\n- Step 1", encoding="utf-8")
+    (notes_dir / "verification_chat_1.json").write_text('{"messages": ["hello"]}', encoding="utf-8")
+
+    snapshot_dir = tmp_path / "notes_snapshot"
+
+    # Test async offloading of create_snapshot (operator pattern)
+    await asyncio.to_thread(create_snapshot, notes_dir, snapshot_dir)
+    assert snapshot_dir.exists()
+    assert (snapshot_dir / "task_plan.md").read_text(encoding="utf-8") == "# Plan\n- Step 1"
+
+    # Modify live task_plan.md
+    (notes_dir / "task_plan.md").write_text("# Plan\n- Step 1 (Modified)", encoding="utf-8")
+
+    # Test async offloading of restore_snapshot (graph pattern when checker fails)
+    def _restore_and_preserve():
+        chat_contents = {}
+        for p in notes_dir.glob("verification_chat_*.json"):
+            chat_contents[p.name] = p.read_text(encoding="utf-8")
+        restore_snapshot(snapshot_dir, notes_dir)
+        for name, content in chat_contents.items():
+            (notes_dir / name).write_text(content, encoding="utf-8")
+
+    await asyncio.to_thread(_restore_and_preserve)
+    assert (notes_dir / "task_plan.md").read_text(encoding="utf-8") == "# Plan\n- Step 1"
+    assert (notes_dir / "verification_chat_1.json").read_text(
+        encoding="utf-8"
+    ) == '{"messages": ["hello"]}'
+
+
+@pytest.mark.asyncio
+async def test_ocr_api_persistent_http_client_singleton_and_tls_reuse():
+    """Verify get_http_client returns a persistent singleton httpx.AsyncClient that reuses TLS connections across OCR calls."""
+    import httpx
+    from artemis.utils import ocr_api
+
+    # Reset singleton to test fresh creation & persistence
+    ocr_api._HTTP_CLIENT = None
+    client1 = ocr_api.get_http_client()
+    client2 = ocr_api.get_http_client()
+
+    assert isinstance(client1, httpx.AsyncClient)
+    assert client1 is client2
+
+    # Verify perform_ocr uses this exact persistent client
+    fake_ocr_resp = MagicMock()
+    fake_ocr_resp.status_code = 200
+    fake_ocr_resp.json.return_value = {
+        "responses": [
+            {
+                "textAnnotations": [
+                    {
+                        "description": "Combined text block",
+                        "boundingPoly": {"vertices": []},
+                    },
+                    {
+                        "description": "Hello World",
+                        "boundingPoly": {"vertices": [{"x": 10, "y": 20}, {"x": 100, "y": 200}]},
+                    },
+                ]
+            }
+        ]
+    }
+
+    with patch.object(client1, "post", new_callable=AsyncMock) as mock_post:
+        mock_post.return_value = fake_ocr_resp
+        result = await ocr_api.perform_ocr(screenshot_b64="ZHVtbXlfYjY0")
+
+        mock_post.assert_called_once()
+        assert len(result) == 1
+        assert result[0]["text"] == "Hello World"
+
+
+@pytest.mark.asyncio
+async def test_ui_filter_and_ui_automator_client_pre_parsed_bounds_o1_lookup():
+    """Verify XML parsing pre-populates parsed_bounds and _parse_bounds hits this cache in O(1) without regex matches."""
+    import re
+    from artemis.clients.ui_automator_client import _parse_hierarchy_xml_to_elements
+    from artemis.utils.ui_filter import _parse_bounds
+
+    # 1. Test ingestion pre-parsing in _parse_hierarchy_xml_to_elements
+    xml_data = '<hierarchy><node text="Login Btn" bounds="[15,25][350,120]"/></hierarchy>'
+    elements = _parse_hierarchy_xml_to_elements(xml_data)
+    assert len(elements) == 1
+    assert elements[0]["parsed_bounds"] == {
+        "left": 15,
+        "top": 25,
+        "right": 350,
+        "bottom": 120,
+    }
+
+    # 2. Test _parse_bounds O(1) hit vs fallback
+    with patch.object(re, "match", wraps=re.match) as spy_match:
+        # Hit pre-parsed bounds -> re.match MUST NOT be called
+        hit_result = _parse_bounds(elements[0])
+        assert hit_result == {
+            "left": 15,
+            "top": 25,
+            "right": 350,
+            "bottom": 120,
+        }
+        spy_match.assert_not_called()
+
+        # Fallback raw string -> re.match is called exactly once
+        fallback_result = _parse_bounds("[50,60][150,160]")
+        assert fallback_result == {
+            "left": 50,
+            "top": 60,
+            "right": 150,
+            "bottom": 160,
+        }
+        assert spy_match.call_count == 1
+
+
+@pytest.mark.asyncio
+async def test_validator_pre_execution_loop_reverted_to_exact_safety_contract(
+    mock_artemis_ctx, tmp_path
+):
+    """Verify Validator pre-execution check runs VLM on XML failure, but retains XML error if VLM fails (reverted contract)."""
+    from artemis.agents.validator.validator import ValidatorNode, ValidationErrorCategory
+
+    # Set mock context properties required by tracing decorators
+    mock_artemis_ctx.data_engine = MagicMock()
+
+    node = ValidatorNode(mock_artemis_ctx)
+    session = MagicMock()
+    session.call_tool = AsyncMock(
+        return_value=MagicMock(content=[MagicMock(text="ZHVtbXlfZGF0YQ==")])
+    )
+
+    decisions = json.dumps(
+        [
+            {
+                "action": "tap",
+                "coordinates": [200, 250],
+                "target_text": "Login",
+                "target_bounds": [100, 200, 300, 300],
+                "target_resource_id": "btn_login",
+            }
+        ]
+    )
+    state = State(
+        messages=[],
+        initial_goal="test",
+        complete_subgoals_by_ids=[],
+        validator_messages=[],
+        latest_ui_hierarchy=[],
+        latest_screenshot="dummy.png",
+        structured_decisions=decisions,
+        focused_app_info=None,
+        device_date=None,
+    )
+
+    with (
+        patch.object(
+            ValidatorNode, "_get_initial_screenshot", new_callable=AsyncMock
+        ) as mock_get_screen,
+        patch.object(
+            ValidatorNode, "_validate_action_precondition", new_callable=AsyncMock
+        ) as mock_xml,
+        patch.object(
+            ValidatorNode,
+            "_validate_action_precondition_pixel",
+            new_callable=AsyncMock,
+        ) as mock_pixel,
+        patch(
+            "artemis.agents.validator.failure_analyzer.FailureAnalyzer.analyze",
+            new_callable=AsyncMock,
+        ) as mock_analyze,
+    ):
+        mock_get_screen.return_value = ("dummy_b64", "dummy_stem")
+        # 1. XML check returns failure (TARGET_OCCUPIED by 'Sign Up' button)
+        mock_xml.return_value = (
+            False,
+            ValidationErrorCategory.TARGET_OCCUPIED,
+            "Occupied by Sign Up button",
+        )
+        # 2. VLM check returns failure (PIXEL_TARGET_DISAPPEARED)
+        mock_pixel.return_value = (
+            False,
+            ValidationErrorCategory.PIXEL_TARGET_DISAPPEARED,
+            "VLM target missing",
+        )
+
+        mock_analyze.return_value = {
+            "status": "cannot_fix",
+            "analysis": "Unrepairable drift",
+        }
+
+        result = await node(state)
+
+        # Verify both XML check and VLM check were called (reverted double-check contract)
+        mock_xml.assert_called_once()
+        mock_pixel.assert_called_once()
+
+        # Verify FailureAnalyzer received the EXACT XML error (TARGET_OCCUPIED), NOT the VLM error
+        mock_analyze.assert_called_once()
+        args, kwargs = mock_analyze.call_args
+        assert kwargs.get("error_category") == ValidationErrorCategory.TARGET_OCCUPIED
+        assert "Occupied by Sign Up button" in args[2]
+
+
+@pytest.mark.asyncio
+async def test_failure_analyzer_prune_intermediate_screenshots(mock_artemis_ctx):
+    """Verify FailureAnalyzer._prune_intermediate_screenshots keeps initial reference images and the latest live image, while pruning intermediate images."""
+    from google.genai import types
+    from artemis.agents.validator.failure_analyzer import FailureAnalyzer
+
+    analyzer = FailureAnalyzer(mock_artemis_ctx)
+
+    # Build simulated contents history:
+    # 0: User (Initial reference screenshots + prompt)
+    # 1: Model
+    # 2: User (Intermediate tool results)
+    # 3: User (Intermediate observation screenshot - SHOULD BE PRUNED)
+    # 4: Model
+    # 5: User (Intermediate tool results 2)
+    # 6: User (Latest observation screenshot - SHOULD BE PRESERVED)
+    contents = [
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_text(text="Reference screenshot 1"),
+                types.Part.from_bytes(data=b"initial_pre_image_bytes", mime_type="image/jpeg"),
+                types.Part.from_bytes(data=b"initial_post_image_bytes", mime_type="image/jpeg"),
+            ],
+        ),
+        types.Content(role="model", parts=[types.Part.from_text(text="Clicked something")]),
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text="Result text")],
+        ),
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_text(text="Intermediate observation"),
+                types.Part.from_bytes(
+                    data=b"intermediate_image_bytes_to_prune",
+                    mime_type="image/jpeg",
+                ),
+            ],
+        ),
+        types.Content(role="model", parts=[types.Part.from_text(text="Clicking again")]),
+        types.Content(
+            role="user",
+            parts=[types.Part.from_text(text="Result text 2")],
+        ),
+        types.Content(
+            role="user",
+            parts=[
+                types.Part.from_text(text="Latest live observation"),
+                types.Part.from_bytes(
+                    data=b"latest_live_image_bytes_to_keep",
+                    mime_type="image/jpeg",
+                ),
+            ],
+        ),
+    ]
+
+    analyzer._prune_intermediate_screenshots(contents)
+
+    # 1. Verify index 0 images are untouched
+    assert contents[0].parts[1].inline_data.data == b"initial_pre_image_bytes"
+    assert contents[0].parts[2].inline_data.data == b"initial_post_image_bytes"
+
+    # 2. Verify index 3 (intermediate image) is replaced by text placeholder
+    assert getattr(contents[3].parts[1], "inline_data", None) is None
+    assert "[Screenshot of intermediate step omitted for performance]" in contents[3].parts[1].text
+
+    # 3. Verify index 6 (latest image) is fully preserved
+    assert contents[6].parts[1].inline_data.data == b"latest_live_image_bytes_to_keep"
+
+
+@pytest.mark.asyncio
+async def test_record_step_pre_allocated_id_reset(mock_artemis_ctx):
+    """Verify that after record_step is called, the pre-allocated current_step_id is reset to None so subsequent steps generate distinct IDs."""
+    from artemis.data_engine.engine import DataEngine
+
+    engine = DataEngine(mock_artemis_ctx)
+    engine.start_session("Pre-allocated ID Reset Test")
+
+    # 1. Pre-allocate ID for step 1
+    step_id_1 = engine.allocate_step_id()
+    assert engine.current_step_id == step_id_1
+
+    # 2. Record step 1 (should use step_id_1 and reset current_step_id to None)
+    returned_id_1 = engine.record_step(summary="Step 1")
+    assert returned_id_1 == step_id_1
+    assert engine.current_step_id is None
+
+    # 3. Pre-allocate ID for step 2 (should get a distinct new ID)
+    step_id_2 = engine.allocate_step_id()
+    assert step_id_2 != step_id_1
+    assert engine.current_step_id == step_id_2
+
+    # 4. Record step 2 (should use step_id_2 and reset current_step_id to None)
+    returned_id_2 = engine.record_step(summary="Step 2")
+    assert returned_id_2 == step_id_2
+    assert returned_id_2 != returned_id_1
+    assert engine.current_step_id is None

tests/unit/data_engine/test_trace_enhancements.py

@@ -0,0 +1,189 @@
+# 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 json
+from unittest.mock import MagicMock
+
+from artemis.context import ArtemisContext
+from artemis.data_engine.trace import (
+    TraceSpan,
+    smart_serialize,
+)
+import pytest
+
+
+def test_smart_serialize():
+    # Test dict
+    d = {"a": 1, "b": "string"}
+    assert smart_serialize(d) == json.dumps(d, ensure_ascii=False)
+
+    # Test list
+    items_list = [1, 2, "three"]
+    assert smart_serialize(items_list) == json.dumps(items_list, ensure_ascii=False)
+
+    # Test bytes
+    b = b"hello world"
+    import hashlib
+
+    expected_sha = hashlib.sha256(b).hexdigest()[:8]
+    assert smart_serialize(b) == f"<Bytes length={len(b)} sha256={expected_sha}>"
+
+    # Test other types
+    assert smart_serialize(123) == "123"
+    assert smart_serialize("just a string") == "just a string"
+
+
+@pytest.mark.asyncio
+async def test_trace_span():
+    mock_ctx = MagicMock(spec=ArtemisContext)
+    mock_ctx.data_engine = MagicMock()
+
+    with TraceSpan(name="test_span", ctx=mock_ctx) as span:
+        span.result = "span_result"
+        # Simulate work
+
+    # TraceSpan writes twice: start (running) and end (success)
+    assert mock_ctx.data_engine.record_trace.call_count == 2
+
+    # Verify the final success call (which is the last call)
+    args, kwargs = mock_ctx.data_engine.record_trace.call_args
+    assert kwargs["type"] == "span"
+    assert kwargs["name"] == "test_span"
+    assert kwargs["status"] == "success"
+    assert kwargs["payload"]["result"] == "span_result"
+
+
+@pytest.mark.asyncio
+async def test_trace_span_failure():
+    mock_ctx = MagicMock(spec=ArtemisContext)
+    mock_ctx.data_engine = MagicMock()
+
+    try:
+        with TraceSpan(name="failed_span", ctx=mock_ctx):
+            raise ValueError("Span error")
+    except ValueError:
+        pass
+
+    # TraceSpan writes twice: start (running) and end (failed)
+    assert mock_ctx.data_engine.record_trace.call_count == 2
+
+    # Verify the final failed call (which is the last call)
+    args, kwargs = mock_ctx.data_engine.record_trace.call_args
+    assert kwargs["type"] == "span"
+    assert kwargs["name"] == "failed_span"
+    assert kwargs["status"] == "failed"
+    assert "error" in kwargs["payload"]
+    assert "Span error" in kwargs["payload"]["error"]
+
+
+def test_data_engine_callback_handler():
+    from uuid import uuid4
+    from langchain_core.messages import HumanMessage, AIMessage
+    from artemis.data_engine.trace import DataEngineCallbackHandler
+
+    mock_ctx = MagicMock(spec=ArtemisContext)
+    mock_ctx.data_engine = MagicMock()
+
+    mock_step_id = uuid4()
+    mock_ctx.data_engine.current_step_id = mock_step_id
+
+    handler = DataEngineCallbackHandler(mock_ctx)
+    run_id = uuid4()
+
+    # 1. Simulate on_chat_model_start
+    serialized = {"name": "test_model"}
+    messages = [[HumanMessage(content="Hello LLM")]]
+    handler.on_chat_model_start(serialized, messages, run_id=run_id)
+
+    assert mock_ctx.data_engine.record_trace.call_count == 1
+    call1_args, call1_kwargs = mock_ctx.data_engine.record_trace.call_args_list[0]
+    assert call1_kwargs["type"] == "llm_call"
+    assert call1_kwargs["name"] == "test_model"
+    assert call1_kwargs["status"] == "running"
+    assert call1_kwargs["step_id"] == mock_step_id
+    assert call1_kwargs["payload"]["messages"][0]["content"] == "Hello LLM"
+
+    # 2. Simulate on_llm_end (successful complete response)
+    class MockGeneration:
+        def __init__(self, message):
+            self.message = message
+
+    class MockLLMResult:
+        def __init__(self, generations):
+            self.generations = generations
+
+    response_message = AIMessage(content="Hello User")
+    mock_response = MockLLMResult(generations=[[MockGeneration(response_message)]])
+
+    handler.on_llm_end(mock_response, run_id=run_id)
+
+    assert mock_ctx.data_engine.record_trace.call_count == 3
+    call2_args, call2_kwargs = mock_ctx.data_engine.record_trace.call_args_list[1]
+    assert call2_kwargs["type"] == "llm_call"
+    assert call2_kwargs["status"] == "success"
+    assert call2_kwargs["step_id"] == mock_step_id
+    assert call2_kwargs["payload"]["messages"][0]["content"] == "Hello LLM"
+    assert call2_kwargs["payload"]["response"][0]["content"] == "Hello User"
+    assert isinstance(call2_kwargs.get("duration"), float)
+    assert call2_kwargs["duration"] >= 0.0
+
+    call3_args, call3_kwargs = mock_ctx.data_engine.record_trace.call_args_list[2]
+    assert call3_kwargs["type"] == "raw_thinking"
+    assert call3_kwargs["parent_trace_id"] == run_id
+    assert call3_kwargs["payload"]["thought"] == "Hello User"
+
+
+def test_callback_handler_chain_parenting():
+    from uuid import uuid4
+    from langchain_core.messages import HumanMessage
+    from artemis.data_engine.trace import DataEngineCallbackHandler, CURRENT_TRACE_ID
+
+    mock_ctx = MagicMock(spec=ArtemisContext)
+    mock_ctx.data_engine = MagicMock()
+
+    handler = DataEngineCallbackHandler(mock_ctx)
+
+    agent_run_id = uuid4()
+    chain1_run_id = uuid4()
+    chain2_run_id = uuid4()
+    llm_run_id = uuid4()
+
+    token = CURRENT_TRACE_ID.set(agent_run_id)
+    try:
+        # Simulate chain start (top level, parent is agent_run_id)
+        handler.on_chain_start({}, {}, run_id=chain1_run_id, parent_run_id=None)
+        # Simulate nested chain start (parent is chain1_run_id)
+        handler.on_chain_start({}, {}, run_id=chain2_run_id, parent_run_id=chain1_run_id)
+
+        # Simulate LLM start inside nested chain (parent is chain2_run_id)
+        messages = [[HumanMessage(content="Hello Nested LLM")]]
+        handler.on_chat_model_start(
+            {"name": "test_nested_model"},
+            messages,
+            run_id=llm_run_id,
+            parent_run_id=chain2_run_id,
+        )
+
+        # Verify parent trace ID resolves to agent_run_id (since chains are not recorded)
+        assert mock_ctx.data_engine.record_trace.call_count == 1
+        args, kwargs = mock_ctx.data_engine.record_trace.call_args
+        assert kwargs["parent_trace_id"] == agent_run_id
+
+        # Cleanup
+        handler.on_chain_end({}, run_id=chain2_run_id)
+        handler.on_chain_end({}, run_id=chain1_run_id)
+
+        assert len(handler.chain_parents) == 0
+    finally:
+        CURRENT_TRACE_ID.reset(token)