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

"""Regression coverage for operator recovery and task pause handling."""

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

from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
import pytest

from artemis.agents.operator.operator import OperatorNode, OperatorNoProgressError
from artemis.config.settings import settings
from artemis.runtime import task_control


@pytest.fixture
def operator(monkeypatch):
    node = object.__new__(OperatorNode)
    node.ctx = SimpleNamespace(transcript_ledger=None)
    node._available_device_actions = lambda: {"click", "long_press", "input_text"}
    node._plan_changed_this_turn = lambda: True
    node._max_burst_actions = lambda: 4
    state = SimpleNamespace(
        operator_raw_data={"width": 1280, "height": 2856},
        indexed_elements=[{"center": [310, 2610], "text": "model menu", "resource_id": "model"}]
        * 13,
        open_incident=None,
    )
    monkeypatch.setattr(settings, "ARTEMIS_OPERATOR_MAX_NO_PROGRESS_ATTEMPTS", 3)
    return node, state


@pytest.mark.parametrize("name", ["click", "long_press", "input_text"])
@pytest.mark.parametrize("target", ["13", "[240, 914]", "240, 914"])
def test_quoted_targets_keep_same_action_and_grounding(operator, name, target):
    node, state = operator
    expected_target = 13 if target == "13" else [240, 914]
    args = {"target": target, "target_description": "model menu", "text": "test"}
    actions, error = node._translate_and_validate_tool({"name": name, "args": args}, state)
    expected, _ = node._translate_and_validate_tool(
        {"name": name, "args": {**args, "target": expected_target}}, state
    )
    assert error is None
    assert actions == expected
    if target == "13":
        assert actions[0]["target_resource_id"] == "model"
    else:
        assert actions[0]["target_description"] == "model menu"


@pytest.mark.parametrize("target", ["0", "14", "model menu", True, "true"])
def test_invalid_index_is_not_guessed_or_dispatched(operator, target):
    node, state = operator
    actions, error = node._translate_and_validate_tool(
        {"name": "click", "args": {"target": target}}, state
    )
    assert not actions
    assert error


def test_quoted_coordinates_still_need_description(operator):
    node, state = operator
    actions, error = node._translate_and_validate_tool(
        {"name": "click", "args": {"target": "[240, 914]"}}, state
    )
    assert not actions
    assert "target_description" in error


def click(target="13"):
    return AIMessage(
        content="Open the model menu.",
        tool_calls=[{"name": "click", "args": {"target": target}, "id": "call_click"}],
    )


async def run_loop(operator, monkeypatch, responses):
    node, state = operator
    invoke = AsyncMock(side_effect=responses)
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    messages = []
    result = await node._invoke_llm_loop(Mock(), messages, [], [], state)
    return result, messages, invoke


@pytest.mark.asyncio
async def test_prose_retry_keeps_context_and_recovers_once(operator, monkeypatch):
    prose = AIMessage(content="I will click the model menu.")
    result, messages, invoke = await run_loop(operator, monkeypatch, [prose, click()])
    assert invoke.await_count == 2
    assert len(result[0]) == 1
    assert result[0][0]["coordinates"] == [310, 2610]
    assert prose in messages
    assert any(isinstance(m, SystemMessage) and "No tool call" in m.content for m in messages)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "scenario", ["complete", "pending", "unfinished", "unchanged", "failed", "continuous", "pro"]
)
async def test_completed_exploration_plan_does_not_force_an_extra_action(
    operator, monkeypatch, scenario
):
    node, state = operator
    node.ctx.exploration_id = None if scenario == "pro" else "exploration"
    node.ctx.data_engine = None
    plan = {"text": "- [ ] Return home"}
    if scenario == "unchanged":
        plan["text"] = "- [x] Return home"
    node._read_task_plan_text = lambda: plan["text"]
    monkeypatch.setattr("artemis.exploration.runtime.publish_response", Mock())
    monkeypatch.setattr(
        "artemis.exploration.runtime.needs_exploration_turn", lambda ctx: scenario == "pending"
    )

    async def save_note(**kwargs):
        plan["text"] = "- [ ] Return home" if scenario == "unfinished" else "- [x] Return home"
        if scenario == "continuous":
            plan["text"] += " [Loop:continuous]"
        return ToolMessage(
            content="Saved plan",
            tool_call_id="plan",
            status="error" if scenario == "failed" else "success",
        )

    monkeypatch.setattr("artemis.agents.operator.operator.invoke_tool_with_injection", save_note)
    response = AIMessage(
        content="Confirmed home; recorded evidence.",
        tool_calls=[{"name": "save_note", "args": {}, "id": "plan"}],
    )
    final = AIMessage(content="# 完整调查报告\n已回主页，证据见步骤 3。未检查其他页面。")
    finishing = scenario in {"complete", "unchanged"}
    invoke = AsyncMock(side_effect=[response, final if finishing else click()])
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    messages = []
    result = await node._invoke_llm_loop(
        Mock(), messages, [SimpleNamespace(name="save_note")], [], state
    )
    if finishing:
        assert invoke.await_count == 2
        assert result[0] is None
        assert final.content in result[1]
        assert state.execution_report == final.content
        assert any("完整报告正文" in m.content for m in messages if isinstance(m, SystemMessage))
        assert result[3] is False
    elif scenario == "pending":
        assert invoke.await_count == 1
        assert result[0] is None
        assert state.execution_report is None
        return
    else:
        assert invoke.await_count == 2
        assert result[0]
        assert state.execution_report is None
    assert any(isinstance(m, ToolMessage) and m.tool_call_id == "plan" for m in messages)


@pytest.mark.asyncio
async def test_final_report_is_discarded_when_guidance_arrives(operator, monkeypatch):
    node, state = operator
    state.execution_report = "旧报告"
    node.ctx.exploration_id = "exploration"
    node._read_task_plan_text = lambda: "- [x] 调查完成"
    pending = {"value": False}
    monkeypatch.setattr(
        "artemis.exploration.runtime.needs_exploration_turn", lambda ctx: pending["value"]
    )
    publish = Mock()
    monkeypatch.setattr("artemis.exploration.runtime.publish_response", publish)

    async def answer(*args):
        pending["value"] = True
        return AIMessage(content="这是过时的最终报告")

    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", answer)
    result = await node._invoke_llm_loop(Mock(), [], [], [], state)
    assert result[0] is None
    assert state.execution_report is None
    publish.assert_not_called()


@pytest.mark.asyncio
async def test_empty_final_report_uses_bounded_recovery(operator, monkeypatch):
    node, state = operator
    node.ctx.exploration_id = "exploration"
    node._read_task_plan_text = lambda: "- [x] 调查完成"
    monkeypatch.setattr("artemis.exploration.runtime.needs_exploration_turn", lambda ctx: False)
    monkeypatch.setattr("artemis.exploration.runtime.publish_response", Mock())
    monkeypatch.setattr("artemis.agents.operator.operator.get_llm", Mock(return_value=Mock()))
    monkeypatch.setattr(settings, "ARTEMIS_TOOL_FALLBACK_MAX_NO_PROGRESS_ATTEMPTS", 2)
    invoke = AsyncMock(return_value=AIMessage(content="  "))
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    with pytest.raises(OperatorNoProgressError):
        await node._invoke_llm_loop(Mock(), [], [], [], state)
    assert state.execution_report is None
    assert invoke.await_count == 5


@pytest.mark.asyncio
async def test_invalid_target_can_be_corrected_before_limit(operator, monkeypatch):
    result, messages, invoke = await run_loop(operator, monkeypatch, [click("14"), click("13")])
    assert invoke.await_count == 2
    assert len(result[0]) == 1
    assert result[0][0]["coordinates"] == [310, 2610]
    outputs = [m for m in messages if isinstance(m, ToolMessage)]
    assert [m.status for m in outputs] == ["error", "success"]


@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [1, 3])
async def test_empty_responses_fail_at_configured_limit(operator, monkeypatch, limit):
    monkeypatch.setattr(settings, "ARTEMIS_OPERATOR_MAX_NO_PROGRESS_ATTEMPTS", limit)
    monkeypatch.setattr(settings, "ARTEMIS_TOOL_FALLBACK_MAX_NO_PROGRESS_ATTEMPTS", 3)
    fallback = Mock()
    get_fallback = Mock(return_value=fallback)
    monkeypatch.setattr("artemis.agents.operator.operator.get_llm", get_fallback)
    invoke = AsyncMock(return_value=AIMessage(content="I will click now."))
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    node, state = operator
    with pytest.raises(OperatorNoProgressError, match="未提供可执行的工具调用"):
        await node._invoke_llm_loop(Mock(), [], [], [], state)
    assert invoke.await_count == limit + 3
    get_fallback.assert_called_once_with(ctx=node.ctx, name="operator", use_fallback=True)
    assert invoke.call_args_list[-1].args[0] is fallback.bind_tools.return_value


@pytest.mark.asyncio
async def test_operator_fallback_is_scoped_to_current_step(operator, monkeypatch):
    node, state = operator
    primary = Mock()
    fallback = Mock()
    get_fallback = Mock(return_value=fallback)
    monkeypatch.setattr("artemis.agents.operator.operator.get_llm", get_fallback)
    invoke = AsyncMock(side_effect=[AIMessage(content="no tool")] * 3 + [click(), click()])
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)

    first = await node._invoke_llm_loop(primary, [], [], [], state)
    second = await node._invoke_llm_loop(primary, [], [], [], state)

    assert first[0] and second[0]
    assert get_fallback.call_count == 1
    assert [call.args[0] for call in invoke.call_args_list] == [
        primary.bind_tools.return_value,
        primary.bind_tools.return_value,
        primary.bind_tools.return_value,
        fallback.bind_tools.return_value,
        primary.bind_tools.return_value,
    ]


@pytest.mark.asyncio
async def test_invalid_calls_and_prose_share_bounded_recovery(operator, monkeypatch):
    invoke = AsyncMock(
        side_effect=[click("missing"), AIMessage(content="Retry."), click("14")] + [click("14")] * 3
    )
    monkeypatch.setattr("artemis.agents.operator.operator.get_llm", Mock(return_value=Mock()))
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    node, state = operator
    messages = []
    with pytest.raises(OperatorNoProgressError, match="校验仍失败"):
        await node._invoke_llm_loop(Mock(), messages, [], [], state)
    assert invoke.await_count == 6
    assert any(isinstance(m, ToolMessage) and m.status == "error" for m in messages)
    assert not any(isinstance(m, ToolMessage) and m.content == "Action Recorded" for m in messages)


@pytest.mark.asyncio
async def test_iteration_cap_cannot_restart_no_progress_next_turn(operator, monkeypatch):
    monkeypatch.setattr("artemis.agents.operator.operator.OPERATOR_MAX_TOOL_ITERATIONS", 2)
    invoke = AsyncMock(return_value=AIMessage(content="I will click now."))
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    node, state = operator
    with pytest.raises(OperatorNoProgressError, match="预算内未恢复"):
        await node._invoke_llm_loop(Mock(), [], [], [], state)
    assert invoke.await_count == 2


@pytest.mark.asyncio
async def test_pause_during_diagnosis_returns_to_fresh_observation(operator, monkeypatch, tmp_path):
    node, state = operator
    monkeypatch.setattr(task_control, "TRACES_PATH", tmp_path)
    started = asyncio.Event()
    cancelled = asyncio.Event()
    events = []

    class Lock:
        def __init__(self, device_id, **kwargs):
            self.device_id = device_id
            self.description = "test"
            self.max_concurrency = None
            self.concurrency_mode = "per_device"
            self.ingress = "web"
            self.lock_scope = "test"

        def release(self):
            events.append("released")

        def acquire(self, cancel_event=None):
            events.append("acquired")

    monkeypatch.setattr(task_control, "DeviceExecutionLock", Lock)
    node.ctx.data_engine = None
    node.ctx.task_control = task_control.TaskControl("pro-task", Lock("device"))
    node.ctx.on_task_pause = AsyncMock()
    node.ctx.on_task_resume = AsyncMock()

    async def diagnose(**kwargs):
        started.set()
        try:
            await asyncio.Event().wait()
        finally:
            cancelled.set()

    monkeypatch.setattr("artemis.agents.operator.operator.invoke_tool_with_injection", diagnose)
    response = AIMessage(
        content="Inspect the failure",
        tool_calls=[
            {"name": "ask_diagnoser", "args": {"query": "why"}, "id": "diagnosis"},
            {"name": "click", "args": {"target": 13}, "id": "tap"},
        ],
    )
    invoke = AsyncMock(return_value=response)
    monkeypatch.setattr("artemis.agents.operator.operator.acomplete", invoke)
    tool = SimpleNamespace(name="ask_diagnoser")
    messages = []
    execution = asyncio.create_task(node._invoke_llm_loop(Mock(), messages, [tool], [], state))
    await asyncio.wait_for(started.wait(), 1)
    task_control.request_pause("pro-task")
    for _ in range(40):
        if task_control.read_state("pro-task").get("status") == "task_paused":
            break
        await asyncio.sleep(0.01)
    assert task_control.read_state("pro-task")["status"] == "task_paused"
    assert cancelled.is_set()
    assert events == ["released"]
    task_control.request_resume("pro-task")
    assert (await asyncio.wait_for(execution, 2))[0] is None
    assert events == ["released", "acquired"]
    assert invoke.await_count == 1
    assert [m.status for m in messages if isinstance(m, ToolMessage)] == ["error", "error"]