diff --git a/tests/test_scratch_redirection.py b/tests/test_scratch_redirection.py new file mode 100644 index 00000000..fcfbe3d2 --- /dev/null +++ b/tests/test_scratch_redirection.py @@ -0,0 +1,48 @@ +from pathlib import Path + +from villani_code.benchmark.policy import filter_meaningful_touched_paths +from villani_code.tools import BashInput, ReadInput, WriteInput, execute_tool + + +def test_new_root_level_scratch_file_is_redirected(tmp_path: Path) -> None: + out = execute_tool("Write", WriteInput(file_path="test_fix.py", content="print('x')").model_dump(), tmp_path) + assert out["is_error"] is False + assert not (tmp_path / "test_fix.py").exists() + assert (tmp_path / ".villani_code" / "scratch" / "default" / "test_fix.py").exists() + + +def test_existing_root_level_scratch_file_is_not_redirected(tmp_path: Path) -> None: + root_file = tmp_path / "test_fix.py" + root_file.write_text("old", encoding="utf-8") + execute_tool("Write", WriteInput(file_path="test_fix.py", content="new").model_dump(), tmp_path) + assert root_file.read_text(encoding="utf-8") == "new" + + +def test_existing_directory_paths_are_not_redirected(tmp_path: Path) -> None: + (tmp_path / "tests").mkdir() + (tmp_path / "src").mkdir() + execute_tool("Write", WriteInput(file_path="tests/test_fix.py", content="a").model_dump(), tmp_path) + execute_tool("Write", WriteInput(file_path="src/test_fix.py", content="b").model_dump(), tmp_path) + assert (tmp_path / "tests" / "test_fix.py").exists() + assert (tmp_path / "src" / "test_fix.py").exists() + + +def test_redirected_scratch_read_resolves_alias(tmp_path: Path) -> None: + execute_tool("Write", WriteInput(file_path="verify.py", content="ok").model_dump(), tmp_path) + out = execute_tool("Read", ReadInput(file_path="verify.py").model_dump(), tmp_path) + assert out["content"] == "ok" + + +def test_redirected_scratch_exec_keeps_repo_imports(tmp_path: Path) -> None: + (tmp_path / "pkg.py").write_text("VALUE = 7\n", encoding="utf-8") + script = "import pkg\nprint(pkg.VALUE)\n" + execute_tool("Write", WriteInput(file_path="verify.py", content=script).model_dump(), tmp_path) + out = execute_tool("Bash", BashInput(command="python verify.py").model_dump(), tmp_path) + assert '"exit_code": 0' in out["content"] + assert "7" in out["content"] + + +def test_final_diff_and_patch_summary_filters_runtime_scratch_paths() -> None: + touched = ["src/app.py", ".villani_code/scratch/default/test_fix.py"] + meaningful = filter_meaningful_touched_paths(touched) + assert meaningful == ["src/app.py"] diff --git a/villani_code/tools.py b/villani_code/tools.py index 423807ed..6ffd03d8 100644 --- a/villani_code/tools.py +++ b/villani_code/tools.py @@ -2,6 +2,8 @@ import glob import json +import os +import shlex import shutil import subprocess from pathlib import Path @@ -114,6 +116,8 @@ class SubmitPlanInput(BaseModel): DENYLIST = ["rm -rf", "del /s", "format ", "mkfs", "dd if=", "curl ", "wget "] +SCRATCH_FILENAMES = {"test_fix.py", "check_fix.py", "verify.py", "verify_fix.py", "repro.py", "debug.py", "tmp_test.py", "test_resolver.py"} + def _error(message: str) -> dict[str, Any]: return {"content": message, "is_error": True} @@ -190,6 +194,70 @@ def _safe_path(repo: Path, raw: str) -> Path: return path + + +def _scratch_root(repo: Path) -> Path: + mission_id = os.getenv("VILLANI_MISSION_ID", "").strip() + if mission_id: + return repo / ".villani_code" / "missions" / mission_id / "scratch" + run_id = os.getenv("VILLANI_RUN_ID", "").strip() or os.getenv("VILLANI_TASK_ID", "").strip() + if run_id: + return repo / ".villani_code" / "runs" / run_id / "scratch" + return repo / ".villani_code" / "scratch" / "default" + + +def _resolve_scratch_alias(repo: Path, raw: str) -> Path: + normalized = raw.strip().replace("\\", "/") + candidate = _safe_path(repo, normalized) + if normalized in {"", "."} or "/" in normalized or candidate.exists(): + return candidate + if Path(normalized).name not in SCRATCH_FILENAMES: + return candidate + scratch_candidate = (_scratch_root(repo) / Path(normalized).name).resolve() + try: + scratch_candidate.relative_to(repo.resolve()) + except ValueError as exc: + raise ValueError("Path escapes repository") from exc + if scratch_candidate.exists(): + return scratch_candidate + return candidate + + +def _resolve_write_target(repo: Path, raw: str) -> tuple[Path, str | None]: + path = _safe_path(repo, raw) + normalized = raw.strip().replace("\\", "/") + if path.exists() or "/" in normalized or Path(normalized).name not in SCRATCH_FILENAMES: + return path, None + redirected = (_scratch_root(repo) / Path(normalized).name).resolve() + try: + redirected.relative_to(repo.resolve()) + except ValueError as exc: + raise ValueError("Path escapes repository") from exc + message = ( + f"Created scratch file {redirected.relative_to(repo)} instead of repo/{Path(normalized).name} " + "because it matched a temporary verification filename." + ) + return redirected, message + + +def _rewrite_bash_scratch_aliases(command: str, cwd: Path, repo: Path) -> str: + tokens = shlex.split(command) + rewritten = list(tokens) + changed = False + for i, token in enumerate(tokens): + if token.startswith("-"): + continue + try: + candidate = _safe_path(repo, str((cwd / token).relative_to(repo)) if not Path(token).is_absolute() else token) + except Exception: + continue + if candidate.exists() or "/" in token or Path(token).name not in SCRATCH_FILENAMES: + continue + scratch = _resolve_scratch_alias(repo, Path(token).name) + if scratch != candidate and scratch.exists(): + rewritten[i] = str(scratch.relative_to(cwd)) + changed = True + return shlex.join(rewritten) if changed else command def _run_ls(data: LsInput, repo: Path) -> str: target = _safe_path(repo, data.path) lines = [] @@ -201,7 +269,7 @@ def _run_ls(data: LsInput, repo: Path) -> str: def _run_read(data: ReadInput, repo: Path, debug_callback: Any | None = None, tool_call_id: str = "") -> str: - path = _safe_path(repo, data.file_path) + path = _resolve_scratch_alias(repo, data.file_path) raw = path.read_bytes()[: data.max_bytes] if callable(debug_callback): debug_callback("file_read", {"file_path": data.file_path, "size_bytes": len(raw), "ok": True, "tool_call_id": tool_call_id}) @@ -242,14 +310,19 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N if bad in lowered: raise ValueError(f"Refusing command: {bad.strip()}") cwd = _safe_path(repo, data.cwd) + command = _rewrite_bash_scratch_aliases(data.command, cwd, repo) if callable(debug_callback): debug_callback("command_started", {"command": data.command, "cwd": data.cwd, "tool_call_id": tool_call_id}) - proc = subprocess.run(data.command, shell=True, cwd=str(cwd), capture_output=True, text=True, timeout=data.timeout_sec) + env = dict(os.environ) + repo_path = str(repo.resolve()) + existing_pythonpath = env.get("PYTHONPATH", "") + env["PYTHONPATH"] = repo_path if not existing_pythonpath else f"{repo_path}{os.pathsep}{existing_pythonpath}" + proc = subprocess.run(command, shell=True, cwd=str(cwd), capture_output=True, text=True, timeout=data.timeout_sec, env=env) if callable(debug_callback): debug_callback( "command_finished", { - "command": data.command, + "command": command, "cwd": data.cwd, "exit_code": proc.returncode, "stdout": proc.stdout, @@ -258,11 +331,11 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N "tool_call_id": tool_call_id, }, ) - return json.dumps({"command": data.command, "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) + return json.dumps({"command": command, "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) def _run_write(data: WriteInput, repo: Path, debug_callback: Any | None = None, tool_call_id: str = "") -> str: - path = _safe_path(repo, data.file_path) + path, redirected_message = _resolve_write_target(repo, data.file_path) if data.mkdirs: path.parent.mkdir(parents=True, exist_ok=True) path.write_text(data.content, encoding="utf-8") @@ -276,6 +349,8 @@ def _run_write(data: WriteInput, repo: Path, debug_callback: Any | None = None, "tool_call_id": tool_call_id, }, ) + if redirected_message: + return f"Wrote {path}\n{redirected_message}" return f"Wrote {path}"