agentic-mobile-control/tests/test_agents_collab_runner.py

from __future__ import annotations

from typing import Any
from unittest.mock import MagicMock, patch

from agents.collab_runner import CollaborativeTaskRunner, CollaborativeTaskRunnerConfig
from agents.config import CollaborationConfig
from agents.models import Observation, ReflectionAction, ReflectionOutcome, VerificationVerdict
from core.models import Bounds, Scene, SceneElement, Task
from runtime.executor import StepResult
from runtime.planner import PlannedStep


def _scene() -> Scene:
    return Scene(
        width=1080,
        height=1920,
        elements=[SceneElement(id="btn1", type="button", bounds=Bounds(10, 20, 100, 50), text="OK")],
    )


def _task() -> Task:
    return Task(goal="Test goal", device_id="dev1")


def _planned_step() -> PlannedStep:
    return PlannedStep(action="describe_screen", description="Observe", args={})


def _success_step_result() -> StepResult:
    return StepResult(step=_planned_step(), success=True, attempts=1, result="ok")


def _observation() -> Observation:
    return Observation(scene_summary="Screen 1080x1920")


def _achieved_verdict() -> VerificationVerdict:
    return VerificationVerdict(result="achieved", reasoning="Effect observed.")


def _not_achieved_verdict() -> VerificationVerdict:
    return VerificationVerdict(result="not_achieved", reasoning="No change.")


def _replan_outcome() -> ReflectionOutcome:
    return ReflectionOutcome(replan=True, reasoning="No recovery.")


def _recovery_outcome() -> ReflectionOutcome:
    return ReflectionOutcome(
        replan=False,
        action=ReflectionAction(action="swipe", description="Scroll", args={"direction": "up"}),
        reasoning="Try scrolling.",
    )


def test_completes_without_verification_failure() -> None:
    """Task completes normally when verifier reports achieved."""
    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.return_value = True

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.return_value = _achieved_verdict()

    reflector = MagicMock()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "completed"
    reflector.reflect.assert_not_called()


def test_recovers_via_reflector_action() -> None:
    """Task recovers when reflector proposes a distinct action."""
    call_count = 0

    def verifier_side_effect(**kwargs: Any) -> VerificationVerdict:
        nonlocal call_count
        call_count += 1
        if call_count == 1:
            return _not_achieved_verdict()
        return _achieved_verdict()

    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.side_effect = [False, True]

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.side_effect = verifier_side_effect

    reflector = MagicMock()
    reflector.reflect.return_value = _recovery_outcome()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=5),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "completed"
    reflector.reflect.assert_called_once()
    executor.execute.assert_called()  # recovery action also executed


def test_exhausts_recovery_ceiling() -> None:
    """Task fails when reflection-recovery ceiling is reached."""
    planner = MagicMock()
    planner.plan.return_value = [_planned_step()]
    planner.goal_reached.return_value = False

    executor = MagicMock()
    executor.execute.return_value = _success_step_result()

    observer = MagicMock()
    observer.observe.return_value = _observation()

    verifier = MagicMock()
    verifier.verify.return_value = _not_achieved_verdict()

    reflector = MagicMock()
    reflector.reflect.return_value = _replan_outcome()

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        observer=observer,
        verifier=verifier,
        reflector=reflector,
        config=CollaborativeTaskRunnerConfig(max_steps=10, max_recovery_attempts=2),
        collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=2),
    )

    with patch("tools.describe_screen.describe_screen", return_value=_scene()):
        task = runner.run(_task())

    assert task.status == "failed"
    assert "ceiling" in (task.failure_reason or "").lower()


def test_disabled_collaboration_runs_plain() -> None:
    """When collaboration is disabled, delegates to plain TaskRunner."""
    planner = MagicMock()
    executor = MagicMock()

    mock_task_runner = MagicMock()
    expected_task = _task()
    expected_task.status = "completed"
    mock_task_runner.run.return_value = expected_task

    runner = CollaborativeTaskRunner(
        planner=planner,
        executor=executor,
        task_runner=mock_task_runner,
        collaboration_config=CollaborationConfig(enabled=False),
    )

    result = runner.run(_task())
    assert result.status == "completed"
    mock_task_runner.run.assert_called_once()