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:
Mathias Gesbert 2026-03-26 10:11:34 +01:00 committed by GitHub
parent 5ea69b547b
commit 6a50d1d521
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
49 changed files with 3726 additions and 99 deletions

View 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"

View 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]