Redirect temporary root-level scratch files to .villani_code scratch storage

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"]

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}"