fix(retain): keep oversized items in one async child to stop FK race (#1795)

hindsight-api-slim/hindsight_api/engine/memory_engine.py

@@ -273,6 +273,12 @@ def _split_contents_into_sub_batches(
     would pass through as a ``1/1`` sub-batch holding the entire
     payload — which contradicts the splitter's log and lets the
     orchestrator OOM under realistic memory limits (see issue #1571).
+
+    Used by the in-process ``retain_batch_async`` path, which processes
+    the returned sub-batches SEQUENTIALLY with ``is_first_batch=(i==1)``.
+    The async submission path uses ``_split_contents_into_async_children``
+    instead, which never fragments a single item across children — see
+    that helper for the reasoning.
     """
     from .retain import fact_extraction
 
@@ -326,6 +332,66 @@ def _flush() -> None:
     return _SubBatchSplit(sub_batches=sub_batches, origin_indices=origin_indices)
 
 
+def _split_contents_into_async_children(
+    contents: list[RetainContentDict],
+    tokens_per_batch: int,
+) -> list[list[RetainContentDict]]:
+    """Pack retain contents into child operations for async submission.
+
+    Unlike ``_split_contents_into_sub_batches`` (used by the in-process
+    path), this NEVER fragments a single input item across multiple
+    children. Items where ``count_tokens(content) > tokens_per_batch``
+    are emitted as their own single-item child holding the FULL
+    un-chunked content; the in-process ``retain_batch_async`` then
+    re-chunks them SEQUENTIALLY inside one worker slot with correct
+    ``is_first_batch=(i==1)`` semantics.
+
+    The previous behavior — chunking oversized items into N independent
+    child async-operations sharing one ``document_id`` — let workers
+    claim siblings concurrently with no per-document gate (the busy-bank
+    guard in ``claim_tasks`` only covers consolidation). Each concurrent
+    child ran ``handle_document_tracking(is_first_batch=True)``, which
+    cascade-deletes the prior winner's ``memory_units`` for that
+    document. The loser's final ANN pass then attempted to insert
+    ``memory_links`` referencing now-deleted units → FK violations on
+    ``fk_memory_links_from_unit_id_memory_units``, partial document
+    state, and worker thread exhaustion from sentence-transformer pools
+    spun up per concurrent child. See issue #1795.
+
+    Items smaller than the budget are still packed together so genuinely
+    independent items keep cross-worker parallelism.
+    """
+    children: list[list[RetainContentDict]] = []
+    current: list[RetainContentDict] = []
+    current_tokens = 0
+
+    def _flush() -> None:
+        nonlocal current, current_tokens
+        if current:
+            children.append(current)
+            current = []
+            current_tokens = 0
+
+    for item in contents:
+        item_tokens = count_tokens(item.get("content", "") or "")
+
+        if item_tokens > tokens_per_batch:
+            # Oversized: flush in-flight items into their own child,
+            # then emit this item AS-IS (un-chunked) into its own child.
+            # The worker will sequentially chunk it inside retain_batch_async.
+            _flush()
+            children.append([item])
+            continue
+
+        if current and current_tokens + item_tokens > tokens_per_batch:
+            _flush()
+        current.append(item)
+        current_tokens += item_tokens
+
+    _flush()
+    return children
+
+
 def _is_invalid_embedding_dimension_error(e: Exception) -> bool:
     """Return True for deterministic embedding-dimension failures.
 
@@ -9655,28 +9721,28 @@ async def submit_async_retain(
         config = get_config()
         tokens_per_batch = config.retain_batch_tokens
 
-        # Split into sub-batches based on token count. Oversized single
-        # items get chunked into per-chunk sub-batches by the helper
-        # (see issue #1571). origin_indices is unused here because
-        # submit_async_retain returns only the parent operation_id, not
-        # per-input results.
-        sub_batches = _split_contents_into_sub_batches(
-            cast(list[RetainContentDict], contents), tokens_per_batch
-        ).sub_batches
+        # Pack items into child operations by token budget. An oversized
+        # single item is emitted as its own un-chunked child rather than
+        # being fragmented across siblings — workers have no
+        # per-document serialization, so concurrent siblings would race
+        # on the same document_id and trigger FK violations in the final
+        # ANN pass (issue #1795). The worker's in-process splitter
+        # handles intra-document chunking sequentially.
+        sub_batches = _split_contents_into_async_children(cast(list[RetainContentDict], contents), tokens_per_batch)
 
         # Log splitting info if we actually split
         if len(sub_batches) > 1:
             sub_batch_sizes = [len(b) for b in sub_batches]
             if len(sub_batches) <= 20:
                 logger.info(
                     f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). "
-                    f"Split into {len(sub_batches)} sub-batches: {sub_batch_sizes} items each"
+                    f"Split into {len(sub_batches)} child operations: {sub_batch_sizes} items each"
                 )
             else:
                 logger.info(
                     f"Large async retain batch ({total_tokens:,} tokens from {len(contents)} items). "
-                    f"Split into {len(sub_batches)} sub-batches "
-                    f"(items per sub-batch: min={min(sub_batch_sizes)}, "
+                    f"Split into {len(sub_batches)} child operations "
+                    f"(items per child: min={min(sub_batch_sizes)}, "
                     f"max={max(sub_batch_sizes)}, total={sum(sub_batch_sizes)})"
                 )
 
@@ -9742,7 +9808,7 @@ async def submit_async_retain(
                     if len(sub_batches) > 1:
                         sub_batch_tokens = sum(count_tokens(item.get("content", "")) for item in sub_batch)
                         logger.info(
-                            f"Submitting sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
+                            f"Submitting child {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_tokens:,} tokens"
                         )
 
                     task_payload: dict[str, Any] = {"contents": sub_batch}

hindsight-api-slim/tests/test_async_batch_retain.py

@@ -832,6 +832,173 @@ async def failing_submit_task(_task_dict):
     assert payload["contents"] == [{"content": "hello", "document_id": "d1"}]
 
 
+# ---------------------------------------------------------------------------
+# Regression tests for issue #1795: submit_async_retain must NOT fragment a
+# single oversized item across multiple child async-operations. Workers have
+# no per-document serialization for retain (the busy-bank guard in
+# claim_tasks only covers consolidation), so concurrent siblings sharing one
+# document_id race on handle_document_tracking(is_first_batch=True), cascade-
+# delete each other's memory_units, and trip FK violations on memory_links in
+# the final ANN pass. The fix: keep the oversized item un-chunked in a single
+# child; the worker's in-process splitter handles intra-document chunking
+# sequentially with correct is_first_batch semantics.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_oversized_single_item_creates_one_child_not_many(memory, request_context, monkeypatch):
+    """One oversized document → one child operation holding the un-chunked content."""
+    from hindsight_api.config import get_config
+    from hindsight_api.engine.memory_engine import count_tokens
+
+    bank_id = f"test_oversized_one_child_{uuid.uuid4().hex[:8]}"
+    pool = await memory._get_pool()
+    await _ensure_bank(pool, bank_id)
+
+    # No-op the worker dispatch: this is a structural assertion on the
+    # async_operations rows that submit_async_retain inserts (those rows are
+    # committed before submit_task is invoked), so we don't need the
+    # SyncTaskBackend to drive the slow LLM pipeline to completion.
+    async def noop_submit_task(_task_dict):
+        return None
+
+    monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task)
+
+    tokens_per_batch = get_config().retain_batch_tokens
+    # Build content that comfortably exceeds the per-batch token budget.
+    # The pre-fix code would split this into many siblings, all sharing
+    # the same document_id and racing on handle_document_tracking.
+    unit = "The quick brown fox jumps over the lazy dog. " * 50  # ~570 tokens
+    repetitions = max(1, (tokens_per_batch * 4) // count_tokens(unit) + 1)
+    big_content = unit * repetitions
+    document_id = f"doc-oversize-{uuid.uuid4().hex[:8]}"
+
+    total_tokens = count_tokens(big_content)
+    assert total_tokens > tokens_per_batch, (
+        f"Test setup error: content has {total_tokens} tokens, must exceed the batch budget of {tokens_per_batch}"
+    )
+
+    result = await memory.submit_async_retain(
+        bank_id=bank_id,
+        contents=[{"content": big_content, "document_id": document_id}],
+        request_context=request_context,
+    )
+    parent_operation_id = result["operation_id"]
+
+    # Parent metadata must say one child — even though token count is huge,
+    # the un-fragmenting splitter keeps the single item in one child.
+    parent_row = await pool.fetchrow(
+        "SELECT result_metadata FROM async_operations WHERE operation_id = $1",
+        uuid.UUID(parent_operation_id),
+    )
+    parent_meta = (
+        json.loads(parent_row["result_metadata"])
+        if isinstance(parent_row["result_metadata"], str)
+        else parent_row["result_metadata"]
+    )
+    assert parent_meta["num_sub_batches"] == 1, (
+        f"Expected 1 child for an oversized single item, got {parent_meta['num_sub_batches']}. "
+        "Issue #1795: per-chunk children race on the shared document_id."
+    )
+
+    # Exactly one retain child row.
+    children = await pool.fetch(
+        """
+        SELECT operation_id, status, task_payload, result_metadata
+        FROM async_operations
+        WHERE bank_id = $1 AND operation_type = 'retain'
+        """,
+        bank_id,
+    )
+    assert len(children) == 1, (
+        f"Expected exactly 1 child retain operation for bank_id={bank_id}, "
+        f"got {len(children)}. Issue #1795 regression: oversized item was "
+        f"fragmented into per-chunk children."
+    )
+
+    child = children[0]
+    payload = child["task_payload"]
+    payload = json.loads(payload) if isinstance(payload, str) else payload
+    assert payload["type"] == "batch_retain"
+    assert payload["bank_id"] == bank_id
+    # Child holds the full un-chunked content — the worker's in-process
+    # splitter will re-chunk it sequentially with is_first_batch=(i==1).
+    assert len(payload["contents"]) == 1, (
+        f"Child payload should hold the original single item, got {len(payload['contents'])} items"
+    )
+    assert payload["contents"][0]["document_id"] == document_id
+    assert payload["contents"][0]["content"] == big_content, (
+        "Child must carry the FULL un-chunked content — chunking happens inside the worker, not at submit time"
+    )
+
+
+@pytest.mark.asyncio
+async def test_oversized_item_among_small_items_keeps_small_items_packed(memory, request_context, monkeypatch):
+    """Mixed batch: small items pack together; oversized item isolates to its own child."""
+    from hindsight_api.config import get_config
+    from hindsight_api.engine.memory_engine import count_tokens
+
+    bank_id = f"test_mixed_pack_{uuid.uuid4().hex[:8]}"
+    pool = await memory._get_pool()
+    await _ensure_bank(pool, bank_id)
+
+    # Structural assertion on async_operations rows only — skip worker dispatch.
+    async def noop_submit_task(_task_dict):
+        return None
+
+    monkeypatch.setattr(memory._task_backend, "submit_task", noop_submit_task)
+
+    tokens_per_batch = get_config().retain_batch_tokens
+    big_unit = "The quick brown fox jumps over the lazy dog. " * 50
+    big_repetitions = max(1, (tokens_per_batch * 3) // count_tokens(big_unit) + 1)
+    big_content = big_unit * big_repetitions
+
+    big_doc = f"doc-big-{uuid.uuid4().hex[:8]}"
+    small_docs = [f"doc-small-{i}-{uuid.uuid4().hex[:6]}" for i in range(3)]
+    contents = [
+        {"content": "Alice works at Google.", "document_id": small_docs[0]},
+        {"content": big_content, "document_id": big_doc},
+        {"content": "Bob loves Python.", "document_id": small_docs[1]},
+        {"content": "Carol writes Rust.", "document_id": small_docs[2]},
+    ]
+
+    result = await memory.submit_async_retain(
+        bank_id=bank_id,
+        contents=contents,
+        request_context=request_context,
+    )
+
+    children = await pool.fetch(
+        """
+        SELECT task_payload
+        FROM async_operations
+        WHERE bank_id = $1 AND operation_type = 'retain'
+        ORDER BY created_at, operation_id
+        """,
+        bank_id,
+    )
+
+    # Find the child holding the big doc — it must hold ONLY the big doc.
+    big_child = None
+    small_doc_ids_seen: list[str] = []
+    for row in children:
+        payload = row["task_payload"]
+        payload = json.loads(payload) if isinstance(payload, str) else payload
+        items = payload["contents"]
+        doc_ids = [item["document_id"] for item in items]
+        if big_doc in doc_ids:
+            assert big_child is None, f"Big doc {big_doc} should appear in exactly one child, found in 2"
+            big_child = items
+            assert doc_ids == [big_doc], f"Child holding the oversized item must hold ONLY that item, got {doc_ids}"
+            assert items[0]["content"] == big_content
+        else:
+            small_doc_ids_seen.extend(doc_ids)
+
+    assert big_child is not None, "Big doc must appear in some child"
+    # Every small input present, exactly once, across the other children.
+    assert sorted(small_doc_ids_seen) == sorted(small_docs)
+
+
 @pytest.mark.asyncio
 async def test_submit_async_batch_retain_rolls_back_parent_on_child_failure(
     memory_no_llm_verify, request_context, monkeypatch