Add event-driven patch-recovery controller for failed Patch/Edit operations

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 == {}

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", [])],
         )
 
 

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",
                     {