Artemis-Eapil/tests/unit/agents/test_flash_cross_target_observation.py

"""Cross-target observations save work without executing stale phone decisions."""

from types import SimpleNamespace
from unittest.mock import AsyncMock, Mock

import pytest
from langchain_core.messages import AIMessage, ToolMessage

from artemis.agents.flash import runner as flash
from artemis.agents.validator.tool_declarations import ToolExecutionResult
from artemis.mcp.observation import DeviceObservationError
from artemis.memory.transcript import TranscriptLedger
from artemis.runtime.task_control import TaskControl


PC = "plug_pc_action"


@pytest.fixture
def cross_target(monkeypatch):
    plugins = SimpleNamespace(
        bindings={PC: ("pc", None)},
        tool_names={PC},
        previous_image=lambda *_: b"pc-before",
    )
    monkeypatch.setattr("artemis.plugins.runtime.get_plugin_runtime", lambda _: plugins)
    runner = flash.FlashRunner.__new__(flash.FlashRunner)
    runner.ctx = SimpleNamespace(execution_target="mixed", data_engine=None, task_control=None)
    runner._android_observation_stale = False
    runner.controller, runner.summarizer, runner._llm = Mock(), None, None
    state = SimpleNamespace(
        latest_screenshot="old-phone.png",
        indexed_elements=[{"index": 1, "text": "Old button"}],
        indexed_points=[[1, 2]],
    )

    async def execute(name, args, tc_id, current, **kwargs):
        if name == PC:
            return ToolExecutionResult(
                tool_name=name,
                tool_call_id=tc_id,
                status="success",
                screenshot_bytes=b"pc-after",
                ui_elements_text="PC controls",
                text_summary="PC result",
                metadata={"plugin_id": "pc"},
                raw_result=[
                    {"type": "text", "text": "PC controls"},
                    {
                        "type": "image_url",
                        "image_url": {"url": "data:image/png;base64,cGM="},
                    },
                ],
            )
        current.indexed_elements = [{"index": 1, "text": "Fresh button"}]
        current.indexed_points = [[10, 20]]
        current.latest_screenshot = "fresh-phone.png"
        return ToolExecutionResult(
            tool_name=name,
            tool_call_id=tc_id,
            status="success",
            text_summary="OK",
            screenshot_bytes=b"fresh-phone",
            ui_elements_text="Fresh button",
        )

    runner.executor = SimpleNamespace(
        action_tool_names={PC, "click", "wait_for_delay"},
        execute=AsyncMock(side_effect=execute),
    )

    async def capture(ctx, current, controller, **kwargs):
        current.indexed_elements = [{"index": 1, "text": "Fresh button"}]
        current.indexed_points = [[10, 20]]
        return "fresh-phone.png", b"fresh-phone", "Fresh button"

    capture_mock = AsyncMock(side_effect=capture)
    monkeypatch.setattr(flash, "capture_screenshot_and_parse_ui", capture_mock)
    return runner, state, capture_mock


def _call(name, call_id, **args):
    return {"name": name, "id": call_id, "args": args}


async def _dispatch(runner, state, calls, image=None, tree=""):
    messages = []
    turn = flash._TurnRecord()
    result = await runner._process_tool_calls(calls, state, messages, "", {}, image, tree, 0, turn)
    return result, messages, turn


@pytest.mark.asyncio
async def test_continuous_pc_calls_do_not_capture_or_resend_the_phone(cross_target):
    runner, state, capture = cross_target
    for i in range(3):
        result, messages, _ = await _dispatch(runner, state, [_call(PC, str(i))])
        assert result[1] is None
        assert "不是安卓画面" in result[2]
        assert "data:image/png;base64,cGM=" in str(messages)
        assert not state.indexed_elements and not state.indexed_points
        assert state.latest_screenshot is None
    capture.assert_not_awaited()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "args", [{"target": 1}, {"target": [100, 200], "target_description": "Button"}]
)
async def test_stale_phone_action_and_remainder_are_deferred_until_new_decision(cross_target, args):
    runner, state, capture = cross_target
    calls = [
        _call(PC, "pc"),
        _call("click", "tap", **args),
        _call("report_task_status", "end", status="completed"),
    ]
    result, messages, turn = await _dispatch(runner, state, calls, b"old-phone", "Old button")
    assert result[0] is None
    assert result[1:3] == (b"fresh-phone", "Fresh button")
    assert [c.args[0] for c in runner.executor.execute.call_args_list] == [PC]
    assert [m.tool_call_id for m in messages if isinstance(m, ToolMessage)] == ["pc", "tap", "end"]
    assert all("未执行" in str(m.content) for m in messages[-2:])
    assert [action[0] for action in turn.actions] == [PC]
    capture.assert_awaited_once()

    # A subsequent model decision sees the new indices; the old index snapshot
    # from the mixed batch must not leak into this new click.
    await _dispatch(runner, state, [_call("click", "fresh-tap", target=1)], *result[1:3])
    assert runner.executor.execute.call_args.kwargs["index_elements"][0]["text"] == "Fresh button"
    capture.assert_awaited_once()


@pytest.mark.asyncio
async def test_explicit_refresh_executes_wait_once_but_defers_same_turn_click(cross_target):
    runner, state, capture = cross_target
    await _dispatch(runner, state, [_call(PC, "pc")])
    result, messages, _ = await _dispatch(
        runner,
        state,
        [
            _call("wait_for_delay", "observe", time_in_ms=0),
            _call("click", "tap", target=1),
        ],
    )
    assert [c.args[0] for c in runner.executor.execute.call_args_list] == [PC, "wait_for_delay"]
    assert result[1:3] == (b"fresh-phone", "Fresh button")
    assert messages[-1].status == "error" and messages[-1].tool_call_id == "tap"
    assert not runner._android_observation_stale
    capture.assert_not_awaited()  # The wait result already contains a fresh phone observation.


@pytest.mark.asyncio
async def test_pc_loading_wait_keeps_phone_stale_and_next_phone_click_is_deferred(
    cross_target, monkeypatch
):
    runner, state, capture = cross_target
    sleep = AsyncMock()
    monkeypatch.setattr(flash.asyncio, "sleep", sleep)
    await _dispatch(runner, state, [_call(PC, "pc")])
    result, messages, turn = await _dispatch(
        runner, state, [_call("wait_for_delay", "wait", time_in_ms=5000)]
    )
    sleep.assert_awaited_once_with(5)
    assert runner._android_observation_stale and result[1] is None
    assert "未采集新画面" in messages[-1].content
    assert turn.actions[0][0] == "wait_for_delay"
    capture.assert_not_awaited()
    assert runner.executor.execute.await_count == 1
    await _dispatch(runner, state, [_call("click", "tap", target=1)])
    capture.assert_awaited_once()
    assert runner.executor.execute.await_count == 1


@pytest.mark.asyncio
async def test_failed_pc_call_also_invalidates_the_phone_without_a_capture(cross_target):
    runner, state, capture = cross_target
    runner.executor.execute.side_effect = RuntimeError("PC connection lost")
    result, messages, _ = await _dispatch(runner, state, [_call(PC, "pc")])
    assert result[1] is None and runner._android_observation_stale
    assert state.latest_screenshot is None and state.indexed_elements == []
    assert messages[0].status == "error"
    capture.assert_not_awaited()


@pytest.mark.asyncio
async def test_failed_phone_refresh_never_dispatches_stale_click(cross_target):
    runner, state, capture = cross_target
    await _dispatch(runner, state, [_call(PC, "pc")])
    capture.side_effect = DeviceObservationError("phone disconnected")
    with pytest.raises(DeviceObservationError):
        await _dispatch(runner, state, [_call("click", "tap", target=1)])
    assert runner.executor.execute.await_count == 1
    assert runner._android_observation_stale


@pytest.mark.asyncio
@pytest.mark.parametrize("intervention", ["guidance", "resume", "resume_with_guidance"])
async def test_user_intervention_refreshes_phone_before_next_model_decision(
    cross_target, tmp_path, intervention
):
    runner, state, capture = cross_target
    runner.max_turns, runner.goal = 3, "cross-target task"
    runner._init_llm = Mock(return_value=Mock())
    runner._get_tools = Mock(return_value=[SimpleNamespace(name="report_task_status")])
    runner._build_ledger = Mock(return_value=TranscriptLedger())
    runner._render_system_prompt = Mock(return_value="test")
    runner._read_injected_instruction = AsyncMock(
        side_effect=[
            (None, None),
            ("按新指导继续", None) if intervention != "resume" else (None, None),
        ]
    )
    if intervention.startswith("resume"):
        control = Mock(spec=TaskControl)
        control.directory = tmp_path
        control.checkpoint = AsyncMock(side_effect=[False, True])

        async def uninterrupted(operation, **kwargs):
            return False, await operation

        control.run_interruptibly = AsyncMock(side_effect=uninterrupted)
        runner.ctx.task_control = control
    seen = []
    answers = [
        AIMessage(content="Observe PC", tool_calls=[_call(PC, "pc")]),
        AIMessage(
            content="Done", tool_calls=[_call("report_task_status", "end", status="completed")]
        ),
    ]

    async def model(llm, tools, messages):
        seen.append(messages[-1])
        return answers.pop(0)

    runner._invoke_model = AsyncMock(side_effect=model)
    result = await runner._run(state)
    assert result["status"] == "completed"
    assert capture.await_count == 2  # Initial phone observation + one refresh for the intervention.
    assert "Fresh button" in str(seen[-1].content)
    assert "不是安卓画面" not in str(seen[-1].content)
    assert not runner._android_observation_stale