fix(reflect): unwrap JSON answer envelopes (#2345)

* fix(reflect): unwrap JSON answer envelopes

* fix(reflect): clarify leaked done argument recovery

Author: xmh1011

hindsight-api-slim/hindsight_api/engine/reflect/agent.py

@@ -90,12 +90,87 @@ def _is_done_tool(name: str) -> bool:
     r'\s*```(?:json)?\s*\{[^}]*(?:"(?:observation_ids|memory_ids|mental_model_ids)"|\})\s*```\s*$',
     re.DOTALL | re.IGNORECASE,
 )
-_LEAKED_JSON_OBJECT = re.compile(
-    r'\s*\{[^{]*"(?:observation_ids|memory_ids|mental_model_ids|answer)"[^}]*\}\s*$', re.DOTALL
-)
 _TRAILING_IDS_PATTERN = re.compile(
     r"\s*(?:observation_ids|memory_ids|mental_model_ids)\s*[=:]\s*\[.*?\]\s*$", re.DOTALL | re.IGNORECASE
 )
+_JSON_CODE_FENCE_PATTERN = re.compile(r"^\s*```(?:json)?\s*(\{.*\})\s*```\s*$", re.DOTALL | re.IGNORECASE)
+
+_DONE_ARGUMENT_KEYS = frozenset(
+    {
+        "answer",
+        "directive_compliance",
+        "memory_ids",
+        "mental_model_ids",
+        "observation_ids",
+        "model_ids",
+    }
+)
+_DONE_ARGUMENT_MARKER_KEYS = _DONE_ARGUMENT_KEYS - {"answer"}
+_LEAKED_JSON_ID_KEYS = frozenset({"memory_ids", "mental_model_ids", "observation_ids", "model_ids"})
+
+
+def _unwrap_leaked_done_arguments(text: str) -> str | None:
+    """Return the answer when a done tool call was rendered as JSON text.
+
+    Some providers leak the done tool's argument object instead of surfacing it
+    as a native tool call, e.g. {"answer": "...", "memory_ids": [...]}. Only
+    unwrap objects that match the done argument shape so normal JSON answers
+    stay intact.
+    """
+    candidate = text.strip()
+    if not candidate:
+        return None
+
+    fenced = _JSON_CODE_FENCE_PATTERN.match(candidate)
+    if fenced:
+        candidate = fenced.group(1).strip()
+
+    try:
+        payload = json.loads(candidate)
+    except json.JSONDecodeError:
+        return None
+
+    if not isinstance(payload, dict):
+        return None
+    answer = payload.get("answer")
+    if not isinstance(answer, str) or not answer.strip():
+        return None
+
+    keys = set(payload)
+    if not keys.intersection(_DONE_ARGUMENT_MARKER_KEYS):
+        return None
+    if not keys.issubset(_DONE_ARGUMENT_KEYS):
+        return None
+
+    for key in ("memory_ids", "mental_model_ids", "observation_ids", "model_ids"):
+        value = payload.get(key)
+        if value is not None and not isinstance(value, list):
+            return None
+
+    return answer.strip()
+
+
+def _strip_trailing_id_json_object(text: str) -> str:
+    stripped = text.rstrip()
+    if not stripped.endswith("}"):
+        return text.strip()
+
+    start = stripped.rfind("{")
+    if start < 0:
+        return text.strip()
+
+    try:
+        payload = json.loads(stripped[start:])
+    except json.JSONDecodeError:
+        return text.strip()
+
+    if not isinstance(payload, dict) or not payload:
+        return text.strip()
+    keys = set(payload)
+    if not keys.issubset(_LEAKED_JSON_ID_KEYS):
+        return text.strip()
+
+    return stripped[:start].strip()
 
 
 def _clean_answer_text(text: str) -> str:
@@ -104,6 +179,10 @@ def _clean_answer_text(text: str) -> str:
     Some LLMs output the done() call as text instead of a proper tool call.
     This strips out patterns like: done({"answer": "...", ...})
     """
+    unwrapped = _unwrap_leaked_done_arguments(text)
+    if unwrapped is not None:
+        return unwrapped
+
     # Remove done() call pattern from the end of the text
     cleaned = _DONE_CALL_PATTERN.sub("", text).strip()
     return cleaned if cleaned else text
@@ -122,13 +201,17 @@ def _clean_done_answer(text: str) -> str:
     if not text:
         return text
 
+    unwrapped = _unwrap_leaked_done_arguments(text)
+    if unwrapped is not None:
+        return unwrapped
+
     cleaned = text
 
     # Remove leaked JSON in code blocks at the end
     cleaned = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
 
     # Remove leaked raw JSON objects at the end
-    cleaned = _LEAKED_JSON_OBJECT.sub("", cleaned).strip()
+    cleaned = _strip_trailing_id_json_object(cleaned)
 
     # Remove trailing ID patterns
     cleaned = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()

hindsight-api-slim/tests/test_reflect_agent.py

@@ -69,6 +69,25 @@ def test_clean_text_multiline_done(self):
         cleaned = _clean_answer_text(text)
         assert cleaned == "Summary of findings."
 
+    def test_clean_text_recovers_leaked_done_arguments(self):
+        """A done tool-call argument object rendered as text should keep only answer."""
+        text = """{
+            "answer": "Use the inbound table for API consumers.",
+            "directive_compliance": "Directive 1 followed.",
+            "memory_ids": ["mem-1"],
+            "mental_model_ids": [],
+            "observation_ids": []
+        }"""
+        cleaned = _clean_answer_text(text)
+        assert cleaned == "Use the inbound table for API consumers."
+        assert "directive_compliance" not in cleaned
+
+    def test_clean_text_leaves_non_done_json_answer_unchanged(self):
+        """Plain JSON answers are valid user-visible content."""
+        text = '{"status": "ok", "items": [1, 2]}'
+        cleaned = _clean_answer_text(text)
+        assert cleaned == text
+
 
 class TestCleanDoneAnswer:
     """Test cleanup of answer field from done() tool call that leaks structured output."""
@@ -142,6 +161,37 @@ def test_clean_answer_multiline_with_markdown(self):
         assert "Point 2" in cleaned
         assert "mental_model_ids" not in cleaned
 
+    def test_clean_answer_recovers_leaked_done_arguments(self):
+        """A done answer that contains leaked done arguments should keep answer content."""
+        text = """{
+            "answer": "Render two markdown tables: inbound and outbound.",
+            "directive_compliance": "All directives followed.",
+            "memory_ids": ["mem-1", "mem-2"],
+            "mental_model_ids": [],
+            "observation_ids": ["obs-1"]
+        }"""
+        cleaned = _clean_done_answer(text)
+        assert cleaned == "Render two markdown tables: inbound and outbound."
+
+    def test_clean_answer_recovers_fenced_leaked_done_arguments(self):
+        """Some providers put leaked done arguments in a JSON code fence."""
+        text = """```json
+{
+  "answer": "The current interface is HTTP only.",
+  "memory_ids": [],
+  "mental_model_ids": [],
+  "observation_ids": []
+}
+```"""
+        cleaned = _clean_done_answer(text)
+        assert cleaned == "The current interface is HTTP only."
+
+    def test_clean_answer_rejects_done_arguments_with_unexpected_keys(self):
+        """Avoid rewriting user-requested JSON that happens to contain answer."""
+        text = '{"answer": "yes", "payload": {"format": "json"}, "memory_ids": []}'
+        cleaned = _clean_done_answer(text)
+        assert cleaned == text
+
 
 class TestToolNameNormalization:
     """Test tool name normalization for various LLM output formats."""