fix(consolidation): refresh mental models every round, not just the final one (#3411)

hindsight-api-slim/hindsight_api/engine/consolidation/consolidator.py

@@ -1088,12 +1088,65 @@ async def _count_unconsolidated_rows(
     )
 
 
+def _as_op_uuid(operation_id: str | uuid.UUID) -> uuid.UUID:
+    return uuid.UUID(operation_id) if isinstance(operation_id, str) else operation_id
+
+
+async def _persist_pending_refresh_tags(conn, operation_id: str, new_tags: list[str]) -> None:
+    """Union ``new_tags`` into the consolidation op's durable ``pending_refresh_tags``.
+
+    Called inside each batch's witness transaction, so the tags of an
+    already-consolidated batch are durable the instant that batch is — a mid-round
+    worker crash no longer loses them. On retry the op re-reads ``task_payload`` and the
+    final round still refreshes those models (#3411); without this, a crash after batch 1
+    committed but before the round finished would drop batch 1's tags, because the retry
+    skips its now-consolidated rows and never re-collects them. ``SELECT ... FOR UPDATE``
+    serialises the concurrent batches of one op so their unions don't clobber each other.
+    """
+    op_uuid = _as_op_uuid(operation_id)
+    row = await conn.fetchrow(
+        f"SELECT task_payload FROM {fq_table('async_operations')} WHERE operation_id = $1 FOR UPDATE",
+        op_uuid,
+    )
+    if row is None:
+        return
+    payload = row["task_payload"]
+    payload = json.loads(payload) if isinstance(payload, str) else (payload or {})
+    existing = set(payload.get("pending_refresh_tags") or [])
+    merged = existing | set(new_tags)
+    if merged == existing:
+        return
+    payload["pending_refresh_tags"] = sorted(merged)
+    await conn.execute(
+        f"UPDATE {fq_table('async_operations')} SET task_payload = $1::jsonb, updated_at = now() "
+        f"WHERE operation_id = $2",
+        json.dumps(payload),
+        op_uuid,
+    )
+
+
+async def _read_pending_refresh_tags(pool, operation_id: str) -> set[str]:
+    """Read the op's durably-accumulated ``pending_refresh_tags`` (crash-safe source of
+    truth for the final-round flush)."""
+    async with acquire_with_retry(pool) as conn:
+        row = await conn.fetchrow(
+            f"SELECT task_payload FROM {fq_table('async_operations')} WHERE operation_id = $1",
+            _as_op_uuid(operation_id),
+        )
+    if row is None:
+        return set()
+    payload = row["task_payload"]
+    payload = json.loads(payload) if isinstance(payload, str) else (payload or {})
+    return set(payload.get("pending_refresh_tags") or [])
+
+
 async def run_consolidation_job(
     memory_engine: "MemoryEngine",
     bank_id: str,
     request_context: "RequestContext",
     operation_id: str | None = None,
     observation_scopes: list[list[str]] | None = None,
+    pending_refresh_tags: list[str] | None = None,
 ) -> dict[str, Any]:
     """
     Run consolidation job for a bank.
@@ -1108,6 +1161,9 @@ async def run_consolidation_job(
         observation_scopes: Optional list of tag scopes. When provided, only
             unconsolidated memories whose tags contain all tags in at least one
             scope are processed.
+        pending_refresh_tags: Tags of memories consolidated by earlier rounds of this
+            round-limited chain, carried through the re-queue so the final round can
+            refresh every affected mental model exactly once (#3411).
 
     Returns:
         Dict with consolidation results
@@ -1127,7 +1183,14 @@ async def run_consolidation_job(
     trace_token = set_trace_context(trace_ctx) if trace_ctx is not None else None
     try:
         return await _run_consolidation_job(
-            memory_engine, bank_id, request_context, config, llm_config, operation_id, observation_scopes
+            memory_engine,
+            bank_id,
+            request_context,
+            config,
+            llm_config,
+            operation_id,
+            observation_scopes,
+            pending_refresh_tags,
         )
     finally:
         if trace_token is not None:
@@ -1145,6 +1208,7 @@ async def _run_consolidation_job(
     llm_config: Any,
     operation_id: str | None = None,
     observation_scopes: list[list[str]] | None = None,
+    pending_refresh_tags: list[str] | None = None,
 ) -> dict[str, Any]:
     """Core consolidation flow. See ``run_consolidation_job`` for the public entrypoint."""
     perf = ConsolidationPerfLog(bank_id)
@@ -1439,6 +1503,22 @@ async def _process_one_llm_batch(llm_batch_local: list[dict[str, Any]], batch_nu
                         )
                     async with conn.transaction():
                         await _txn_provider.write_txn_witness(_batch_txn, conn=conn, fq_table=fq_table)
+                        # Persist this batch's mental-model refresh tags atomically with the
+                        # witness, so they share the batch's fate: durable iff the batch is
+                        # (#3411). Only the succeeded source facts — the ones just marked
+                        # consolidated — contribute a tag.
+                        if operation_id and succeeded_ids:
+                            succeeded_set = {str(mem_id) for mem_id in succeeded_ids}
+                            batch_tags = sorted(
+                                {
+                                    t
+                                    for m in llm_batch_local
+                                    if str(m["id"]) in succeeded_set
+                                    for t in (m.get("tags") or [])
+                                }
+                            )
+                            if batch_tags:
+                                await _persist_pending_refresh_tags(conn, operation_id, batch_tags)
             except BaseException:
                 # The witness row was never committed, so this batch's writes are invisible;
                 # discard the write-group rather than leaving it pending for the recovery
@@ -1653,6 +1733,19 @@ async def _run_group(
     # execute_task's retry handler means the op is retried with backoff; on retry the
     # consolidator skips already-consolidated rows via the consolidated_at filter and
     # picks up the remainder. Issue #1842.
+    # The affected-tag union for the whole round-limited chain. Refresh fires once, when
+    # the backlog has fully drained (the final round), not once per round — a model's
+    # memories can straddle rounds, and gating on the final round alone (the prior
+    # behaviour) dropped every model consolidated earlier because the final round's tags
+    # no longer named them (#3411). The union is durable: each batch writes its tags into
+    # the op's ``task_payload`` inside the batch's own witness txn (crash-safe), and the
+    # re-queue threads the accumulated set forward to the next round. Prefer that durable
+    # value; fall back to the in-memory union when there is no backing op (a direct
+    # ``run_consolidation_job`` call, e.g. in tests).
+    all_refresh_tags = set(pending_refresh_tags or []) | consolidated_tags
+    if operation_id:
+        all_refresh_tags |= await _read_pending_refresh_tags(pool, operation_id)
+
     if hit_round_limit:
         remaining = total_count - stats["memories_processed"]
         logger.info(
@@ -1663,6 +1756,7 @@ async def _run_group(
             bank_id=bank_id,
             request_context=request_context,
             observation_scopes=observation_scopes,
+            pending_refresh_tags=sorted(all_refresh_tags) or None,
         )
 
     # Build summary
@@ -1699,11 +1793,19 @@ def _fmt(key: str) -> str:
     if timing_parts:
         perf.log(f"[4] Timing breakdown: {', '.join(timing_parts)}")
 
-    # Trigger mental model refreshes only on the final round (when all memories are processed).
-    # If we hit the round limit and re-queued, skip MM refresh — the next round will handle it.
+    # Trigger mental-model refreshes once, when the chain has fully drained. On a
+    # round-limited round we skip and carry the affected tags forward (above); the
+    # final round flushes the accumulated union, so a model whose memories were
+    # consolidated in ANY round is refreshed exactly once — deduplicated, not dropped
+    # (#3411). Each model is still refreshed at most once per drain: a strict tagged
+    # model appears once in the trigger's candidate query regardless of how many rounds
+    # its tag spanned.
     if hit_round_limit:
         stats["mental_models_refreshed"] = 0
-        logger.info(f"[CONSOLIDATION] bank={bank_id} skipping mental model refresh (round limit hit, re-queued)")
+        logger.info(
+            f"[CONSOLIDATION] bank={bank_id} deferring mental model refresh to the final round "
+            f"(round limit hit; carrying {len(all_refresh_tags)} tags forward)"
+        )
     else:
         set_stage("consolidation.refreshing_mental_models")
         await memory_engine._write_operation_progress(
@@ -1717,7 +1819,7 @@ def _fmt(key: str) -> str:
             memory_engine=memory_engine,
             bank_id=bank_id,
             request_context=request_context,
-            consolidated_tags=list(consolidated_tags) if consolidated_tags else None,
+            consolidated_tags=sorted(all_refresh_tags) or None,
             perf=perf,
         )
         stats["mental_models_refreshed"] = mental_models_refreshed
@@ -1827,10 +1929,14 @@ async def _trigger_mental_model_refreshes(
     for row in rows:
         mental_model_id = row["id"]
         try:
+            # skip_if_in_flight: a consolidation chain fires this every round and
+            # overlapping consolidations can run on the same bank, so a model still
+            # pending/processing a refresh must not be enqueued a second time (#3411).
             await memory_engine.submit_async_refresh_mental_model(
                 bank_id=bank_id,
                 mental_model_id=mental_model_id,
                 request_context=request_context,
+                skip_if_in_flight=True,
             )
             refreshed_count += 1
             logger.info(

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

@@ -2448,6 +2448,7 @@ async def _handle_consolidation(self, task_dict: dict[str, Any]):
             request_context=internal_context,
             operation_id=task_dict.get("operation_id"),
             observation_scopes=task_dict.get("observation_scopes"),
+            pending_refresh_tags=task_dict.get("pending_refresh_tags"),
         )
 
         logger.info(f"[CONSOLIDATION] bank={bank_id} completed: {result.get('memories_processed', 0)} processed")
@@ -15920,6 +15921,30 @@ async def _submit_async_operation(
                         row_payload = row["task_payload"]
                         row_dict = json.loads(row_payload) if isinstance(row_payload, str) else (row_payload or {})
                         if row_dict.get("observation_scopes") is None:
+                            # A round-limit consolidation chain carries its accumulated
+                            # mental-model refresh tags in ``pending_refresh_tags`` (#3411).
+                            # If its re-queue is deduped into an unrelated pending
+                            # consolidation (e.g. one a retain enqueued mid-drain), fold
+                            # those tags into the surviving op so its final round still
+                            # refreshes every affected model — otherwise the accumulated
+                            # set is silently lost and the models stay stale. Safe to
+                            # UPDATE here: dedupe_by_bank holds FOR NO KEY UPDATE on the
+                            # bank row, serialising concurrent submits for this bank.
+                            incoming_tags = full_payload.get("pending_refresh_tags")
+                            if incoming_tags:
+                                existing_tags = row_dict.get("pending_refresh_tags") or []
+                                merged_tags = sorted(set(existing_tags) | set(incoming_tags))
+                                if merged_tags != existing_tags:
+                                    row_dict["pending_refresh_tags"] = merged_tags
+                                    await conn.execute(
+                                        f"""
+                                        UPDATE {fq_table("async_operations")}
+                                        SET task_payload = $1::jsonb, updated_at = now()
+                                        WHERE operation_id = $2
+                                        """,
+                                        json.dumps(row_dict, default=_json_default),
+                                        row["operation_id"],
+                                    )
                             logger.debug(
                                 f"{operation_type} task already pending for bank_id={bank_id}, "
                                 f"skipping duplicate (existing operation_id={row['operation_id']})"
@@ -16418,6 +16443,7 @@ async def submit_async_consolidation(
         *,
         request_context: "RequestContext",
         observation_scopes: list[list[str]] | None = None,
+        pending_refresh_tags: list[str] | None = None,
     ) -> dict[str, Any]:
         """Submit a consolidation operation to run asynchronously.
 
@@ -16429,6 +16455,9 @@ async def submit_async_consolidation(
             request_context: Request context for authentication
             observation_scopes: Optional list of tag scopes to consolidate. When provided,
                 only unconsolidated memories matching at least one scope are processed.
+            pending_refresh_tags: Set by the round-limit re-queue only — the union of tags
+                consolidated by earlier rounds of this chain, so the final round refreshes
+                every affected mental model exactly once (#3411). Not a caller-facing knob.
 
         Returns:
             Dict with operation_id
@@ -16454,6 +16483,8 @@ async def submit_async_consolidation(
             task_payload["_api_key_id"] = request_context.api_key_id
         if observation_scopes is not None:
             task_payload["observation_scopes"] = observation_scopes
+        if pending_refresh_tags is not None:
+            task_payload["pending_refresh_tags"] = pending_refresh_tags
 
         # Skip bank-level deduplication when scoped — the caller wants a
         # targeted run that should not be merged into a pending full-bank sweep.