perf: bound large retrieval snapshot reads

apps/api/tests/contract/test_retrieval_snapshot_large_corpus_contract.py

@@ -6,6 +6,7 @@
 from uuid import uuid4
 
 from httpx import AsyncClient
+import pytest
 
 from shared.models.database.document import Document, DocumentChunk, DocumentSection
 from shared.models.database.job_result import JobResult
@@ -15,6 +16,7 @@
     _REVISION_GROUP_SIZE,
     load_nav_snapshot,
 )
+import shared.services.retrieval.nav_snapshot as nav_snapshot_module
 from sqlalchemy import Executable, Result, select
 from sqlalchemy.engine import Row
 from sqlalchemy.sql.selectable import Select
@@ -252,10 +254,19 @@
     developer_api_client_factory: Callable[
         [], AbstractAsyncContextManager[AsyncClient]
     ],
+    monkeypatch: pytest.MonkeyPatch,
 ) -> None:
     namespace = f"large-corpus-{uuid4().hex[:8]}"
     async with developer_api_client_factory():
         await _seed_large_retrieval_corpus(namespace)
+        async def unexpected_manifest_load(*_args: object, **_kwargs: object) -> None:
+            raise AssertionError("large snapshots must use normalized retrieval rows")
+
+        monkeypatch.setattr(
+            nav_snapshot_module,
+            "_load_manifest_sections",
+            unexpected_manifest_load,
+        )
         legacy_rows = await _load_legacy_rows(namespace)
         async with contract_db_session() as db:
             counting_db = _CountingSession(db)

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

@@ -23,7 +23,6 @@
 )
 from shared.services.retrieval.search.ranking import rank_retrieval_candidates
 from shared.services.retrieval.search.scoped_corpus import (
-    count_manifest_chunks,
     count_scoped_chunks,
     load_all_scoped_chunks,
 )
@@ -58,27 +57,15 @@ async def _try_run_small_corpus_route(
     context: RetrievalRouteContext,
 ) -> RetrievalRouteOutcome | None:
     total_chunk_count: int | None = None
-    if (
-        context.use_agentic is not False
-        and context.revision_pins is not None
-        and not context.exclude_document_ids
-        and not context.exclude_sections
-        and context.allowed_chunk_types is None
-        and not context.signal_paths
-    ):
-        total_chunk_count = await count_manifest_chunks(
-            context.db,
-            revision_pins=context.revision_pins,
-        )
-    if total_chunk_count is None:
-        total_chunk_count = await count_scoped_chunks(
-            context.db,
-            user_id=context.user_id,
-            namespace=context.namespace,
-            exclude_document_ids=context.exclude_document_ids,
-            allowed_chunk_types=context.allowed_chunk_types,
-            revision_pins=context.revision_pins,
-        )
+    total_chunk_count = await count_scoped_chunks(
+        context.db,
+        user_id=context.user_id,
+        namespace=context.namespace,
+        exclude_document_ids=context.exclude_document_ids,
+        allowed_chunk_types=context.allowed_chunk_types,
+        revision_pins=context.revision_pins,
+        max_count=context.top_k + 1,
+    )
 
     logger.info(f"\n  Total chunks in scope: {total_chunk_count}")
     if total_chunk_count > context.top_k:

packages/shared-python/shared/services/retrieval/nav_snapshot.py

@@ -55,6 +55,12 @@
 # Keep revision predicates bounded while reducing round trips for large
 # namespaces. Keyset paging still caps each payload query at 10,000 rows.
 _REVISION_GROUP_SIZE = 64
+# Compressed manifests are efficient for small namespaces, but transferring a
+# large set of binary payloads can exceed the asyncpg statement timeout before
+# PostgreSQL has done meaningful work. For larger snapshots, the normalized
+# section/chunk index loaders below transfer only the fields map-nav needs and
+# page them by revision.
+_MANIFEST_MAX_REVISION_COUNT = 32
 _logger = logging.getLogger(__name__)
 
 
@@ -206,12 +212,14 @@ async def load_nav_snapshot(
         if job_result_id and job_id
     }
 
-    manifest_sections = await _load_manifest_sections(
-        db,
-        document_revisions=document_revisions,
-        exclude_sections=excluded_secs,
-        job_id_by_result_id=job_id_by_result_id,
-    )
+    manifest_sections = None
+    if len(document_revisions) <= _MANIFEST_MAX_REVISION_COUNT:
+        manifest_sections = await _load_manifest_sections(
+            db,
+            document_revisions=document_revisions,
+            exclude_sections=excluded_secs,
+            job_id_by_result_id=job_id_by_result_id,
+        )
     if manifest_sections is None:
         sections_by_doc, section_path_by_id = await _load_sections(
             db,
@@ -358,26 +366,31 @@ async def _load_manifest_sections(
         if len(manifest_entries) != len(document_revisions):
             return None
     else:
-        statement = select(
-            RetrievalServingRevisionManifest.document_id,
-            RetrievalServingRevisionManifest.job_result_id,
-            RetrievalServingRevisionManifest.payload_zlib,
-            RetrievalServingRevisionManifest.checksum,
-            RetrievalServingRevisionManifest.format_version,
-        ).where(
-            tuple_(
+        manifest_entries = []
+        for group_start in range(0, len(document_revisions), _REVISION_GROUP_SIZE):
+            revision_group = document_revisions[
+                group_start : group_start + _REVISION_GROUP_SIZE
+            ]
+            statement = select(
                 RetrievalServingRevisionManifest.document_id,
                 RetrievalServingRevisionManifest.job_result_id,
-            ).in_(document_revisions)
-        )
-        try:
-            rows = (await db.execute(statement)).all()
-        except Exception:
-            await db.rollback()
-            return None
-        if len(rows) != len(document_revisions):
-            return None
-        manifest_entries = [tuple(row) for row in rows]
+                RetrievalServingRevisionManifest.payload_zlib,
+                RetrievalServingRevisionManifest.checksum,
+                RetrievalServingRevisionManifest.format_version,
+            ).where(
+                tuple_(
+                    RetrievalServingRevisionManifest.document_id,
+                    RetrievalServingRevisionManifest.job_result_id,
+                ).in_(revision_group)
+            )
+            try:
+                rows = (await db.execute(statement)).all()
+            except Exception:
+                await db.rollback()
+                return None
+            if len(rows) != len(revision_group):
+                return None
+            manifest_entries.extend(tuple(row) for row in rows)
     by_doc: dict[str, list[SectionRow]] = {}
     path_by_id: dict[str, str] = {}
     ids_by_doc: dict[str, list[str]] = {}

packages/shared-python/shared/services/retrieval/search/scoped_corpus.py

@@ -77,10 +77,11 @@ async def count_scoped_chunks(
     exclude_document_ids: list[str],
     allowed_chunk_types: set[str] | None,
     revision_pins: Mapping[str, str] | None = None,
+    max_count: int | None = None,
 ) -> int:
     if revision_pins is None:
         stmt = (
-            select(func.count(DocumentChunk.id))
+            select(DocumentChunk.id)
             .join(
                 Document,
                 (Document.document_id == DocumentChunk.document_id)
@@ -92,7 +93,7 @@ async def count_scoped_chunks(
         )
     else:
         stmt = (
-            select(func.count(DocumentChunk.id))
+            select(DocumentChunk.id)
             .join(Document, Document.document_id == DocumentChunk.document_id)
             .where(Document.user_id == user_id)
             .where(Document.namespace == namespace)
@@ -106,7 +107,13 @@ async def count_scoped_chunks(
         stmt = stmt.where(Document.document_id.notin_(list(exclude_document_ids)))
     if allowed_chunk_types is not None:
         stmt = stmt.where(func.lower(DocumentChunk.chunk_type).in_(list(allowed_chunk_types)))
-    result = await db.execute(stmt)
+
+    if max_count is not None:
+        stmt = stmt.limit(max_count)
+
+    result = await db.execute(
+        select(func.count()).select_from(stmt.order_by(None).subquery())
+    )
     return result.scalar() or 0