fix(reflect): safely strip leaked done-tool siblings

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

@@ -101,6 +101,14 @@ def _is_done_tool(name: str) -> bool:
 )
 _JSON_CODE_FENCE_PATTERN = re.compile(r"^\s*```(?:json)?\s*(\{.*\})\s*```\s*$", re.DOTALL | re.IGNORECASE)
 
+# Candidate boundary between a truncated answer string and sibling done-tool
+# arguments, e.g. ...end of prose", "memory_ids": ["a", "b"]\n}.
+# Values are deliberately not parsed by regex; the suffix is reconstructed as
+# JSON and validated against IDs actually retrieved during this reflect run.
+_DONE_SIBLING_BOUNDARY_PATTERN = re.compile(
+    r'"\s*,\s*"(?:directive_compliance|memory_ids|mental_model_ids|observation_ids|model_ids)"\s*:'
+)
+
 _DONE_ARGUMENT_KEYS = frozenset(
     {
         "answer",
@@ -194,7 +202,104 @@ def _clean_answer_text(text: str) -> str:
     return cleaned if cleaned else text
 
 
-def _clean_done_answer(text: str) -> str:
+def _is_complete_json_document(text: str) -> bool:
+    """Return whether all answer text is one JSON document, optionally fenced."""
+    candidate = text.strip()
+    fenced = _JSON_CODE_FENCE_PATTERN.match(candidate)
+    if fenced:
+        candidate = fenced.group(1).strip()
+    try:
+        json.loads(candidate)
+    except (json.JSONDecodeError, ValueError):
+        return False
+    return True
+
+
+def _strip_leaked_done_sibling_suffix(
+    text: str,
+    *,
+    available_memory_ids: set[str] | frozenset[str],
+    available_mental_model_ids: set[str] | frozenset[str],
+    available_observation_ids: set[str] | frozenset[str],
+) -> str | None:
+    """Strip a terminal sibling-argument fragment proven to be a done-tool leak.
+
+    The provider can occasionally put the closing quote and remaining done-tool
+    arguments inside ``answer``. Shape alone is not enough evidence to rewrite
+    user prose, so a candidate with an answer prefix is removed only when it is
+    valid JSON and contains at least one reference retrieved during this reflect
+    run. A metadata-only fragment has no answer to preserve and is removed even
+    when its ID lists are empty. ``None`` means no proven leak; an empty string
+    means a proven leak contained no answer and must remain empty for the
+    caller's failure handling.
+    """
+    allowed_ids = {
+        "memory_ids": available_memory_ids,
+        "mental_model_ids": available_mental_model_ids,
+        "model_ids": available_mental_model_ids,
+        "observation_ids": available_observation_ids,
+    }
+
+    for boundary in _DONE_SIBLING_BOUNDARY_PATTERN.finditer(text):
+        # A quote escaped by an odd run of backslashes belongs to answer prose,
+        # not to the JSON boundary that closed the answer argument.
+        backslash_count = 0
+        cursor = boundary.start() - 1
+        while cursor >= 0 and text[cursor] == "\\":
+            backslash_count += 1
+            cursor -= 1
+        if backslash_count % 2:
+            continue
+
+        comma = text.find(",", boundary.start(), boundary.end())
+        if comma < 0:
+            continue
+        suffix = text[comma + 1 :].strip()
+        candidate = "{" + suffix
+        if not candidate.rstrip().endswith("}"):
+            candidate += "}"
+
+        try:
+            payload = json.loads(candidate)
+        except (json.JSONDecodeError, ValueError):
+            continue
+        if not isinstance(payload, dict) or not payload:
+            continue
+
+        keys = set(payload)
+        if not keys.issubset(_DONE_ARGUMENT_MARKER_KEYS):
+            continue
+        directive_compliance = payload.get("directive_compliance")
+        if directive_compliance is not None and not isinstance(directive_compliance, str):
+            continue
+
+        id_keys = keys.intersection(_LEAKED_JSON_ID_KEYS)
+        known_reference_count = 0
+        references_are_valid = True
+        for key in id_keys:
+            value = payload[key]
+            if not isinstance(value, list) or not all(isinstance(reference_id, str) for reference_id in value):
+                references_are_valid = False
+                break
+            known_reference_count += sum(reference_id in allowed_ids[key] for reference_id in value)
+
+        if not references_are_valid:
+            continue
+
+        answer_prefix = text[: boundary.start()].rstrip()
+        if not answer_prefix or known_reference_count:
+            return answer_prefix
+
+    return None
+
+
+def _clean_done_answer(
+    text: str,
+    *,
+    available_memory_ids: set[str] | frozenset[str] = frozenset(),
+    available_mental_model_ids: set[str] | frozenset[str] = frozenset(),
+    available_observation_ids: set[str] | frozenset[str] = frozenset(),
+) -> str:
     """Clean up the answer field from a done() tool call.
 
     Some LLMs leak structured output patterns into the answer text, such as:
@@ -210,19 +315,45 @@ def _clean_done_answer(text: str) -> str:
     unwrapped = _unwrap_leaked_done_arguments(text)
     if unwrapped is not None:
         return unwrapped
+    if _is_complete_json_document(text):
+        # A complete JSON document is user-visible answer content, not the
+        # truncated remainder of a done-tool argument object. This check must
+        # precede every suffix cleaner, including the legacy raw-object path.
+        return text.strip()
 
     cleaned = text
+    removed_leak = False
 
     # Remove leaked JSON in code blocks at the end
-    cleaned = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
+    stripped = _LEAKED_JSON_SUFFIX.sub("", cleaned).strip()
+    removed_leak = removed_leak or stripped != cleaned.strip()
+    cleaned = stripped
 
     # Remove leaked raw JSON objects at the end
-    cleaned = _strip_trailing_id_json_object(cleaned)
+    stripped = _strip_trailing_id_json_object(cleaned)
+    removed_leak = removed_leak or stripped != cleaned.strip()
+    cleaned = stripped
 
     # Remove trailing ID patterns
-    cleaned = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()
+    stripped = _TRAILING_IDS_PATTERN.sub("", cleaned).strip()
+    removed_leak = removed_leak or stripped != cleaned.strip()
+    cleaned = stripped
+
+    # Remove sibling done-argument fields left after a closed answer string.
+    stripped_sibling = _strip_leaked_done_sibling_suffix(
+        cleaned,
+        available_memory_ids=available_memory_ids,
+        available_mental_model_ids=available_mental_model_ids,
+        available_observation_ids=available_observation_ids,
+    )
+    if stripped_sibling is not None:
+        removed_leak = True
+        cleaned = stripped_sibling
 
-    return cleaned if cleaned else text
+    # Do not turn a proven empty result back into metadata-shaped content. The
+    # caller converts empty to NO_ANSWER_TEXT, which the refresh guard treats as
+    # a failed render and refuses to persist over healthy content.
+    return cleaned if cleaned or removed_leak else text
 
 
 async def _generate_structured_output(
@@ -1435,7 +1566,16 @@ async def _process_done_tool(
 
     # Extract and clean the answer - some LLMs leak structured output into the answer text
     raw_answer = args.get("answer", "").strip()
-    answer = _clean_done_answer(raw_answer) if raw_answer else ""
+    answer = (
+        _clean_done_answer(
+            raw_answer,
+            available_memory_ids=available_memory_ids,
+            available_mental_model_ids=available_mental_model_ids,
+            available_observation_ids=available_observation_ids,
+        )
+        if raw_answer
+        else ""
+    )
     if not answer:
         answer = NO_ANSWER_TEXT
 

hindsight-api-slim/tests/test_reflect_agent.py

@@ -9,12 +9,14 @@
 """
 
 import asyncio
+from itertools import permutations
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
 
 from hindsight_api.engine.llm_interface import LLM_TOOL_CHOICE_AUTO, LLMToolChoice
 from hindsight_api.engine.reflect.agent import (
+    NO_ANSWER_TEXT,
     _all_mental_models_are_usable_and_fresh,
     _cache_cleanup_tasks,
     _clean_answer_text,
@@ -24,8 +26,10 @@
     _is_context_overflow_error,
     _is_done_tool,
     _normalize_tool_name,
+    _process_done_tool,
     run_reflect_agent,
 )
+from hindsight_api.engine.reflect.models import TokenUsageSummary
 from hindsight_api.engine.response_models import LLMToolCall, LLMToolCallResult, TokenUsage
 from tests.llm_judge import assert_meets_criteria
 
@@ -194,6 +198,135 @@ def test_clean_answer_rejects_done_arguments_with_unexpected_keys(self):
         cleaned = _clean_done_answer(text)
         assert cleaned == text
 
+        report = '{"report": "yes", "memory_ids": ["mem-1"]}'
+        assert _clean_done_answer(report, available_memory_ids={"mem-1"}) == report
+
+        id_only = '{"memory_ids": ["mem-1"], "observation_ids": []}'
+        assert _clean_done_answer(id_only, available_memory_ids={"mem-1"}) == id_only
+
+        fenced_id_only = '```json\n{"memory_ids": ["mem-1"]}\n```'
+        assert _clean_done_answer(fenced_id_only, available_memory_ids={"mem-1"}) == fenced_id_only
+
+    def test_clean_answer_strips_sibling_fields_after_closed_answer(self):
+        """A truncated done-argument object leaks its sibling fields into the answer.
+
+        Observed in production: the answer string was closed and its siblings
+        followed, so 17 memory UUIDs were persisted into a mental model's stored
+        content. Nothing caught it -- the text does not parse as JSON (the opening
+        {"answer": " was consumed as the answer's start), carries no ``` fence,
+        does not begin at "{", and its key is quoted with a "}" after the "]".
+        """
+        text = (
+            '# Home Network\n\n- OPNsense router\n- Admin via tunnel", '
+            '"memory_ids": ["97424a2f-f166-4ab3-9c5a-baf3dae4ea5a", '
+            '"fb0d8a6d-c55d-450f-a679-e8715bdfa4d4"]\n}'
+        )
+        cleaned = _clean_done_answer(
+            text,
+            available_memory_ids={
+                "97424a2f-f166-4ab3-9c5a-baf3dae4ea5a",
+                "fb0d8a6d-c55d-450f-a679-e8715bdfa4d4",
+            },
+        )
+        assert "memory_ids" not in cleaned
+        assert "97424a2f" not in cleaned
+        assert cleaned.endswith("Admin via tunnel")
+
+    def test_clean_answer_strips_multiple_sibling_id_fields(self):
+        """Several id fields may follow the closed answer, with or without a brace."""
+        for tail, available_memory_ids, available_mental_model_ids, available_observation_ids in (
+            ('", "memory_ids": ["a"], "observation_ids": ["b"]}', {"a"}, set(), {"b"}),
+            ('", "mental_model_ids": ["c"]', set(), {"c"}, set()),
+            ('", "model_ids": ["d"]}', set(), {"d"}, set()),
+        ):
+            cleaned = _clean_done_answer(
+                "Real prose" + tail,
+                available_memory_ids=available_memory_ids,
+                available_mental_model_ids=available_mental_model_ids,
+                available_observation_ids=available_observation_ids,
+            )
+            assert cleaned == "Real prose", f"tail {tail!r} -> {cleaned!r}"
+
+    def test_clean_answer_parses_sibling_fields_in_any_order(self):
+        """Field order and escaped JSON string content must not affect cleanup."""
+        fields = (
+            '"directive_compliance": "Use \\\\server\\\\share and \\"quoted\\" text with ]"',
+            '"memory_ids": ["mem-1"]',
+            '"observation_ids": ["obs-1"]',
+        )
+        for ordered_fields in permutations(fields):
+            text = 'Real prose", ' + ", ".join(ordered_fields) + "}"
+            cleaned = _clean_done_answer(
+                text,
+                available_memory_ids={"mem-1"},
+                available_observation_ids={"obs-1"},
+            )
+            assert cleaned == "Real prose", f"field order {ordered_fields!r} -> {cleaned!r}"
+
+    def test_clean_answer_preserves_terminal_schema_fragment_without_retrieval_provenance(self):
+        """A prose example that resembles done arguments is not itself proof of a leak."""
+        text = 'Document this fragment", "memory_ids": ["example-id"]}'
+        assert _clean_done_answer(text) == text
+        assert _clean_done_answer(text, available_memory_ids={"different-id"}) == text
+
+        mixed_ids = 'Document this fragment", "memory_ids": ["known-id", "example-id"]}'
+        assert _clean_done_answer(mixed_ids, available_memory_ids={"known-id"}) == "Document this fragment"
+
+    def test_clean_answer_preserves_intended_terminal_quote(self):
+        """Removing the JSON boundary must not remove a quote belonging to the answer."""
+        text = 'The operator called it "safe"", "memory_ids": ["mem-1"]}'
+        cleaned = _clean_done_answer(text, available_memory_ids={"mem-1"})
+        assert cleaned == 'The operator called it "safe"'
+
+    def test_clean_answer_returns_empty_for_metadata_only_sibling_leak(self):
+        """A proven leak with no answer must reach the caller as an empty failure."""
+        assert _clean_done_answer('", "memory_ids": ["mem-1"]}', available_memory_ids={"mem-1"}) == ""
+        assert _clean_done_answer('", "memory_ids": []}') == ""
+        assert _clean_done_answer('", "directive_compliance": "No answer rendered"}') == ""
+
+    @pytest.mark.asyncio
+    async def test_process_done_tool_maps_metadata_only_sibling_leak_to_no_answer(self):
+        """The persist-side guard sees the established failure sentinel, not leaked metadata."""
+        result = await _process_done_tool(
+            LLMToolCall(id="done-1", name="done", arguments={"answer": '", "memory_ids": ["mem-1"]}'}),
+            available_memory_ids={"mem-1"},
+            available_mental_model_ids=set(),
+            available_observation_ids=set(),
+            iterations=1,
+            total_tools_called=1,
+            tool_trace=[],
+            llm_trace=[],
+            usage=TokenUsageSummary(),
+            log_completion=MagicMock(),
+            reflect_id="reflect-1",
+            directives_applied=[],
+        )
+        assert result.text == NO_ANSWER_TEXT
+
+        empty_ids_result = await _process_done_tool(
+            LLMToolCall(id="done-2", name="done", arguments={"answer": '", "memory_ids": []}'}),
+            available_memory_ids=set(),
+            available_mental_model_ids=set(),
+            available_observation_ids=set(),
+            iterations=1,
+            total_tools_called=1,
+            tool_trace=[],
+            llm_trace=[],
+            usage=TokenUsageSummary(),
+            log_completion=MagicMock(),
+            reflect_id="reflect-2",
+            directives_applied=[],
+        )
+        assert empty_ids_result.text == NO_ANSWER_TEXT
+
+    def test_clean_answer_keeps_prose_that_merely_mentions_id_fields(self):
+        """The stripper must not eat ordinary text or an inline JSON example."""
+        for text in (
+            "We store memory_ids in the database for audit purposes.",
+            'Example:\n\n```json\n{"memory_ids": ["x"]}\n```\n\nMore prose follows.',
+        ):
+            assert _clean_done_answer(text) == text
+
 
 class TestToolNameNormalization:
     """Test tool name normalization for various LLM output formats."""