v2.7.0 (#534)
Co-authored-by: Quentin Torroba <quentin.torroba@mistral.ai> Co-authored-by: Clément Drouin <clement.drouin@mistral.ai> Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai> Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
parent
5ea69b547b
commit
6a50d1d521
49 changed files with 3726 additions and 99 deletions
308
tests/core/test_rewind_integration.py
Normal file
308
tests/core/test_rewind_integration.py
Normal file
|
|
@ -0,0 +1,308 @@
|
|||
"""Integration tests for the rewind feature through agent_loop.act().
|
||||
|
||||
These tests drive the full pipeline:
|
||||
act() → create_checkpoint() → _process_one_tool_call() →
|
||||
get_file_snapshot() → add_snapshot() → tool writes file →
|
||||
rewind_to_message() → verify file restoration.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.conftest import build_test_agent_loop, build_test_vibe_config
|
||||
from tests.mock.utils import mock_llm_chunk
|
||||
from tests.stubs.fake_backend import FakeBackend
|
||||
from vibe.core.agents.models import BuiltinAgentName
|
||||
from vibe.core.types import BaseEvent, FunctionCall, ToolCall
|
||||
|
||||
|
||||
async def _act_and_collect(agent_loop, prompt: str) -> list[BaseEvent]:
|
||||
return [ev async for ev in agent_loop.act(prompt)]
|
||||
|
||||
|
||||
def _write_file_tool_call(
|
||||
path: str, content: str, *, call_id: str = "call_1", overwrite: bool = False
|
||||
) -> ToolCall:
|
||||
args = json.dumps({"path": path, "content": content, "overwrite": overwrite})
|
||||
return ToolCall(
|
||||
id=call_id, index=0, function=FunctionCall(name="write_file", arguments=args)
|
||||
)
|
||||
|
||||
|
||||
def _search_replace_tool_call(
|
||||
file_path: str, search: str, replace: str, *, call_id: str = "call_1"
|
||||
) -> ToolCall:
|
||||
content = f"<<<<<<< SEARCH\n{search}\n=======\n{replace}\n>>>>>>> REPLACE"
|
||||
args = json.dumps({"file_path": file_path, "content": content})
|
||||
return ToolCall(
|
||||
id=call_id,
|
||||
index=0,
|
||||
function=FunctionCall(name="search_replace", arguments=args),
|
||||
)
|
||||
|
||||
|
||||
def _bash_tool_call(command: str, *, call_id: str = "call_1") -> ToolCall:
|
||||
args = json.dumps({"command": command})
|
||||
return ToolCall(
|
||||
id=call_id, index=0, function=FunctionCall(name="bash", arguments=args)
|
||||
)
|
||||
|
||||
|
||||
def _make_agent_loop(backend: FakeBackend):
|
||||
config = build_test_vibe_config(
|
||||
enabled_tools=["write_file", "search_replace", "bash"],
|
||||
tools={
|
||||
"write_file": {"permission": "always"},
|
||||
"search_replace": {"permission": "always"},
|
||||
"bash": {"permission": "always"},
|
||||
},
|
||||
system_prompt_id="tests",
|
||||
include_project_context=False,
|
||||
include_prompt_detail=False,
|
||||
)
|
||||
return build_test_agent_loop(
|
||||
config=config, agent_name=BuiltinAgentName.AUTO_APPROVE, backend=backend
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestRewindIntegration:
|
||||
async def test_write_file_rewind_restores_original(
|
||||
self, tmp_working_directory: Path
|
||||
) -> None:
|
||||
"""Write a file in turn 1, rewind to turn 1 → file should not exist."""
|
||||
target = tmp_working_directory / "hello.txt"
|
||||
|
||||
backend = FakeBackend([
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Creating file.",
|
||||
tool_calls=[_write_file_tool_call(str(target), "hello world")],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Done.")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "create hello.txt")
|
||||
assert target.read_text() == "hello world"
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
assert len(rewindable) == 1
|
||||
await rm.rewind_to_message(rewindable[0][0], restore_files=True)
|
||||
|
||||
assert not target.exists()
|
||||
|
||||
async def test_search_replace_rewind_restores_previous_version(
|
||||
self, tmp_working_directory: Path
|
||||
) -> None:
|
||||
"""Edit a pre-existing file with search_replace, rewind restores original."""
|
||||
target = tmp_working_directory / "config.yaml"
|
||||
target.write_text("key: original\n", encoding="utf-8")
|
||||
|
||||
backend = FakeBackend([
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Updating config.",
|
||||
tool_calls=[
|
||||
_search_replace_tool_call(
|
||||
str(target), "key: original", "key: modified"
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Updated.")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "update config")
|
||||
assert target.read_text() == "key: modified\n"
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
await rm.rewind_to_message(rewindable[0][0], restore_files=True)
|
||||
|
||||
assert target.read_text() == "key: original\n"
|
||||
|
||||
async def test_write_then_search_replace_rewind_to_middle(
|
||||
self, tmp_working_directory: Path
|
||||
) -> None:
|
||||
"""Turn 1 creates a file with write_file, turn 2 patches it with
|
||||
search_replace. Rewind to turn 2 restores the turn 1 version.
|
||||
"""
|
||||
target = tmp_working_directory / "app.py"
|
||||
|
||||
backend = FakeBackend([
|
||||
# Turn 1: create the file
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Creating.",
|
||||
tool_calls=[
|
||||
_write_file_tool_call(str(target), "def main():\n pass\n")
|
||||
],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Created.")],
|
||||
# Turn 2: patch with search_replace
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Updating.",
|
||||
tool_calls=[
|
||||
_search_replace_tool_call(
|
||||
str(target),
|
||||
" pass",
|
||||
' print("hello")',
|
||||
call_id="call_2",
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Updated.")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "create app.py")
|
||||
assert "pass" in target.read_text()
|
||||
|
||||
await _act_and_collect(agent_loop, "update app.py")
|
||||
assert 'print("hello")' in target.read_text()
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
assert len(rewindable) == 2
|
||||
|
||||
await rm.rewind_to_message(rewindable[1][0], restore_files=True)
|
||||
assert target.read_text() == "def main():\n pass\n"
|
||||
|
||||
async def test_rewind_without_restore_keeps_files(
|
||||
self, tmp_working_directory: Path
|
||||
) -> None:
|
||||
"""Rewind with restore_files=False keeps the file as-is."""
|
||||
target = tmp_working_directory / "data.json"
|
||||
|
||||
backend = FakeBackend([
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Writing.",
|
||||
tool_calls=[_write_file_tool_call(str(target), '{"a": 1}')],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Done.")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "write data")
|
||||
assert target.read_text() == '{"a": 1}'
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
await rm.rewind_to_message(rewindable[0][0], restore_files=False)
|
||||
|
||||
assert target.read_text() == '{"a": 1}'
|
||||
|
||||
async def test_rewind_then_new_turn(self, tmp_working_directory: Path) -> None:
|
||||
"""After rewind, a new turn creates fresh checkpoints that work correctly."""
|
||||
target = tmp_working_directory / "code.py"
|
||||
|
||||
backend = FakeBackend([
|
||||
# Turn 1
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="v1.", tool_calls=[_write_file_tool_call(str(target), "v1")]
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="ok")],
|
||||
# Turn 2
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="v2.",
|
||||
tool_calls=[
|
||||
_write_file_tool_call(
|
||||
str(target), "v2", call_id="call_2", overwrite=True
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="ok")],
|
||||
# Turn 3 (after rewind, new turn)
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="v2bis.",
|
||||
tool_calls=[
|
||||
_write_file_tool_call(
|
||||
str(target), "v2bis", call_id="call_3", overwrite=True
|
||||
)
|
||||
],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="ok")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "turn1")
|
||||
await _act_and_collect(agent_loop, "turn2")
|
||||
assert target.read_text() == "v2"
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
await rm.rewind_to_message(rewindable[1][0], restore_files=True)
|
||||
assert target.read_text() == "v1"
|
||||
|
||||
await _act_and_collect(agent_loop, "turn2bis")
|
||||
assert target.read_text() == "v2bis"
|
||||
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
await rm.rewind_to_message(rewindable[1][0], restore_files=True)
|
||||
assert target.read_text() == "v1"
|
||||
|
||||
async def test_rewind_restores_file_deleted_by_bash(
|
||||
self, tmp_working_directory: Path
|
||||
) -> None:
|
||||
"""A tracked file deleted via bash in turn 2 is restored on rewind.
|
||||
|
||||
Turn 1: create the file with write_file (file becomes tracked).
|
||||
Turn 2: delete it with bash (bash has no snapshot, but
|
||||
create_checkpoint re-reads known files).
|
||||
Rewind to turn 2 → file restored to its turn-1 content.
|
||||
"""
|
||||
target = tmp_working_directory / "important.txt"
|
||||
|
||||
backend = FakeBackend([
|
||||
# Turn 1: create the file
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Creating.",
|
||||
tool_calls=[_write_file_tool_call(str(target), "precious data")],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Created.")],
|
||||
# Turn 2: delete it via bash
|
||||
[
|
||||
mock_llm_chunk(
|
||||
content="Deleting.",
|
||||
tool_calls=[_bash_tool_call(f"rm {target}", call_id="call_2")],
|
||||
)
|
||||
],
|
||||
[mock_llm_chunk(content="Deleted.")],
|
||||
])
|
||||
agent_loop = _make_agent_loop(backend)
|
||||
|
||||
await _act_and_collect(agent_loop, "create file")
|
||||
assert target.read_text() == "precious data"
|
||||
|
||||
await _act_and_collect(agent_loop, "delete file")
|
||||
assert not target.exists()
|
||||
|
||||
rm = agent_loop.rewind_manager
|
||||
rewindable = rm.get_rewindable_messages()
|
||||
assert len(rewindable) == 2
|
||||
|
||||
# Rewind to turn 2 → restores the file to its state before turn 2
|
||||
await rm.rewind_to_message(rewindable[1][0], restore_files=True)
|
||||
assert target.exists()
|
||||
assert target.read_text() == "precious data"
|
||||
594
tests/core/test_rewind_manager.py
Normal file
594
tests/core/test_rewind_manager.py
Normal file
|
|
@ -0,0 +1,594 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vibe.core.rewind.manager import FileSnapshot, RewindError, RewindManager
|
||||
from vibe.core.types import LLMMessage, MessageList, Role
|
||||
|
||||
|
||||
def _make_messages(*contents: str) -> MessageList:
|
||||
"""Create a MessageList with a system message followed by user/assistant pairs."""
|
||||
msgs = MessageList([LLMMessage(role=Role.system, content="system")])
|
||||
for content in contents:
|
||||
msgs.append(LLMMessage(role=Role.user, content=content))
|
||||
msgs.append(LLMMessage(role=Role.assistant, content=f"reply to {content}"))
|
||||
return msgs
|
||||
|
||||
|
||||
def _snap(path: Path) -> FileSnapshot:
|
||||
"""Create a FileSnapshot by reading a file (or None if missing)."""
|
||||
resolved = str(path.resolve())
|
||||
try:
|
||||
content: bytes | None = path.read_bytes()
|
||||
except FileNotFoundError:
|
||||
content = None
|
||||
return FileSnapshot(path=resolved, content=content)
|
||||
|
||||
|
||||
def _make_manager(
|
||||
messages: MessageList,
|
||||
) -> tuple[RewindManager, list[bool], list[bool]]:
|
||||
save_calls: list[bool] = []
|
||||
reset_calls: list[bool] = []
|
||||
|
||||
async def save_messages() -> None:
|
||||
save_calls.append(True)
|
||||
|
||||
def reset_session() -> None:
|
||||
reset_calls.append(True)
|
||||
|
||||
mgr = RewindManager(
|
||||
messages=messages, save_messages=save_messages, reset_session=reset_session
|
||||
)
|
||||
return mgr, save_calls, reset_calls
|
||||
|
||||
|
||||
class TestCheckpoints:
|
||||
def test_create_checkpoint_carries_forward_snapshots(self, tmp_path: Path) -> None:
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("v1", encoding="utf-8")
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr.add_snapshot(_snap(f))
|
||||
f.write_text("v2", encoding="utf-8")
|
||||
|
||||
mgr.create_checkpoint()
|
||||
|
||||
# Second checkpoint should have re-read the file
|
||||
assert len(mgr.checkpoints) == 2
|
||||
assert mgr.checkpoints[1].files[0].content == b"v2"
|
||||
|
||||
def test_add_snapshot_to_all_checkpoints(self, tmp_path: Path) -> None:
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr.create_checkpoint()
|
||||
|
||||
f = tmp_path / "late.txt"
|
||||
f.write_text("content", encoding="utf-8")
|
||||
mgr.add_snapshot(_snap(f))
|
||||
|
||||
resolved = str(f.resolve())
|
||||
assert any(s.path == resolved for s in mgr.checkpoints[0].files)
|
||||
assert any(s.path == resolved for s in mgr.checkpoints[1].files)
|
||||
|
||||
def test_add_snapshot_no_duplicate(self, tmp_path: Path) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
mgr.create_checkpoint()
|
||||
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("content", encoding="utf-8")
|
||||
mgr.add_snapshot(_snap(f))
|
||||
mgr.add_snapshot(_snap(f))
|
||||
|
||||
resolved = str(f.resolve())
|
||||
matches = [s for s in mgr.checkpoints[0].files if s.path == resolved]
|
||||
assert len(matches) == 1
|
||||
|
||||
def test_has_changes_detects_new_file(self, tmp_path: Path) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
f = tmp_path / "new.txt"
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr.add_snapshot(FileSnapshot(path=str(f.resolve()), content=None))
|
||||
assert not mgr.has_file_changes_at(len(messages))
|
||||
|
||||
f.write_text("created", encoding="utf-8")
|
||||
assert mgr.has_file_changes_at(len(messages))
|
||||
|
||||
def test_has_changes_false_when_unchanged(self, tmp_path: Path) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
f = tmp_path / "f.txt"
|
||||
f.write_text("content", encoding="utf-8")
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr.add_snapshot(_snap(f))
|
||||
assert not mgr.has_file_changes_at(len(messages))
|
||||
|
||||
def test_has_file_changes_at_no_checkpoint(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
assert not mgr.has_file_changes_at(1)
|
||||
|
||||
|
||||
class TestRewind:
|
||||
def test_get_rewindable_messages(self) -> None:
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
result = mgr.get_rewindable_messages()
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0] == (1, "hello")
|
||||
assert result[1] == (3, "world")
|
||||
|
||||
def test_get_rewindable_messages_excludes_injected(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
# Insert an injected middleware message between turns
|
||||
messages.append(
|
||||
LLMMessage(role=Role.user, content="plan mode reminder", injected=True)
|
||||
)
|
||||
messages.append(LLMMessage(role=Role.user, content="world"))
|
||||
messages.append(LLMMessage(role=Role.assistant, content="reply to world"))
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
result = mgr.get_rewindable_messages()
|
||||
|
||||
assert len(result) == 2
|
||||
assert result[0] == (1, "hello")
|
||||
# Index 3 is the injected message — it must be skipped
|
||||
assert result[1] == (4, "world")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewind_to_message(self) -> None:
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, save_calls, reset_calls = _make_manager(messages)
|
||||
|
||||
content, errors = await mgr.rewind_to_message(3, restore_files=False)
|
||||
|
||||
assert content == "world"
|
||||
assert errors == []
|
||||
assert len(save_calls) == 1
|
||||
assert len(reset_calls) == 1
|
||||
assert len(messages) == 3
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewind_to_message_invalid_index(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
with pytest.raises(RewindError, match="Invalid message index"):
|
||||
await mgr.rewind_to_message(99, restore_files=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewind_to_message_not_user(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
with pytest.raises(RewindError, match="not a user message"):
|
||||
await mgr.rewind_to_message(2, restore_files=False)
|
||||
|
||||
def test_messages_reset_clears_checkpoints(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
mgr.create_checkpoint()
|
||||
|
||||
assert len(mgr.checkpoints) == 1
|
||||
|
||||
messages.reset([LLMMessage(role=Role.system, content="system")])
|
||||
|
||||
assert len(mgr.checkpoints) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_rewind_preserves_earlier_checkpoints(self) -> None:
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr._checkpoints[-1].message_index = 1
|
||||
mgr.create_checkpoint()
|
||||
mgr._checkpoints[-1].message_index = 3
|
||||
|
||||
assert len(mgr.checkpoints) == 2
|
||||
|
||||
await mgr.rewind_to_message(3, restore_files=False)
|
||||
|
||||
assert len(mgr.checkpoints) == 1
|
||||
|
||||
def test_update_system_prompt_preserves_checkpoints(self) -> None:
|
||||
"""Switching agents via shift+tab calls update_system_prompt which must
|
||||
NOT clear rewind checkpoints (unlike a full reset).
|
||||
"""
|
||||
messages = _make_messages("hello", "world")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
mgr.create_checkpoint()
|
||||
mgr._checkpoints[-1].message_index = 1
|
||||
mgr.create_checkpoint()
|
||||
mgr._checkpoints[-1].message_index = 3
|
||||
|
||||
assert len(mgr.checkpoints) == 2
|
||||
|
||||
# Simulate shift+tab agent switch: only the system prompt changes
|
||||
messages.update_system_prompt("new agent system prompt")
|
||||
|
||||
assert len(mgr.checkpoints) == 2
|
||||
assert messages[0].content == "new agent system prompt"
|
||||
|
||||
def test_create_checkpoint_uses_current_message_count(self) -> None:
|
||||
messages = _make_messages("hello")
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
|
||||
mgr.create_checkpoint()
|
||||
|
||||
assert mgr.checkpoints[0].message_index == len(messages)
|
||||
|
||||
|
||||
class _Turn:
|
||||
"""Helper that simulates one conversation turn."""
|
||||
|
||||
def __init__(self, mgr: RewindManager, messages: MessageList) -> None:
|
||||
self._mgr = mgr
|
||||
self._messages = messages
|
||||
|
||||
def begin(self, user_msg: str) -> None:
|
||||
"""Start a new turn: create checkpoint then append user message.
|
||||
|
||||
This mirrors agent_loop.act() which calls create_checkpoint()
|
||||
*before* the user message is added to the message list.
|
||||
"""
|
||||
self._mgr.create_checkpoint()
|
||||
self._messages.append(LLMMessage(role=Role.user, content=user_msg))
|
||||
|
||||
def tool_write(self, path: Path, content: str) -> None:
|
||||
"""Simulate a tool writing to a file (snapshot → write)."""
|
||||
self._mgr.add_snapshot(_snap(path))
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
|
||||
def tool_delete(self, path: Path) -> None:
|
||||
"""Simulate a tool deleting a file (snapshot → unlink)."""
|
||||
self._mgr.add_snapshot(_snap(path))
|
||||
path.unlink()
|
||||
|
||||
def end(self, assistant_reply: str = "ok") -> None:
|
||||
"""End the turn: append assistant reply."""
|
||||
self._messages.append(LLMMessage(role=Role.assistant, content=assistant_reply))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
class TestRewindScenarios:
|
||||
@staticmethod
|
||||
def _setup() -> tuple[RewindManager, MessageList, _Turn]:
|
||||
messages = MessageList([LLMMessage(role=Role.system, content="system")])
|
||||
mgr, _, _ = _make_manager(messages)
|
||||
return mgr, messages, _Turn(mgr, messages)
|
||||
|
||||
async def test_edit_file_across_turns_rewind_to_middle(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
"""A file is edited every turn. Rewinding restores the version from
|
||||
*before* the target turn.
|
||||
"""
|
||||
mgr, messages, turn = self._setup()
|
||||
f = tmp_path / "app.py"
|
||||
f.write_text("v0", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "v1")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "v2")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn3")
|
||||
turn.tool_write(f, "v3")
|
||||
turn.end()
|
||||
|
||||
# Rewind to turn2 (message index 3) → file should be v1
|
||||
rewindable = mgr.get_rewindable_messages()
|
||||
turn2_idx = rewindable[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
|
||||
assert f.read_text(encoding="utf-8") == "v1"
|
||||
|
||||
async def test_file_created_then_rewind_before_creation(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
"""A file that didn't exist before should be deleted on rewind."""
|
||||
mgr, messages, turn = self._setup()
|
||||
new_file = tmp_path / "generated.py"
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(new_file, "print('hello')")
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
await mgr.rewind_to_message(turn1_idx, restore_files=True)
|
||||
|
||||
assert not new_file.exists()
|
||||
|
||||
async def test_file_deleted_by_tool_rewind_restores(self, tmp_path: Path) -> None:
|
||||
"""Rewinding past a deletion restores the file."""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "config.yaml"
|
||||
f.write_text("key: value", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "key: value") # touch to start tracking
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_delete(f)
|
||||
turn.end()
|
||||
|
||||
assert not f.exists()
|
||||
|
||||
turn2_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
|
||||
assert f.exists()
|
||||
assert f.read_text(encoding="utf-8") == "key: value"
|
||||
|
||||
async def test_mixed_create_and_edit(self, tmp_path: Path) -> None:
|
||||
"""Multiple files: one pre-existing and edited, one created mid-session."""
|
||||
mgr, messages, turn = self._setup()
|
||||
existing = tmp_path / "main.py"
|
||||
existing.write_text("original", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(existing, "modified")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
new_file = tmp_path / "utils.py"
|
||||
turn.tool_write(new_file, "def helper(): ...")
|
||||
turn.tool_write(existing, "modified again")
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
await mgr.rewind_to_message(turn1_idx, restore_files=True)
|
||||
|
||||
assert existing.read_text(encoding="utf-8") == "original"
|
||||
assert not new_file.exists()
|
||||
|
||||
async def test_user_manual_edit_between_turns(self, tmp_path: Path) -> None:
|
||||
"""If the user edits a file between turns, the checkpoint at the next
|
||||
turn captures the user's version, so rewind restores it.
|
||||
"""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "readme.md"
|
||||
f.write_text("initial", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "tool wrote this")
|
||||
turn.end()
|
||||
|
||||
# User manually edits the file outside the tool loop
|
||||
f.write_text("user edited this", encoding="utf-8")
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "tool overwrote user")
|
||||
turn.end()
|
||||
|
||||
# Rewind to turn2 → should restore the user's manual edit
|
||||
turn2_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
|
||||
assert f.read_text(encoding="utf-8") == "user edited this"
|
||||
|
||||
async def test_rewind_without_restore(self, tmp_path: Path) -> None:
|
||||
"""Rewinding with restore_files=False truncates messages but keeps
|
||||
files as they are.
|
||||
"""
|
||||
mgr, messages, turn = self._setup()
|
||||
f = tmp_path / "data.json"
|
||||
f.write_text("{}", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, '{"a": 1}')
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, '{"a": 1, "b": 2}')
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
await mgr.rewind_to_message(turn1_idx, restore_files=False)
|
||||
|
||||
# File untouched
|
||||
assert f.read_text(encoding="utf-8") == '{"a": 1, "b": 2}'
|
||||
# But messages were truncated
|
||||
assert len(messages) == 1 # only system message
|
||||
|
||||
async def test_rewind_then_new_turns_then_rewind_again(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
"""After a rewind, new turns create new checkpoints. A second rewind
|
||||
should work correctly with the new history.
|
||||
"""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "code.py"
|
||||
f.write_text("v0", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "v1")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "v2")
|
||||
turn.end()
|
||||
|
||||
# Rewind to turn2
|
||||
turn2_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
assert f.read_text(encoding="utf-8") == "v1"
|
||||
|
||||
# New turn after rewind
|
||||
turn.begin("turn2-bis")
|
||||
turn.tool_write(f, "v2-bis")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn3-bis")
|
||||
turn.tool_write(f, "v3-bis")
|
||||
turn.end()
|
||||
|
||||
# Rewind to turn2-bis
|
||||
turn2bis_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2bis_idx, restore_files=True)
|
||||
assert f.read_text(encoding="utf-8") == "v1"
|
||||
|
||||
async def test_agent_switch_between_turns_preserves_rewind(
|
||||
self, tmp_path: Path
|
||||
) -> None:
|
||||
"""Pressing shift+tab between two messages switches agents, which calls
|
||||
update_system_prompt. Checkpoints must survive so a subsequent rewind
|
||||
restores files correctly.
|
||||
"""
|
||||
mgr, messages, turn = self._setup()
|
||||
f = tmp_path / "main.py"
|
||||
f.write_text("v0", encoding="utf-8")
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "v1")
|
||||
turn.end()
|
||||
|
||||
# User presses shift+tab → agent switch → system prompt replaced
|
||||
messages.update_system_prompt("switched agent prompt")
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "v2")
|
||||
turn.end()
|
||||
|
||||
# Rewind to turn2 should restore "v1"
|
||||
turn2_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
|
||||
assert f.read_text(encoding="utf-8") == "v1"
|
||||
|
||||
async def test_binary_file_snapshot_and_restore(self, tmp_path: Path) -> None:
|
||||
"""Binary files (non-UTF-8) are snapshotted and restored correctly."""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "image.bin"
|
||||
original = bytes(range(256))
|
||||
f.write_bytes(original)
|
||||
|
||||
turn.begin("turn1")
|
||||
mgr.add_snapshot(_snap(f))
|
||||
f.write_bytes(b"\x00" * 256)
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
await mgr.rewind_to_message(turn1_idx, restore_files=True)
|
||||
|
||||
assert f.read_bytes() == original
|
||||
|
||||
async def test_create_edit_delete_full_lifecycle(self, tmp_path: Path) -> None:
|
||||
"""File goes through create → edit → delete. Rewind to each point
|
||||
restores the correct state.
|
||||
"""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "temp.txt"
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(f, "created")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "edited")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn3")
|
||||
turn.tool_delete(f)
|
||||
turn.end()
|
||||
|
||||
assert not f.exists()
|
||||
|
||||
# Rewind to turn3 → file should be "edited" (state before deletion)
|
||||
turn3_idx = mgr.get_rewindable_messages()[2][0]
|
||||
await mgr.rewind_to_message(turn3_idx, restore_files=True)
|
||||
assert f.read_text(encoding="utf-8") == "edited"
|
||||
|
||||
async def test_user_creates_file_tool_overwrites(self, tmp_path: Path) -> None:
|
||||
"""User creates a file manually before a turn. The tool overwrites it.
|
||||
Rewind restores the user's version.
|
||||
"""
|
||||
mgr, _, turn = self._setup()
|
||||
f = tmp_path / "notes.txt"
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.end()
|
||||
|
||||
# User creates the file manually between turns
|
||||
f.write_text("user notes", encoding="utf-8")
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(f, "overwritten by tool")
|
||||
turn.end()
|
||||
|
||||
turn2_idx = mgr.get_rewindable_messages()[1][0]
|
||||
await mgr.rewind_to_message(turn2_idx, restore_files=True)
|
||||
|
||||
assert f.read_text(encoding="utf-8") == "user notes"
|
||||
|
||||
async def test_nested_directory_files(self, tmp_path: Path) -> None:
|
||||
"""Files in nested directories are restored including parent dirs."""
|
||||
mgr, _, turn = self._setup()
|
||||
deep = tmp_path / "src" / "pkg" / "module.py"
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.tool_write(deep, "def foo(): pass")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
turn.tool_write(deep, "def foo(): return 42")
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
|
||||
# Delete everything
|
||||
deep.unlink()
|
||||
(tmp_path / "src" / "pkg").rmdir()
|
||||
(tmp_path / "src").rmdir()
|
||||
|
||||
await mgr.rewind_to_message(turn1_idx, restore_files=True)
|
||||
|
||||
# File didn't exist before turn1 → should be deleted
|
||||
assert not deep.exists()
|
||||
|
||||
async def test_rewind_restores_errors_collected(self, tmp_path: Path) -> None:
|
||||
"""When removing a file during rewind fails, errors are returned in the tuple."""
|
||||
mgr, _, turn = self._setup()
|
||||
created_file = tmp_path / "locked.txt"
|
||||
|
||||
turn.begin("turn1")
|
||||
turn.end()
|
||||
|
||||
turn.begin("turn2")
|
||||
# Snapshot runs before write → earlier checkpoints record content=None
|
||||
turn.tool_write(created_file, "created in turn2")
|
||||
turn.end()
|
||||
|
||||
turn1_idx = mgr.get_rewindable_messages()[0][0]
|
||||
with patch(
|
||||
"vibe.core.rewind.manager.os.remove",
|
||||
side_effect=OSError("mocked removal failure"),
|
||||
):
|
||||
_, errors = await mgr.rewind_to_message(turn1_idx, restore_files=True)
|
||||
|
||||
assert len(errors) == 1
|
||||
assert "Failed to delete file" in errors[0]
|
||||
assert "locked.txt" in errors[0]
|
||||
Loading…
Add table
Add a link
Reference in a new issue