fix(retrieval): bound map-nav snapshot memory and SQL loads

apps/api/alembic/versions/0c1d2e3f4a5b_add_chunk_revision_section_order_index.py

@@ -0,0 +1,43 @@
+"""Add the index used by lazy map-nav section loads."""
+
+from __future__ import annotations
+
+from alembic import op
+
+
+revision = "0c1d2e3f4a5b"
+down_revision = "fbf0c1d2e3f4"
+branch_labels = None
+depends_on = None
+
+_INDEX_NAME = "idx_document_chunks_revision_section_order"
+
+
+def upgrade() -> None:
+    external_transaction = bool(
+        op.get_context().opts.get("knowhere_external_transaction", False)
+    )
+    if external_transaction:
+        op.execute(
+            f"CREATE INDEX IF NOT EXISTS {_INDEX_NAME} "
+            "ON document_chunks "
+            "(document_id, job_result_id, section_id, sort_order, chunk_id, id)"
+        )
+        return
+    with op.get_context().autocommit_block():
+        op.execute(
+            f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {_INDEX_NAME} "
+            "ON document_chunks "
+            "(document_id, job_result_id, section_id, sort_order, chunk_id, id)"
+        )
+
+
+def downgrade() -> None:
+    external_transaction = bool(
+        op.get_context().opts.get("knowhere_external_transaction", False)
+    )
+    if external_transaction:
+        op.execute(f"DROP INDEX IF EXISTS {_INDEX_NAME}")
+        return
+    with op.get_context().autocommit_block():
+        op.execute(f"DROP INDEX CONCURRENTLY IF EXISTS {_INDEX_NAME}")

apps/api/tests/contract/test_retrieval_lazy_snapshot_quality_contract.py

@@ -0,0 +1,176 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from collections.abc import Sequence
+from typing import Any
+
+from shared.services.retrieval.nav.nav_hierarchy import ProviderToolSpace
+from shared.services.retrieval.nav.nav_knowhere import (
+    KnowhereProvider,
+    LazyKnowhereProvider,
+    NamespaceKnowhereProvider,
+    SectionRow,
+    UnitRow,
+)
+from shared.services.retrieval.nav.nav_map_scores import (
+    build_score_units,
+    compute_corpus_map_and_unit_scores,
+)
+from shared.services.retrieval.nav.knowhere_hybrid import (
+    ScoreUnitRow,
+    score_rows_hybrid_all,
+    score_unit_stream_hybrid_all,
+)
+
+
+@dataclass
+class _FakeChunkStore:
+    units_by_section: dict[str, list[UnitRow]]
+
+    def load_section_units(
+        self,
+        document_id: str,
+        section_id: str,
+        extra_chunk_ids: Sequence[str] = (),
+    ) -> list[UnitRow]:
+        del document_id
+        units = list(self.units_by_section.get(section_id, ()))
+        known = {unit.chunk_id for unit in units}
+        for units_in_section in self.units_by_section.values():
+            for unit in units_in_section:
+                if unit.chunk_id in extra_chunk_ids and unit.chunk_id not in known:
+                    units.append(unit)
+        return units
+
+    def close(self) -> None:
+        return None
+
+
+def _providers() -> tuple[ProviderToolSpace, ProviderToolSpace]:
+    sections = [
+        SectionRow("root", None, "Root", "Root", 0, "", 0),
+        SectionRow("section", "root", "Root / Section", "Section", 1, "", 1),
+        SectionRow("leaf", "section", "Root / Section / Leaf", "Leaf", 2, "", 2),
+    ]
+    text = UnitRow(
+        "duplicate-chunk",
+        "leaf",
+        "text",
+        "alpha retrieval evidence",
+        1,
+        metadata={"connect_to": [{"target": "asset-1", "relation": "embeds"}]},
+    )
+    asset = UnitRow(
+        "asset-1",
+        "root",
+        "image",
+        "",
+        2,
+        file_path="images/asset.png",
+        metadata={"summary": "supporting image"},
+    )
+    eager = NamespaceKnowhereProvider(
+        [KnowhereProvider(doc_id="doc", sections=sections, units=[text, asset])],
+        titles={"doc": "document"},
+    )
+    store = _FakeChunkStore({"root": [asset], "leaf": [text]})
+    lazy = NamespaceKnowhereProvider(
+        [
+            LazyKnowhereProvider(
+                doc_id="doc",
+                sections=sections,
+                chunk_store=store,
+                known_chunk_ids=[text.chunk_id, asset.chunk_id],
+                root_asset_ids=[asset.chunk_id],
+                remounted_assets_by_section={"leaf": [asset.chunk_id]},
+            )
+        ],
+        titles={"doc": "document"},
+        chunk_owner_by_id={"duplicate-chunk": "doc", "asset-1": "doc"},
+    )
+    return ProviderToolSpace(eager), ProviderToolSpace(lazy)
+
+
+def test_lazy_provider_preserves_score_units_and_scores() -> None:
+    eager, lazy = _providers()
+
+    assert build_score_units(eager, "doc") == build_score_units(lazy, "doc")
+    assert compute_corpus_map_and_unit_scores(
+        eager, doc_ids=["doc"], query="alpha retrieval"
+    ) == compute_corpus_map_and_unit_scores(
+        lazy, doc_ids=["doc"], query="alpha retrieval"
+    )
+
+    lazy_provider = lazy._provider
+    self_units = getattr(lazy_provider, "self_units")
+    assert [unit.chunk_id for unit in self_units("leaf")] == [
+        "duplicate-chunk",
+        "asset-1",
+    ]
+
+
+def test_streaming_scorer_preserves_exact_eager_scores() -> None:
+    rows: list[ScoreUnitRow] = [
+        {
+            "chunk_id": "unit-a",
+            "path_search_text": "root alpha",
+            "content_search_text": "alpha alpha evidence",
+            "term_search_text": "alpha alpha evidence root",
+        },
+        {
+            "chunk_id": "unit-b",
+            "path_search_text": "root beta",
+            "content_search_text": "beta evidence",
+            "term_search_text": "beta evidence root",
+        },
+        {
+            "chunk_id": "unit-c",
+            "path_search_text": "root common",
+            "content_search_text": "common evidence",
+            "term_search_text": "common evidence root",
+        },
+    ]
+    eager_rows: list[dict[str, Any]] = [dict(row) for row in rows]
+    eager_scores = {
+        str(row["chunk_id"]): float(row["score"])
+        for row in score_rows_hybrid_all(eager_rows, "alpha evidence")
+    }
+    replay_count: int = 0
+
+    def unit_factory() -> Sequence[ScoreUnitRow]:
+        nonlocal replay_count
+        replay_count += 1
+        return rows
+
+    assert score_unit_stream_hybrid_all(unit_factory, "alpha evidence") == eager_scores
+    assert replay_count == 1
+
+
+def test_streaming_scorer_preserves_duplicate_id_eager_semantics() -> None:
+    rows: list[ScoreUnitRow] = [
+        {
+            "chunk_id": "duplicate",
+            "path_search_text": "alpha",
+            "content_search_text": "alpha",
+            "term_search_text": "alpha",
+        },
+        {
+            "chunk_id": "duplicate",
+            "path_search_text": "beta",
+            "content_search_text": "beta",
+            "term_search_text": "beta",
+        },
+        {
+            "chunk_id": "other",
+            "path_search_text": "alpha beta",
+            "content_search_text": "alpha beta",
+            "term_search_text": "alpha beta",
+        },
+    ]
+    eager_rows: list[dict[str, Any]] = [dict(row) for row in rows]
+    eager_scores = {
+        str(row["chunk_id"]): float(row["score"])
+        for row in score_rows_hybrid_all(eager_rows, "alpha beta")
+    }
+
+    assert score_unit_stream_hybrid_all(lambda: rows, "alpha beta") == eager_scores

apps/api/tests/migrations/test_schema_contract.py

@@ -233,6 +233,28 @@ def test_should_index_document_chunks_in_snapshot_pagination_order(
     )
 
 
+def test_should_index_document_chunks_in_lazy_section_order(
+    migrated_head_engine: Engine,
+) -> None:
+    with migrated_head_engine.begin() as connection:
+        index_definition = connection.execute(
+            text(
+                """
+                SELECT indexdef
+                FROM pg_indexes
+                WHERE schemaname = current_schema()
+                  AND tablename = 'document_chunks'
+                  AND indexname = 'idx_document_chunks_revision_section_order'
+                """
+            )
+        ).scalar_one()
+
+    assert (
+        "(document_id, job_result_id, section_id, sort_order, chunk_id, id)"
+        in str(index_definition)
+    )
+
+
 def test_should_upgrade_with_a_caller_owned_connection(
     alembic_engine: Engine,
 ) -> None:

packages/shared-python/shared/models/database/document.py

@@ -210,6 +210,15 @@ class DocumentChunk(Base):
             "chunk_id",
             "id",
         ),
+        Index(
+            "idx_document_chunks_revision_section_order",
+            "document_id",
+            "job_result_id",
+            "section_id",
+            "sort_order",
+            "chunk_id",
+            "id",
+        ),
         Index("idx_document_chunks_section", "section_id"),
         Index(
             "idx_chunk_content_search_tsv",

packages/shared-python/shared/services/retrieval/execution/routes.py

@@ -186,6 +186,7 @@ async def _run_mapnav_route(
         namespace=context.namespace,
         exclude_document_ids=context.exclude_document_ids,
         exclude_sections=context.exclude_sections,
+        lazy=True,
     )
 
     # Small-corpus count / snapshot reads may leave a checkout; drop it before
@@ -196,19 +197,22 @@ async def _run_mapnav_route(
     cfg = build_nav_config()
     toolspace = ProviderToolSpace(snapshot.provider)
 
-    episode = await asyncio.to_thread(
-        run_nav_episode,
-        None,
-        context.query,
-        corpus_doc_ids=list(snapshot.document_ids),
-        budget_chars=budget,
-        compose_answer=False,
-        policy="llm",
-        config=cfg,
-        toolspace=toolspace,
-    )
+    try:
+        episode = await asyncio.to_thread(
+            run_nav_episode,
+            None,
+            context.query,
+            corpus_doc_ids=list(snapshot.document_ids),
+            budget_chars=budget,
+            compose_answer=False,
+            policy="llm",
+            config=cfg,
+            toolspace=toolspace,
+        )
 
-    refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot)
+        refs, score_by_chunk_id = build_referenced_chunks(episode, snapshot)
+    finally:
+        snapshot.close()
 
     async with open_fresh_database_context() as final_db:
         resolved = await resolve_workflow_references(