From 6e3c2336db53f62aa45ad18e8424084793dba3bd Mon Sep 17 00:00:00 2001 From: mmprotest Date: Thu, 9 Apr 2026 18:09:59 +1000 Subject: [PATCH] Add malformed tool recovery and shell preflight guards --- tests/test_loop.py | 89 +++++++++++++++++++++++++++++++++++ tests/test_tools_hardening.py | 24 +++++++++- villani_code/state.py | 49 +++++++++++++++++++ villani_code/tools.py | 42 +++++++++++++++++ 4 files changed, 203 insertions(+), 1 deletion(-) diff --git a/tests/test_loop.py b/tests/test_loop.py index 6ad25ecf..a52bfff7 100644 --- a/tests/test_loop.py +++ b/tests/test_loop.py @@ -214,6 +214,59 @@ def create_message(self, payload, stream): } +class FakeClientMalformedWriteRecovery: + def __init__(self) -> None: + self.calls = 0 + self.payloads: list[dict] = [] + + def create_message(self, payload, stream): + self.calls += 1 + self.payloads.append(payload) + if self.calls <= 3: + return { + "id": str(self.calls), + "role": "assistant", + "content": [{"type": "tool_use", "id": f"tool-{self.calls}", "name": "Write", "input": {}}], + } + return {"id": "4", "role": "assistant", "content": [{"type": "text", "text": "done"}]} + + +class FakeClientMalformedThenValidWrite: + def __init__(self) -> None: + self.calls = 0 + self.payloads: list[dict] = [] + + def create_message(self, payload, stream): + self.calls += 1 + self.payloads.append(payload) + if self.calls <= 3: + return { + "id": f"bad-{self.calls}", + "role": "assistant", + "content": [{"type": "tool_use", "id": f"tool-bad-{self.calls}", "name": "Write", "input": {}}], + } + if self.calls == 4: + return { + "id": "good", + "role": "assistant", + "content": [ + { + "type": "tool_use", + "id": "tool-good", + "name": "Write", + "input": {"file_path": "ok.txt", "content": "ok"}, + } + ], + } + if self.calls == 5: + return { + "id": "bad-again", + "role": "assistant", + "content": [{"type": "tool_use", "id": "tool-bad-again", "name": "Write", "input": {}}], + } + return {"id": "done", "role": "assistant", "content": [{"type": "text", "text": "done"}]} + + def test_loop_retries_on_empty_assistant_turn(tmp_path: Path): client = FakeClientEmptyThenDone() runner = Runner(client=client, repo=tmp_path, model="m", stream=False) @@ -233,6 +286,42 @@ def test_loop_retries_on_empty_assistant_turn(tmp_path: Path): assert continuation_messages +def test_repeated_malformed_tool_calls_inject_compact_recovery_message(tmp_path: Path) -> None: + client = FakeClientMalformedWriteRecovery() + runner = Runner(client=client, repo=tmp_path, model="m", stream=False) + + runner.run("write a file") + + recovery_messages = [ + block.get("text", "") + for payload in client.payloads[1:] + for message in payload.get("messages", []) + if message.get("role") == "user" + for block in message.get("content", []) + if block.get("type") == "text" and "Recent Write calls were invalid" in block.get("text", "") + ] + assert recovery_messages + assert "Issue one valid Write call using the tool schema, or inspect relevant context/files before retrying." in recovery_messages[-1] + + +def test_malformed_tool_recovery_state_clears_after_valid_call(tmp_path: Path) -> None: + client = FakeClientMalformedThenValidWrite() + runner = Runner(client=client, repo=tmp_path, model="m", stream=False) + runner.run("write a file") + + recovery_injections = [ + payload + for payload in client.payloads + if payload.get("messages") + and payload["messages"][-1].get("role") == "user" + and payload["messages"][-1].get("content") + and payload["messages"][-1]["content"][0].get("type") == "text" + and "Recent Write calls were invalid" + in payload["messages"][-1]["content"][0].get("text", "") + ] + assert len(recovery_injections) == 1 + + def test_tool_result_followup_is_pure_tool_result_message(tmp_path: Path): client = FakeClientToolUseThenDone() runner = Runner(client=client, repo=tmp_path, model="m", stream=False) diff --git a/tests/test_tools_hardening.py b/tests/test_tools_hardening.py index d2f0f020..870580ce 100644 --- a/tests/test_tools_hardening.py +++ b/tests/test_tools_hardening.py @@ -2,7 +2,7 @@ import pytest -from villani_code.tools import GitSimpleInput, LsInput, _safe_path +from villani_code.tools import GitSimpleInput, LsInput, _preflight_shell_command, _safe_path def test_safe_path_accepts_in_repo_path(tmp_path: Path) -> None: @@ -29,3 +29,25 @@ def test_model_defaults_do_not_share_list_state() -> None: git_two = GitSimpleInput() git_one.args.append("status") assert git_two.args == [] + + +def test_shell_preflight_rejects_windows_heredoc() -> None: + reason = _preflight_shell_command("python < None: + assert _preflight_shell_command("python < None: + reason = _preflight_shell_command("x" * 9000, env=("windows", "cmd")) + assert reason is not None + assert "command-length limits" in reason + + +def test_shell_preflight_rejects_clearly_bash_specific_windows_form() -> None: + reason = _preflight_shell_command("set -euo pipefail && echo ok", env=("windows", "cmd")) + assert reason is not None + assert "bash-specific" in reason diff --git a/villani_code/state.py b/villani_code/state.py index 1a2c80d4..5a702d77 100644 --- a/villani_code/state.py +++ b/villani_code/state.py @@ -498,6 +498,9 @@ def __init__( self._no_progress_cycles = 0 self._recovery_count = 0 self._last_failed_tool_sig = "" + self._tool_attempt_counter = 0 + self._malformed_tool_attempts: dict[str, list[int]] = {} + self._malformed_tool_last_recovery_turn: dict[str, int] = {} self._repo_map = "" self._retriever: Retriever | None = None self._context_budget = ( @@ -1355,6 +1358,7 @@ def _budget_reason( } tool_results: list[dict[str, Any]] = [] + protocol_recovery_messages: list[str] = [] for block in tool_uses: tool_name = block.get("name", "") tool_input = dict(block.get("input", {})) @@ -1391,6 +1395,9 @@ def _budget_reason( result = self._execute_tool_with_policy( tool_name, tool_input, tool_use_id, len(messages) ) + recovery_message = self._track_malformed_tool_call_and_recovery(tool_name, result) + if recovery_message: + protocol_recovery_messages.append(recovery_message) tool_calls_used += 1 if self.small_model: result = self._truncate_tool_result(tool_name, result) @@ -1529,6 +1536,13 @@ def _budget_reason( ) self._pending_verification = "" messages.append({"role": "user", "content": next_user_content}) + for recovery_message in _dedupe_preserve(protocol_recovery_messages): + messages.append( + { + "role": "user", + "content": [{"type": "text", "text": recovery_message}], + } + ) reason = _budget_reason() if reason: @@ -1640,6 +1654,41 @@ def _execute_tool_with_policy( message_count, ) + def _is_malformed_tool_result(self, tool_name: str, result: dict[str, Any]) -> bool: + if not result.get("is_error"): + return False + content = str(result.get("content", "")) + return f"Invalid input for {tool_name}:" in content + + def _track_malformed_tool_call_and_recovery(self, tool_name: str, result: dict[str, Any]) -> str: + self._tool_attempt_counter += 1 + malformed = self._is_malformed_tool_result(tool_name, result) + attempts = self._malformed_tool_attempts.setdefault(tool_name, []) + window = 6 + threshold = 3 + if malformed: + attempts.append(self._tool_attempt_counter) + cutoff = self._tool_attempt_counter - window + 1 + self._malformed_tool_attempts[tool_name] = [idx for idx in attempts if idx >= cutoff] + if len(self._malformed_tool_attempts[tool_name]) >= threshold: + last_recovery_turn = self._malformed_tool_last_recovery_turn.get(tool_name, -1) + if self._tool_attempt_counter != last_recovery_turn: + self._malformed_tool_last_recovery_turn[tool_name] = self._tool_attempt_counter + return ( + f"Recent {tool_name} calls were invalid because required fields were missing or malformed. " + f"Issue one valid {tool_name} call using the tool schema, or inspect relevant context/files before retrying." + ) + return "" + + if not result.get("is_error"): + self._malformed_tool_attempts.pop(tool_name, None) + self._malformed_tool_last_recovery_turn.pop(tool_name, None) + return "" + + cutoff = self._tool_attempt_counter - window + 1 + self._malformed_tool_attempts[tool_name] = [idx for idx in attempts if idx >= cutoff] + return "" + def _build_tool_result_event_payload( self, tool_name: str, tool_use_id: str, result: dict[str, Any] ) -> dict[str, Any]: diff --git a/villani_code/tools.py b/villani_code/tools.py index 423807ed..19e0a045 100644 --- a/villani_code/tools.py +++ b/villani_code/tools.py @@ -2,6 +2,9 @@ import glob import json +import os +import platform +import re import shutil import subprocess from pathlib import Path @@ -113,6 +116,10 @@ class SubmitPlanInput(BaseModel): } DENYLIST = ["rm -rf", "del /s", "format ", "mkfs", "dd if=", "curl ", "wget "] +_WINDOWS_HEREDOC_RE = re.compile(r"<<\s*[A-Za-z0-9_'-]+") +_WINDOWS_BASH_ONLY_RE = re.compile(r"(^|\s)(set -euo pipefail|\[\[|2>/dev/null|\|\s*tail\b)") +_WINDOWS_COMMAND_LENGTH_LIMIT = 8000 +_GENERIC_COMMAND_LENGTH_LIMIT = 32768 def _error(message: str) -> dict[str, Any]: @@ -236,6 +243,9 @@ def _run_search(data: SearchInput, repo: Path) -> str: def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | None = None, tool_call_id: str = "") -> str: + guard_error = _preflight_shell_command(data.command) + if guard_error: + raise ValueError(guard_error) lowered = data.command.lower() if not unsafe: for bad in DENYLIST: @@ -261,6 +271,38 @@ def _run_bash(data: BashInput, repo: Path, unsafe: bool, debug_callback: Any | N return json.dumps({"command": data.command, "exit_code": proc.returncode, "stdout": proc.stdout, "stderr": proc.stderr}, indent=2) +def _detect_shell_environment() -> tuple[str, str]: + shell_hint = str(os.environ.get("COMSPEC") or os.environ.get("SHELL") or "").lower() + platform_hint = platform.system().lower() + if os.name == "nt" or platform_hint.startswith("win") or "powershell" in shell_hint or "cmd" in shell_hint: + if "powershell" in shell_hint or "pwsh" in shell_hint: + return "windows", "powershell" + return "windows", "cmd" + return "posix", "bash" + + +def _preflight_shell_command(command: str, env: tuple[str, str] | None = None) -> str | None: + family, shell_name = env or _detect_shell_environment() + raw = str(command or "") + if len(raw) > _GENERIC_COMMAND_LENGTH_LIMIT: + return ( + "Rejected command: this command form is likely to exceed shell command-length limits. " + "Use a smaller command or another method." + ) + if family != "windows": + return None + if len(raw) > _WINDOWS_COMMAND_LENGTH_LIMIT: + return ( + "Rejected command: this command form is likely to exceed shell command-length limits. " + "Use a smaller command or another method." + ) + if _WINDOWS_HEREDOC_RE.search(raw): + return "Rejected command: heredoc-style redirection is incompatible with this Windows shell environment." + if shell_name in {"cmd", "powershell"} and _WINDOWS_BASH_ONLY_RE.search(raw): + return "Rejected command: this bash-specific form is incompatible with this Windows shell environment." + return None + + def _run_write(data: WriteInput, repo: Path, debug_callback: Any | None = None, tool_call_id: str = "") -> str: path = _safe_path(repo, data.file_path) if data.mkdirs: