diff --git a/tests/test_state_runtime.py b/tests/test_state_runtime.py index 16e604a7..51dfd820 100644 --- a/tests/test_state_runtime.py +++ b/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 diff --git a/villani_code/state.py b/villani_code/state.py index 1a2c80d4..f8be3b24 100644 --- a/villani_code/state.py +++ b/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} diff --git a/villani_code/state_runtime.py b/villani_code/state_runtime.py index 9e03824d..4e9b2ec0 100644 --- a/villani_code/state_runtime.py +++ b/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 []