Harden AI plan parsing against truncated LLM JSON responses

qa_agent/ai_planner.py

@@ -9,6 +9,7 @@
 """
 
 import json
+import logging
 import time
 
 from .llm_client import (
@@ -255,6 +256,56 @@
     "visible", "hidden", "text_contains", "url_contains", "element_count"
 })
 
+logger = logging.getLogger(__name__)
+
+
+def _repair_truncated_json(text: str) -> str | None:
+    """Attempt to repair truncated JSON by closing open strings/containers."""
+    if not text.strip().startswith("{"):
+        return None
+
+    stack: list[str] = []
+    in_string = False
+    escaped = False
+
+    for ch in text:
+        if in_string:
+            if escaped:
+                escaped = False
+                continue
+            if ch == "\\":
+                escaped = True
+            elif ch == '"':
+                in_string = False
+            continue
+
+        if ch == '"':
+            in_string = True
+        elif ch in "{[":
+            stack.append(ch)
+        elif ch in "}]":
+            if not stack:
+                return None
+            opener = stack.pop()
+            if (opener == "{" and ch != "}") or (opener == "[" and ch != "]"):
+                return None
+
+    if not in_string and not stack:
+        return None
+
+    # Can't safely repair mid-escape sequence - unclear what the intended escape was
+    if escaped:
+        return None
+
+    repaired = text
+    if in_string:
+        repaired += '"'
+    for opener in reversed(stack):
+        repaired += "}" if opener == "{" else "]"
+
+    return repaired
+
+
 def validate_plan(plan: "TestPlan") -> list[str]:
     """Return rule-based reliability warnings for a generated TestPlan.
 
@@ -420,6 +471,20 @@ def _parse_json(self, text: str) -> dict:
         try:
             data = json.loads(stripped)
         except json.JSONDecodeError as exc:
+            repaired = _repair_truncated_json(stripped)
+            if repaired is not None:
+                try:
+                    data = json.loads(repaired)
+                except json.JSONDecodeError:
+                    pass  # Repair failed, fall through to original error
+                else:
+                    if isinstance(data, dict):
+                        logger.warning(
+                            "Recovered from truncated LLM response (%d chars repaired). "
+                            "Original length: %d, repaired length: %d",
+                            len(repaired) - len(stripped), len(stripped), len(repaired)
+                        )
+                        return data
             preview = text[:_MAX_RAW_RESPONSE_IN_ERROR]
             suffix = "…" if len(text) > _MAX_RAW_RESPONSE_IN_ERROR else ""
             raise ValueError(

tests/test_ai_planner.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import json
 import os
 from unittest.mock import MagicMock, patch
 
@@ -89,6 +90,94 @@ def test_malformed_json_raises_value_error(self):
         with pytest.raises(ValueError, match="invalid JSON"):
             planner.plan("test", "https://example.com")
 
+    def test_truncated_json_warning_string_is_repaired(self):
+        data = json.loads(VALID_PLAN_JSON)
+        data["warnings"] = [
+            "CSS checks cannot be verified via Playwright computed-style assertions; use visual regression."
+        ]
+        truncated = json.dumps(data)[:-3]  # chop trailing quote/bracket/brace
+        planner = self._planner(truncated)
+        plan = planner.plan("test", "https://example.com")
+        assert plan.warnings
+        assert "computed-style assertions" in plan.warnings[0]
+
+    def test_truncated_after_backslash_raises_error(self):
+        """Truncation mid-escape should fail safely - can't infer intended escape."""
+        data = json.loads(VALID_PLAN_JSON)
+        data["notes"] = "Line 1\nLine 2"  # Contains actual newline character
+        full_json = json.dumps(data)
+        # In JSON, newline becomes \n escape sequence. Find it and truncate after backslash
+        idx = full_json.index("\\n")  # Find the \n in the JSON string
+        truncated = full_json[:idx + 1]  # Keep backslash, remove the 'n'
+
+        planner = self._planner(truncated)
+        # Should raise because repair returns None for mid-escape truncation
+        with pytest.raises(ValueError, match="invalid JSON"):
+            planner.plan("test", "https://example.com")
+
+    def test_truncated_nested_structures_repaired(self):
+        """Multiple unclosed containers should all be closed."""
+        # Create JSON with nested structures and truncate mid-way
+        data = json.loads(VALID_PLAN_JSON)
+        data["custom_steps"][0]["actions"].append({"type": "hover", "selector": "#menu"})
+        full_json = json.dumps(data)
+        # Truncate in the middle of the nested structure
+        # Find a point deep in the nesting and truncate there
+        truncate_at = full_json.index('"hover"') + len('"hover"')
+        truncated = full_json[:truncate_at]
+
+        planner = self._planner(truncated)
+        # Should successfully repair by closing all open containers
+        plan = planner.plan("test", "https://example.com")
+        assert isinstance(plan, TestPlan)
+        # At minimum should have the summary from before truncation
+        assert plan.summary == "Test the login flow"
+
+    def test_already_valid_json_returns_none_from_repair_fn(self):
+        """_repair_truncated_json should return None for already-valid JSON."""
+        from qa_agent.ai_planner import _repair_truncated_json
+        assert _repair_truncated_json('{"key": "value"}') is None
+        assert _repair_truncated_json(VALID_PLAN_JSON) is None
+
+    def test_repair_closes_unclosed_string(self):
+        """Single unclosed string should be closed."""
+        partial = '{"summary": "test in progress'
+        from qa_agent.ai_planner import _repair_truncated_json
+        repaired = _repair_truncated_json(partial)
+        assert repaired is not None
+        # Should add closing quote and brace
+        assert repaired == partial + '"}'
+        # Verify it's valid JSON
+        parsed = json.loads(repaired)
+        assert parsed["summary"] == "test in progress"
+
+    def test_repair_closes_multiple_containers(self):
+        """Multiple unclosed objects/arrays should all be closed."""
+        partial = '{"custom_steps": [{"actions": [{"type": "click"'
+        from qa_agent.ai_planner import _repair_truncated_json
+        repaired = _repair_truncated_json(partial)
+        assert repaired is not None
+        # Should close: string, object (action), array (actions), object (step), array (steps), object (root)
+        assert repaired.endswith('"}]}]}')
+        # Verify it's valid JSON
+        parsed = json.loads(repaired)
+        assert "custom_steps" in parsed
+
+    def test_repair_rejects_mismatched_brackets(self):
+        """Malformed JSON with bracket mismatches should return None."""
+        from qa_agent.ai_planner import _repair_truncated_json
+        # Opening { but closing ]
+        assert _repair_truncated_json('{"key": [}') is None
+        # Closing without opening
+        assert _repair_truncated_json('{"key": "value"}}') is None
+
+    def test_repair_rejects_non_object_json(self):
+        """Repair only works for truncated objects, not arrays or primitives."""
+        from qa_agent.ai_planner import _repair_truncated_json
+        assert _repair_truncated_json('["array"') is None
+        assert _repair_truncated_json('"string') is None
+        assert _repair_truncated_json('123') is None
+
     def test_no_text_content_raises_llm_error(self):
         from qa_agent.llm_client import LLMError
         client = MagicMock()