Add finalization-only cleanup gate for provisional scratch artifacts

tests/test_scratch_cleanup_gate.py

@@ -0,0 +1,74 @@
+from pathlib import Path
+import subprocess
+
+from villani_code import state_runtime
+
+
+class Dummy:
+    def __init__(self, repo: Path):
+        self.repo = repo
+        self._provisional_scratch_candidates = set()
+        self._explicit_mutation_created_paths = set()
+        self._cleanup_summary = {}
+        self._execution_plan = None
+        self.event_callback = lambda *_args, **_kwargs: None
+
+
+def _init_repo(repo: Path) -> None:
+    subprocess.run(["git", "init"], cwd=repo, check=True, capture_output=True)
+
+
+def test_bash_created_untracked_removed_when_validation_passes(tmp_path: Path, monkeypatch):
+    _init_repo(tmp_path)
+    (tmp_path / "README.md").write_text("ok\n")
+    (tmp_path / "src").mkdir()
+    runner = Dummy(tmp_path)
+    f = tmp_path / "probe.txt"
+    f.write_text("x")
+    runner._provisional_scratch_candidates.add("probe.txt")
+
+    monkeypatch.setattr(state_runtime, "run_validation", lambda *a, **k: type("R", (), {"passed": True})())
+    state_runtime.cleanup_provisional_scratch_artifacts(runner, ["README.md"])
+    assert not f.exists()
+    assert runner._cleanup_summary["cleanup_kept"] is True
+
+
+def test_bash_created_untracked_restored_when_validation_fails(tmp_path: Path, monkeypatch):
+    _init_repo(tmp_path)
+    (tmp_path / "README.md").write_text("ok\n")
+    (tmp_path / "src").mkdir()
+    runner = Dummy(tmp_path)
+    f = tmp_path / "probe.txt"
+    f.write_text("x")
+    runner._provisional_scratch_candidates.add("probe.txt")
+    monkeypatch.setattr(state_runtime, "run_validation", lambda *a, **k: type("R", (), {"passed": False})())
+    state_runtime.cleanup_provisional_scratch_artifacts(runner, ["README.md"])
+    assert f.exists()
+    assert runner._cleanup_summary["cleanup_restored"] is True
+
+
+def test_write_created_never_candidate(tmp_path: Path, monkeypatch):
+    _init_repo(tmp_path)
+    (tmp_path / "src").mkdir()
+    (tmp_path / "README.md").write_text("ok\n")
+    runner = Dummy(tmp_path)
+    f = tmp_path / "probe.txt"
+    f.write_text("x")
+    runner._provisional_scratch_candidates.add("probe.txt")
+    runner._explicit_mutation_created_paths.add("probe.txt")
+    monkeypatch.setattr(state_runtime, "run_validation", lambda *a, **k: type("R", (), {"passed": True})())
+    state_runtime.cleanup_provisional_scratch_artifacts(runner, ["README.md"])
+    assert f.exists()
+
+
+def test_src_and_tests_never_removed(tmp_path: Path, monkeypatch):
+    _init_repo(tmp_path)
+    (tmp_path / "src").mkdir(); (tmp_path / "tests").mkdir()
+    (tmp_path / "README.md").write_text("ok\n")
+    sf = tmp_path / "src" / "probe.txt"; sf.write_text("x")
+    tf = tmp_path / "tests" / "probe.txt"; tf.write_text("x")
+    runner = Dummy(tmp_path)
+    runner._provisional_scratch_candidates.update({"src/probe.txt", "tests/probe.txt"})
+    monkeypatch.setattr(state_runtime, "run_validation", lambda *a, **k: type("R", (), {"passed": True})())
+    state_runtime.cleanup_provisional_scratch_artifacts(runner, ["README.md"])
+    assert sf.exists() and tf.exists()

villani_code/debug_recorder.py

@@ -381,7 +381,7 @@ def write_prompt_rendered(self, text: str) -> None:
     def write_working_context(self, text: str) -> None:
         self._safe(append_text, self.artifacts.path("working_context.txt"), text)
 
-    def write_final_summary(self, *, status: str, termination_reason: str, total_turns: int, mission_id: str = "") -> Path:
+    def write_final_summary(self, *, status: str, termination_reason: str, total_turns: int, mission_id: str = "", cleanup_summary: dict[str, Any] | None = None) -> Path:
         if status == "completed":
             self._emit("run_completed", {"termination_reason": termination_reason, "mission_id": mission_id, "total_turns": total_turns})
         elif status == "failed":
@@ -398,6 +398,8 @@ def write_final_summary(self, *, status: str, termination_reason: str, total_tur
             summary_payload["changed_files"] = sorted(self._changed_files)
             summary_payload["last_failed_command"] = self._last_failed_command
             summary_payload["last_failed_validation"] = self._last_failed_validation
+            if isinstance(cleanup_summary, dict):
+                summary_payload.update(cleanup_summary)
             final_path = self.artifacts.path("final_summary.json")
             self._safe(write_json, final_path, summary_payload)
             return final_path

villani_code/state.py

@@ -509,6 +509,9 @@ def __init__(
         self._pending_verification = ""
         self._intended_targets: set[str] = set()
         self._before_contents: dict[str, str] = {}
+        self._provisional_scratch_candidates: set[str] = set()
+        self._explicit_mutation_created_paths: set[str] = set()
+        self._cleanup_summary = {}
         self._current_verification_targets: set[str] = set()
         self._current_verification_before_contents: dict[str, str] = {}
         self._verification_baseline_changed: set[str] = set()
@@ -887,6 +890,9 @@ def run(
         self._verification_baseline_changed = set(baseline_changed)
         self._intended_targets: set[str] = set()
         self._before_contents: dict[str, str] = {}
+        self._provisional_scratch_candidates: set[str] = set()
+        self._explicit_mutation_created_paths: set[str] = set()
+        self._cleanup_summary: dict[str, Any] = {}
         self._current_verification_targets: set[str] = set()
         self._current_verification_before_contents: dict[str, str] = {}
         self._last_verification_fingerprint = ""
@@ -1021,6 +1027,9 @@ def _finish_bounded(
             if not self._planning_read_only:
                 transcript_path = self._save_transcript_and_link(transcript)
             post = self._run_post_execution_validation(_change_summary()[2])
+            if completed and (not post or "passed" in post.lower()):
+                from villani_code import state_runtime
+                state_runtime.cleanup_provisional_scratch_artifacts(self, _change_summary()[2])
             if post:
                 response.setdefault("content", []).append({"type": "text", "text": post})
             self._save_session_snapshot(messages)
@@ -1034,6 +1043,7 @@ def _finish_bounded(
                     termination_reason=reason,
                     total_turns=turns_used,
                     mission_id=self._mission_id,
+                        cleanup_summary=getattr(self, "_cleanup_summary", None),
                 )
             return {
                 "response": response,

villani_code/state_runtime.py

@@ -1378,3 +1378,64 @@ def run_post_execution_validation(runner: Any, changed_files: list[str]) -> str:
         runner._mission_state.last_checkpoint_id = checkpoint.checkpoint_id
         save_mission_state(runner.repo, runner._mission_state)
     return outcome.message
+
+
+def cleanup_provisional_scratch_artifacts(runner: Any, changed_files: list[str]) -> None:
+    summary = {"cleanup_candidates_seen": 0, "cleanup_candidates_eligible": 0, "cleanup_files_removed": [], "cleanup_verification_rerun": False, "cleanup_kept": False, "cleanup_restored": False, "cleanup_skipped_reason": ""}
+    candidates = set(getattr(runner, "_provisional_scratch_candidates", set()))
+    summary["cleanup_candidates_seen"] = len(candidates)
+    if not candidates:
+        summary["cleanup_skipped_reason"] = "no_candidates"
+        runner._cleanup_summary = summary
+        return
+    source_roots = ["src", "lib", "app"]
+    test_roots = ["tests", "test"]
+    pkg_roots = ["pkg"]
+    confident = any((runner.repo / root).exists() for root in source_roots + test_roots + pkg_roots)
+    if not confident:
+        summary["cleanup_skipped_reason"] = "root_detection_not_confident"
+        runner._cleanup_summary = summary
+        return
+    eligible=[]
+    for rel in sorted(candidates):
+        if rel in getattr(runner, "_explicit_mutation_created_paths", set()):
+            continue
+        path=(runner.repo/rel).resolve()
+        if not path.exists() or not path.is_file():
+            continue
+        try:
+            path.relative_to(runner.repo.resolve())
+        except ValueError:
+            continue
+        if rel in changed_files and rel not in candidates:
+            continue
+        if any(rel == r or rel.startswith(r + "/") for r in source_roots + test_roots + pkg_roots if (runner.repo / r).exists()):
+            continue
+        eligible.append(rel)
+    summary["cleanup_candidates_eligible"] = len(eligible)
+    if not eligible:
+        summary["cleanup_skipped_reason"] = "no_eligible_candidates"
+        runner._cleanup_summary = summary
+        return
+    backups=[]
+    try:
+        for rel in eligible:
+            p=runner.repo/rel
+            backups.append((rel,p.read_bytes(),p.stat().st_mode))
+            p.unlink()
+        summary["cleanup_files_removed"] = [b[0] for b in backups]
+        summary["cleanup_verification_rerun"] = True
+        result = run_validation(runner.repo, changed_files, event_callback=runner.event_callback, repo_map=load_repo_map(runner.repo), change_impact=getattr(getattr(runner,'_execution_plan',None),'change_impact',None), action_classes=list(getattr(getattr(runner,'_execution_plan',None),'action_classes',[])), task_mode=str(getattr(getattr(runner,'_execution_plan',None),'task_mode',TaskMode.GENERAL.value)))
+        if result.passed:
+            summary["cleanup_kept"] = True
+        else:
+            raise RuntimeError("verification_failed")
+    except Exception as exc:
+        for rel, data, mode in backups:
+            p=runner.repo/rel
+            p.parent.mkdir(parents=True, exist_ok=True)
+            p.write_bytes(data)
+            p.chmod(mode)
+        summary["cleanup_restored"] = bool(backups)
+        summary["cleanup_skipped_reason"] = f"restored:{exc}"
+    runner._cleanup_summary = summary

villani_code/state_tooling.py

@@ -269,7 +269,10 @@ def _benchmark_post_write_python_validation(
                 "message": message,
             }
             runner.event_callback(event_payload)
-            runner.event_callback(
+            if tool_name == "Bash":
+        files_after = _snapshot_repo_files(runner.repo)
+        runner._provisional_scratch_candidates.update(sorted(files_after - files_before))
+    runner.event_callback(
                 {
                     "type": "failure_classified",
                     "category": "benchmark_post_write_validation_failed",
@@ -492,6 +495,11 @@ def execute_tool_with_policy(
         if runner.small_model:
             runner._tighten_tool_input(tool_name, tool_input)
     if tool_name in {"Write", "Patch"}:
+        repo_path = runner.repo / str(tool_input.get("file_path", ""))
+        if not repo_path.exists():
+            normalized = str(tool_input.get("file_path", "")).replace("\\", "/").lstrip("./")
+            if normalized:
+                runner._explicit_mutation_created_paths.add(normalized)
         _normalize_mutation_payload(tool_name, tool_input)
 
     policy = runner.permissions.evaluate_with_reason(
@@ -593,6 +601,11 @@ def execute_tool_with_policy(
                 return {"content": msg, "is_error": True}
 
     if tool_name in {"Write", "Patch"}:
+        repo_path = runner.repo / str(tool_input.get("file_path", ""))
+        if not repo_path.exists():
+            normalized = str(tool_input.get("file_path", "")).replace("\\", "/").lstrip("./")
+            if normalized:
+                runner._explicit_mutation_created_paths.add(normalized)
         targets = _benchmark_mutation_targets(tool_name, tool_input)
         normalized_targets = sorted(
             {
@@ -625,6 +638,18 @@ def execute_tool_with_policy(
     return _benchmark_post_write_python_validation(runner, tool_name, tool_input, result)
 
 
+
+
+def _snapshot_repo_files(repo: Path) -> set[str]:
+    files: set[str] = set()
+    for path in repo.rglob("*"):
+        if path.is_file():
+            try:
+                files.add(path.relative_to(repo).as_posix())
+            except ValueError:
+                continue
+    return files
+
 def execute_tool_with_lifecycle(
     *,
     runner: Any,
@@ -658,6 +683,7 @@ def _debug_callback_with_turn(event_type: str, payload: dict[str, Any]) -> None:
         if callable(debug_callback):
             debug_callback(event_type, callback_payload)
 
+    files_before = _snapshot_repo_files(runner.repo) if tool_name == "Bash" else set()
     result = execute_tool(
         tool_name,
         tool_input,