151 lines
5.7 KiB
Python
151 lines
5.7 KiB
Python
# 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.
|
||
|
||
"""replay_steps: range replay, per-call step cap, whole-step token budget,
|
||
``[Screen]`` status lines, and the no-history degradation."""
|
||
|
||
from types import SimpleNamespace
|
||
from unittest.mock import MagicMock
|
||
|
||
import pytest
|
||
|
||
from artemis.tools.history import ReplayStepsTool, replay_steps, replay_steps_text
|
||
|
||
|
||
def _friendly_step(number, tool_result="The toggle is ON", **overrides):
|
||
step = {
|
||
"step_id": f"step_{number}",
|
||
"step_number": number,
|
||
"relative_time": f"{number * 3}.0s",
|
||
"summary": f"Step {number} summary.",
|
||
"action_taken": {"action": "click", "target": [0.5, 0.5], "target_text": "Save"},
|
||
"last_execution_result": {"status": "success"},
|
||
"interleaved_events": [
|
||
{"type": "thought", "content": f"Thinking at step {number}."},
|
||
{
|
||
"type": "tool_call",
|
||
"name": "ask_explorer",
|
||
"args": {"question": "toggle?"},
|
||
"result": tool_result,
|
||
},
|
||
],
|
||
}
|
||
step.update(overrides)
|
||
return step
|
||
|
||
|
||
def _reader(steps):
|
||
reader = MagicMock()
|
||
reader.get_agent_friendly_steps_in_range.side_effect = lambda s, e: [
|
||
st for st in steps if s <= st["step_number"] <= e
|
||
]
|
||
return reader
|
||
|
||
|
||
def test_replay_range_renders_every_step_in_full():
|
||
reader = _reader([_friendly_step(3), _friendly_step(4, tool_result="The toggle is OFF " * 100)])
|
||
out = replay_steps_text(reader, 3, 4)
|
||
|
||
reader.get_agent_friendly_steps_in_range.assert_called_once_with(3, 4)
|
||
assert "- **Step 3 (Start: 9.0s)**" in out
|
||
assert "- **Step 4 (Start: 12.0s)**" in out
|
||
assert "[Screen]: Step 3 summary." in out
|
||
assert "Thinking at step 4." in out
|
||
assert "`ask_explorer(" in out
|
||
# Loose clamp: a ~1.8k-char tool result is replayed in full.
|
||
assert "The toggle is OFF " * 100 in out
|
||
assert "[Planned Action]: Tapped 'Save' at [0.5, 0.5]" in out
|
||
|
||
|
||
def test_replay_single_step_reversed_bounds_and_missing_step():
|
||
reader = _reader([_friendly_step(2)])
|
||
assert "- **Step 2 (Start: 6.0s)**" in replay_steps_text(reader, 2)
|
||
reader.get_agent_friendly_steps_in_range.assert_called_once_with(2, 2)
|
||
|
||
reader = _reader([_friendly_step(2), _friendly_step(3)])
|
||
out = replay_steps_text(reader, 3, 2)
|
||
reader.get_agent_friendly_steps_in_range.assert_called_once_with(2, 3)
|
||
assert out.index("**Step 2") < out.index("**Step 3")
|
||
|
||
assert replay_steps_text(_reader([]), 9) == "Error: step 9 not found."
|
||
assert replay_steps_text(_reader([]), 9, 12) == "Error: no recorded steps in range 9–12."
|
||
assert "must be integers" in replay_steps_text(_reader([]), "x")
|
||
|
||
|
||
def test_replay_caps_the_range_per_call():
|
||
reader = _reader([_friendly_step(n) for n in range(1, 30)])
|
||
out = replay_steps_text(reader, 1, 20, max_steps=5)
|
||
reader.get_agent_friendly_steps_in_range.assert_called_once_with(1, 5)
|
||
assert "**Step 5" in out
|
||
assert "**Step 6" not in out
|
||
assert "Call again from Step 6" in out
|
||
|
||
|
||
def test_replay_token_budget_drops_whole_trailing_steps():
|
||
big = "x" * 4000 # ~1k tokens per step
|
||
reader = _reader([_friendly_step(n, tool_result=big) for n in range(1, 6)])
|
||
out = replay_steps_text(reader, 1, 5, max_steps=5, max_tokens=2500)
|
||
|
||
# Two steps fit; the rest are dropped as whole steps, never cut mid-step.
|
||
assert "**Step 1" in out and "**Step 2" in out
|
||
assert "**Step 3" not in out
|
||
assert out.count(big) == 2
|
||
assert "Token budget reached: Steps 1–2 replayed; 3 more step(s)" in out
|
||
assert "Call again from Step 3" in out
|
||
|
||
|
||
def test_replay_always_shows_at_least_the_first_step():
|
||
reader = _reader([_friendly_step(1, tool_result="y" * 4000), _friendly_step(2)])
|
||
out = replay_steps_text(reader, 1, 2, max_tokens=10)
|
||
assert "**Step 1" in out
|
||
assert "y" * 4000 in out
|
||
assert "**Step 2" not in out
|
||
|
||
|
||
def test_replay_screen_line_states():
|
||
with_text = replay_steps_text(_reader([_friendly_step(1)]), 1)
|
||
assert "[Screen]: Step 1 summary." in with_text
|
||
assert "[Summary]" not in with_text
|
||
|
||
pending = replay_steps_text(
|
||
_reader([_friendly_step(1, summary=None, extra_metadata={"summary_status": "pending"})]),
|
||
1,
|
||
)
|
||
assert "[Screen]: (screen description pending)" in pending
|
||
|
||
failed = replay_steps_text(
|
||
_reader([_friendly_step(1, summary="", extra_metadata={"summary_status": "failed"})]),
|
||
1,
|
||
)
|
||
assert "[Screen]: (screen description unavailable)" in failed
|
||
|
||
no_status = replay_steps_text(_reader([_friendly_step(1, summary=None)]), 1)
|
||
assert "[Screen]" not in no_status
|
||
|
||
|
||
def test_replay_without_history():
|
||
assert "no execution history" in replay_steps_text(None, 1)
|
||
|
||
|
||
@pytest.mark.asyncio
|
||
async def test_tool_execute_reads_the_context_engine_and_degrades_without_one():
|
||
reader = _reader([_friendly_step(2)])
|
||
out = await ReplayStepsTool().execute(ctx=SimpleNamespace(data_engine=reader), start_step=2)
|
||
assert "- **Step 2 (Start: 6.0s)**" in out
|
||
|
||
out = await replay_steps.execute(ctx=SimpleNamespace(data_engine=None), start_step=1)
|
||
assert "no execution history" in out
|
||
assert replay_steps.is_available(SimpleNamespace(data_engine=None)) is False
|
||
assert replay_steps.is_available(SimpleNamespace(data_engine=reader)) is True
|