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

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Operator prompt snapshots and transcript mode.

Flag off (``agent.memory.transcript.enabled=false``): the legacy 2-message
build must stay byte-for-byte identical — pinned by SHA-256 goldens recorded
for the current template.

Flag on: ``_build_prompt`` renders ``S + F + A + tail`` from the session
transcript ledger; the turn's messages are staged after the tool loop and
committed (with step key and validator result) at the next build.
"""

import base64
import hashlib
import re
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

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

from artemis.agents.operator.operator import OperatorNode
from artemis.agents.operator.prompts import load_operator_prompts
from artemis.agents.operator.prompts import (
    PLAN_HISTORY_STATIC_POINTER,
    PLAN_HISTORY_TEMPLATE_SECTION,
    PromptBuilder,
    TemplatePromptComponent,
    render_transcript_static_system,
)
from artemis.config.agent import MemoryTranscriptConfig
from artemis.context import ArtemisContext
from artemis.memory.step_memory import StepMemoryService
from artemis.memory.transcript import (
    EPHEMERAL_BLOCKS_KEY,
    EXECUTION_RESULT_MARKER,
    PLAN_RECITATION_MARKER,
    RESTORED_HISTORY_HEADER,
    TranscriptLedger,
)

# SHA-256 snapshots of the legacy system message with the fixed inputs below.
# Update these when an intentional template change alters the rendered prompt.
GOLDEN_EMPTY_PLAN = "e83c2c6cc1e28b04898e6602194cb2fbac67fe5a3aec8e109872c375f4e59944"
GOLDEN_SENTINEL_PLAN = "dba4eb2cdcd9389148ac3228c6e9e69fdbe8d407dec76b58aa65509ad768c392"

SCREENSHOT_B64 = base64.b64encode(b"fake-jpeg-bytes").decode("utf-8")


async def _render_legacy_system(plan_and_history: str) -> str:
    builder = PromptBuilder()
    component = TemplatePromptComponent()
    ctx = SimpleNamespace(execution_setup=None, actuator=None, data_engine=None)
    state = SimpleNamespace(initial_goal="GOLDEN_FIXED_GOAL")
    await component(
        builder,
        state,
        ctx,
        prompts=load_operator_prompts(),
        template_name="main_template",
        plan_and_history=plan_and_history,
    )
    return "".join(builder.system_parts)


@pytest.mark.asyncio
async def test_legacy_prompt_bytes_match_chinese_prompt_golden():
    empty = await _render_legacy_system("No plan or history yet.")
    sentinel = await _render_legacy_system("PLAN_AND_HISTORY_SENTINEL\nline2")
    assert hashlib.sha256(empty.encode("utf-8")).hexdigest() == GOLDEN_EMPTY_PLAN
    assert hashlib.sha256(sentinel.encode("utf-8")).hexdigest() == GOLDEN_SENTINEL_PLAN


def test_main_template_contains_exactly_one_plan_history_section():
    template = load_operator_prompts()["main_template"]
    assert template.count(PLAN_HISTORY_TEMPLATE_SECTION) == 1


def test_transcript_static_system_is_stable_and_carries_no_history():
    ctx = SimpleNamespace(execution_setup=None, actuator=None, data_engine=None)
    state = SimpleNamespace(initial_goal="STATIC_GOAL")
    prompts = load_operator_prompts()

    first = render_transcript_static_system(prompts, ctx, state)
    second = render_transcript_static_system(prompts, ctx, state)
    assert first == second  # byte-stable S region

    assert "STATIC_GOAL" in first
    # The volatile section is replaced by the static pointer...
    assert PLAN_HISTORY_STATIC_POINTER.splitlines()[1] in first
    # ...and no unrendered slot or trailing observation header survives.
    assert "{{" not in first
    assert not first.endswith("# CURRENT OBSERVATION\n")
    # The history guide names the marker a resolved screenshot renders as.
    assert "--- Historical Visual Transition ---" in first


def _transcript_ctx():
    ctx = MagicMock(spec=ArtemisContext)
    ctx.execution_setup = None
    ctx.actuator = None
    ctx.data_engine = None
    ctx.step_memory = StepMemoryService(ctx=None)
    ctx.transcript_ledger = None
    return ctx


def _transcript_state(**overrides):
    state = MagicMock()
    state.subagent_calls = []
    state.initial_goal = "Transcript goal"
    state.operator_feedback = None
    state.injected_instruction = None
    state.operator_tool_limit_exceeded = False
    state.structured_decisions = None
    state.last_execution_result = None
    state.current_step_id = None
    state.latest_ui_hierarchy = []
    state.operator_raw_data = {
        "screenshot_b64": SCREENSHOT_B64,
        "xml_hierarchy": [],
        "ocr_results": None,
        "width": 1080,
        "height": 2400,
    }
    for key, value in overrides.items():
        setattr(state, key, value)
    return state


@pytest.mark.asyncio
@pytest.mark.parametrize("exploration", [False, True])
@pytest.mark.parametrize("transcript", [False, True])
async def test_plugin_result_is_delivered_with_pc_image_and_separate_phone_observation(
    tmp_path, exploration, transcript, monkeypatch
):
    from langchain_openai import ChatOpenAI
    from PIL import Image

    ctx = _transcript_ctx()
    ctx.exploration_id = "explore" if exploration else None
    monkeypatch.setattr(
        "artemis.exploration.runtime.exploration_messages",
        lambda _: [SystemMessage(content="Explore"), HumanMessage(content="Guidance")],
    )
    ctx.data_engine = MagicMock()
    ctx.data_engine.base_dir = str(tmp_path)
    pc = tmp_path / "pc.jpg"
    Image.new("RGB", (200, 100), "blue").save(pc)
    ctx.data_engine.get_image_path.return_value = pc
    node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=transcript))
    ledger = node._ensure_transcript_ledger()
    tool = "plug_test_get_app_state"
    ledger.stage_turn(
        [
            AIMessage(content="", tool_calls=[{"name": tool, "args": {}, "id": "pc_call"}]),
            ToolMessage(content="Action Recorded", tool_call_id="pc_call"),
        ]
    )
    report = {
        "status": "dispatched",
        "execution": [
            {
                "action": tool,
                "tool_call_id": "pc_call",
                "plugin_id": "test",
                "plugin_status": "success",
                "result": "PC Address bar",
                "post_image_name": "pc",
            }
        ],
    }
    state = _transcript_state(structured_decisions="[]", last_execution_result=report)
    messages = await node._build_prompt(state, SCREENSHOT_B64, [], "PHONE_ELEMENTS", 2, [], "")
    if transcript:
        output = next(m for m in messages if isinstance(m, ToolMessage))
        assert output.tool_call_id == "pc_call" and "PC Address bar" in str(output.content)
    else:
        assert "PC Address bar" in str(messages)
    assert base64.b64encode(pc.read_bytes()).decode() in str(messages)
    assert SCREENSHOT_B64 not in str(messages) and "PHONE_ELEMENTS" in str(messages)
    wire = ChatOpenAI(model="fixture", api_key="fixture", use_responses_api=True, store=False)
    payload = wire._get_request_payload(messages)["input"]
    assert len([x for x in payload if x.get("type") == "function_call_output"]) == int(transcript)
    guided = await node._build_prompt(
        _transcript_state(
            structured_decisions="[]", last_execution_result=report, injected_instruction="回到手机"
        ),
        SCREENSHOT_B64,
        [],
        "PHONE_ELEMENTS",
        3,
        [],
        "",
    )
    assert SCREENSHOT_B64 in str(guided)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "exploration,plugins", [(False, False), (True, False), (False, True), (True, True)]
)
async def test_injected_prefix_and_compression_preserve_tool_pairs(
    monkeypatch, exploration, plugins
):
    from langchain_openai import ChatOpenAI

    ctx = _transcript_ctx()
    ctx.exploration_id = "test-exploration" if exploration else None
    ctx.plugins = SimpleNamespace(prompt=lambda: "Plugin instructions") if plugins else None
    monkeypatch.setattr(
        "artemis.exploration.runtime.exploration_messages",
        lambda _: [
            SystemMessage(content="Exploration instructions"),
            HumanMessage(content="User guidance"),
        ],
    )
    monkeypatch.setattr("artemis.agents.operator.operator.get_llm", lambda **_: MagicMock())
    node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))
    requests = []
    wire_model = ChatOpenAI(model="test", api_key="test", use_responses_api=True, store=False)

    async def tool_loop(current_messages, **_kwargs):
        messages = current_messages
        payload = wire_model._get_request_payload(messages)
        requests.append(payload["input"])
        call_id = f"call_{len(requests)}"
        messages.extend(
            [
                AIMessage(
                    content="", tool_calls=[{"name": "click", "args": {"target": 1}, "id": call_id}]
                ),
                ToolMessage(content="Action Recorded", tool_call_id=call_id),
            ]
        )
        return [{"action": "tap"}], None, None, False

    monkeypatch.setattr(node, "_invoke_llm_loop", tool_loop)
    for turn in range(8):
        await node(_transcript_state(current_step_id=f"step-{turn}"))
        if turn == 4:
            ctx.transcript_ledger.freeze_turns(2, [HumanMessage(content="Summarized first turns")])

    for items in requests:
        pending, seen = set(), set()
        for item in items:
            if item.get("type") == "function_call":
                assert item["call_id"] not in seen
                pending.add(item["call_id"])
                seen.add(item["call_id"])
            elif item.get("type") == "function_call_output":
                assert item["call_id"] in pending
                pending.remove(item["call_id"])
        assert not pending
    assert {item["call_id"] for item in requests[-1] if item.get("type") == "function_call"} == {
        f"call_{n}" for n in range(3, 8)
    }


def _action_llm(captured: list):
    mock_llm = MagicMock()
    mock_response = MagicMock()
    mock_response.tool_calls = [
        {
            "name": "click",
            "args": {"target": [50, 50], "target_description": "button"},
            "id": "call_click",
        }
    ]
    mock_response.content = "Click the button."

    async def mock_ainvoke(*args, **kwargs):
        captured.append(list(args[0]))
        return mock_response

    mock_llm.ainvoke = AsyncMock(side_effect=mock_ainvoke)
    mock_llm.bind_tools.return_value = mock_llm
    return mock_llm


@pytest.mark.asyncio
async def test_transcript_mode_two_turns_build_four_regions(tmp_path):
    ctx = _transcript_ctx()
    # Turn 2 has a recorded step but no committed ledger turn yet. This is
    # an active session, so it must not trigger cold-start history seeding.
    ctx.data_engine = MagicMock()
    ctx.data_engine.base_dir = str(tmp_path)
    ctx.data_engine.global_base_dir = str(tmp_path)
    ctx.data_engine.get_agent_friendly_steps.return_value = []
    captured: list = []
    mock_llm = _action_llm(captured)

    with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
        node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))

        # Turn 1
        await node(_transcript_state())
        ledger = ctx.transcript_ledger
        assert isinstance(ledger, TranscriptLedger)
        assert ledger.has_static_prefix
        assert ledger.has_staged_turn
        assert ledger.turn_count == 0

        turn1_messages = captured[0]
        assert isinstance(turn1_messages[0], SystemMessage)
        assert isinstance(turn1_messages[-1], HumanMessage)
        tail1_text = str(turn1_messages[-1].content)
        assert "# CURRENT OBSERVATION [T+" in tail1_text
        assert PLAN_RECITATION_MARKER in tail1_text

        # Turn 2: step 1 is now recorded, and the previous turn is committed
        # with its step id and result.
        ctx.data_engine.get_agent_friendly_steps.return_value = [
            {
                "step_id": "step-1",
                "step_number": 1,
                "summary": "Step 1 summary",
                "relative_time": "1.0s",
                "extra_metadata": {"subgoal_hash": "default"},
            }
        ]
        await node(
            _transcript_state(
                current_step_id="step-1",
                structured_decisions='[{"action": "tap"}]',
                last_execution_result={"status": "dispatched"},
            )
        )

        # The staged-but-uncommitted turn 2 state never triggers cold-start
        # seeding (the transcript path stayed live instead of falling back).
        assert not ledger.has_restored_history
        assert ledger.turn_count == 1
        turn2_messages = captured[1]
        # S region is the same object with the same bytes.
        assert turn2_messages[0] is turn1_messages[0]

        active = list(ledger.active_messages)
        # Committed turn-1 observation followed by its validator result.
        assert any(EXECUTION_RESULT_MARKER in str(m.content) for m in active)
        committed_obs = active[0]
        # Text edge at depth 1: the committed observation has lost its plan
        # recitation as soon as the next tail exists (its screenshot stays
        # until the screenshot edge).
        assert PLAN_RECITATION_MARKER in str(committed_obs.content)  # retained until compaction
        # The live tail still recites the plan.
        assert PLAN_RECITATION_MARKER in str(turn2_messages[-1].content)
        # Message layout is S + A + tail.
        assert turn2_messages[-1] is not turn1_messages[-1]
        assert re.search(r"T\+\d{2,}:\d{2}", str(turn2_messages[-1].content))


@pytest.mark.asyncio
async def test_transcript_tail_marks_per_turn_notices_ephemeral():
    """The hand-built tail carries the builder's ephemeral indices: per-turn
    notices are deleted by the scrub edge; header, plan recitation and the
    observation itself are not."""
    from artemis.agents.operator.prompts import USER_GUIDANCE_MARKER

    ctx = _transcript_ctx()
    captured: list = []
    mock_llm = _action_llm(captured)
    with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
        node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))
        await node(
            _transcript_state(
                injected_instruction="Skip the login step",
                user_stop_requested=False,
                operator_tool_limit_exceeded=True,
            )
        )

    tail = captured[0][-1]
    assert isinstance(tail, HumanMessage)
    texts = {
        i: block.get("text", "")
        for i, block in enumerate(tail.content)
        if isinstance(block, dict) and block.get("type") == "text"
    }
    guidance = next(i for i, t in texts.items() if t.startswith(USER_GUIDANCE_MARKER))
    warning = next(i for i, t in texts.items() if "tool-call budget" in t)
    assert set(tail.additional_kwargs[EPHEMERAL_BLOCKS_KEY]) == {guidance, warning}
    # The user's stop bit is relayed on the guidance block itself.
    assert "authorized stopping" not in texts[guidance]
    # The verbatim instruction is its own persistent block right after the
    # wrapper (stamped with the ledger's session offset), never ephemeral:
    # a standing instruction must outlive the turn it arrived on.
    body = guidance + 1
    assert re.fullmatch(r'User instruction \(T\+\d{2,}:\d{2}\): "Skip the login step"', texts[body])
    assert body not in tail.additional_kwargs[EPHEMERAL_BLOCKS_KEY]
    assert "Skip the login step" not in texts[guidance]
    # Header, plan recitation and the observation stay in the history.
    assert texts[0].startswith("# CURRENT OBSERVATION")
    plan = next(i for i, t in texts.items() if t.startswith(PLAN_RECITATION_MARKER))
    screenshot = next(i for i, t in texts.items() if t == "--- Current Screenshot ---")
    assert not {0, plan, screenshot} & set(tail.additional_kwargs[EPHEMERAL_BLOCKS_KEY])


@pytest.mark.asyncio
async def test_transcript_actionless_turn_commits_without_validator_message():
    ctx = _transcript_ctx()
    captured: list = []
    mock_llm = _action_llm(captured)

    with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
        node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))
        await node(_transcript_state())
        # Planner-rejected turns clear structured_decisions: no validator ran.
        await node(
            _transcript_state(
                current_step_id="step-1",
                structured_decisions="",
                last_execution_result={"status": "dispatched"},
            )
        )

    ledger = ctx.transcript_ledger
    assert ledger.turn_count == 1
    assert not any(EXECUTION_RESULT_MARKER in str(m.content) for m in ledger.active_messages)


@pytest.mark.asyncio
async def test_transcript_cold_start_builds_restored_history_block(tmp_path):
    ctx = _transcript_ctx()
    ctx.data_engine = MagicMock()
    ctx.data_engine.base_dir = str(tmp_path)
    ctx.data_engine.global_base_dir = str(tmp_path)
    ctx.data_engine.get_agent_friendly_steps.return_value = [
        {
            "step_id": f"step_{i}",
            "step_number": i,
            "summary": f"Step {i} summary",
            "relative_time": f"{i}.0s",
            "extra_metadata": {"subgoal_hash": "default"},
        }
        for i in range(1, 4)
    ]

    captured: list = []
    mock_llm = _action_llm(captured)

    with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
        node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))
        await node(_transcript_state())

    ledger = ctx.transcript_ledger
    assert ledger.has_restored_history

    messages = captured[0]
    restored_text = str(messages[1].content)
    assert RESTORED_HISTORY_HEADER in restored_text
    assert "--- Execution History ---" in restored_text
    assert "Step 1 summary" in restored_text
    # The live tail still follows normally and A stays appendable.
    assert "# CURRENT OBSERVATION [T+" in str(messages[-1].content)
    assert ledger.has_staged_turn


@pytest.mark.asyncio
async def test_flag_off_keeps_legacy_two_message_build_and_no_ledger():
    ctx = _transcript_ctx()
    captured: list = []
    mock_llm = _action_llm(captured)

    with patch("artemis.agents.operator.operator.get_llm", return_value=mock_llm):
        node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=False))
        await node(_transcript_state())

    assert len(captured[0]) == 2
    assert isinstance(captured[0][0], SystemMessage)
    assert isinstance(captured[0][1], HumanMessage)
    assert ctx.transcript_ledger is None


@pytest.mark.asyncio
async def test_exploration_guidance_appends_without_rewriting_static_prefix(monkeypatch):
    ctx = _transcript_ctx()
    ctx.exploration_id = "explore"
    guidance = ["initial guidance"]
    monkeypatch.setattr(
        "artemis.exploration.runtime.exploration_messages",
        lambda _: [SystemMessage(content="rules"), HumanMessage(content=guidance[0])],
    )
    node = OperatorNode(ctx, transcript_config=MemoryTranscriptConfig(enabled=True))
    state = _transcript_state()
    first = await node._build_prompt_transcript(state, "", "", 1, [], "")
    prefix = first[0].content
    assert "initial guidance" not in prefix
    assert "initial guidance" in str(first[-1].content)
    node._stage_transcript_turn(first)
    second = await node._build_prompt_transcript(state, "", "", 2, [], "")
    assert second[0].content == prefix
    assert "initial guidance" not in str(second[-1].content)
    node._stage_transcript_turn(second)
    guidance[0] = "new direction"
    third = await node._build_prompt_transcript(state, "", "", 3, [], "")
    assert third[0].content == prefix
    assert "new direction" in str(third[-1].content)
    # An interrupted/uncommitted build must not acknowledge away the guidance.
    retried = await node._build_prompt_transcript(state, "", "", 3, [], "")
    assert "new direction" in str(retried[-1].content)
    node._stage_transcript_turn(retried)
    await node._build_prompt_transcript(state, "", "", 4, [], "")
    ctx.transcript_ledger.freeze_turns(3, [HumanMessage(content="compressed history")])
    compacted = await node._build_prompt_transcript(state, "", "", 4, [], "")
    assert compacted[0].content == prefix
    assert "new direction" in str(compacted[-1].content)