Improve protocol recovery for malformed tool calls and add Windows-aware shell guardrails

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)

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 <<EOF\nprint('x')\nEOF", env=("windows", "powershell"))
+    assert reason is not None
+    assert "heredoc-style redirection" in reason
+
+
+def test_shell_preflight_does_not_overblock_posix_heredoc() -> None:
+    assert _preflight_shell_command("python <<EOF\nprint('x')\nEOF", env=("posix", "bash")) is None
+
+
+def test_shell_preflight_rejects_excessive_command_length() -> 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

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

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: