fix(retain): preserve JSON chunks during output retry

hindsight-api-slim/hindsight_api/engine/retain/fact_extraction.py

@@ -232,6 +232,55 @@ class FactExtractionResponse(BaseModel):
     facts: list[ExtractedFact] = Field(description="List of extracted factual statements")
 
 
+def _split_chunk_for_output_retry(chunk: str) -> tuple[str, str] | None:
+    """Split an oversized extraction chunk without corrupting structured input."""
+    stripped = chunk.strip()
+    if len(stripped) <= 1:
+        return None
+
+    try:
+        parsed = json.loads(stripped)
+    except (TypeError, ValueError, json.JSONDecodeError):
+        parsed = None
+
+    if isinstance(parsed, list):
+        if len(parsed) >= 2:
+            mid = len(parsed) // 2
+            return json.dumps(parsed[:mid]), json.dumps(parsed[mid:])
+
+        if len(parsed) == 1 and isinstance(parsed[0], dict):
+            turn = parsed[0]
+            content = turn.get("content")
+            if isinstance(content, str) and len(content) > 1:
+                cut = len(content) // 2
+                first_turn = dict(turn)
+                second_turn = dict(turn)
+                first_turn["content"] = content[:cut]
+                second_turn["content"] = content[cut:]
+                return json.dumps([first_turn]), json.dumps([second_turn])
+
+        return None
+
+    # Split plain text at the midpoint, preferring sentence boundaries nearby.
+    mid_point = len(stripped) // 2
+    search_range = int(len(stripped) * 0.2)
+    search_start = max(0, mid_point - search_range)
+    search_end = min(len(stripped), mid_point + search_range)
+
+    best_split = mid_point
+    for ending in [". ", "! ", "? ", "\n\n"]:
+        pos = stripped.rfind(ending, search_start, search_end)
+        if pos != -1:
+            best_split = pos + len(ending)
+            break
+
+    first_half = stripped[:best_split].strip()
+    second_half = stripped[best_split:].strip()
+    if not first_half or not second_half or first_half == stripped or second_half == stripped:
+        return None
+    return first_half, second_half
+
+
 class ExtractedFactVerbose(BaseModel):
     """A single extracted fact with verbose field descriptions for detailed extraction."""
 
@@ -1664,33 +1713,22 @@ async def _extract_facts_with_auto_split(
             metadata=metadata,
         )
     except OutputTooLongError:
-        # Output exceeded token limits - split the chunk in half and retry
+        # Output exceeded token limits - split the chunk and retry. Conversation
+        # chunks are JSON arrays, so preserve array/turn boundaries when possible.
         logger.warning(
             f"Output too long for chunk {chunk_index + 1}/{total_chunks} "
-            f"({len(chunk)} chars). Splitting in half and retrying..."
+            f"({len(chunk)} chars). Splitting and retrying..."
         )
 
-        # Split at the midpoint, preferring sentence boundaries
-        mid_point = len(chunk) // 2
-
-        # Try to find a sentence boundary near the midpoint
-        # Look for ". ", "! ", "? " within 20% of midpoint
-        search_range = int(len(chunk) * 0.2)
-        search_start = max(0, mid_point - search_range)
-        search_end = min(len(chunk), mid_point + search_range)
-
-        sentence_endings = [". ", "! ", "? ", "\n\n"]
-        best_split = mid_point
-
-        for ending in sentence_endings:
-            pos = chunk.rfind(ending, search_start, search_end)
-            if pos != -1:
-                best_split = pos + len(ending)
-                break
+        split_chunks = _split_chunk_for_output_retry(chunk)
+        if split_chunks is None:
+            logger.warning(
+                f"Cannot make progress splitting chunk {chunk_index + 1}/{total_chunks} "
+                f"({len(chunk)} chars); dropping this sub-chunk."
+            )
+            return [], TokenUsage()
 
-        # Split the chunk
-        first_half = chunk[:best_split].strip()
-        second_half = chunk[best_split:].strip()
+        first_half, second_half = split_chunks
 
         logger.info(
             f"Split chunk {chunk_index + 1} into two sub-chunks: {len(first_half)} chars and {len(second_half)} chars"

hindsight-api-slim/tests/test_fact_extraction_retry.py

@@ -9,12 +9,87 @@
 BadRequestError handler and not for non-dict JSON responses.
 """
 
+import json
 from datetime import datetime, timezone
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
 
+def test_output_retry_split_preserves_conversation_array_boundaries():
+    """OutputTooLong retry splitting must keep conversation chunks valid JSON arrays."""
+    from hindsight_api.engine.retain.fact_extraction import _split_chunk_for_output_retry
+
+    turns = [
+        {"role": "user", "content": "alpha"},
+        {"role": "assistant", "content": "bravo"},
+        {"role": "user", "content": "charlie"},
+        {"role": "assistant", "content": "delta"},
+    ]
+
+    split = _split_chunk_for_output_retry(json.dumps(turns))
+
+    assert split is not None
+    first, second = split
+    assert json.loads(first) == turns[:2]
+    assert json.loads(second) == turns[2:]
+
+
+def test_output_retry_split_divides_single_oversized_turn_content():
+    """A lone oversized conversation turn is split inside content and rewrapped."""
+    from hindsight_api.engine.retain.fact_extraction import _split_chunk_for_output_retry
+
+    turn = {"role": "user", "content": "abcdefghijklmnopqrstuvwxyz", "name": "casey"}
+
+    split = _split_chunk_for_output_retry(json.dumps([turn]))
+
+    assert split is not None
+    first, second = split
+    first_turn = json.loads(first)[0]
+    second_turn = json.loads(second)[0]
+    assert first_turn["role"] == "user"
+    assert second_turn["role"] == "user"
+    assert first_turn["name"] == "casey"
+    assert second_turn["name"] == "casey"
+    assert first_turn["content"] + second_turn["content"] == turn["content"]
+
+
+def test_output_retry_split_returns_none_when_no_progress_possible():
+    """Pathological tiny chunks should be dropped instead of recursively retried."""
+    from hindsight_api.engine.retain.fact_extraction import _split_chunk_for_output_retry
+
+    assert _split_chunk_for_output_retry("x") is None
+    assert _split_chunk_for_output_retry(json.dumps([{"role": "user", "content": ""}])) is None
+
+
+@pytest.mark.asyncio
+async def test_output_too_long_drops_unsplittable_subchunk_without_recursing():
+    """If a chunk cannot be reduced further, auto-split exits gracefully."""
+    from hindsight_api.engine.llm_wrapper import OutputTooLongError
+    from hindsight_api.engine.retain.fact_extraction import _extract_facts_with_auto_split
+
+    config = _make_config(llm_max_retries=1)
+    llm_config = _make_llm_config(mock_response={})
+
+    with patch(
+        "hindsight_api.engine.retain.fact_extraction._extract_facts_from_chunk",
+        side_effect=OutputTooLongError("too long"),
+    ) as extract:
+        facts, usage = await _extract_facts_with_auto_split(
+            chunk="x",
+            chunk_index=0,
+            total_chunks=1,
+            event_date=datetime(2023, 1, 1, tzinfo=timezone.utc),
+            context="",
+            llm_config=llm_config,
+            config=config,
+            agent_name="agent",
+        )
+
+    assert facts == []
+    assert extract.call_count == 1
+
+
 def _make_config(llm_max_retries: int = 3, retain_llm_max_retries: int | None = None):
     """Build a minimal HindsightConfig for fact extraction tests."""
     from hindsight_api.config import HindsightConfig