agentic-mobile-control/tests/test_ai_planner.py

from __future__ import annotations

from typing import Any

import pytest

from core.errors import TaskFailedError
from core.models import Bounds, Scene, SceneElement
from runtime.ai_planner import AIPlanner
from runtime.context import TaskContext
from runtime.executor import StepResult
from runtime.planner_config import PlannerConfig
from runtime.tool_calling_client import ToolCallDecision
from runtime.tool_specs import ALL_TOOL_SPECS


class FakeToolCallingClient:
    def __init__(self, decision: ToolCallDecision) -> None:
        self.decision = decision
        self.calls: list[dict[str, Any]] = []

    def decide(
        self,
        *,
        system_prompt: str,
        user_prompt: str,
        screenshot: bytes | None,
        tools: list[Any],
        timeout: float,
        history: list[dict[str, Any]] | None = None,
    ) -> ToolCallDecision:
        self.calls.append(
            {
                "system_prompt": system_prompt,
                "user_prompt": user_prompt,
                "screenshot": screenshot,
                "tools": tools,
                "timeout": timeout,
                "history": list(history) if history is not None else None,
            }
        )
        return self.decision


def _scene() -> Scene:
    return Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="send", type="button", text="Send", bounds=Bounds(1, 2, 3, 4)
            )
        ],
    )


def _context() -> TaskContext:
    return TaskContext(task_id="task-1", goal="send a message")


def test_ai_planner_returns_single_planned_step_for_action_decision() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(
            tool_name="tap",
            arguments={"x": 1, "y": 2},
            purpose="Open the send control.",
            expected_outcome="The message composer is focused.",
        )
    )
    planner = AIPlanner(client=client)

    steps = planner.plan(goal="send a message", scene=_scene(), context=_context())

    assert len(steps) == 1
    step = steps[0]
    assert step.action == "tap"
    assert step.description == "AI planner: Open the send control."
    assert step.args == {"x": 1, "y": 2}
    assert step.purpose == "Open the send control."
    assert step.expected_outcome == "The message composer is focused."
    assert step.expected_text == "The message composer is focused."


def test_ai_planner_finish_task_success_returns_empty_plan() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(
            tool_name="finish_task", arguments={"success": True, "reason": "done"}
        )
    )
    planner = AIPlanner(client=client)

    steps = planner.plan(goal="send a message", scene=_scene(), context=_context())

    assert steps == []


def test_ai_planner_finish_task_failure_raises_task_failed_error_with_reason() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(
            tool_name="finish_task",
            arguments={"success": False, "reason": "stuck on login"},
        )
    )
    planner = AIPlanner(client=client)

    with pytest.raises(TaskFailedError, match="stuck on login"):
        planner.plan(goal="send a message", scene=_scene(), context=_context())


def test_ai_planner_finish_task_failure_without_reason_uses_default_message() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="finish_task", arguments={"success": False})
    )
    planner = AIPlanner(client=client)

    with pytest.raises(TaskFailedError, match="task failed"):
        planner.plan(goal="send a message", scene=_scene(), context=_context())


def test_ai_planner_goal_reached_is_always_false() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client)

    assert (
        planner.goal_reached(goal="anything", scene=_scene(), context=_context())
        is False
    )


def test_ai_planner_forwards_tools_screenshot_and_timeout_to_client() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client, config=PlannerConfig(timeout=12.5))

    planner.plan(
        goal="send a message",
        scene=_scene(),
        context=_context(),
        screenshot=b"fake-bytes",
    )

    call = client.calls[0]
    assert call["tools"] == ALL_TOOL_SPECS
    assert call["screenshot"] == b"fake-bytes"
    assert call["timeout"] == 12.5
    assert "send a message" in call["user_prompt"]


def test_ai_planner_includes_device_platform_from_task_context() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client)
    context = TaskContext(
        task_id="task-1",
        goal="send a message",
        device_platform="android",
    )

    planner.plan(goal="send a message", scene=_scene(), context=context)

    assert "Device type: android" in client.calls[0]["user_prompt"]


def test_ai_planner_populates_step_prompt_from_user_prompt() -> None:
    """PlannedStep.prompt should carry the actual user prompt sent to the LLM,
    not the bare task goal."""
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client)

    steps = planner.plan(goal="send a message", scene=_scene(), context=_context())

    assert len(steps) == 1
    assert steps[0].prompt is not None
    # The per-step prompt contains the goal but also scene JSON and instruction text
    assert "send a message" in steps[0].prompt
    assert "Current Scene (JSON)" in steps[0].prompt
    assert "Call exactly one tool" in steps[0].prompt


def test_ai_planner_step_prompt_reflects_scene_changes() -> None:
    """Per-step prompts differ when the scene changes, proving they are not
    just the repeated task goal."""
    from runtime.context import TaskContext

    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client)

    scene_a = Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="btn_a", type="button", text="Alpha", bounds=Bounds(1, 2, 3, 4)
            )
        ],
    )
    scene_b = Scene(
        width=10,
        height=20,
        elements=[
            SceneElement(
                id="btn_b", type="button", text="Beta", bounds=Bounds(5, 6, 7, 8)
            )
        ],
    )

    steps_a = planner.plan(
        goal="test", scene=scene_a, context=TaskContext(task_id="t", goal="test")
    )
    steps_b = planner.plan(
        goal="test", scene=scene_b, context=TaskContext(task_id="t", goal="test")
    )

    assert steps_a[0].prompt != steps_b[0].prompt
    assert "Alpha" in steps_a[0].prompt
    assert "Beta" in steps_b[0].prompt


def test_ai_planner_propagates_rationale_and_thinking_to_planned_step() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(
            tool_name="tap",
            arguments={"x": 1, "y": 2},
            text_output="Previous step opened settings. Now tapping account.",
            thinking="I need to navigate deeper.",
            purpose="Open account settings.",
            expected_outcome="The account settings page is visible.",
        )
    )
    planner = AIPlanner(client=client)

    steps = planner.plan(goal="open account", scene=_scene(), context=_context())

    assert steps[0].rationale == "Previous step opened settings. Now tapping account."
    assert steps[0].thinking == "I need to navigate deeper."
    assert steps[0].purpose == "Open account settings."
    assert steps[0].expected_outcome == "The account settings page is visible."


def test_ai_planner_carries_completed_turn_into_the_next_llm_call() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(
            tool_name="tap",
            arguments={"x": 1, "y": 2},
            text_output="Opening the send control.",
            purpose="Open the send control.",
            expected_outcome="The composer is focused.",
        )
    )
    planner = AIPlanner(client=client)
    context = _context()

    first_step = planner.plan(goal=context.goal, scene=_scene(), context=context)[0]
    context.add_step_result(
        StepResult(
            step=first_step,
            success=True,
            attempts=1,
            result={"ok": True},
        )
    )
    planner.plan(goal=context.goal, scene=_scene(), context=context)

    assert client.calls[0]["history"] == []
    history = client.calls[1]["history"]
    assert history is not None
    assert history[0]["tool_name"] == "tap"
    assert history[0]["arguments"] == {
        "x": 1,
        "y": 2,
        "purpose": "Open the send control.",
        "expected_outcome": "The composer is focused.",
    }
    assert history[0]["tool_result"]["success"] is True
    assert history[0]["tool_result"]["result"] == {"ok": True}


def test_ai_planner_limits_history_sent_to_the_llm() -> None:
    client = FakeToolCallingClient(
        ToolCallDecision(tool_name="tap", arguments={"x": 1, "y": 2})
    )
    planner = AIPlanner(client=client, config=PlannerConfig(history_max_turns=2))
    context = _context()
    context.planner_history.extend(
        {
            "user_prompt": f"turn-{index}",
            "tool_name": "tap",
            "arguments": {},
            "rationale": None,
            "tool_result": {"success": True},
        }
        for index in range(3)
    )

    planner.plan(goal=context.goal, scene=_scene(), context=context)

    assert [turn["user_prompt"] for turn in client.calls[0]["history"]] == [
        "turn-1",
        "turn-2",
    ]