From bf8e697b37571f721a13e388c20017552d095d8b Mon Sep 17 00:00:00 2001 From: mmprotest Date: Sat, 25 Apr 2026 19:51:48 +1000 Subject: [PATCH] Add generic patch-recovery controller for failed patch loops --- tests/test_patch_recovery_controller.py | 211 ++++++++++++++++++++++++ villani_code/mission_state.py | 6 + villani_code/state.py | 22 +++ villani_code/state_runtime.py | 196 ++++++++++++++++++++++ villani_code/state_tooling.py | 97 +++++++++++ 5 files changed, 532 insertions(+) create mode 100644 tests/test_patch_recovery_controller.py diff --git a/tests/test_patch_recovery_controller.py b/tests/test_patch_recovery_controller.py new file mode 100644 index 00000000..49de472e --- /dev/null +++ b/tests/test_patch_recovery_controller.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from pathlib import Path + +from villani_code.state import Runner + + +class _PatchFailThenDoneClient: + def __init__(self) -> None: + self.calls = 0 + self.second_payload: dict | None = None + + def create_message(self, payload, stream): + self.calls += 1 + if self.calls == 1: + return { + "id": "1", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-1", + "name": "Patch", + "input": { + "file_path": "src/a.py", + "unified_diff": "--- a/src/a.py\n+++ b/src/a.py\n@@ -1 +1 @@\n-missing\n+patched\n", + }, + } + ], + } + self.second_payload = payload + return {"id": "2", "role": "assistant", "content": [{"type": "text", "text": "done"}]} + + +class _PatchSuccessThenDoneClient: + def __init__(self) -> None: + self.calls = 0 + self.second_payload: dict | None = None + + def create_message(self, payload, stream): + self.calls += 1 + if self.calls == 1: + return { + "id": "1", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-1", + "name": "Patch", + "input": { + "file_path": "src/a.py", + "unified_diff": "--- a/src/a.py\n+++ b/src/a.py\n@@ -1 +1 @@\n-old\n+new\n", + }, + } + ], + } + self.second_payload = payload + return {"id": "2", "role": "assistant", "content": [{"type": "text", "text": "done"}]} + + +class _PatchFailThenBashThenDoneClient: + def __init__(self, command: str) -> None: + self.calls = 0 + self.command = command + self.third_payload: dict | None = None + + def create_message(self, payload, stream): + self.calls += 1 + if self.calls == 1: + return { + "id": "1", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-1", + "name": "Patch", + "input": { + "file_path": "src/a.py", + "unified_diff": "--- a/src/a.py\n+++ b/src/a.py\n@@ -1 +1 @@\n-missing\n+patched\n", + }, + } + ], + } + if self.calls == 2: + return { + "id": "2", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-2", + "name": "Bash", + "input": {"command": self.command}, + } + ], + } + self.third_payload = payload + return {"id": "3", "role": "assistant", "content": [{"type": "text", "text": "done"}]} + + +def _find_injected_recovery_text(payload: dict) -> str: + for message in payload.get("messages", []): + if message.get("role") != "user": + continue + for block in message.get("content", []): + if not isinstance(block, dict) or block.get("type") != "text": + continue + text = str(block.get("text", "")) + if "Patch recovery: the previous patch failed" in text: + return text + return "" + + +def test_patch_failure_records_structured_state_and_refreshes_context(tmp_path: Path) -> None: + target = tmp_path / "src" / "a.py" + target.parent.mkdir(parents=True) + target.write_text("old\nsecond\n", encoding="utf-8") + + client = _PatchFailThenDoneClient() + runner = Runner(client=client, repo=tmp_path, model="m", stream=False, plan_mode="off") + runner.run("fix whatever task shape") + + assert runner._mission_state is not None + patch_state = runner._mission_state.last_patch_failure + assert patch_state["target_file"] == "src/a.py" + assert patch_state["next_required_action"] == "retry_patch_against_refreshed_context" + assert "src/a.py (total_lines=" in patch_state["refreshed_context"] + assert "old" in patch_state["refreshed_context"] + assert patch_state["attempts_for_file"] == 1 + assert runner._mission_state.recent_tool_failures + + +def test_patch_failure_injects_single_turn_recovery_instruction(tmp_path: Path) -> None: + target = tmp_path / "src" / "a.py" + target.parent.mkdir(parents=True) + target.write_text("old\n", encoding="utf-8") + + client = _PatchFailThenDoneClient() + runner = Runner(client=client, repo=tmp_path, model="m", stream=False, plan_mode="off") + runner.run("some benchmark fastapi language-specific task") + + assert client.second_payload is not None + injected = _find_injected_recovery_text(client.second_payload) + assert "target_file: src/a.py" in injected + assert "patch_error:" in injected + assert "refreshed_context:" in injected + + +def test_patch_recovery_blocks_immediate_bash_rewrite_of_same_target(tmp_path: Path) -> None: + target = tmp_path / "src" / "a.py" + target.parent.mkdir(parents=True) + target.write_text("old\n", encoding="utf-8") + events: list[dict] = [] + command = "python -c \"from pathlib import Path; Path('src/a.py').write_text('rewritten\\n')\"" + client = _PatchFailThenBashThenDoneClient(command) + runner = Runner( + client=client, + repo=tmp_path, + model="m", + stream=False, + plan_mode="off", + event_callback=events.append, + ) + runner.run("generic task") + + assert client.third_payload is not None + tool_result_messages = [ + message + for message in client.third_payload["messages"] + if message.get("role") == "user" + and message.get("content") + and isinstance(message["content"][0], dict) + and message["content"][0].get("type") == "tool_result" + ] + assert any("Patch recovery active" in str(block.get("content", "")) for msg in tool_result_messages for block in msg["content"]) + assert any(event.get("type") == "patch_recovery_bash_blocked" for event in events) + + +def test_patch_recovery_allows_read_only_and_test_bash_commands(tmp_path: Path) -> None: + target = tmp_path / "src" / "a.py" + target.parent.mkdir(parents=True) + target.write_text("old\n", encoding="utf-8") + + client_read = _PatchFailThenBashThenDoneClient("cat src/a.py") + runner_read = Runner(client=client_read, repo=tmp_path, model="m", stream=False, plan_mode="off") + runner_read.run("generic task") + assert client_read.third_payload is not None + assert "Patch recovery active" not in str(client_read.third_payload) + + client_test = _PatchFailThenBashThenDoneClient("pytest -q") + runner_test = Runner(client=client_test, repo=tmp_path, model="m", stream=False, plan_mode="off") + runner_test.run("generic task") + assert client_test.third_payload is not None + assert "Patch recovery active" not in str(client_test.third_payload) + + +def test_patch_recovery_does_not_trigger_when_patch_succeeds(tmp_path: Path) -> None: + target = tmp_path / "src" / "a.py" + target.parent.mkdir(parents=True) + target.write_text("old\n", encoding="utf-8") + + client = _PatchSuccessThenDoneClient() + runner = Runner(client=client, repo=tmp_path, model="m", stream=False, plan_mode="off") + runner.run("generic task") + + assert client.second_payload is not None + assert _find_injected_recovery_text(client.second_payload) == "" + assert runner._mission_state is not None + assert runner._mission_state.last_patch_failure == {} diff --git a/villani_code/mission_state.py b/villani_code/mission_state.py index 8552f94d..6b761605 100644 --- a/villani_code/mission_state.py +++ b/villani_code/mission_state.py @@ -61,6 +61,9 @@ class MissionState: autonomous_satisfied_keys_summary: list[str] = field(default_factory=list) autonomous_blockers_summary: list[str] = field(default_factory=list) autonomous_stop_reason: str = "" + last_patch_failure: dict[str, Any] = field(default_factory=dict) + recent_tool_failures: list[str] = field(default_factory=list) + dirty_temp_files: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: return asdict(self) @@ -92,6 +95,9 @@ def from_dict(cls, payload: dict[str, Any]) -> "MissionState": autonomous_satisfied_keys_summary=[str(v) for v in payload.get("autonomous_satisfied_keys_summary", [])], autonomous_blockers_summary=[str(v) for v in payload.get("autonomous_blockers_summary", [])], autonomous_stop_reason=str(payload.get("autonomous_stop_reason", "")), + last_patch_failure=dict(payload.get("last_patch_failure", {}) or {}), + recent_tool_failures=[str(v) for v in payload.get("recent_tool_failures", [])], + dirty_temp_files=[str(v) for v in payload.get("dirty_temp_files", [])], ) diff --git a/villani_code/state.py b/villani_code/state.py index 1a2c80d4..17beee34 100644 --- a/villani_code/state.py +++ b/villani_code/state.py @@ -528,6 +528,9 @@ def __init__( self._patch_sanity_retry_pending = False self._first_attempt_write_lock_active = False self._first_attempt_locked_target = "" + self._patch_recovery_state: dict[str, Any] | None = None + self._patch_recovery_pending_turn = False + self._patch_recovery_attempts_by_file: dict[str, int] = {} self._context_governance = ContextGovernanceManager(self.repo) self._planning_read_only = False self._runtime_mode: Literal["execution", "planning"] = "execution" @@ -1409,6 +1412,18 @@ def _budget_reason( if result.get("is_error"): result_text = str(result.get("content", "")) self._update_mission_state(last_failed_command=f"{tool_name} {tool_input}", last_failed_summary=result_text[:500]) + if tool_name in {"Patch", "Edit"}: + from villani_code import state_runtime + + state_runtime.record_patch_failure( + self, + tool_name=tool_name, + tool_input=tool_input, + error_text=result_text, + turn_index=self._current_turn_index + if isinstance(self._current_turn_index, int) + else turns_used, + ) if ( self.benchmark_config.enabled and tool_name in {"Write", "Patch"} @@ -1446,6 +1461,13 @@ def _budget_reason( } ) + if tool_name in {"Patch", "Edit"} and not result.get("is_error"): + from villani_code import state_runtime + + state_runtime.clear_patch_recovery_if_target_succeeded( + self, tool_input=tool_input + ) + self.hooks.run_event( "PostToolUse", { diff --git a/villani_code/state_runtime.py b/villani_code/state_runtime.py index 9e03824d..27cd178e 100644 --- a/villani_code/state_runtime.py +++ b/villani_code/state_runtime.py @@ -26,6 +26,7 @@ from villani_code.mission_state import save_mission_state from villani_code.summarizer import summarize_mission_state from villani_code.utils import ensure_dir +from villani_code.patch_apply import extract_unified_diff_targets, parse_unified_diff _DIAGNOSIS_KEYS = ("target_file", "bug_class", "fix_intent") @@ -371,6 +372,7 @@ def inject_diagnosis_hint(messages: list[dict[str, Any]], diagnosis: dict[str, s def prepare_messages_for_model(runner: Any, messages: list[dict[str, Any]]) -> list[dict[str, Any]]: prepared = deepcopy(messages) + _inject_patch_recovery_instruction(runner, prepared) if runner.small_model: inject_retrieval_briefing(runner, prepared) if runner._context_budget: @@ -458,6 +460,200 @@ def inject_retrieval_briefing(runner: Any, messages: list[dict[str, Any]]) -> No content.insert(0, {"type": "text", "text": f"\n{briefing}\n"}) +def _extract_patch_target_file(tool_input: dict[str, Any]) -> str: + explicit = _normalize_repo_path(str(tool_input.get("file_path", ""))) + if explicit: + return explicit + diff = str(tool_input.get("unified_diff", "")) + try: + targets = extract_unified_diff_targets(diff, default_file_path=None) + except Exception: + targets = [] + for candidate in targets: + normalized = _normalize_repo_path(candidate) + if normalized: + return normalized + return "" + + +def _extract_failed_hunk_snippet(tool_input: dict[str, Any], target_file: str) -> str: + diff = str(tool_input.get("unified_diff", "")) + if not diff.strip(): + return "" + try: + parsed = parse_unified_diff(diff) + except Exception: + return diff[:1200] + normalized_target = _normalize_repo_path(target_file) + for file_patch in parsed: + candidate = file_patch.new_path if file_patch.new_path != "/dev/null" else file_patch.old_path + if _normalize_repo_path(candidate) != normalized_target: + continue + if not file_patch.hunks: + continue + lines = file_patch.hunks[0].lines[:40] + return "\n".join(lines)[:1200] + return diff[:1200] + + +def _anchor_candidates_from_hunk(hunk_text: str) -> list[str]: + lines = [] + for raw in str(hunk_text or "").splitlines(): + if not raw: + continue + if raw.startswith(("-", " ")): + candidate = raw[1:].rstrip() + if candidate.strip(): + lines.append(candidate) + unique: list[str] = [] + for line in lines: + if line not in unique: + unique.append(line) + return unique[:12] + + +def _render_context_slice(file_lines: list[str], start: int, end: int, total_lines: int) -> str: + rendered: list[str] = [] + for idx in range(start, min(end, total_lines)): + rendered.append(f"{idx + 1:>5}: {file_lines[idx]}") + return "\n".join(rendered) + + +def _extract_refreshed_patch_context(repo: Path, target_file: str, failed_hunk: str) -> str: + normalized = _normalize_repo_path(target_file) + if not normalized: + return "No target file was identified from the failed patch." + target_path = (repo / normalized).resolve() + try: + target_path.relative_to(repo.resolve()) + except Exception: + return f"Target file is outside workspace: {normalized}" + if not target_path.exists() or not target_path.is_file(): + return f"Target file missing or unreadable: {normalized}" + text = target_path.read_text(encoding="utf-8", errors="replace") + lines = text.splitlines() + total_lines = len(lines) + if total_lines == 0: + return f"{normalized} is empty (0 lines)." + anchors = _anchor_candidates_from_hunk(failed_hunk) + anchor_line_index = None + for anchor in anchors: + for idx, line in enumerate(lines): + if anchor in line: + anchor_line_index = idx + break + if anchor_line_index is not None: + break + context_radius = 20 + hard_cap = 120 + if anchor_line_index is None: + end = min(total_lines, hard_cap) + body = _render_context_slice(lines, 0, end, total_lines) + return f"{normalized} (total_lines={total_lines}, anchor=not_found)\n{body}" + start = max(0, anchor_line_index - context_radius) + end = min(total_lines, anchor_line_index + context_radius + 1) + if end - start > hard_cap: + end = start + hard_cap + body = _render_context_slice(lines, start, end, total_lines) + return ( + f"{normalized} (total_lines={total_lines}, anchor_line={anchor_line_index + 1}, " + f"slice={start + 1}-{end})\n{body}" + ) + + +def record_patch_failure( + runner: Any, + *, + tool_name: str, + tool_input: dict[str, Any], + error_text: str, + turn_index: int | None, +) -> None: + target_file = _extract_patch_target_file(tool_input) + failed_hunk = _extract_failed_hunk_snippet(tool_input, target_file) + refreshed_context = _extract_refreshed_patch_context(runner.repo, target_file, failed_hunk) + attempts = 1 + if target_file: + attempts = int(getattr(runner, "_patch_recovery_attempts_by_file", {}).get(target_file, 0)) + 1 + runner._patch_recovery_attempts_by_file[target_file] = attempts + state = { + "tool": tool_name, + "target_file": target_file, + "error": str(error_text or "")[:800], + "failed_hunk": failed_hunk[:1200], + "refreshed_context": refreshed_context[:5000], + "attempts_for_file": attempts, + "turn_index": int(turn_index or 0), + "next_required_action": "retry_patch_against_refreshed_context", + } + runner._patch_recovery_state = state + runner._patch_recovery_pending_turn = True + runner.event_callback( + { + "type": "patch_failure_detected", + "target_file": target_file, + "error": state["error"], + "attempts_for_file": attempts, + } + ) + runner.event_callback( + { + "type": "patch_recovery_context_refreshed", + "target_file": target_file, + "context_preview": refreshed_context[:400], + } + ) + mission = getattr(runner, "_mission_state", None) + if mission is not None: + mission.last_patch_failure = dict(state) + summary = f"Patch failure: {target_file or 'unknown'}: {state['error']}" + recent = list(getattr(mission, "recent_tool_failures", [])) + recent.append(summary[:280]) + mission.recent_tool_failures = recent[-8:] + validation_failures = list(getattr(mission, "validation_failures", [])) + validation_failures.append(summary[:280]) + mission.validation_failures = validation_failures[-8:] + save_mission_state(runner.repo, mission) + + +def clear_patch_recovery_if_target_succeeded(runner: Any, *, tool_input: dict[str, Any]) -> None: + active = getattr(runner, "_patch_recovery_state", None) + if not isinstance(active, dict): + return + active_target = _normalize_repo_path(str(active.get("target_file", ""))) + if not active_target: + return + success_target = _extract_patch_target_file(tool_input) + if _normalize_repo_path(success_target) != active_target: + return + runner._patch_recovery_state = None + runner._patch_recovery_pending_turn = False + + +def _inject_patch_recovery_instruction(runner: Any, messages: list[dict[str, Any]]) -> None: + if not bool(getattr(runner, "_patch_recovery_pending_turn", False)): + return + state = getattr(runner, "_patch_recovery_state", None) + if not isinstance(state, dict): + return + text = ( + "Patch recovery: the previous patch failed. Do not use Bash/Python scripts to rewrite this target file yet. " + "Use the refreshed context to produce a smaller patch against current contents. " + "If patching is impossible, explain why in one sentence.\n" + f"target_file: {state.get('target_file', 'unknown')}\n" + f"patch_error: {state.get('error', '')}\n" + f"refreshed_context:\n{state.get('refreshed_context', '')}" + ) + if prepend_text_to_latest_safe_user_message(messages, text): + runner._patch_recovery_pending_turn = False + runner.event_callback( + { + "type": "patch_recovery_instruction_injected", + "target_file": state.get("target_file", ""), + } + ) + + def init_small_model_support(runner: Any) -> None: index_path = runner.repo / ".villani_code" / "index" / "index.json" if index_path.exists(): diff --git a/villani_code/state_tooling.py b/villani_code/state_tooling.py index c84c64ca..28473907 100644 --- a/villani_code/state_tooling.py +++ b/villani_code/state_tooling.py @@ -443,6 +443,98 @@ def _prepare_global_mutation_policy( return tool_name, tool_input, None +def _is_read_only_or_test_shell_command(command: str) -> bool: + lowered = str(command or "").strip().lower() + if not lowered: + return True + read_only_prefixes = ( + "pwd", + "ls", + "cat ", + "rg ", + "grep ", + "find ", + "head ", + "tail ", + "wc ", + "git status", + "git diff", + "git log", + "git show", + "git branch", + "git rev-parse", + "git ls-files", + "pytest", + "python -m pytest", + "uv run pytest", + "poetry run pytest", + "npm test", + "pnpm test", + "yarn test", + ) + if any(lowered.startswith(prefix) for prefix in read_only_prefixes): + return True + if "pytest" in lowered or " test" in lowered: + return True + return False + + +def _looks_like_shell_rewrite_of_target(command: str, target_file: str) -> bool: + lowered = str(command or "").strip().lower() + target = str(target_file or "").replace("\\", "/").lstrip("./").lower() + if not lowered or not target: + return False + target_basename = target.split("/")[-1] + if target not in lowered and target_basename not in lowered: + return False + rewrite_markers = ( + ">", + ">>", + "| tee", + "sed -i", + "perl -pi", + ".write_text(", + ".write_bytes(", + "open(", + " mv ", + " cp ", + " truncate ", + " dd ", + ) + return any(marker in lowered for marker in rewrite_markers) + + +def _validate_patch_recovery_bash_guard( + runner: Any, tool_name: str, tool_input: dict[str, Any] +) -> str | None: + if tool_name != "Bash": + return None + active = getattr(runner, "_patch_recovery_state", None) + if not isinstance(active, dict): + return None + target_file = str(active.get("target_file", "") or "").replace("\\", "/").lstrip("./") + attempts = int(active.get("attempts_for_file", 0) or 0) + if not target_file or attempts >= 2: + return None + command = str(tool_input.get("command", "")) + if _is_read_only_or_test_shell_command(command): + return None + if not _looks_like_shell_rewrite_of_target(command, target_file): + return None + runner.event_callback( + { + "type": "patch_recovery_bash_blocked", + "target_file": target_file, + "attempts_for_file": attempts, + } + ) + return ( + "Patch recovery active: previous patch failed for " + f"{target_file}. Retry a smaller Patch against refreshed context before " + "using Bash/Python to rewrite this file." + ) + + def execute_tool_with_policy( runner: Any, tool_name: str, @@ -491,6 +583,11 @@ def execute_tool_with_policy( return {"content": policy_error, "is_error": True} if runner.small_model: runner._tighten_tool_input(tool_name, tool_input) + patch_recovery_bash_error = _validate_patch_recovery_bash_guard( + runner, tool_name, tool_input + ) + if patch_recovery_bash_error: + return {"content": patch_recovery_bash_error, "is_error": True} if tool_name in {"Write", "Patch"}: _normalize_mutation_payload(tool_name, tool_input)