diff --git a/coworker/agent.py b/coworker/agent.py index 9479718c8..71d843656 100644 --- a/coworker/agent.py +++ b/coworker/agent.py @@ -267,7 +267,11 @@ def build_engine( executor = LocalExecutor(cwd=ws) if ws is not None else None todo = TodoList() context = AgentContext( - workspace=ws, executor=executor, todo=todo, roots=root_list or None + workspace=ws, + executor=executor, + todo=todo, + roots=root_list or None, + session_id=session_id, ) registry = ToolRegistry() @@ -533,6 +537,7 @@ def context_provider() -> str: model=model, instructions=instructions, approver=approver, + session_id=session_id or "default", # Stop kills the in-flight foreground shell command, not just the loop. interrupt_hooks=[executor.interrupt_now] if executor is not None else None, max_iterations=( diff --git a/coworker/agents/base.py b/coworker/agents/base.py index 43ac03de8..bfa61ee55 100644 --- a/coworker/agents/base.py +++ b/coworker/agents/base.py @@ -23,6 +23,7 @@ class AgentContext: # When None, tools fall back to the single `workspace` root. Held by reference so runtime # add/remove of folders is seen by the file tools built from it. roots: Optional[list] = None + session_id: Optional[str] = None @dataclass diff --git a/coworker/catalog.py b/coworker/catalog.py index cdf4fae37..86fe4d579 100644 --- a/coworker/catalog.py +++ b/coworker/catalog.py @@ -92,7 +92,8 @@ def _files(context: AgentContext) -> list: def _git(context: AgentContext) -> list: ws = str(context.workspace) - return [*ai.toolkits.git(root=ws), *git_tools(ws)] # git_status, git_diff, git_log + sid = getattr(context, "session_id", None) or "" + return [*ai.toolkits.git(root=ws), *git_tools(ws, session_id=sid)] def _search(context: AgentContext) -> list: diff --git a/coworker/engine.py b/coworker/engine.py index e492443c4..1596c204b 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -97,6 +97,7 @@ def __init__( model: str, instructions: Optional[str] = None, approver: Optional[Approver] = None, + session_id: str = "default", max_iterations: int = 12, model_settings: Optional[dict[str, Any]] = None, messages: Optional[list[dict[str, Any]]] = None, @@ -129,6 +130,9 @@ def __init__( self.permissions = permissions self.model = model self.approver = approver or _deny_all + self.session_id = session_id or "default" + self.turn_index = 0 + self._turn_checkpoint_created = False self.max_iterations = max_iterations self.model_settings = dict(model_settings or {}) self.messages: list[dict[str, Any]] = list(messages or []) @@ -286,6 +290,25 @@ def queue_steering( ) -> None: self._steering.append((text, source)) + def revert_turn(self, turn: Optional[int] = None) -> dict[str, Any]: + """Revert workspace files to the checkpoint taken before turn `turn`.""" + from .tools.git import list_checkpoints, restore_checkpoint + + target = turn + if target is None or target <= 0: + ckpts = list_checkpoints( + self.permissions.workspace_root, session_id=self.session_id + ) + if not ckpts: + return { + "ok": False, + "error": "No checkpoints available to revert.", + } + target = ckpts[-1]["turn"] + return restore_checkpoint( + self.permissions.workspace_root, self.session_id, target + ) + # -- main loop -------------------------------------------------------------- async def run( self, @@ -313,6 +336,12 @@ async def run( message["_display"] = display self.messages.append(message) self._cancel.clear() + if self.turn_index == 0: + from .tools.git import list_checkpoints + previous = list_checkpoints(self.permissions.workspace_root, self.session_id) + self.turn_index = max((c["turn"] for c in previous), default=0) + self.turn_index += 1 + self._turn_checkpoint_created = False if self.session_facts is not None: self.session_facts.begin_turn() # §8.4 retry guard resets per user turn: two reviewer denials in one turn route @@ -817,6 +846,21 @@ async def _handle_tool_calls( if allowed: cleared.append(tool_call) + if cleared and not self._turn_checkpoint_created: + from .risk import WRITE_TOOLS + + if any(tc.name in WRITE_TOOLS for tc in cleared): + try: + from .tools.git import create_checkpoint + + self._turn_checkpoint_created = create_checkpoint( + self.permissions.workspace_root, + self.session_id, + self.turn_index, + ) is not None + except Exception: + pass + concurrent = ( [tc for tc in cleared if self._parallel_safe(tc)] if len(cleared) > 1 diff --git a/coworker/risk.py b/coworker/risk.py index 4873271be..1e78b1dfb 100644 --- a/coworker/risk.py +++ b/coworker/risk.py @@ -55,6 +55,7 @@ class RiskClass(str, Enum): _BASE: dict[str, RiskClass] = { **{name: RiskClass.WRITE_LOCAL for name in WRITE_TOOLS}, + "revert_turn": RiskClass.WRITE_LOCAL, SHELL_TOOL: RiskClass.EXEC, **{name: RiskClass.EGRESS for name in EGRESS_TOOLS}, } diff --git a/coworker/tools/git.py b/coworker/tools/git.py index 169b8fd3a..1f4727a9c 100644 --- a/coworker/tools/git.py +++ b/coworker/tools/git.py @@ -1,27 +1,33 @@ -"""`git_log` — recent commit history for context (read-only). +"""`git_log` — recent commit history for context (read-only), and workspace +turn checkpoints via lightweight git shadow refs for safe rollback. -aisuite's git toolkit gives `git_status`/`git_diff`; this adds history so the agent can see how -a file came to be the way it is before changing it. Read-only; no commit/push here (the prompt -forbids those without explicit ask, and they'd go through run_shell anyway). +Checkpoints capture the exact working tree state (including untracked and modified +files) before write tools mutate the repository, enabling safe `revert_turn` +rollbacks without modifying git history or HEAD. """ from __future__ import annotations +import os +import re import subprocess +import uuid from pathlib import Path -from typing import Any, Optional +from typing import Any import aisuite as ai _SEP = "\x1f" +CHECKPOINT_REF_PREFIX = "refs/openworker/checkpoints" _SCHEMA = { "type": "function", "function": { "name": "git_log", "description": ( - "Recent git commit history (hash, author, date, subject). Optionally scope to a path. " - "Use it to understand how code evolved before editing. Read-only." + "Recent git commit history (hash, author, date, subject). Optionally " + "scope to a path. Use it to understand how code evolved before editing. " + "Read-only." ), "parameters": { "type": "object", @@ -39,11 +45,242 @@ }, } +_REVERT_TURN_SCHEMA = { + "type": "function", + "function": { + "name": "revert_turn", + "description": ( + "Revert workspace changes to the git shadow checkpoint captured " + "before a specific turn began. If turn is omitted or 0, reverts to " + "the checkpoint taken before the latest turn." + ), + "parameters": { + "type": "object", + "properties": { + "turn": { + "type": "integer", + "description": ( + "The turn number to revert to (default 0 for latest turn " + "checkpoint)." + ), + }, + }, + }, + }, +} + + +def _git_env(extra: dict[str, str] | None = None) -> dict[str, str]: + env = { + **os.environ, + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_SYSTEM": os.devnull, + "GIT_CONFIG_NOSYSTEM": "1", + "GIT_AUTHOR_NAME": "OpenWorker", + "GIT_AUTHOR_EMAIL": "checkpoint@openworker.invalid", + "GIT_COMMITTER_NAME": "OpenWorker", + "GIT_COMMITTER_EMAIL": "checkpoint@openworker.invalid", + } + if extra: + env.update(extra) + return env + + +def _sanitize_session_id(session_id: str) -> str: + cleaned = re.sub(r"[^a-zA-Z0-9_-]", "_", session_id or "") + return cleaned or "default" + + +def is_git_repo(workspace: str | Path) -> bool: + """Return True if workspace is inside a git work tree.""" + root = Path(workspace).expanduser().resolve() + try: + out = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--is-inside-work-tree"], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + return out.returncode == 0 and out.stdout.strip() == "true" + except (OSError, subprocess.SubprocessError): + return False + + +def _git_dir(workspace: str | Path) -> Path | None: + root = Path(workspace).expanduser().resolve() + try: + out = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--git-dir"], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + if out.returncode == 0 and out.stdout.strip(): + raw = Path(out.stdout.strip()) + return raw if raw.is_absolute() else (root / raw).resolve() + except (OSError, subprocess.SubprocessError): + pass + return None -def git_tools(workspace: str) -> list: + +def _run_git(root: Path, *args: str, env=None, input=None) -> bytes: + return subprocess.run( + ["git", "-C", str(root), *args], env=env or _git_env(), + input=input, capture_output=True, check=True, timeout=20, + ).stdout + + +def _checkpoint_root(root: Path) -> bool: + # A subfolder checkout must not snapshot or restore its parent repository. + top = os.fsdecode(_run_git(root, "rev-parse", "--show-toplevel")).strip() + return Path(top).resolve() == root + + +def create_checkpoint( + workspace: str | Path, session_id: str, turn_index: int +) -> str | None: + """Snapshot working files and the index separately without changing HEAD.""" + root = Path(workspace).expanduser().resolve() + tmp_idx = None + try: + if not _checkpoint_root(root): + return None + git_dir = _git_dir(root) + if git_dir is None: + return None + sid = _sanitize_session_id(session_id) + ref = f"{CHECKPOINT_REF_PREFIX}/{sid}/{turn_index}" + index_ref = f"refs/openworker/checkpoint-index/{sid}/{turn_index}" + # write-tree fails closed for an unmerged index. + index_tree = _run_git(root, "write-tree").decode().strip() + tmp_idx = git_dir / f"ow_ckpt_{uuid.uuid4().hex}" + env = _git_env({"GIT_INDEX_FILE": str(tmp_idx)}) + _run_git(root, "read-tree", index_tree, env=env) + _run_git(root, "add", "-A", env=env) + tree = _run_git(root, "write-tree", env=env).decode().strip() + commit = _run_git(root, "commit-tree", tree, "-m", + f"openworker checkpoint {sid} turn {turn_index}").decode().strip() + # Publish both snapshots atomically; never overwrite an earlier turn. + commands = f"start\ncreate {ref} {commit}\ncreate {index_ref} {index_tree}\nprepare\ncommit\n" + _run_git(root, "update-ref", "--stdin", input=commands.encode()) + return ref + except (OSError, subprocess.SubprocessError): + return None + finally: + if tmp_idx is not None: + tmp_idx.unlink(missing_ok=True) + + +def list_checkpoints( + workspace: str | Path, session_id: str | None = None +) -> list[dict[str, Any]]: + """List available turn checkpoints for the workspace.""" + if not is_git_repo(workspace): + return [] + root = Path(workspace).expanduser().resolve() + prefix = CHECKPOINT_REF_PREFIX + if session_id: + sid = _sanitize_session_id(session_id) + prefix = f"{CHECKPOINT_REF_PREFIX}/{sid}" + + try: + out = subprocess.run( + [ + "git", + "-C", + str(root), + "for-each-ref", + "--format=%(refname) %(objectname) %(creatordate:iso8601)", + f"{prefix}/", + ], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=10, + ) + if out.returncode != 0: + return [] + results = [] + for line in out.stdout.splitlines(): + parts = line.strip().split(maxsplit=2) + if len(parts) >= 2: + refname = parts[0] + commit = parts[1] + date_str = parts[2] if len(parts) > 2 else "" + ref_parts = refname.split("/") + if len(ref_parts) >= 5: + ckpt_sid = ref_parts[3] + try: + turn = int(ref_parts[4]) + except ValueError: + turn = 0 + results.append( + { + "ref": refname, + "session_id": ckpt_sid, + "turn": turn, + "commit": commit, + "date": date_str, + } + ) + results.sort(key=lambda c: c["turn"]) + return results + except (OSError, subprocess.SubprocessError): + return [] + + +def restore_checkpoint( + workspace: str | Path, session_id: str, turn_index: int +) -> dict[str, Any]: + """Restore exact filenames and the captured staging state; keep HEAD unchanged.""" + root = Path(workspace).expanduser().resolve() + if not is_git_repo(root): + return {"ok": False, "error": "workspace is not a git repository"} + sid = _sanitize_session_id(session_id) + ref = f"{CHECKPOINT_REF_PREFIX}/{sid}/{turn_index}" + index_ref = f"refs/openworker/checkpoint-index/{sid}/{turn_index}" + try: + if not _checkpoint_root(root): + return {"ok": False, "error": "checkpoint restore requires the repository root"} + # Resolve and enumerate everything before changing a file. Old checkpoints + # without an index snapshot cannot promise a safe staging-state restore. + tree = _run_git(root, "rev-parse", "--verify", f"{ref}^{{tree}}").decode().strip() + index_tree = _run_git(root, "rev-parse", "--verify", f"{index_ref}^{{tree}}").decode().strip() + cp_files = {os.fsdecode(f) for f in _run_git(root, "ls-tree", "-rz", "--name-only", tree).split(b"\0") if f} + current = {os.fsdecode(f) for f in _run_git(root, "ls-files", "-z", "--cached", "--others", "--exclude-standard").split(b"\0") if f} + removed = sorted(current - cp_files) + # Worktree-only restore never stages a formerly untracked/unstaged file. + if cp_files: + _run_git(root, "restore", f"--source={tree}", "--worktree", "--", ".") + for name in removed: + path = root / name + if path.is_file() or path.is_symlink(): + path.unlink() + # Only prune empty parents of files removed by this restore. + for name in removed: + parent = (root / name).parent + while parent != root: + try: + parent.rmdir() + except OSError: + break + parent = parent.parent + _run_git(root, "read-tree", index_tree) + return {"ok": True, "ref": ref, "turn": turn_index, "removed_files": removed, + "message": f"Restored workspace and staging state before turn {turn_index}."} + except (OSError, subprocess.SubprocessError) as exc: + return {"ok": False, "error": f"checkpoint not found or restore failed: {exc}"} + + +def git_tools(workspace: str, session_id: str = "") -> list: root = str(Path(workspace).resolve()) - def git_log(path: Optional[str] = None, max_count: int = 20) -> dict[str, Any]: + def git_log(path: str | None = None, max_count: int = 20) -> dict[str, Any]: n = max_count if isinstance(max_count, int) and max_count > 0 else 20 n = min(n, 200) cmd = [ @@ -58,8 +295,15 @@ def git_log(path: Optional[str] = None, max_count: int = 20) -> dict[str, Any]: if path: cmd += ["--", path] try: - out = subprocess.run(cmd, capture_output=True, text=True, timeout=15) - except Exception as exc: + out = subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=15, + ) + except (OSError, subprocess.SubprocessError) as exc: return {"error": f"git log failed: {exc}"} if out.returncode != 0: return {"error": (out.stderr or "git log failed").strip()[:300]} @@ -77,6 +321,22 @@ def git_log(path: Optional[str] = None, max_count: int = 20) -> dict[str, Any]: ) return {"count": len(commits), "commits": commits} + def revert_turn(turn: int = 0) -> dict[str, Any]: + """Revert workspace to the git shadow checkpoint captured before a turn.""" + target_turn = turn + if target_turn <= 0: + ckpts = list_checkpoints(root, session_id=session_id) + if not ckpts: + return { + "ok": False, + "error": "No checkpoints available to revert.", + } + target_turn = ckpts[-1]["turn"] + res = restore_checkpoint(root, session_id or "default", target_turn) + if not res.get("ok"): + return {"ok": False, "error": res.get("error", "Revert failed")} + return res + git_log.__name__ = "git_log" git_log.__doc__ = _SCHEMA["function"]["description"] git_log.__aisuite_tool_metadata__ = ai.ToolMetadata( @@ -87,4 +347,16 @@ def git_log(path: Optional[str] = None, max_count: int = 20) -> dict[str, Any]: requires_approval=False, ) git_log.__coworker_schema__ = _SCHEMA - return [git_log] + + revert_turn.__name__ = "revert_turn" + revert_turn.__doc__ = _REVERT_TURN_SCHEMA["function"]["description"] + revert_turn.__aisuite_tool_metadata__ = ai.ToolMetadata( + name="revert_turn", + category="git", + risk_level="high", + capabilities=["git"], + requires_approval=True, + ) + revert_turn.__coworker_schema__ = _REVERT_TURN_SCHEMA + + return [git_log, revert_turn] diff --git a/tests/test_catalog.py b/tests/test_catalog.py index e9ebcbb12..3181fd737 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -27,6 +27,7 @@ "git_status", "git_diff", "git_log", + "revert_turn", "grep", "run_shell", "shell_task_output", diff --git a/tests/test_git_checkpoints.py b/tests/test_git_checkpoints.py new file mode 100644 index 000000000..ce6e3e54f --- /dev/null +++ b/tests/test_git_checkpoints.py @@ -0,0 +1,340 @@ +"""Tests for workspace turn checkpointing via git shadow refs (Issue #614).""" + +import subprocess +from pathlib import Path + +import pytest + +from coworker.engine import TurnEngine +from coworker.permissions import Mode, PermissionEngine +from coworker.providers import AssistantTurn, ProviderClient, ToolCall +from coworker.tools import ToolRegistry +from coworker.tools.git import ( + _git_env, + create_checkpoint, + git_tools, + is_git_repo, + list_checkpoints, + restore_checkpoint, +) + + +def _init_git_repo(path: Path) -> None: + env = _git_env() + subprocess.run( + ["git", "init"], + cwd=path, + check=True, + capture_output=True, + env=env, + ) + subprocess.run( + [ + "git", + "-c", + "user.name=Test", + "-c", + "user.email=test@test.local", + "commit", + "--allow-empty", + "-m", + "initial", + ], + cwd=path, + check=True, + capture_output=True, + env=env, + ) + + +def test_is_git_repo(tmp_path): + assert not is_git_repo(tmp_path) + _init_git_repo(tmp_path) + assert is_git_repo(tmp_path) + + +def test_create_checkpoint_non_git_returns_none(tmp_path): + assert create_checkpoint(tmp_path, "sess-1", 1) is None + assert list_checkpoints(tmp_path) == [] + res = restore_checkpoint(tmp_path, "sess-1", 1) + assert not res["ok"] + assert "not a git repository" in res["error"] + + +def test_create_and_list_checkpoints(tmp_path): + _init_git_repo(tmp_path) + (tmp_path / "hello.txt").write_text("v1") + + ref1 = create_checkpoint(tmp_path, "session-a", 1) + assert ref1 == "refs/openworker/checkpoints/session-a/1" + + (tmp_path / "hello.txt").write_text("v2") + ref2 = create_checkpoint(tmp_path, "session-a", 2) + assert ref2 == "refs/openworker/checkpoints/session-a/2" + + ckpts = list_checkpoints(tmp_path, session_id="session-a") + assert len(ckpts) == 2 + assert ckpts[0]["turn"] == 1 + assert ckpts[1]["turn"] == 2 + assert ckpts[0]["session_id"] == "session-a" + + +def test_create_and_restore_checkpoint(tmp_path): + _init_git_repo(tmp_path) + + # Setup tracked file + (tmp_path / "app.py").write_text("print('original')") + env = _git_env() + subprocess.run( + ["git", "add", "app.py"], + cwd=tmp_path, + check=True, + capture_output=True, + env=env, + ) + subprocess.run( + [ + "git", + "-c", + "user.name=Test", + "-c", + "user.email=test@test.local", + "commit", + "-m", + "add app.py", + ], + cwd=tmp_path, + check=True, + capture_output=True, + env=env, + ) + + # Pre-existing untracked file and gitignored file + (tmp_path / "untracked_pre.txt").write_text("pre-existing untracked") + (tmp_path / ".gitignore").write_text("*.log\n") + (tmp_path / "build.log").write_text("log line 1") + + # Capture checkpoint before turn 1 + ref = create_checkpoint(tmp_path, "sess-1", 1) + assert ref is not None + + # Agent executes turn: modifies app.py, deletes untracked_pre.txt, creates new file + (tmp_path / "app.py").write_text("print('corrupted by agent')") + (tmp_path / "untracked_pre.txt").unlink() + sub = tmp_path / "new_dir" + sub.mkdir() + (sub / "generated.py").write_text("bad code") + (tmp_path / "build.log").write_text("log line 2") + + # Restore checkpoint + res = restore_checkpoint(tmp_path, "sess-1", 1) + assert res["ok"] + assert "new_dir/generated.py" in res["removed_files"] + + # Assertions + assert (tmp_path / "app.py").read_text() == "print('original')" + assert (tmp_path / "untracked_pre.txt").read_text() == "pre-existing untracked" + assert not (sub / "generated.py").exists() + assert not sub.exists() + # gitignored file untouched + assert (tmp_path / "build.log").read_text() == "log line 2" + + +def test_revert_turn_tool(tmp_path): + _init_git_repo(tmp_path) + (tmp_path / "main.py").write_text("def run(): pass") + + create_checkpoint(tmp_path, "session-test", 1) + (tmp_path / "main.py").write_text("syntax error !!!") + + tools = git_tools(str(tmp_path), session_id="session-test") + assert len(tools) == 2 + revert_fn = tools[1] + assert revert_fn.__name__ == "revert_turn" + + # Call revert_turn without argument -> reverts latest turn (turn 1) + res = revert_fn() + assert res["ok"] + assert (tmp_path / "main.py").read_text() == "def run(): pass" + + # Call on nonexistent turn + err = revert_fn(turn=99) + assert not err["ok"] + assert "not found" in err["error"] + + +class DummyProvider(ProviderClient): + def __init__(self, responses: list[AssistantTurn]) -> None: + self.responses = list(responses) + + def complete(self, *, model, messages, tools=None, **settings): + if self.responses: + return self.responses.pop(0) + return AssistantTurn(text="Done") + + def capabilities(self, model): + from coworker.providers.base import ModelCapabilities + + return ModelCapabilities() + + +@pytest.mark.asyncio +async def test_engine_automatic_checkpoint_before_writes(tmp_path): + _init_git_repo(tmp_path) + (tmp_path / "target.txt").write_text("initial state") + + written_files = [] + + def write_file(path: str, content: str) -> str: + p = tmp_path / path + p.write_text(content) + written_files.append(path) + return f"Wrote {path}" + + write_file.__name__ = "write_file" + write_file.__aisuite_tool_metadata__ = None + + registry = ToolRegistry() + registry.register(write_file) + + permissions = PermissionEngine(workspace_root=tmp_path, mode=Mode.BYPASS_APPROVALS) + + # Provider will request write_file + call = ToolCall( + id="c1", + name="write_file", + arguments={"path": "target.txt", "content": "agent mutated state"}, + ) + provider = DummyProvider( + [ + AssistantTurn(text="Writing file", tool_calls=[call]), + AssistantTurn(text="Finished write"), + ] + ) + + engine = TurnEngine( + provider=provider, + registry=registry, + permissions=permissions, + model="mock-model", + session_id="test-session", + ) + + events = [] + async for event in engine.run("Please write target.txt"): + events.append(event) + + assert (tmp_path / "target.txt").read_text() == "agent mutated state" + + # A checkpoint should have been created for turn 1 + ckpts = list_checkpoints(tmp_path, session_id="test-session") + assert len(ckpts) == 1 + assert ckpts[0]["turn"] == 1 + + # Reverting via engine.revert_turn restores target.txt + revert_res = engine.revert_turn() + assert revert_res["ok"] + assert (tmp_path / "target.txt").read_text() == "initial state" + + +@pytest.mark.asyncio +async def test_engine_skips_checkpoint_gracefully_in_non_git_workspace(tmp_path): + assert not is_git_repo(tmp_path) + (tmp_path / "file.txt").write_text("initial") + + def write_file(path: str, content: str) -> str: + (tmp_path / path).write_text(content) + return "ok" + + write_file.__name__ = "write_file" + write_file.__aisuite_tool_metadata__ = None + + registry = ToolRegistry() + registry.register(write_file) + + permissions = PermissionEngine(workspace_root=tmp_path, mode=Mode.BYPASS_APPROVALS) + call = ToolCall( + id="c1", + name="write_file", + arguments={"path": "file.txt", "content": "updated"}, + ) + provider = DummyProvider( + [ + AssistantTurn(text="Write", tool_calls=[call]), + AssistantTurn(text="Done"), + ] + ) + + engine = TurnEngine( + provider=provider, + registry=registry, + permissions=permissions, + model="mock-model", + session_id="non-git-session", + ) + + # Must complete cleanly without errors + events = [] + async for event in engine.run("Update file"): + events.append(event) + + assert (tmp_path / "file.txt").read_text() == "updated" + assert list_checkpoints(tmp_path) == [] + + +def test_restore_preserves_special_names_and_staged_state(tmp_path): + _init_git_repo(tmp_path) + def git(*args): + return subprocess.check_output(["git", "-C", str(tmp_path), *args], env=_git_env()) + names = ["with ü.txt", "tab\tname", "line\nname"] + for name in names: + (tmp_path / name).write_text("original") + f = tmp_path / "staged.txt" + f.write_text("staged") + git("add", "--", ".") + f.write_text("unstaged") + (tmp_path / "untracked.txt").write_text("untracked") + before_index = git("write-tree") + assert create_checkpoint(tmp_path, "special", 1) + f.write_text("agent edit") + assert restore_checkpoint(tmp_path, "special", 1)["ok"] + assert git("write-tree") == before_index + assert f.read_text() == "unstaged" + for name in names: + assert (tmp_path / name).read_text() == "original" + assert b"untracked.txt" in git("ls-files", "--others", "--exclude-standard") + + +def test_restore_enumeration_failure_changes_nothing(tmp_path, monkeypatch): + import coworker.tools.git as mod + _init_git_repo(tmp_path) + f = tmp_path / "file.txt" + f.write_text("before") + assert create_checkpoint(tmp_path, "s", 1) + f.write_text("after") + original = mod._run_git + def fail(root, *args, **kwargs): + if args[0] == "ls-tree": + raise OSError("tree unavailable") + return original(root, *args, **kwargs) + monkeypatch.setattr(mod, "_run_git", fail) + assert not restore_checkpoint(tmp_path, "s", 1)["ok"] + assert f.read_text() == "after" + + +def test_checkpoint_does_not_overwrite_an_existing_turn(tmp_path): + _init_git_repo(tmp_path) + f = tmp_path / "file.txt" + f.write_text("first") + assert create_checkpoint(tmp_path, "s", 1) + f.write_text("second") + assert create_checkpoint(tmp_path, "s", 1) is None + assert restore_checkpoint(tmp_path, "s", 1)["ok"] + assert f.read_text() == "first" + + +def test_checkpoint_skips_repository_subfolder(tmp_path): + _init_git_repo(tmp_path) + sub = tmp_path / "sub" + sub.mkdir() + assert create_checkpoint(sub, "s", 1) is None