agentic-mobile-control/tests/test_task_loop.py

from __future__ import annotations

from pathlib import Path

from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import Executor, ExecutorConfig
from runtime.planner import PlannedStep, Planner
from runtime.task import TaskRunner, TaskRunnerConfig
from storage.artifact_store import ArtifactStore
from storage.task_metadata import TaskMetadataStore
from storage.timeline import Timeline
from tests.fakes import PNG_10X20


class ScriptedPlanner(Planner):
    def __init__(self, steps: list[PlannedStep]) -> None:
        self.steps = steps

    def plan(self, *, goal, scene, context):
        if len(context.step_results) >= len(self.steps):
            return []
        return [self.steps[len(context.step_results)]]

    def goal_reached(self, *, goal, scene, context):
        return len(context.step_results) >= len(self.steps) and all(
            result.success for result in context.step_results
        )


def test_task_runner_executes_loop_and_writes_timeline(tmp_path) -> None:
    scene = Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="search",
                type="input",
                text="Search",
                bounds=Bounds(1, 2, 4, 4),
            )
        ],
    )
    planner = ScriptedPlanner(
        [
            PlannedStep(action="tap", description="tap search", args={"x": 3, "y": 4}),
            PlannedStep(
                action="input_text",
                description="type query",
                args={"text": "Mac mini"},
            ),
        ]
    )
    executor = Executor(
        tools={
            "tap": lambda **kwargs: {"ok": True, **kwargs},
            "input_text": lambda **kwargs: {"ok": True, **kwargs},
        },
        config=ExecutorConfig(max_retries=1, backoff_seconds=0),
    )
    metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3")
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    task = Task(goal="open app and search", device_id="iphone-1")

    runner = TaskRunner(
        planner=planner,
        executor=executor,
        metadata_store=metadata,
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=5),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    result = runner.run(task)

    assert result.status == "completed"
    assert len(timeline.read(task.id)) == 2
    assert metadata.get_task(task.id)["status"] == "completed"


def test_task_runner_stops_before_the_next_planned_action() -> None:
    scene = Scene(width=10, height=20, elements=[])
    stop_requested = False
    actions: list[str] = []

    def record_action(**kwargs):
        nonlocal stop_requested
        actions.append("tap")
        stop_requested = True
        return {"ok": True}

    runner = TaskRunner(
        planner=ScriptedPlanner(
            [
                PlannedStep(action="tap", description="first", args={}),
                PlannedStep(action="tap", description="second", args={}),
            ]
        ),
        executor=Executor(
            tools={"tap": record_action},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        config=TaskRunnerConfig(max_steps=5),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    result = runner.run(
        Task(goal="perform two actions", device_id="phone"),
        should_stop=lambda: stop_requested,
    )

    assert result.status == "failed"
    assert result.failure_reason == "execution interrupted"
    assert actions == ["tap"]


def test_task_runner_persists_action_evidence_and_raw_ocr(tmp_path) -> None:
    scene = Scene(
        width=10,
        height=20,
        elements=[],
        ocr_elements=[
            SceneElement(
                id="ocr-001",
                type="text",
                text="Search",
                bounds=Bounds(1, 2, 3, 4),
                confidence=0.98,
                source="ocr",
            )
        ],
    )
    screenshots = iter(
        [
            b"planning",
            b"before-action",
            b"after-action",
            b"completion-check",
        ]
    )
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    runner = TaskRunner(
        planner=ScriptedPlanner(
            [PlannedStep(action="tap", description="tap search", args={})]
        ),
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=2),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: next(screenshots),
    )

    result = runner.run(Task(goal="tap search", device_id="iphone-1"))

    assert result.status == "completed"
    record = timeline.read(result.id)[0]
    assert Path(record["before_screenshot_path"]).read_bytes() == b"before-action"
    assert Path(record["after_screenshot_path"]).read_bytes() == b"after-action"
    assert record["tool_call"]["description"] == "tap search"
    assert record["ocr_results"][0]["text"] == "Search"


def test_task_runner_persists_ui_tree_elements_from_fused_scene(tmp_path) -> None:
    scene = Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="ui-000",
                type="button",
                text="Search",
                bounds=Bounds(1, 2, 3, 4),
                source="ui",
            ),
            SceneElement(
                id="ocr-000",
                type="text",
                text="Unrelated label",
                bounds=Bounds(5, 6, 3, 4),
                source="ocr",
            ),
        ],
    )
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    runner = TaskRunner(
        planner=ScriptedPlanner(
            [PlannedStep(action="tap", description="tap search", args={})]
        ),
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=2),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )

    result = runner.run(Task(goal="tap search", device_id="iphone-1"))

    record = timeline.read(result.id)[0]
    assert [element["text"] for element in record["ui_tree_results"]] == ["Search"]
    assert [element["text"] for element in record["ocr_results"]] == ["Unrelated label"]


def test_task_runner_persists_empty_ui_tree_results_when_scene_has_no_ui_elements(
    tmp_path,
) -> None:
    scene = Scene(width=10, height=20, elements=[])
    timeline = Timeline(ArtifactStore(tmp_path / "history"))
    runner = TaskRunner(
        planner=ScriptedPlanner(
            [PlannedStep(action="tap", description="tap search", args={})]
        ),
        executor=Executor(
            tools={"tap": lambda **kwargs: {"ok": True}},
            config=ExecutorConfig(max_retries=1, backoff_seconds=0),
        ),
        timeline=timeline,
        config=TaskRunnerConfig(max_steps=2),
        observer=lambda device_id: scene,
        screenshot_provider=lambda device_id: PNG_10X20,
    )
    result = runner.run(Task(goal="tap search", device_id="iphone-1"))
    record = timeline.read(result.id)[0]
    assert record["ui_tree_results"] == []