diff --git a/tests/test_scratch_cleanup_gate.py b/tests/test_scratch_cleanup_gate.py new file mode 100644 index 00000000..aa3dc693 --- /dev/null +++ b/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() diff --git a/villani_code/debug_recorder.py b/villani_code/debug_recorder.py index 0628962f..9a8a536b 100644 --- a/villani_code/debug_recorder.py +++ b/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 diff --git a/villani_code/state.py b/villani_code/state.py index 1a2c80d4..6234b285 100644 --- a/villani_code/state.py +++ b/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, diff --git a/villani_code/state_runtime.py b/villani_code/state_runtime.py index 9e03824d..91199331 100644 --- a/villani_code/state_runtime.py +++ b/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 diff --git a/villani_code/state_tooling.py b/villani_code/state_tooling.py index c84c64ca..6c8146fd 100644 --- a/villani_code/state_tooling.py +++ b/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,