336 lines
11 KiB
Python
336 lines
11 KiB
Python
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
|
|
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
|
|
from world.config import WorldConfig
|
|
|
|
|
|
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_collaborative_run_shares_bookkeeping_with_task_runner(tmp_path) -> None:
|
|
"""A completed collaborative task records Timeline/TaskMetadataStore/
|
|
on_task_succeeded exactly like a plain TaskRunner.run() would for the
|
|
same scripted single-step scenario, and populates context.world along
|
|
the way (verified indirectly via what the Observer is called with)."""
|
|
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.return_value = _achieved_verdict()
|
|
|
|
reflector = MagicMock()
|
|
|
|
metadata = TaskMetadataStore(tmp_path / "tasks.sqlite3")
|
|
timeline = Timeline(ArtifactStore(tmp_path / "history"))
|
|
on_task_succeeded = MagicMock()
|
|
|
|
shared_task_runner = TaskRunner(
|
|
metadata_store=metadata,
|
|
timeline=timeline,
|
|
on_task_succeeded=on_task_succeeded,
|
|
world_config=WorldConfig(enabled=True),
|
|
config=TaskRunnerConfig(max_steps=5),
|
|
screenshot_provider=lambda device_id: PNG_10X20,
|
|
)
|
|
|
|
runner = CollaborativeTaskRunner(
|
|
planner=planner,
|
|
executor=executor,
|
|
observer=observer,
|
|
verifier=verifier,
|
|
reflector=reflector,
|
|
task_runner=shared_task_runner,
|
|
config=CollaborativeTaskRunnerConfig(max_steps=5),
|
|
collaboration_config=CollaborationConfig(enabled=True, max_recovery_attempts=3),
|
|
)
|
|
|
|
task = _task()
|
|
metadata.create_task(task)
|
|
|
|
with patch("tools.describe_screen.describe_screen", return_value=_scene()):
|
|
result = runner.run(task)
|
|
|
|
assert result.status == "completed"
|
|
assert len(timeline.read(task.id)) == 1
|
|
record = timeline.read(task.id)[0]
|
|
assert record["before_screenshot_path"]
|
|
assert record["after_screenshot_path"]
|
|
assert metadata.get_task(task.id)["status"] == "completed"
|
|
on_task_succeeded.assert_called_once_with(task.id, task.goal, timeline)
|
|
|
|
# context.world was populated (not None) throughout the run: every
|
|
# Observer.observe call received a non-None `world` kwarg.
|
|
assert observer.observe.call_args_list
|
|
for call in observer.observe.call_args_list:
|
|
assert call.kwargs["world"] is not None
|
|
|
|
|
|
def test_second_step_verified_against_post_step_one_observation() -> None:
|
|
"""When a single Planner.plan() call returns a 2-element plan, the
|
|
second step's Verifier.verify() call uses the first step's post
|
|
observation as its pre_observation, not the stale pre-plan observation."""
|
|
step_one = _planned_step()
|
|
step_two = PlannedStep(action="tap", description="tap next", args={"x": 1, "y": 2})
|
|
|
|
planner = MagicMock()
|
|
planner.plan.return_value = [step_one, step_two]
|
|
planner.goal_reached.side_effect = [False, True]
|
|
|
|
executor = MagicMock()
|
|
executor.execute.return_value = _success_step_result()
|
|
|
|
obs_pre_plan = Observation(scene_summary="obs0-pre-plan")
|
|
obs_after_step_one = Observation(scene_summary="obs1-after-step-one")
|
|
obs_after_step_two = Observation(scene_summary="obs2-after-step-two")
|
|
obs_next_outer_iter = Observation(scene_summary="obs3-next-outer-iter")
|
|
|
|
observer = MagicMock()
|
|
observer.observe.side_effect = [
|
|
obs_pre_plan,
|
|
obs_after_step_one,
|
|
obs_after_step_two,
|
|
obs_next_outer_iter,
|
|
]
|
|
|
|
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()
|
|
assert verifier.verify.call_count == 2
|
|
|
|
first_call_kwargs = verifier.verify.call_args_list[0].kwargs
|
|
second_call_kwargs = verifier.verify.call_args_list[1].kwargs
|
|
|
|
assert first_call_kwargs["pre_observation"] is obs_pre_plan
|
|
assert first_call_kwargs["post_observation"] is obs_after_step_one
|
|
|
|
# The bug being guarded against: the second step must not be verified
|
|
# against the stale pre-whole-plan observation.
|
|
assert second_call_kwargs["pre_observation"] is obs_after_step_one
|
|
assert second_call_kwargs["pre_observation"] is not obs_pre_plan
|
|
assert second_call_kwargs["post_observation"] is obs_after_step_two
|
|
|
|
|
|
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()
|