agentic-mobile-control/tests/test_world_model.py

from __future__ import annotations

from core.models import Bounds, Scene, SceneElement
from runtime.executor import StepResult
from runtime.planner import PlannedStep
from semantic.models import SemanticScene, SemanticWidget
from world.config import WorldConfig, load_config
from world.model import WorldModel


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


def _semantic_scene(page: str = "Chat") -> SemanticScene:
    return SemanticScene(
        page=page,
        intents=["send a message"],
        widgets=[SemanticWidget(element_id="send", purpose="send message")],
    )


def _result(step: PlannedStep, *, success: bool = True) -> StepResult:
    return StepResult(step=step, success=success, attempts=1)


def test_world_config_defaults_to_enabled_with_history_size_bound() -> None:
    config = load_config({})

    assert config.enabled is True
    assert config.history_size == 10


def test_observe_updates_current_page_from_semantic_scene() -> None:
    step = PlannedStep(action="tap", description="tap send")
    model = WorldModel(config=WorldConfig(history_size=2))

    model.observe(_scene(), _semantic_scene("Chat"), step, _result(step))

    assert model.state.current_page == "Chat"


def test_observe_leaves_current_page_unchanged_without_semantic_scene() -> None:
    step = PlannedStep(action="tap", description="tap send")
    model = WorldModel(config=WorldConfig(history_size=2))
    model.state.current_page = "Chat"

    model.observe(_scene(), None, step, _result(step))

    assert model.state.current_page == "Chat"


def test_observe_updates_current_app_for_successful_launch_and_clear_for_terminate() -> None:
    model = WorldModel(config=WorldConfig(history_size=2))
    launch = PlannedStep(
        action="launch_app",
        description="launch chat",
        args={"app_id": "com.example.chat"},
    )
    terminate = PlannedStep(
        action="terminate_app",
        description="close chat",
        args={"app_id": "com.example.chat"},
    )

    model.observe(_scene(), None, launch, _result(launch))
    assert model.state.current_app == "com.example.chat"

    model.observe(_scene(), None, terminate, _result(terminate))
    assert model.state.current_app is None


def test_observe_does_not_update_current_app_on_failed_or_unrelated_step() -> None:
    model = WorldModel(config=WorldConfig(history_size=3))
    model.state.current_app = "com.example.chat"
    failed_launch = PlannedStep(
        action="launch_app",
        description="launch other",
        args={"app_id": "com.example.other"},
    )
    tap = PlannedStep(action="tap", description="tap", args={"x": 1, "y": 2})

    model.observe(_scene(), None, failed_launch, _result(failed_launch, success=False))
    model.observe(_scene(), None, tap, _result(tap))

    assert model.state.current_app == "com.example.chat"


def test_observe_merges_explicit_remember_variables_only() -> None:
    model = WorldModel(config=WorldConfig(history_size=3))
    remember_step = PlannedStep(
        action="tap",
        description="select chat",
        args={"remember": {"contact": "Zhang San"}},
    )
    plain_step = PlannedStep(action="tap", description="tap send", args={})

    model.observe(_scene(), None, remember_step, _result(remember_step))
    model.observe(_scene(), None, plain_step, _result(plain_step))

    assert model.state.variables == {"contact": "Zhang San"}


def test_observe_skips_malformed_fields_without_raising() -> None:
    model = WorldModel(config=WorldConfig(history_size=2))
    missing_app = PlannedStep(action="launch_app", description="launch", args={})
    bad_remember = PlannedStep(
        action="tap",
        description="tap",
        args={"remember": "not-a-dict"},
    )

    model.observe(_scene(), None, missing_app, _result(missing_app))
    model.observe(_scene(), None, bad_remember, _result(bad_remember))

    assert model.state.current_app is None
    assert model.state.variables == {}
    assert len(model.state.history) == 2


def test_concurrent_tasks_on_same_world_model_stay_isolated() -> None:
    model = WorldModel(config=WorldConfig(history_size=5))
    launch_a = PlannedStep(
        action="launch_app",
        description="launch a's app",
        args={"app_id": "com.example.a"},
    )
    launch_b = PlannedStep(
        action="launch_app",
        description="launch b's app",
        args={"app_id": "com.example.b"},
    )

    task_a = model.start_task("task-a")
    task_a.observe(_scene(), _semantic_scene("A-Home"), launch_a, _result(launch_a))

    # Task B starts on the SAME WorldModel instance *after* task A has
    # already observed an event but *before* task A observes its next one --
    # this is the interleaving that corrupted state when WorldModel tracked
    # "the current task" via a single shared mutable attribute.
    task_b = model.start_task("task-b")
    task_b.observe(_scene(), _semantic_scene("B-Home"), launch_b, _result(launch_b))

    remember_a = PlannedStep(
        action="tap",
        description="a remembers",
        args={"remember": {"contact": "Alice"}},
    )
    task_a.observe(_scene(), _semantic_scene("A-Chat"), remember_a, _result(remember_a))

    assert task_a.state.current_app == "com.example.a"
    assert task_a.state.current_page == "A-Chat"
    assert task_a.state.variables == {"contact": "Alice"}
    assert [event.action for event in task_a.state.history] == ["launch_app", "tap"]

    assert task_b.state.current_app == "com.example.b"
    assert task_b.state.current_page == "B-Home"
    assert task_b.state.variables == {}
    assert [event.action for event in task_b.state.history] == ["launch_app"]

    # External readers keep working via state_for().
    assert model.state_for("task-a") is task_a.state
    assert model.state_for("task-b") is task_b.state
    assert model.state_for("task-a") is not model.state_for("task-b")


def test_observe_appends_semantic_or_raw_scene_history_with_eviction() -> None:
    model = WorldModel(config=WorldConfig(history_size=2))
    first = PlannedStep(action="first", description="first")
    second = PlannedStep(action="second", description="second")
    third = PlannedStep(action="third", description="third")

    model.observe(_scene(), _semantic_scene("Chat"), first, _result(first))
    model.observe(_scene(), None, second, _result(second))
    model.observe(_scene(), None, third, _result(third))

    assert [event.action for event in model.state.history] == ["second", "third"]
    assert model.state.history[0].scene_summary.to_dict() == _scene().to_dict()
    assert model.state.history.maxlen == 2