Conservative cleanup of provisional scratch files after successful verification

tests/test_state_runtime.py

@@ -664,3 +664,85 @@ def test_diagnosis_confidence_weak_without_file_evidence(tmp_path: Path) -> None
     }
     confidence = state_runtime.classify_diagnosis_target_confidence(runner, diagnosis, failure_evidence=None)
     assert confidence == "weak"
+
+def _cleanup_runner(tmp_path: Path) -> SimpleNamespace:
+    return SimpleNamespace(
+        repo=tmp_path,
+        _provisional_scratch_candidates=set(),
+        benchmark_config=SimpleNamespace(visible_verification=["python -c 'print(\"ok\")'"]),
+        event_callback=lambda _e: None,
+    )
+
+
+def test_cleanup_removes_root_helper_when_verify_still_passes(tmp_path: Path) -> None:
+    runner = _cleanup_runner(tmp_path)
+    (tmp_path / "helper_tmp.py").write_text("print('x')\n", encoding="utf-8")
+    (tmp_path / "src").mkdir()
+    runner._provisional_scratch_candidates = {"helper_tmp.py"}
+    cleaned = state_runtime._cleanup_provisional_scratch_after_success(runner)
+    assert cleaned == ["helper_tmp.py"]
+    assert not (tmp_path / "helper_tmp.py").exists()
+
+
+def test_cleanup_restores_when_verify_fails(tmp_path: Path) -> None:
+    runner = _cleanup_runner(tmp_path)
+    runner.benchmark_config.visible_verification = ["python -c \"import pathlib; raise SystemExit(1 if not pathlib.Path('helper_tmp.py').exists() else 0)\""]
+    (tmp_path / "helper_tmp.py").write_text("print('x')\n", encoding="utf-8")
+    (tmp_path / "src").mkdir()
+    runner._provisional_scratch_candidates = {"helper_tmp.py"}
+    cleaned = state_runtime._cleanup_provisional_scratch_after_success(runner)
+    assert cleaned == []
+    assert (tmp_path / "helper_tmp.py").exists()
+
+
+def test_cleanup_keeps_existing_file_modifications(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    runner = _cleanup_runner(tmp_path)
+    (tmp_path / "helper_tmp.py").write_text("print('x')\n", encoding="utf-8")
+    (tmp_path / "src").mkdir()
+    runner._provisional_scratch_candidates = {"helper_tmp.py"}
+    monkeypatch.setattr(state_runtime, "git_changed_files", lambda _repo: ["helper_tmp.py"])
+    cleaned = state_runtime._cleanup_provisional_scratch_after_success(runner)
+    assert cleaned == []
+    assert (tmp_path / "helper_tmp.py").exists()
+
+
+def test_cleanup_keeps_files_inside_source_roots(tmp_path: Path) -> None:
+    runner = _cleanup_runner(tmp_path)
+    (tmp_path / "src").mkdir()
+    (tmp_path / "src" / "helper_tmp.py").write_text("print('x')\n", encoding="utf-8")
+    runner._provisional_scratch_candidates = {"src/helper_tmp.py"}
+    cleaned = state_runtime._cleanup_provisional_scratch_after_success(runner)
+    assert cleaned == []
+    assert (tmp_path / "src" / "helper_tmp.py").exists()
+
+
+def test_cleanup_skips_when_root_detection_uncertain(tmp_path: Path) -> None:
+    runner = _cleanup_runner(tmp_path)
+    (tmp_path / "helper_tmp.py").write_text("print('x')\n", encoding="utf-8")
+    runner._provisional_scratch_candidates = {"helper_tmp.py"}
+    cleaned = state_runtime._cleanup_provisional_scratch_after_success(runner)
+    assert cleaned == []
+    assert (tmp_path / "helper_tmp.py").exists()
+
+def test_explicit_write_created_source_file_not_marked_scratch(tmp_path: Path) -> None:
+    _seed_repo(tmp_path)
+    runner = Runner(client=DummyClient(), repo=tmp_path, model="m", stream=False)
+    before = runner._snapshot_repo_files()
+    (tmp_path / "src").mkdir(exist_ok=True)
+    (tmp_path / "src" / "new_module.py").write_text("x=1\n", encoding="utf-8")
+    created = runner._new_files_since(before)
+    runner._explicit_tool_created_files.update(created)
+    runner._provisional_scratch_candidates.update(created - runner._explicit_tool_created_files)
+    assert "src/new_module.py" not in runner._provisional_scratch_candidates
+
+
+def test_explicit_write_created_test_file_not_marked_scratch(tmp_path: Path) -> None:
+    _seed_repo(tmp_path)
+    runner = Runner(client=DummyClient(), repo=tmp_path, model="m", stream=False)
+    before = runner._snapshot_repo_files()
+    (tmp_path / "tests").mkdir(exist_ok=True)
+    (tmp_path / "tests" / "test_new.py").write_text("def test_x():\n    assert True\n", encoding="utf-8")
+    created = runner._new_files_since(before)
+    runner._explicit_tool_created_files.update(created)
+    runner._provisional_scratch_candidates.update(created - runner._explicit_tool_created_files)
+    assert "tests/test_new.py" not in runner._provisional_scratch_candidates

villani_code/state.py

@@ -518,6 +518,9 @@ def __init__(
         self._last_verification_fingerprint = ""
         self._repeated_stale_verification_count = 0
         self._last_verification_intentional: set[str] = set()
+        self._provisional_scratch_candidates: set[str] = set()
+        self._explicit_tool_created_files: set[str] = set()
+        self._cleaned_scratch_files: list[str] = []
         self._last_verification_artifact_count = 0
         self._last_validation_target = ""
         self._last_validation_summary = ""
@@ -1388,9 +1391,18 @@ def _budget_reason(
                         "transcript": transcript,
                     }
 
+                pre_repo_files: set[str] = set()
+                if tool_name in {"Bash", "Write", "Patch"}:
+                    pre_repo_files = set(self._snapshot_repo_files())
                 result = self._execute_tool_with_policy(
                     tool_name, tool_input, tool_use_id, len(messages)
                 )
+                if tool_name in {"Write", "Patch"}:
+                    self._explicit_tool_created_files.update(self._new_files_since(pre_repo_files) if pre_repo_files else [])
+                if tool_name == "Bash":
+                    self._provisional_scratch_candidates.update(
+                        self._new_files_since(pre_repo_files) - self._explicit_tool_created_files
+                    )
                 tool_calls_used += 1
                 if self.small_model:
                     result = self._truncate_tool_result(tool_name, result)
@@ -1555,6 +1567,21 @@ def _is_mutating_tool_call(
             return not command.startswith(readonly_prefixes)
         return False
 
+
+    def _snapshot_repo_files(self) -> set[str]:
+        files: set[str] = set()
+        for path in self.repo.rglob("*"):
+            if not path.is_file():
+                continue
+            rel = path.relative_to(self.repo).as_posix()
+            if rel.startswith(".git/"):
+                continue
+            files.add(rel)
+        return files
+
+    def _new_files_since(self, before: set[str]) -> set[str]:
+        return self._snapshot_repo_files() - set(before)
+
     def _dispatch_event(self, event: dict[str, Any]) -> None:
         if "turn_index" not in event and isinstance(self._current_turn_index, int):
             event = {**event, "turn_index": self._current_turn_index}

villani_code/state_runtime.py

@@ -161,6 +161,11 @@ def run_pre_edit_failure_localization(runner: Any) -> dict[str, Any] | None:
     visible_commands = list(getattr(cfg, "visible_verification", []) if cfg else [])
     visible_command = str(visible_commands[0]).strip() if visible_commands else ""
     expected_file = _single_clear_file(list(getattr(cfg, "expected_files", []) if cfg else []))
+    cleaned = _cleanup_provisional_scratch_after_success(runner)
+    if cleaned:
+        runner._cleaned_scratch_files = cleaned
+        runner.event_callback({"type": "scratch_cleanup_applied", "files": cleaned})
+        changed_files = [f for f in changed_files if f not in set(cleaned)]
     plan = getattr(runner, "_execution_plan", None)
     relevant_file = _single_clear_file(list(getattr(plan, "relevant_files", []) if plan else []))
     has_traceback = bool(getattr(runner, "_pending_verification", "").strip())
@@ -252,6 +257,11 @@ def classify_diagnosis_target_confidence(
         if excerpt_path and excerpt_path == target_file:
             return "strong"
 
+    cleaned = _cleanup_provisional_scratch_after_success(runner)
+    if cleaned:
+        runner._cleaned_scratch_files = cleaned
+        runner.event_callback({"type": "scratch_cleanup_applied", "files": cleaned})
+        changed_files = [f for f in changed_files if f not in set(cleaned)]
     plan = getattr(runner, "_execution_plan", None)
     relevant_file = _single_clear_file(list(getattr(plan, "relevant_files", []) if plan else []))
     if relevant_file and relevant_file == target_file:
@@ -291,6 +301,11 @@ def run_pre_edit_diagnosis(
 ) -> dict[str, str] | None:
     runner.event_callback({"type": "diagnosis_attempted"})
     evidence_lines = [f"Objective: {instruction.strip()}"]
+    cleaned = _cleanup_provisional_scratch_after_success(runner)
+    if cleaned:
+        runner._cleaned_scratch_files = cleaned
+        runner.event_callback({"type": "scratch_cleanup_applied", "files": cleaned})
+        changed_files = [f for f in changed_files if f not in set(cleaned)]
     plan = getattr(runner, "_execution_plan", None)
     if plan is not None:
         if getattr(plan, "validation_steps", None):
@@ -628,6 +643,11 @@ def _is_pytest_based_verification(runner: Any) -> bool:
     visible = list(getattr(cfg, "visible_verification", []) if cfg else [])
     if any("pytest" in str(cmd).lower() for cmd in visible):
         return True
+    cleaned = _cleanup_provisional_scratch_after_success(runner)
+    if cleaned:
+        runner._cleaned_scratch_files = cleaned
+        runner.event_callback({"type": "scratch_cleanup_applied", "files": cleaned})
+        changed_files = [f for f in changed_files if f not in set(cleaned)]
     plan = getattr(runner, "_execution_plan", None)
     steps = list(getattr(plan, "validation_steps", []) if plan else [])
     return any("pytest" in str(step).lower() for step in steps)
@@ -1278,11 +1298,69 @@ def ensure_project_memory_and_plan(runner: Any, instruction: str) -> None:
 
 
 
+
+
+def _conservative_source_roots(repo: Path) -> set[str] | None:
+    roots = set()
+    for candidate in ("src", "tests", "test", "package"):
+        if (repo / candidate).exists():
+            roots.add(candidate)
+    return roots or None
+
+
+def _cleanup_provisional_scratch_after_success(runner: Any) -> list[str]:
+    candidates = sorted(set(getattr(runner, "_provisional_scratch_candidates", set())))
+    if not candidates:
+        return []
+    roots = _conservative_source_roots(runner.repo)
+    if roots is None:
+        return []
+    changed = set(git_changed_files(runner.repo))
+    eligible: list[str] = []
+    for rel in candidates:
+        p = runner.repo / rel
+        if not p.exists() or not p.is_file():
+            continue
+        if rel in changed:
+            continue
+        if any(rel == r or rel.startswith(r + "/") for r in roots):
+            continue
+        eligible.append(rel)
+    if not eligible:
+        return []
+    visible = list(getattr(getattr(runner, "benchmark_config", None), "visible_verification", []) or [])
+    if len(visible) != 1:
+        return []
+    cmd = visible[0]
+    backups: dict[str, bytes] = {}
+    for rel in eligible:
+        fp = runner.repo / rel
+        backups[rel] = fp.read_bytes()
+        fp.unlink(missing_ok=True)
+    ok = True
+    try:
+        proc = subprocess.run(cmd, shell=True, cwd=runner.repo, capture_output=True, text=True)
+        ok = proc.returncode == 0
+    except Exception:
+        ok = False
+    if not ok:
+        for rel, data in backups.items():
+            fp = runner.repo / rel
+            fp.parent.mkdir(parents=True, exist_ok=True)
+            fp.write_bytes(data)
+        return []
+    return eligible
+
 def run_post_execution_validation(runner: Any, changed_files: list[str]) -> str:
     if getattr(runner, "_planning_read_only", False):
         return ""
     if not changed_files:
         return ""
+    cleaned = _cleanup_provisional_scratch_after_success(runner)
+    if cleaned:
+        runner._cleaned_scratch_files = cleaned
+        runner.event_callback({"type": "scratch_cleanup_applied", "files": cleaned})
+        changed_files = [f for f in changed_files if f not in set(cleaned)]
     plan = getattr(runner, "_execution_plan", None)
     plan_impact = getattr(plan, "change_impact", None)
     plan_actions = list(getattr(plan, "action_classes", [])) if plan else []