From 6bf678953044bb7c1e6ab61151a1d8391816adad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= <boschi1997@gmail.com>
Date: Thu, 28 May 2026 11:18:15 +0200
Subject: [PATCH 1/3] fix(retain): keep oversized items in one async child to
 stop FK race (#1795)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

submit_async_retain split oversized retain payloads into N independent
async_operations rows that all shared one document_id. Workers have no
per-document gate for retain (claim_tasks only guards consolidation),
so siblings ran concurrently — each entered handle_document_tracking
with is_first_batch=True, cascade-deleting the previous winner's
memory_units. The loser's final ANN pass then inserted memory_links
referencing now-deleted units, tripping
fk_memory_links_from_unit_id_memory_units. Concurrent siblings also
exhausted OS thread budgets via per-child sentence-transformer pools
(libgomp resource-unavailable failures) and left partial document
state visible to dry-run skip checks.

Add _split_contents_into_async_children for the async submit path: it
packs items into children by token budget but never fragments a single
item across children. Oversized items go into their own one-item child
holding the full un-chunked content; the worker's existing in-process
splitter (retain_batch_async → _split_contents_into_sub_batches)
re-chunks them sequentially inside one worker slot with correct
is_first_batch=(i==1) semantics — the same path that already enforces
SELECT … FOR UPDATE + content-hash gating between batches of one call.

Small items still pack together so genuinely independent inputs keep
cross-worker parallelism. Metadata field names (num_sub_batches,
sub_batch_index, total_sub_batches) are unchanged.

Tests:
- 8 pure-Python tests for the new helper covering single oversized,
  metadata preservation, packing by budget, mixed inputs, multiple
  oversized, boundary positioning, empty input.
- 3 integration tests against the real DB:
  - test_oversized_single_item_creates_one_child_not_many asserts the
    async_operations table has exactly one retain row with the
    un-chunked content (fails on pre-fix code: "got 7" children).
  - test_oversized_single_item_drains_without_fk_violation drives a
    worker drain and asserts no memory_links rows have orphan FKs in
    either direction — the exact invariant pre-fix code violated.
  - test_oversized_item_among_small_items_keeps_small_items_packed
    confirms the parallelism optimization isn't lost.
---
 .../hindsight_api/engine/memory_engine.py     |  90 ++++++-
 .../tests/test_async_batch_retain.py          | 248 ++++++++++++++++++
 .../tests/test_batch_chunking.py              | 153 +++++++++++
 3 files changed, 479 insertions(+), 12 deletions(-)

diff --git a/hindsight-api-slim/hindsight_api/engine/memory_engine.py b/hindsight-api-slim/hindsight_api/engine/memory_engine.py
index 08876f7b92..b47a9c47f9 100644
--- a/hindsight-api-slim/hindsight_api/engine/memory_engine.py
+++ b/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,14 +9721,14 @@ 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:
@@ -9670,13 +9736,13 @@ async def submit_async_retain(
             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}
diff --git a/hindsight-api-slim/tests/test_async_batch_retain.py b/hindsight-api-slim/tests/test_async_batch_retain.py
index 2771e992d0..dda8ef0f8c 100644
--- a/hindsight-api-slim/tests/test_async_batch_retain.py
+++ b/hindsight-api-slim/tests/test_async_batch_retain.py
@@ -832,6 +832,254 @@ 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):
+    """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)
+
+    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_single_item_drains_without_fk_violation(memory, request_context):
+    """End-to-end: oversized doc retains successfully — no FK violation, no orphaned units."""
+    from hindsight_api.config import get_config
+    from hindsight_api.engine.memory_engine import count_tokens
+
+    bank_id = f"test_oversized_drain_{uuid.uuid4().hex[:8]}"
+    pool = await memory._get_pool()
+    await _ensure_bank(pool, bank_id)
+
+    tokens_per_batch = get_config().retain_batch_tokens
+    # Use natural-language sentences so the in-process splitter actually
+    # produces multiple chunks during worker execution (this is what
+    # exercises the is_first_batch=(i==1) sequencing inside retain_batch_async).
+    paragraph = (
+        "Alice met Bob at the coffee shop in Paris. "
+        "They discussed the new project deadline. "
+        "Carol joined later with research notes from Berlin. "
+        "The team agreed to ship the prototype by Friday. "
+    )
+    big_content = paragraph * max(1, (tokens_per_batch * 3) // count_tokens(paragraph) + 1)
+    document_id = f"doc-drain-{uuid.uuid4().hex[:8]}"
+    assert count_tokens(big_content) > 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,
+    )
+
+    # SyncTaskBackend runs the child inline; give it a moment to settle.
+    await asyncio.sleep(0.5)
+
+    parent_status = await memory.get_operation_status(
+        bank_id=bank_id,
+        operation_id=result["operation_id"],
+        request_context=request_context,
+    )
+    assert parent_status["status"] == "completed", parent_status
+    assert len(parent_status["child_operations"]) == 1
+    assert parent_status["child_operations"][0]["status"] == "completed", parent_status
+
+    # Document row exists and is owned by this retain (content_hash is a
+    # real SHA, not the __pending__ placeholder used during in-flight
+    # ownership-locking).
+    doc_row = await pool.fetchrow(
+        "SELECT content_hash FROM documents WHERE id = $1 AND bank_id = $2",
+        document_id,
+        bank_id,
+    )
+    assert doc_row is not None, "Document row must exist after successful retain"
+    assert doc_row["content_hash"] != "__pending__", (
+        f"Document content_hash is still the placeholder: {doc_row['content_hash']!r}. "
+        "Retain finished without writing the real hash — partial state."
+    )
+
+    # Memory units were committed for this document.
+    unit_count = await pool.fetchval(
+        "SELECT COUNT(*) FROM memory_units WHERE document_id = $1 AND bank_id = $2",
+        document_id,
+        bank_id,
+    )
+    assert unit_count > 0, (
+        f"Expected memory_units for document {document_id} after successful retain, "
+        f"got 0. A cascade-delete race (issue #1795) would wipe them."
+    )
+
+    # Every memory_link FK is satisfied — no rows pointing at deleted units.
+    # This is the exact invariant the pre-fix code violated: child A would
+    # commit units {u1..uN}, child B cascade-deletes them, then A's final
+    # ANN pass tries to insert links from u1..uN → FK violation.
+    orphaned_from = await pool.fetchval(
+        """
+        SELECT COUNT(*)
+        FROM memory_links ml
+        LEFT JOIN memory_units mu ON mu.id = ml.from_unit_id
+        WHERE ml.bank_id = $1 AND mu.id IS NULL
+        """,
+        bank_id,
+    )
+    assert orphaned_from == 0, (
+        f"Found {orphaned_from} memory_links with from_unit_id pointing at "
+        f"non-existent memory_units. This is the issue #1795 race symptom."
+    )
+    orphaned_to = await pool.fetchval(
+        """
+        SELECT COUNT(*)
+        FROM memory_links ml
+        LEFT JOIN memory_units mu ON mu.id = ml.to_unit_id
+        WHERE ml.bank_id = $1 AND mu.id IS NULL
+        """,
+        bank_id,
+    )
+    assert orphaned_to == 0, f"Found {orphaned_to} memory_links with to_unit_id pointing at non-existent memory_units."
+
+
+@pytest.mark.asyncio
+async def test_oversized_item_among_small_items_keeps_small_items_packed(memory, request_context):
+    """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)
+
+    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
diff --git a/hindsight-api-slim/tests/test_batch_chunking.py b/hindsight-api-slim/tests/test_batch_chunking.py
index 60674b5c75..f0f94d2827 100644
--- a/hindsight-api-slim/tests/test_batch_chunking.py
+++ b/hindsight-api-slim/tests/test_batch_chunking.py
@@ -7,6 +7,7 @@
 
 from hindsight_api import MemoryEngine
 from hindsight_api.engine.memory_engine import (
+    _split_contents_into_async_children,
     _split_contents_into_sub_batches,
     count_tokens,
 )
@@ -118,6 +119,158 @@ def test_split_small_batch_returns_single_sub_batch():
     assert split.origin_indices == [[0, 1]]
 
 
+# ---------------------------------------------------------------------------
+# Regression tests for issue #1795: the async-submit splitter must NEVER
+# fragment a single oversized item across multiple children. Each child
+# becomes an independent async_operations row claimed by workers in
+# parallel with no per-document gate, so siblings would race on the same
+# document_id, cascade-delete each other's memory_units, and trip FK
+# violations on memory_links in the final ANN pass.
+# ---------------------------------------------------------------------------
+
+
+def test_async_children_oversized_single_item_stays_in_one_child():
+    """An oversized single item must NOT be split across children."""
+    tokens_per_batch = 1_000
+    big_content = "The quick brown fox jumps over the lazy dog. " * 1_000
+    item = {"content": big_content, "document_id": "doc-oversize"}
+    assert count_tokens(big_content) > tokens_per_batch
+
+    children = _split_contents_into_async_children([item], tokens_per_batch)
+
+    # Exactly one child holding the full un-chunked item — the worker
+    # will sequentially chunk it inside retain_batch_async.
+    assert len(children) == 1, (
+        f"Oversized single item must become exactly one child, got {len(children)}. "
+        "Regressing to per-chunk children causes the issue #1795 race."
+    )
+    assert children[0] == [item]
+    assert children[0][0]["content"] == big_content
+
+
+def test_async_children_oversized_item_preserves_metadata():
+    """The single-child payload must carry every original key untouched."""
+    tokens_per_batch = 500
+    item = {
+        "content": "Alice met Bob at the coffee shop. " * 500,
+        "document_id": "doc-42",
+        "context": "shared-context",
+        "tags": ["t1", "t2"],
+        "metadata": {"source": "test"},
+    }
+
+    children = _split_contents_into_async_children([item], tokens_per_batch)
+
+    assert len(children) == 1
+    assert children[0] == [item]
+
+
+def test_async_children_packs_small_items_by_budget():
+    """Small items pack into shared children up to the token budget."""
+    tokens_per_batch = 100
+    item_text = "Alice works at Google and Bob loves Python programming language."
+    item_tokens = count_tokens(item_text)
+    # Pick enough items to force at least 2 children at this budget.
+    num_items = max(4, (tokens_per_batch // max(item_tokens, 1)) * 3)
+    contents = [{"content": item_text, "document_id": f"doc-{i}"} for i in range(num_items)]
+    total = sum(count_tokens(c["content"]) for c in contents)
+    assert total > tokens_per_batch, (
+        f"Test setup error: {total} tokens does not exceed budget {tokens_per_batch}"
+    )
+
+    children = _split_contents_into_async_children(contents, tokens_per_batch)
+
+    assert len(children) >= 2
+    # Every child stays within budget.
+    for child in children:
+        child_tokens = sum(count_tokens(item["content"]) for item in child)
+        assert child_tokens <= tokens_per_batch, (
+            f"Child holding {child_tokens} tokens exceeds budget {tokens_per_batch}"
+        )
+    # Every input item appears exactly once across all children.
+    flat = [item for child in children for item in child]
+    assert flat == contents
+
+
+def test_async_children_mixed_small_and_oversized():
+    """Small items pack together; the oversized item gets its own child as-is."""
+    tokens_per_batch = 1_000
+    small_a = {"content": "Alice works at Google. " * 5, "document_id": "doc-a"}
+    big = {"content": "The quick brown fox jumps over the lazy dog. " * 1_000, "document_id": "doc-big"}
+    small_b = {"content": "Bob loves Python. " * 5, "document_id": "doc-b"}
+    small_c = {"content": "Carol writes Rust. " * 5, "document_id": "doc-c"}
+
+    children = _split_contents_into_async_children([small_a, big, small_b, small_c], tokens_per_batch)
+
+    # The big item must be in its own child, un-chunked.
+    big_children = [c for c in children if any(item is big or item.get("document_id") == "doc-big" for item in c)]
+    assert len(big_children) == 1, (
+        f"Oversized item must occupy exactly one child, found {len(big_children)} containing it"
+    )
+    assert big_children[0] == [big]
+    assert big_children[0][0]["content"] == big["content"]
+
+    # Every small item is present somewhere.
+    flat_doc_ids = [item["document_id"] for child in children for item in child]
+    for small in (small_a, small_b, small_c):
+        assert small["document_id"] in flat_doc_ids
+
+    # No child fragments the big item.
+    for child in children:
+        big_count = sum(1 for item in child if item.get("document_id") == "doc-big")
+        assert big_count <= 1
+
+
+def test_async_children_each_oversized_item_gets_own_child():
+    """Multiple oversized items → one child per item, each un-chunked."""
+    tokens_per_batch = 500
+    big_a = {"content": "Alpha sentence. " * 500, "document_id": "doc-a"}
+    big_b = {"content": "Bravo sentence. " * 500, "document_id": "doc-b"}
+    big_c = {"content": "Charlie sentence. " * 500, "document_id": "doc-c"}
+
+    children = _split_contents_into_async_children([big_a, big_b, big_c], tokens_per_batch)
+
+    assert len(children) == 3
+    assert children == [[big_a], [big_b], [big_c]]
+
+
+def test_async_children_single_small_item_returns_one_child():
+    """A single small item is one child of one item — no splitting."""
+    tokens_per_batch = 10_000
+    item = {"content": "Alice works at Google", "document_id": "doc-1"}
+
+    children = _split_contents_into_async_children([item], tokens_per_batch)
+
+    assert children == [[item]]
+
+
+def test_async_children_empty_returns_empty():
+    """Empty input → empty output (no spurious child)."""
+    assert _split_contents_into_async_children([], 1_000) == []
+
+
+def test_async_children_oversized_at_boundaries():
+    """Oversized items at the start, middle, and end all isolate correctly."""
+    tokens_per_batch = 1_000
+    small_a = {"content": "Small A. " * 5, "document_id": "doc-sa"}
+    small_b = {"content": "Small B. " * 5, "document_id": "doc-sb"}
+    big_start = {"content": "First big. " * 1_000, "document_id": "doc-bs"}
+    big_end = {"content": "Last big. " * 1_000, "document_id": "doc-be"}
+
+    children = _split_contents_into_async_children([big_start, small_a, small_b, big_end], tokens_per_batch)
+
+    # First child holds big_start alone; last child holds big_end alone.
+    assert children[0] == [big_start]
+    assert children[-1] == [big_end]
+    # Middle children contain only small items.
+    middle_items = [item for child in children[1:-1] for item in child]
+    assert all(item["document_id"] in {"doc-sa", "doc-sb"} for item in middle_items)
+    # Every small input is present.
+    flat_ids = [item["document_id"] for child in children for item in child]
+    assert "doc-sa" in flat_ids
+    assert "doc-sb" in flat_ids
+
+
 @pytest.mark.asyncio
 async def test_large_batch_auto_chunks(memory, request_context):
     bank_id = "test_chunking_agent"

From 309bf3b6b517b362462f91dcd8b3c25b26a0c9f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= <boschi1997@gmail.com>
Date: Thu, 28 May 2026 11:53:17 +0200
Subject: [PATCH 2/3] test(retain): no-op worker dispatch in structural tests
 for #1795
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The two structural assertions (test_oversized_single_item_creates_one_child_not_many
and test_oversized_item_among_small_items_keeps_small_items_packed) only need to
verify the async_operations rows that submit_async_retain inserts — those rows
commit before submit_task is called. The previous version let SyncTaskBackend
drive the full LLM-based retain pipeline synchronously, which timed out at
CI's 300s per-test limit even though it ran in ~5s locally.

Monkeypatch _task_backend.submit_task to a no-op so the structural assertions
fire in ~30ms without running the worker.

Also slim the drain test's payload from ~3x to ~1.2x the per-batch token budget.
That still triggers in-process splitting (~2 sub-batches → the path that
exercises is_first_batch=(i==1) sequencing) but cuts LLM extraction work from
~5 chunks to ~2, keeping wall time comfortably under 300s on slower runners.

The structural regression assertions still fail without the engine fix —
verified by temporarily reverting hindsight_api/engine/memory_engine.py and
re-running: "Expected 1 child for an oversized single item, got 7. Issue #1795:
per-chunk children race on the shared document_id."
---
 .../tests/test_async_batch_retain.py          | 31 +++++++++++++++----
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/hindsight-api-slim/tests/test_async_batch_retain.py b/hindsight-api-slim/tests/test_async_batch_retain.py
index dda8ef0f8c..c4a0529aa4 100644
--- a/hindsight-api-slim/tests/test_async_batch_retain.py
+++ b/hindsight-api-slim/tests/test_async_batch_retain.py
@@ -846,7 +846,7 @@ async def failing_submit_task(_task_dict):
 
 
 @pytest.mark.asyncio
-async def test_oversized_single_item_creates_one_child_not_many(memory, request_context):
+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
@@ -855,6 +855,15 @@ async def test_oversized_single_item_creates_one_child_not_many(memory, request_
     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
@@ -934,16 +943,20 @@ async def test_oversized_single_item_drains_without_fk_violation(memory, request
     await _ensure_bank(pool, bank_id)
 
     tokens_per_batch = get_config().retain_batch_tokens
-    # Use natural-language sentences so the in-process splitter actually
-    # produces multiple chunks during worker execution (this is what
-    # exercises the is_first_batch=(i==1) sequencing inside retain_batch_async).
+    # Use natural-language sentences so the in-process splitter produces
+    # multiple chunks during worker execution (the path that exercises
+    # is_first_batch=(i==1) sequencing inside retain_batch_async).
+    # Size kept small (~1.2x budget → 2 sub-batches) so LLM extraction
+    # finishes within CI's 300s per-test timeout — locally this runs in
+    # ~25s but slower CI runners can take an order of magnitude longer.
     paragraph = (
         "Alice met Bob at the coffee shop in Paris. "
         "They discussed the new project deadline. "
         "Carol joined later with research notes from Berlin. "
         "The team agreed to ship the prototype by Friday. "
     )
-    big_content = paragraph * max(1, (tokens_per_batch * 3) // count_tokens(paragraph) + 1)
+    paragraph_tokens = count_tokens(paragraph)
+    big_content = paragraph * max(1, int(tokens_per_batch * 1.2) // paragraph_tokens + 1)
     document_id = f"doc-drain-{uuid.uuid4().hex[:8]}"
     assert count_tokens(big_content) > tokens_per_batch
 
@@ -1020,7 +1033,7 @@ async def test_oversized_single_item_drains_without_fk_violation(memory, request
 
 
 @pytest.mark.asyncio
-async def test_oversized_item_among_small_items_keeps_small_items_packed(memory, request_context):
+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
@@ -1029,6 +1042,12 @@ async def test_oversized_item_among_small_items_keeps_small_items_packed(memory,
     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)

From 84d0e53dc17e9c7df8d7044a4d88c3e8c84678b5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Nicol=C3=B2=20Boschi?= <boschi1997@gmail.com>
Date: Thu, 28 May 2026 12:22:15 +0200
Subject: [PATCH 3/3] =?UTF-8?q?test(retain):=20drop=20end-to-end=20drain?=
 =?UTF-8?q?=20test=20for=20#1795=20=E2=80=94=20too=20CI-flaky?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

test_oversized_single_item_drains_without_fk_violation drives the full
retain pipeline (LLM extraction + embeddings + ANN + consolidation)
synchronously through SyncTaskBackend. Even with the payload trimmed
to ~1.2x the batch budget (~2 sub-batches), Gemini API latency in CI
varies enough that the 300s per-test timeout fires intermittently.

The fix is already covered without it:
- test_oversized_single_item_creates_one_child_not_many is the direct
  regression test for #1795. It asserts on the async_operations rows
  submit_async_retain inserts and was empirically shown to fail on
  the pre-fix engine ("Expected 1 child for an oversized single item,
  got 7"). No worker execution needed.
- test_oversized_item_among_small_items_keeps_small_items_packed
  covers the mixed-batch case structurally.
- 8 unit tests in test_batch_chunking.py cover the helper directly.
- The FK constraint fk_memory_links_from_unit_id_memory_units is
  enforced by Postgres itself; any orphan write would error at insert
  time, so the engine cannot silently regress without other tests
  noticing.
---
 .../tests/test_async_batch_retain.py          | 100 ------------------
 1 file changed, 100 deletions(-)

diff --git a/hindsight-api-slim/tests/test_async_batch_retain.py b/hindsight-api-slim/tests/test_async_batch_retain.py
index c4a0529aa4..78f421bd28 100644
--- a/hindsight-api-slim/tests/test_async_batch_retain.py
+++ b/hindsight-api-slim/tests/test_async_batch_retain.py
@@ -932,106 +932,6 @@ async def noop_submit_task(_task_dict):
     )
 
 
-@pytest.mark.asyncio
-async def test_oversized_single_item_drains_without_fk_violation(memory, request_context):
-    """End-to-end: oversized doc retains successfully — no FK violation, no orphaned units."""
-    from hindsight_api.config import get_config
-    from hindsight_api.engine.memory_engine import count_tokens
-
-    bank_id = f"test_oversized_drain_{uuid.uuid4().hex[:8]}"
-    pool = await memory._get_pool()
-    await _ensure_bank(pool, bank_id)
-
-    tokens_per_batch = get_config().retain_batch_tokens
-    # Use natural-language sentences so the in-process splitter produces
-    # multiple chunks during worker execution (the path that exercises
-    # is_first_batch=(i==1) sequencing inside retain_batch_async).
-    # Size kept small (~1.2x budget → 2 sub-batches) so LLM extraction
-    # finishes within CI's 300s per-test timeout — locally this runs in
-    # ~25s but slower CI runners can take an order of magnitude longer.
-    paragraph = (
-        "Alice met Bob at the coffee shop in Paris. "
-        "They discussed the new project deadline. "
-        "Carol joined later with research notes from Berlin. "
-        "The team agreed to ship the prototype by Friday. "
-    )
-    paragraph_tokens = count_tokens(paragraph)
-    big_content = paragraph * max(1, int(tokens_per_batch * 1.2) // paragraph_tokens + 1)
-    document_id = f"doc-drain-{uuid.uuid4().hex[:8]}"
-    assert count_tokens(big_content) > 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,
-    )
-
-    # SyncTaskBackend runs the child inline; give it a moment to settle.
-    await asyncio.sleep(0.5)
-
-    parent_status = await memory.get_operation_status(
-        bank_id=bank_id,
-        operation_id=result["operation_id"],
-        request_context=request_context,
-    )
-    assert parent_status["status"] == "completed", parent_status
-    assert len(parent_status["child_operations"]) == 1
-    assert parent_status["child_operations"][0]["status"] == "completed", parent_status
-
-    # Document row exists and is owned by this retain (content_hash is a
-    # real SHA, not the __pending__ placeholder used during in-flight
-    # ownership-locking).
-    doc_row = await pool.fetchrow(
-        "SELECT content_hash FROM documents WHERE id = $1 AND bank_id = $2",
-        document_id,
-        bank_id,
-    )
-    assert doc_row is not None, "Document row must exist after successful retain"
-    assert doc_row["content_hash"] != "__pending__", (
-        f"Document content_hash is still the placeholder: {doc_row['content_hash']!r}. "
-        "Retain finished without writing the real hash — partial state."
-    )
-
-    # Memory units were committed for this document.
-    unit_count = await pool.fetchval(
-        "SELECT COUNT(*) FROM memory_units WHERE document_id = $1 AND bank_id = $2",
-        document_id,
-        bank_id,
-    )
-    assert unit_count > 0, (
-        f"Expected memory_units for document {document_id} after successful retain, "
-        f"got 0. A cascade-delete race (issue #1795) would wipe them."
-    )
-
-    # Every memory_link FK is satisfied — no rows pointing at deleted units.
-    # This is the exact invariant the pre-fix code violated: child A would
-    # commit units {u1..uN}, child B cascade-deletes them, then A's final
-    # ANN pass tries to insert links from u1..uN → FK violation.
-    orphaned_from = await pool.fetchval(
-        """
-        SELECT COUNT(*)
-        FROM memory_links ml
-        LEFT JOIN memory_units mu ON mu.id = ml.from_unit_id
-        WHERE ml.bank_id = $1 AND mu.id IS NULL
-        """,
-        bank_id,
-    )
-    assert orphaned_from == 0, (
-        f"Found {orphaned_from} memory_links with from_unit_id pointing at "
-        f"non-existent memory_units. This is the issue #1795 race symptom."
-    )
-    orphaned_to = await pool.fetchval(
-        """
-        SELECT COUNT(*)
-        FROM memory_links ml
-        LEFT JOIN memory_units mu ON mu.id = ml.to_unit_id
-        WHERE ml.bank_id = $1 AND mu.id IS NULL
-        """,
-        bank_id,
-    )
-    assert orphaned_to == 0, f"Found {orphaned_to} memory_links with to_unit_id pointing at non-existent memory_units."
-
-
 @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."""