fix(entity-resolver): keep every co-occurrence pair when canonicalising order (#2750)

The pair canonicalisation in _link_units_to_entities_batch_impl swapped
entity_id_1 and entity_id_2 in place, but entity_id_1 is the outer loop's
iterate:

    for i, entity_id_1 in enumerate(entity_list):
        for entity_id_2 in entity_list[i + 1:]:
            if entity_id_1 > entity_id_2:
                entity_id_1, entity_id_2 = entity_id_2, entity_id_1

Once a swap happens, entity_id_1 stays swapped for the rest of that inner
loop, so every later pair in the same outer iteration is built from the
wrong first element. Those pairs collide with ones already emitted, so the
effect is silently missing edges rather than wrong ones.

entity_list comes from a set, so the ordering (and the bug) varies per run.

Move the canonicalisation into a _canonical_cooccurrence_pairs() helper that
orders each pair into fresh locals, leaving the iterate untouched, and cover
it with order-pinned unit tests that need no database.

Author: ebarkhordar

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

@@ -9,6 +9,7 @@
 import json
 import logging
 from collections import defaultdict
+from collections.abc import Iterator
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from difflib import SequenceMatcher
@@ -75,6 +76,22 @@ def _later_date(a: datetime | None, b: datetime | None) -> datetime | None:
     return a if a > b else b
 
 
+def _canonical_cooccurrence_pairs(entity_list: list[str]) -> Iterator[tuple[str, str]]:
+    """Yield each distinct pair of ``entity_list`` as ``(a, b)`` with ``a < b``.
+
+    Canonical ordering matches the entity_cooccurrences PK and check constraint.
+    The pair is ordered into fresh locals rather than by swapping the loop
+    variables: ``entity_id_1`` is the outer iterate, so swapping it would leak
+    into the remaining inner iterations and build later pairs off the wrong
+    element.
+    """
+    for i, entity_id_1 in enumerate(entity_list):
+        for entity_id_2 in entity_list[i + 1 :]:
+            if entity_id_1 == entity_id_2:
+                continue
+            yield (entity_id_1, entity_id_2) if entity_id_1 < entity_id_2 else (entity_id_2, entity_id_1)
+
+
 @dataclass
 class _CooccurrencePair:
     """A (entity_id_1, entity_id_2) pair observed in a retain batch (for post-txn flush)."""
@@ -853,20 +870,12 @@ async def _link_units_to_entities_batch_impl(self, conn, unit_entity_pairs: list
         for unit_id, entity_ids in unit_to_entities.items():
             entity_list = list(entity_ids)
             event_date = unit_event_date.get(unit_id)
-            for i, entity_id_1 in enumerate(entity_list):
-                for entity_id_2 in entity_list[i + 1 :]:
-                    if entity_id_1 == entity_id_2:
-                        continue
-                    # Canonical ordering (entity_id_1 < entity_id_2) matches the
-                    # entity_cooccurrences PK and check constraint.
-                    if entity_id_1 > entity_id_2:
-                        entity_id_1, entity_id_2 = entity_id_2, entity_id_1
-                    key = (entity_id_1, entity_id_2)
-                    prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
-                    if prev is _SENTINEL_MISSING:
-                        cooccurrence_pairs[key] = event_date
-                    else:
-                        cooccurrence_pairs[key] = _later_date(prev, event_date)
+            for key in _canonical_cooccurrence_pairs(entity_list):
+                prev = cooccurrence_pairs.get(key, _SENTINEL_MISSING)
+                if prev is _SENTINEL_MISSING:
+                    cooccurrence_pairs[key] = event_date
+                else:
+                    cooccurrence_pairs[key] = _later_date(prev, event_date)
 
         # Accumulate co-occurrence pairs for post-transaction flush.
         # The actual INSERT/UPDATE is deferred to flush_pending_stats() to avoid

hindsight-api-slim/tests/test_entity_resolver.py

@@ -2,6 +2,7 @@
 Tests for EntityResolver edge cases.
 """
 
+import itertools
 import json
 import uuid
 from datetime import datetime, timezone
@@ -11,7 +12,7 @@
 
 from hindsight_api.engine.db import create_database_backend
 from hindsight_api.engine.db.result import DictResultRow as ResultRow
-from hindsight_api.engine.entity_resolver import EntityResolver
+from hindsight_api.engine.entity_resolver import EntityResolver, _canonical_cooccurrence_pairs
 from hindsight_api.pg0 import resolve_database_url
 
 # ---------------------------------------------------------------------------
@@ -59,6 +60,46 @@ async def test_discard_pending_stats_does_not_affect_other_task_keys():
     assert other_key in resolver._pending_cooccurrences, "other task's cooccurrences must be preserved"
 
 
+# ---------------------------------------------------------------------------
+# Unit tests for _canonical_cooccurrence_pairs() — no database required
+# ---------------------------------------------------------------------------
+
+
+def test_canonical_cooccurrence_pairs_keeps_every_pair_when_input_is_unsorted():
+    """Every distinct pair must survive, whatever order the entities arrive in.
+
+    ["C", "A", "B"] is the minimal witness: ordering the second pair used to
+    swap the outer loop variable, so ("A", "C") was silently dropped.
+    """
+    pairs = list(_canonical_cooccurrence_pairs(["C", "A", "B"]))
+
+    assert set(pairs) == {("A", "B"), ("A", "C"), ("B", "C")}
+    assert len(pairs) == 3
+
+
+def test_canonical_cooccurrence_pairs_is_input_order_invariant():
+    """entity_list is built from a set, so its order is arbitrary per run: the
+    emitted pairs must not depend on it."""
+    expected = {("e1", "e2"), ("e1", "e3"), ("e1", "e4"), ("e2", "e3"), ("e2", "e4"), ("e3", "e4")}
+
+    for order in itertools.permutations(["e1", "e2", "e3", "e4"]):
+        assert set(_canonical_cooccurrence_pairs(list(order))) == expected, f"lost a pair for input order {order}"
+
+
+def test_canonical_cooccurrence_pairs_emits_canonical_order():
+    """Each pair must come out as (a, b) with a < b, the ordering the
+    entity_cooccurrences PK and check constraint require."""
+    for order in itertools.permutations(["e1", "e2", "e3"]):
+        for a, b in _canonical_cooccurrence_pairs(list(order)):
+            assert a < b, f"non-canonical pair {(a, b)} for input order {order}"
+
+
+def test_canonical_cooccurrence_pairs_skips_duplicates_and_short_input():
+    assert list(_canonical_cooccurrence_pairs(["e1", "e1"])) == []
+    assert list(_canonical_cooccurrence_pairs(["e1"])) == []
+    assert list(_canonical_cooccurrence_pairs([])) == []
+
+
 @pytest.mark.asyncio
 async def test_resolve_entities_batch_handles_unicode_lower_conflicts(pg0_db_url):
     """