fix(api): consolidation discovery probe, mental-model edit race, LlamaParse v2

hindsight-api-slim/hindsight_api/alembic/versions/6ab9509eeb56_banks_needing_consolidation_exists_probe.py

@@ -0,0 +1,142 @@
+"""Make ``banks_needing_consolidation()`` probe each bank instead of scanning facts.
+
+The ``c8b4e2a71f95`` body joined ``memory_units`` to ``banks`` and grouped by
+``bank_id``. When a large share of facts is unconsolidated (a backlog after
+consolidation fell behind), ``GROUP BY`` must visit every matching row, so the
+planner skips ``idx_memory_units_unconsolidated`` and seq-scans the whole table
+on every reconcile pass (#5425).
+
+The rewrite starts from ``banks`` and asks "is there at least one eligible
+fact?" with ``EXISTS``, which stops at the first index hit per bank. Results are
+identical: ``banks.bank_id`` is the primary key (one row per bank, like the
+``GROUP BY``), facts without a ``banks`` row are excluded by both forms, and the
+config / in-flight-op predicates are unchanged. The wrapper (lock_timeout, skip
+arms, install-run gating, ``execute_unless_owned``) is copied unchanged from
+``c8b4e2a71f95``.
+
+Downgrade is a no-op: same signature and same results as the previous body, so
+leaving it in place is harmless; downgrading past earlier migrations restores or
+drops the routine as they define.
+
+Oracle slot intentionally absent: this routine is PostgreSQL-only, and the Oracle
+worker keeps its per-schema sweep.
+
+Revision ID: 6ab9509eeb56
+Revises: c8e2f4a6b1d3
+Create Date: 2026-10-12
+"""
+
+from collections.abc import Sequence
+
+from alembic import context
+
+from hindsight_api.alembic._dialect import run_for_dialect
+from hindsight_api.alembic._owned import execute_unless_owned
+from hindsight_api.config import get_config
+
+revision: str = "6ab9509eeb56"
+down_revision: str | Sequence[str] | None = "c8e2f4a6b1d3"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+# Same values as c8b4e2a71f95 — see that migration for the reasoning.
+_LOCK_TIMEOUT = "250ms"
+
+_SKIP_ARMS = """
+                EXCEPTION
+                    -- Schema or its tables vanished between the pg_class
+                    -- snapshot and this query (tenant dropped or migrating).
+                    WHEN undefined_table OR invalid_schema_name OR undefined_column THEN
+                        CONTINUE;
+                    -- Schema is mid-DDL and holds (or has queued) an
+                    -- AccessExclusiveLock. Skip it rather than wait: waiting is
+                    -- what closes the deadlock cycle. deadlock_detected is
+                    -- belt-and-braces for a cycle formed before lock_timeout.
+                    WHEN lock_not_available OR deadlock_detected THEN
+                        CONTINUE;
+"""
+
+
+def _configured_schema() -> str:
+    """The one schema this deployment's routines live in and are called from."""
+    return get_config().database_schema or "public"
+
+
+def _target_schema() -> str | None:
+    return context.config.get_main_option("target_schema")
+
+
+def _is_install_run() -> bool:
+    """True for the single run that owns the routines (mirrors c8b4e2a71f95)."""
+    target = _target_schema()
+    return not target or target == _configured_schema()
+
+
+def _prefix(schema: str | None) -> str:
+    """Qualifier for ``schema``, or ``""`` to fall back to ``search_path``."""
+    return f'"{schema}".' if schema else ""
+
+
+def _pg_upgrade() -> None:
+    if not _is_install_run():
+        return
+    schema = _prefix(_target_schema())
+
+    execute_unless_owned(
+        "banks_needing_consolidation",
+        f"""
+        CREATE OR REPLACE FUNCTION {schema}banks_needing_consolidation()
+        RETURNS TABLE(schema_name text, bank_id text)
+        LANGUAGE plpgsql STABLE
+        AS $fn$
+        DECLARE
+            sch text;
+            prev_lock_timeout text;
+        BEGIN
+            prev_lock_timeout := current_setting('lock_timeout');
+            PERFORM set_config('lock_timeout', '{_LOCK_TIMEOUT}', true);
+            FOR sch IN
+                SELECT n.nspname
+                FROM pg_class c
+                JOIN pg_namespace n ON n.oid = c.relnamespace
+                WHERE c.relname = 'memory_units' AND c.relkind = 'r'
+            LOOP
+                BEGIN
+                    RETURN QUERY EXECUTE format($q$
+                        SELECT %1$L::text, b.bank_id
+                        FROM %1$I.banks b
+                        WHERE COALESCE(b.config -> 'enable_auto_consolidation', 'true'::jsonb) <> 'false'::jsonb
+                          AND NOT EXISTS (
+                              SELECT 1 FROM %1$I.async_operations o
+                              WHERE o.bank_id = b.bank_id
+                                AND o.operation_type = 'consolidation'
+                                AND o.status IN ('pending', 'processing')
+                          )
+                          AND EXISTS (
+                              SELECT 1 FROM %1$I.memory_units m
+                              WHERE m.bank_id = b.bank_id
+                                AND m.consolidated_at IS NULL
+                                AND m.consolidation_failed_at IS NULL
+                                AND m.fact_type IN ('experience', 'world')
+                          )
+                    $q$, sch);
+{_SKIP_ARMS}                END;
+            END LOOP;
+            PERFORM set_config('lock_timeout', prev_lock_timeout, true);
+        END;
+        $fn$;
+        """,
+    )
+
+
+def _pg_downgrade() -> None:
+    # No-op by design — see the module docstring.
+    return
+
+
+def upgrade() -> None:
+    run_for_dialect(pg=_pg_upgrade)
+
+
+def downgrade() -> None:
+    run_for_dialect(pg=_pg_downgrade)

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

@@ -1855,6 +1855,13 @@ def _resolve_refresh_tag_filtering(
     return RefreshTagFiltering(tags=None, tags_match="any", tag_groups=groups)
 
 
+def _parse_reflect_response(raw: Any) -> dict[str, Any] | None:
+    """A stored ``reflect_response`` as a dict: drivers hand it back as JSON text or decoded."""
+    if isinstance(raw, str):
+        return json.loads(raw) if raw else None
+    return raw
+
+
 def _scope_mental_model_trigger(
     trigger: dict[str, Any] | None,
     tag_scope: list[TagGroup] | None,
@@ -18597,6 +18604,7 @@ async def update_mental_model(
         # The vector as well as its literal: a store-owned index takes the floats, and asking for
         # the string alone would mean embedding the same document twice on every page write.
         new_embedding_vec: list[float] | None = None
+        current_row: Any = None
         document_changed = name is not None or content is not None
         if document_changed:
             async with use_or_acquire(backend, conn) as read_conn:
@@ -18612,20 +18620,16 @@ async def update_mental_model(
             if content is not None:
                 # Snapshot for history, taken from the same read.
                 previous_content = current_row["content"]
-                raw_rr = current_row["reflect_response"]
-                if isinstance(raw_rr, str):
-                    previous_reflect_response = json.loads(raw_rr) if raw_rr else None
-                else:
-                    previous_reflect_response = raw_rr
+                previous_reflect_response = _parse_reflect_response(current_row["reflect_response"])
 
             new_embedding_vec = await self._mental_model_embedding_vector(
                 name if name is not None else (current_row["name"] or ""),
                 content if content is not None else (current_row["content"] or ""),
             )
             new_embedding_str = str(new_embedding_vec) if new_embedding_vec else None
 
-        # The exit stack carries the row lock a trigger patch takes below; it stays
-        # empty (and free) on every other update.
+        # The exit stack carries the row lock a trigger patch or document edit takes
+        # below; it stays empty (and free) on every other update.
         async with use_or_acquire(backend, conn) as conn, AsyncExitStack() as stack:
             # A supplied trigger PATCHES the stored one (see _merge_trigger): callers
             # send only the fields they set, so the rest have to be read back rather
@@ -18634,19 +18638,36 @@ async def update_mental_model(
             # Read-modify-write, so the row is locked for the rest of the statement:
             # two concurrent patches (an agent setting a cron while another flips
             # fact_types) would otherwise both merge onto the same pre-image and the
-            # later write would drop the earlier one. Only this path takes the lock —
-            # a refresh passes no trigger and stays lock-free.
-            if trigger is not None:
+            # later write would drop the earlier one.
+            #
+            # A document edit takes the same lock for the same reason: the embedding above was
+            # computed from a pre-image read off the lock, and a concurrent edit may have moved
+            # the other half since (#5297). Re-checked under the lock, and re-embedded only in
+            # that rare race, so a slow embedder still never pins a connection in the common case.
+            if trigger is not None or document_changed:
                 await stack.enter_async_context(conn.transaction())
-                trigger_row = await conn.fetchrow(
-                    f"SELECT trigger, tags FROM {fq_table('mental_models')} WHERE bank_id = $1 AND id = $2 FOR UPDATE",
+                locked_row = await conn.fetchrow(
+                    f"SELECT trigger, tags, name, content, reflect_response FROM {fq_table('mental_models')} "
+                    "WHERE bank_id = $1 AND id = $2 FOR UPDATE",
                     bank_id,
                     mental_model_id,
                 )
-                if trigger_row is None:
+                if locked_row is None:
                     return None
-                trigger = self._merge_trigger(trigger, base=self._stored_trigger(trigger_row["trigger"]))
-                trigger = _scope_mental_model_trigger(trigger, tag_scope)
+                if trigger is not None:
+                    trigger = self._merge_trigger(trigger, base=self._stored_trigger(locked_row["trigger"]))
+                    trigger = _scope_mental_model_trigger(trigger, tag_scope)
+                if document_changed and (
+                    locked_row["name"] != current_row["name"] or locked_row["content"] != current_row["content"]
+                ):
+                    if content is not None:
+                        previous_content = locked_row["content"]
+                        previous_reflect_response = _parse_reflect_response(locked_row["reflect_response"])
+                    new_embedding_vec = await self._mental_model_embedding_vector(
+                        name if name is not None else (locked_row["name"] or ""),
+                        content if content is not None else (locked_row["content"] or ""),
+                    )
+                    new_embedding_str = str(new_embedding_vec) if new_embedding_vec else None
 
             # Build dynamic update
             updates = []

hindsight-api-slim/hindsight_api/engine/parsers/llama_parse.py

@@ -8,13 +8,17 @@
 from typing import Any
 
 import aiohttp
+from pydantic import BaseModel
 
 from ..aiohttp_session import per_phase_timeout, proxy_env_is_set
 from .base import FileParser, UnsupportedFileTypeError
 
 logger = logging.getLogger(__name__)
 
-_LLAMA_PARSE_BASE_URL = "https://api.cloud.llamaindex.ai/api/parsing"
+_LLAMA_PARSE_BASE_URL = "https://api.cloud.llamaindex.ai/api/v2/parse"
+# v2 requires an explicit tier and version (v1 had a default mode). cost_effective
+# is the closest to v1's default parse at a comparable per-page cost.
+_PARSE_CONFIGURATION = json.dumps({"tier": "cost_effective", "version": "latest"})
 _DEFAULT_POLL_INTERVAL = 2.0  # seconds
 _DEFAULT_TIMEOUT = 300.0  # seconds
 
@@ -23,6 +27,18 @@
 _UNSUPPORTED_FILE_STATUS_CODES = {400, 415, 422}
 
 
+class _ParseJob(BaseModel):
+    status: str | None = None
+    error_message: str | None = None
+
+
+class _ParseResult(BaseModel):
+    """``GET /api/v2/parse/{id}``: ``job`` always, ``markdown_full`` when expanded."""
+
+    job: _ParseJob = _ParseJob()
+    markdown_full: str | None = None
+
+
 class LlamaParseParser(FileParser):
     """
     LlamaParse file parser using LlamaIndex's hosted parsing service.
@@ -73,21 +89,22 @@ async def _convert(self, session: aiohttp.ClientSession, file_data: bytes, filen
         form = aiohttp.FormData()
         # Ensure file_data is plain bytes (storage backends may return obstore.Bytes)
         form.add_field("file", bytes(file_data), filename=filename, content_type=content_type)
+        form.add_field("configuration", _PARSE_CONFIGURATION)
         async with session.post(f"{_LLAMA_PARSE_BASE_URL}/upload", headers=self._auth_headers, data=form) as resp:
             upload_data = await _json_or_raise(resp, filename, "upload")
         job_id: str = upload_data["id"]
 
-        # Step 2: Poll job status until SUCCESS or ERROR
+        # Step 2: Poll job status until COMPLETED, FAILED or CANCELLED
         deadline = time.monotonic() + self._timeout
         while True:
-            async with session.get(f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}", headers=self._auth_headers) as resp:
-                status_data = await _json_or_raise(resp, filename, "poll job status")
-            status = status_data.get("status")
+            async with session.get(f"{_LLAMA_PARSE_BASE_URL}/{job_id}", headers=self._auth_headers) as resp:
+                job = _ParseResult.model_validate(await _json_or_raise(resp, filename, "poll job status")).job
+            status = job.status
 
-            if status == "SUCCESS":
+            if status == "COMPLETED":
                 break
-            if status in ("ERROR", "CANCELLED"):
-                error = status_data.get("error_code") or status_data.get("error") or "unknown error"
+            if status in ("FAILED", "CANCELLED"):
+                error = job.error_message or status.lower()
                 raise RuntimeError(f"LlamaParse job failed for '{filename}': {error}")
 
             if time.monotonic() >= deadline:
@@ -97,10 +114,10 @@ async def _convert(self, session: aiohttp.ClientSession, file_data: bytes, filen
 
         # Step 3: Fetch the markdown result
         async with session.get(
-            f"{_LLAMA_PARSE_BASE_URL}/job/{job_id}/result/markdown", headers=self._auth_headers
+            f"{_LLAMA_PARSE_BASE_URL}/{job_id}", headers=self._auth_headers, params={"expand": "markdown_full"}
         ) as resp:
-            result_data = await _json_or_raise(resp, filename, "fetch markdown result")
-        markdown = result_data.get("markdown")
+            result = _ParseResult.model_validate(await _json_or_raise(resp, filename, "fetch markdown result"))
+        markdown = result.markdown_full
         if not markdown:
             raise RuntimeError(f"No content extracted from '{filename}'")
         return markdown