From 1c10de7727042b84c98e2f23c3dbfdc1121d8d35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=C3=87a=C4=9Fda=C5=9F=20Y=C3=BCrekli?= <25122236+cagdasyurekli@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:32:20 +0200 Subject: [PATCH 1/2] feat(git): workspace turn checkpointing via git shadow refs for safe rollback (#614) --- coworker/agent.py | 7 +- coworker/agents/base.py | 1 + coworker/catalog.py | 3 +- coworker/engine.py | 41 ++++ coworker/risk.py | 1 + coworker/tools/git.py | 402 +++++++++++++++++++++++++++++++++- tests/test_catalog.py | 1 + tests/test_git_checkpoints.py | 282 ++++++++++++++++++++++++ 8 files changed, 724 insertions(+), 14 deletions(-) create mode 100644 tests/test_git_checkpoints.py diff --git a/coworker/agent.py b/coworker/agent.py index 9479718c89..71d843656f 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 43ac03de82..bfa61ee558 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 cdf4fae37f..86fe4d5794 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 e492443c42..5409275728 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,8 @@ async def run( message["_display"] = display self.messages.append(message) self._cancel.clear() + 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 +842,22 @@ 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 + + create_checkpoint( + self.permissions.workspace_root, + self.session_id, + self.turn_index, + ) + self._turn_checkpoint_created = True + 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 4873271bee..1e78b1dfb1 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 169b8fd3ab..0106a3b4bf 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,348 @@ }, } +_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": "/dev/null", + "GIT_CONFIG_SYSTEM": "/dev/null", + "GIT_CONFIG_NOSYSTEM": "1", + "HOME": "/tmp", + } + 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 create_checkpoint( + workspace: str | Path, session_id: str, turn_index: int +) -> str | None: + """Capture workspace working tree as a git shadow ref before writes apply. + + Saves tracked, modified, and untracked files into a temporary index without + affecting the repository's real index, HEAD, or branch. Returns the ref name + on success, or None if the workspace is not a git repo or checkpointing fails. + """ + if not is_git_repo(workspace): + return None + root = Path(workspace).expanduser().resolve() + git_dir = _git_dir(root) + if not git_dir or not git_dir.is_dir(): + return None + + sid = _sanitize_session_id(session_id) + ref = f"{CHECKPOINT_REF_PREFIX}/{sid}/{turn_index}" + tmp_name = f"ow_ckpt_{sid}_{turn_index}_{uuid.uuid4().hex[:8]}" + tmp_idx = git_dir / tmp_name + + try: + env = _git_env({"GIT_INDEX_FILE": str(tmp_idx)}) + add_res = subprocess.run( + ["git", "-C", str(root), "--work-tree", str(root), "add", "-A"], + capture_output=True, + text=True, + check=False, + env=env, + timeout=15, + ) + if add_res.returncode != 0: + return None -def git_tools(workspace: str) -> list: + wt_res = subprocess.run( + ["git", "-C", str(root), "write-tree"], + capture_output=True, + text=True, + check=False, + env=env, + timeout=15, + ) + if wt_res.returncode != 0 or not wt_res.stdout.strip(): + return None + tree_sha = wt_res.stdout.strip() + + parent = None + head_res = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--verify", "HEAD"], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + if head_res.returncode == 0 and head_res.stdout.strip(): + parent = head_res.stdout.strip() + + commit_cmd = [ + "git", + "-C", + str(root), + "commit-tree", + tree_sha, + "-m", + f"openworker checkpoint {sid} turn {turn_index}", + ] + if parent: + commit_cmd.extend(["-p", parent]) + ct_res = subprocess.run( + commit_cmd, + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=15, + ) + if ct_res.returncode != 0 or not ct_res.stdout.strip(): + return None + commit_sha = ct_res.stdout.strip() + + up_res = subprocess.run( + ["git", "-C", str(root), "update-ref", ref, commit_sha], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + if up_res.returncode != 0: + return None + return ref + except (OSError, subprocess.SubprocessError): + return None + finally: + if tmp_idx.exists(): + try: + tmp_idx.unlink() + except OSError: + pass + + +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 workspace files to the checkpoint captured before the turn began.""" + if not is_git_repo(workspace): + return {"ok": False, "error": "workspace is not a git repository"} + root = Path(workspace).expanduser().resolve() + sid = _sanitize_session_id(session_id) + ref = f"{CHECKPOINT_REF_PREFIX}/{sid}/{turn_index}" + + try: + chk = subprocess.run( + ["git", "-C", str(root), "rev-parse", "--verify", ref], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + if chk.returncode != 0: + return { + "ok": False, + "error": f"checkpoint not found for turn {turn_index} ({ref})", + } + + checkout = subprocess.run( + ["git", "-C", str(root), "checkout", ref, "--", "."], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=20, + ) + if checkout.returncode != 0: + return { + "ok": False, + "error": (checkout.stderr or "git checkout failed").strip()[:300], + } + + tree_out = subprocess.run( + ["git", "-C", str(root), "ls-tree", "-r", "--name-only", ref], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=10, + ) + cp_files = set(tree_out.stdout.splitlines()) + + removed_files: list[str] = [] + for dirpath, dirnames, filenames in os.walk(root, topdown=True): + if ".git" in dirnames: + dirnames.remove(".git") + for filename in filenames: + file_path = Path(dirpath) / filename + rel = str(file_path.relative_to(root)) + if rel.startswith(".git") or ".git" in file_path.parts: + continue + if rel not in cp_files: + ign = subprocess.run( + ["git", "-C", str(root), "check-ignore", rel], + capture_output=True, + text=True, + check=False, + env=_git_env(), + timeout=5, + ) + if ign.returncode != 0: + try: + file_path.unlink() + removed_files.append(rel) + except OSError: + pass + for dirpath, dirnames, _ in os.walk(root, topdown=False): + if ".git" in Path(dirpath).parts: + continue + for dirname in dirnames: + dpath = Path(dirpath) / dirname + if dpath.name != ".git" and ".git" not in dpath.parts: + try: + dpath.rmdir() + except OSError: + pass + + return { + "ok": True, + "ref": ref, + "turn": turn_index, + "removed_files": removed_files, + "message": ( + f"Successfully reverted workspace to turn {turn_index} checkpoint " + f"({len(removed_files)} post-turn file(s) removed)." + ), + } + except (OSError, subprocess.SubprocessError) as exc: + return {"ok": False, "error": f"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 +401,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 +427,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 +453,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 e9ebcbb127..3181fd737d 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 0000000000..b5bff24a92 --- /dev/null +++ b/tests/test_git_checkpoints.py @@ -0,0 +1,282 @@ +"""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) == [] From 805791dfa7833002b5bef3586badee4cde7bd91a Mon Sep 17 00:00:00 2001 From: cyurekli Date: Tue, 15 Sep 2026 21:21:39 +0200 Subject: [PATCH 2/2] fix: preserve filenames and staging during checkpoint restore --- coworker/engine.py | 9 +- coworker/tools/git.py | 254 ++++++++++------------------------ tests/test_git_checkpoints.py | 58 ++++++++ 3 files changed, 138 insertions(+), 183 deletions(-) diff --git a/coworker/engine.py b/coworker/engine.py index 5409275728..1596c204bb 100644 --- a/coworker/engine.py +++ b/coworker/engine.py @@ -336,6 +336,10 @@ 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: @@ -849,12 +853,11 @@ async def _handle_tool_calls( try: from .tools.git import create_checkpoint - create_checkpoint( + self._turn_checkpoint_created = create_checkpoint( self.permissions.workspace_root, self.session_id, self.turn_index, - ) - self._turn_checkpoint_created = True + ) is not None except Exception: pass diff --git a/coworker/tools/git.py b/coworker/tools/git.py index 0106a3b4bf..1f4727a9c8 100644 --- a/coworker/tools/git.py +++ b/coworker/tools/git.py @@ -73,10 +73,13 @@ def _git_env(extra: dict[str, str] | None = None) -> dict[str, str]: env = { **os.environ, - "GIT_CONFIG_GLOBAL": "/dev/null", - "GIT_CONFIG_SYSTEM": "/dev/null", + "GIT_CONFIG_GLOBAL": os.devnull, + "GIT_CONFIG_SYSTEM": os.devnull, "GIT_CONFIG_NOSYSTEM": "1", - "HOME": "/tmp", + "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) @@ -124,106 +127,52 @@ def _git_dir(workspace: str | Path) -> Path | None: return None +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: - """Capture workspace working tree as a git shadow ref before writes apply. - - Saves tracked, modified, and untracked files into a temporary index without - affecting the repository's real index, HEAD, or branch. Returns the ref name - on success, or None if the workspace is not a git repo or checkpointing fails. - """ - if not is_git_repo(workspace): - return None + """Snapshot working files and the index separately without changing HEAD.""" root = Path(workspace).expanduser().resolve() - git_dir = _git_dir(root) - if not git_dir or not git_dir.is_dir(): - return None - - sid = _sanitize_session_id(session_id) - ref = f"{CHECKPOINT_REF_PREFIX}/{sid}/{turn_index}" - tmp_name = f"ow_ckpt_{sid}_{turn_index}_{uuid.uuid4().hex[:8]}" - tmp_idx = git_dir / tmp_name - + tmp_idx = None try: - env = _git_env({"GIT_INDEX_FILE": str(tmp_idx)}) - add_res = subprocess.run( - ["git", "-C", str(root), "--work-tree", str(root), "add", "-A"], - capture_output=True, - text=True, - check=False, - env=env, - timeout=15, - ) - if add_res.returncode != 0: - return None - - wt_res = subprocess.run( - ["git", "-C", str(root), "write-tree"], - capture_output=True, - text=True, - check=False, - env=env, - timeout=15, - ) - if wt_res.returncode != 0 or not wt_res.stdout.strip(): - return None - tree_sha = wt_res.stdout.strip() - - parent = None - head_res = subprocess.run( - ["git", "-C", str(root), "rev-parse", "--verify", "HEAD"], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=5, - ) - if head_res.returncode == 0 and head_res.stdout.strip(): - parent = head_res.stdout.strip() - - commit_cmd = [ - "git", - "-C", - str(root), - "commit-tree", - tree_sha, - "-m", - f"openworker checkpoint {sid} turn {turn_index}", - ] - if parent: - commit_cmd.extend(["-p", parent]) - ct_res = subprocess.run( - commit_cmd, - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=15, - ) - if ct_res.returncode != 0 or not ct_res.stdout.strip(): + if not _checkpoint_root(root): return None - commit_sha = ct_res.stdout.strip() - - up_res = subprocess.run( - ["git", "-C", str(root), "update-ref", ref, commit_sha], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=5, - ) - if up_res.returncode != 0: + 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.exists(): - try: - tmp_idx.unlink() - except OSError: - pass + if tmp_idx is not None: + tmp_idx.unlink(missing_ok=True) def list_checkpoints( @@ -288,99 +237,44 @@ def list_checkpoints( def restore_checkpoint( workspace: str | Path, session_id: str, turn_index: int ) -> dict[str, Any]: - """Restore workspace files to the checkpoint captured before the turn began.""" - if not is_git_repo(workspace): - return {"ok": False, "error": "workspace is not a git repository"} + """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: - chk = subprocess.run( - ["git", "-C", str(root), "rev-parse", "--verify", ref], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=5, - ) - if chk.returncode != 0: - return { - "ok": False, - "error": f"checkpoint not found for turn {turn_index} ({ref})", - } - - checkout = subprocess.run( - ["git", "-C", str(root), "checkout", ref, "--", "."], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=20, - ) - if checkout.returncode != 0: - return { - "ok": False, - "error": (checkout.stderr or "git checkout failed").strip()[:300], - } - - tree_out = subprocess.run( - ["git", "-C", str(root), "ls-tree", "-r", "--name-only", ref], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=10, - ) - cp_files = set(tree_out.stdout.splitlines()) - - removed_files: list[str] = [] - for dirpath, dirnames, filenames in os.walk(root, topdown=True): - if ".git" in dirnames: - dirnames.remove(".git") - for filename in filenames: - file_path = Path(dirpath) / filename - rel = str(file_path.relative_to(root)) - if rel.startswith(".git") or ".git" in file_path.parts: - continue - if rel not in cp_files: - ign = subprocess.run( - ["git", "-C", str(root), "check-ignore", rel], - capture_output=True, - text=True, - check=False, - env=_git_env(), - timeout=5, - ) - if ign.returncode != 0: - try: - file_path.unlink() - removed_files.append(rel) - except OSError: - pass - for dirpath, dirnames, _ in os.walk(root, topdown=False): - if ".git" in Path(dirpath).parts: - continue - for dirname in dirnames: - dpath = Path(dirpath) / dirname - if dpath.name != ".git" and ".git" not in dpath.parts: - try: - dpath.rmdir() - except OSError: - pass - - return { - "ok": True, - "ref": ref, - "turn": turn_index, - "removed_files": removed_files, - "message": ( - f"Successfully reverted workspace to turn {turn_index} checkpoint " - f"({len(removed_files)} post-turn file(s) removed)." - ), - } + 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"restore failed: {exc}"} + return {"ok": False, "error": f"checkpoint not found or restore failed: {exc}"} def git_tools(workspace: str, session_id: str = "") -> list: diff --git a/tests/test_git_checkpoints.py b/tests/test_git_checkpoints.py index b5bff24a92..ce6e3e54f0 100644 --- a/tests/test_git_checkpoints.py +++ b/tests/test_git_checkpoints.py @@ -280,3 +280,61 @@ def write_file(path: str, content: str) -> str: 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