feat: add JSON repair to tool-call pipeline (closes #781)

run_agent.py

@@ -106,7 +106,7 @@ from agent.trajectory import (
     convert_scratchpad_to_think, has_incomplete_scratchpad,
     save_trajectory as _save_trajectory_to_file,
 )
-from utils import atomic_json_write, env_var_enabled
+from utils import atomic_json_write, env_var_enabled, repair_and_load_json
 
 
 
@@ -279,12 +279,14 @@ def _should_parallelize_tool_batch(tool_calls) -> bool:
         try:
             function_args = json.loads(tool_call.function.arguments)
         except Exception:
-            logging.debug(
-                "Could not parse args for %s — defaulting to sequential; raw=%s",
-                tool_name,
-                tool_call.function.arguments[:200],
-            )
-            return False
+            function_args = repair_and_load_json(tool_call.function.arguments)
+            if function_args is None:
+                logging.debug(
+                    "Could not parse args for %s — defaulting to sequential; raw=%s",
+                    tool_name,
+                    tool_call.function.arguments[:200],
+                )
+                return False
         if not isinstance(function_args, dict):
             logging.debug(
                 "Non-dict args for %s (%s) — defaulting to sequential",
@@ -5255,7 +5257,12 @@ class AIAgent:
                         try:
                             json.loads(arguments)
                         except json.JSONDecodeError:
-                            has_truncated_tool_args = True
+                            # Try repair before flagging truncation
+                            repaired = repair_and_load_json(arguments)
+                            if repaired is None:
+                                has_truncated_tool_args = True
+                            else:
+                                tc["function"]["arguments"] = json.dumps(repaired)
                     mock_tool_calls.append(SimpleNamespace(
                         id=tc["id"],
                         type=tc["type"],
@@ -7067,7 +7074,7 @@ class AIAgent:
             try:
                 function_args = json.loads(tool_call.function.arguments)
             except json.JSONDecodeError:
-                function_args = {}
+                function_args = repair_and_load_json(tool_call.function.arguments, default={})
             if not isinstance(function_args, dict):
                 function_args = {}
 
@@ -7264,8 +7271,9 @@ class AIAgent:
             try:
                 function_args = json.loads(tool_call.function.arguments)
             except json.JSONDecodeError as e:
-                logging.warning(f"Unexpected JSON error after validation: {e}")
-                function_args = {}
+                function_args = repair_and_load_json(tool_call.function.arguments, default={})
+                if not function_args:
+                    logging.warning(f"Unexpected JSON error after validation: {e}")
             if not isinstance(function_args, dict):
                 function_args = {}
 
@@ -10013,7 +10021,12 @@ class AIAgent:
                         try:
                             json.loads(args)
                         except json.JSONDecodeError as e:
-                            invalid_json_args.append((tc.function.name, str(e)))
+                            # Attempt repair before declaring invalid
+                            repaired = repair_and_load_json(args)
+                            if repaired is not None and isinstance(repaired, dict):
+                                tc.function.arguments = json.dumps(repaired)
+                            else:
+                                invalid_json_args.append((tc.function.name, str(e)))
                     
                     if invalid_json_args:
                         # Check if the invalid JSON is due to truncation rather

tests/test_json_repair.py

@@ -0,0 +1,127 @@
+"""Tests for JSON repair integration in tool-call pipeline."""
+
+import json
+import sys
+import os
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+
+from utils import repair_and_load_json
+
+
+class TestRepairAndLoadJson:
+    """Unit tests for repair_and_load_json helper."""
+
+    def test_valid_json_passthrough(self):
+        """Valid JSON parses without repair."""
+        result = repair_and_load_json('{"key": "value"}')
+        assert result == {"key": "value"}
+
+    def test_truncated_string(self):
+        """Truncated closing quote/brace gets repaired."""
+        result = repair_and_load_json('{"path": "test.py", "content": "hello')
+        assert result is not None
+        assert result["path"] == "test.py"
+
+    def test_missing_closing_brace(self):
+        """Missing closing brace gets repaired."""
+        result = repair_and_load_json('{"name": "foo", "args": {"x": 1}')
+        assert result is not None
+        assert result["name"] == "foo"
+
+    def test_trailing_comma(self):
+        """Trailing comma gets repaired."""
+        result = repair_and_load_json('{"key": "value",}')
+        assert result == {"key": "value"}
+
+    def test_single_quotes(self):
+        """Single-quoted JSON gets repaired."""
+        result = repair_and_load_json("{'key': 'value'}")
+        assert result == {"key": "value"}
+
+    def test_unquoted_keys(self):
+        """Unquoted keys get repaired."""
+        result = repair_and_load_json('{key: "value"}')
+        assert result == {"key": "value"}
+
+    def test_markdown_wrapped(self):
+        """Markdown-fenced JSON gets repaired."""
+        result = repair_and_load_json('```json\n{"key": "value"}\n```')
+        assert result == {"key": "value"}
+
+    def test_extra_trailing_text(self):
+        """Extra text after JSON gets stripped."""
+        result = repair_and_load_json('{"result": true} some extra text')
+        assert result == {"result": True}
+
+    def test_empty_string_returns_default(self):
+        """Empty string returns default."""
+        assert repair_and_load_json("", default={}) == {}
+
+    def test_none_returns_default(self):
+        """None returns default."""
+        assert repair_and_load_json(None, default=[]) == []
+
+    def test_completely_invalid_never_raises(self):
+        """Completely garbled input never raises — always returns something."""
+        # json_repair is aggressive — it may return [] or {} for nonsensical input
+        # The key guarantee: it never raises an exception
+        result = repair_and_load_json("not json at all {{{", default=None)
+        # Just verify it returned without exception — result can be anything
+        assert True  # Reached here means no exception
+
+    def test_default_is_none_by_default(self):
+        """Default parameter defaults to None."""
+        assert repair_and_load_json("totally broken") is None
+
+    def test_list_json(self):
+        """Valid list JSON parses correctly."""
+        result = repair_and_load_json('[1, 2, 3]')
+        assert result == [1, 2, 3]
+
+    def test_nested_truncation(self):
+        """Deeply nested truncated JSON gets repaired."""
+        result = repair_and_load_json('{"a": {"b": {"c": "hello')
+        assert result is not None
+        assert result["a"]["b"]["c"] == "hello"
+
+    def test_tool_call_arguments_pattern(self):
+        """Real-world tool call argument pattern with missing close."""
+        raw = '{"command": "python3 -c \\"print(1)\\"", "timeout": 30'
+        result = repair_and_load_json(raw, default={})
+        assert result is not None
+        assert "command" in result
+
+
+class TestRepairInToolCallContext:
+    """Tests simulating actual tool-call repair scenarios."""
+
+    def test_write_file_truncated_content(self):
+        """Simulates a write_file call truncated mid-content."""
+        raw = '{"path": "/tmp/test.py", "content": "import os\\nprint(os.getcwd())'
+        result = repair_and_load_json(raw, default={})
+        assert isinstance(result, dict)
+        assert result.get("path") == "/tmp/test.py"
+
+    def test_terminal_truncated_command(self):
+        """Simulates a terminal call truncated mid-command."""
+        raw = '{"command": "find . -name \'*.py\' | head -2'
+        result = repair_and_load_json(raw, default={})
+        assert isinstance(result, dict)
+        assert "command" in result
+
+    def test_patch_with_trailing_comma(self):
+        """Simulates a patch call with trailing comma."""
+        raw = '{"path": "file.py", "old_string": "foo", "new_string": "bar",}'
+        result = repair_and_load_json(raw)
+        assert result == {"path": "file.py", "old_string": "foo", "new_string": "bar"}
+
+    def test_empty_arguments_returns_empty_dict(self):
+        """Empty string should be handled gracefully."""
+        result = repair_and_load_json("", default={})
+        assert result == {}
+
+
+if __name__ == "__main__":
+    import pytest
+    pytest.main([__file__, "-v"])

utils.py

@@ -162,3 +162,33 @@ def env_int(key: str, default: int = 0) -> int:
 def env_bool(key: str, default: bool = False) -> bool:
     """Read an environment variable as a boolean."""
     return is_truthy_value(os.getenv(key, ""), default=default)
+
+
+def repair_and_load_json(text: str, default: Any = None) -> Any:
+    """Parse JSON with automatic repair for common LLM malformation patterns.
+
+    Tries json.loads first (fast path). On failure, attempts repair via
+    json_repair for: truncated strings, missing braces, trailing commas,
+    single quotes, unquoted keys, markdown wrapping, extra trailing text.
+
+    Args:
+        text: Raw JSON string (possibly malformed).
+        default: Value to return on complete parse failure.
+
+    Returns:
+        Parsed object on success, *default* on unrecoverable failure.
+    """
+    if not text or not isinstance(text, str):
+        return default
+    try:
+        return json.loads(text)
+    except (json.JSONDecodeError, TypeError, ValueError):
+        pass
+    try:
+        from json_repair import repair_json
+        repaired = repair_json(text)
+        if not repaired:
+            return default
+        return json.loads(repaired)
+    except Exception:
+        return default