fix(core): reuse retrieval for stable rerank pagination

pyproject.toml

@@ -192,6 +192,10 @@ style = "pep440"
 bump = true
 fallback-version = "0.0.0"
 
+[tool.ty.environment]
+# Shared test helpers use the same import root as pytest's pythonpath above.
+extra-paths = ["tests"]
+
 [tool.ty.rules]
 all = "error"
 

src/basic_memory/repository/semantic_vector_index.py

@@ -53,6 +53,9 @@ class VectorMatch:
 
     key: VectorKey
     similarity: float
+    # Zero-based position before adapter-side filtering, when filtering consumes
+    # top-k slots (sqlite-vec). Other adapters use their returned match order.
+    candidate_rank: int | None = None
 
 
 @runtime_checkable

src/basic_memory/repository/sqlite_vec_index.py

@@ -327,25 +327,34 @@ async def search(
         query: Sequence[float],
         *,
         limit: int,
+        stable_prefix: bool = False,
     ) -> list[VectorMatch]:
         if not query or limit <= 0:
             return []
         validate_query_dimensions(self.scope, query)
         await self.initialize()
         vector_k = min(limit, SQLITE_VEC_MAX_K)
+        # vec0's equal-distance membership changes with k. Reranking needs one
+        # repeatable universe to reconstruct smaller prefixes without a second query.
+        # Reuse the adapter's existing ceiling; only the requested slots survive SQL.
+        query_k = str(SQLITE_VEC_MAX_K) if stable_prefix else ":vector_k"
         async with db.scoped_session(self._session_maker) as session:
             await self._ensure_loaded(session)
             result = await session.execute(
                 text(
                     "WITH vector_matches AS MATERIALIZED ("
                     " SELECT rowid, distance, source_hash FROM search_vector_embeddings "
-                    " WHERE embedding MATCH :query AND k = :vector_k"
-                    ") "
-                    "SELECT c.entity_id, c.chunk_key, vector_matches.distance "
-                    "FROM vector_matches "
+                    f" WHERE embedding MATCH :query AND k = {query_k}"
+                    "), ranked_matches AS MATERIALIZED ("
+                    " SELECT *, ROW_NUMBER() OVER (ORDER BY distance, rowid) - 1 AS candidate_rank"
+                    " FROM vector_matches) "
+                    "SELECT c.entity_id, c.chunk_key, vector_matches.distance, "
+                    "vector_matches.candidate_rank "
+                    "FROM ranked_matches AS vector_matches "
                     "JOIN search_vector_chunks c ON c.id = vector_matches.rowid "
                     "AND c.source_hash = vector_matches.source_hash "
                     "WHERE c.project_id = :project_id "
+                    "AND vector_matches.candidate_rank < :vector_k "
                     "AND c.vector_index = 'sqlite-vec' "
                     "AND c.embedding_status = 'ready' "
                     "AND c.embedding_model = :embedding_identity "
@@ -366,6 +375,7 @@ async def search(
                     entity_id=int(row["entity_id"]),
                     chunk_key=str(row["chunk_key"]),
                 ),
+                candidate_rank=int(row["candidate_rank"]),
                 similarity=max(
                     0.0,
                     min(1.0, 1.0 - (float(row["distance"]) ** 2) / 2.0),

test-int/semantic/test_semantic_coverage.py

@@ -211,7 +211,7 @@ async def record_vector_query(
 
     assert growing_prefix_results
     assert reranker.calls == 2
-    assert candidate_limits == [90, 80]
+    assert candidate_limits == [90]
 
     candidate_limits.clear()
     large_page_with_probe = await search_service.search(
@@ -226,7 +226,7 @@ async def record_vector_query(
 
     assert len(large_page_with_probe) == 101
     assert reranker.calls == 3
-    assert candidate_limits == [890, 80]
+    assert candidate_limits == [890]
 
     # A larger retrieval window must extend the same sequence rather than
     # reordering rows already exposed by an earlier deep page.
@@ -243,7 +243,7 @@ async def record_vector_query(
     )
 
     assert len(stable_window) == 40
-    assert candidate_limits == [280, 80]
+    assert candidate_limits == [280]
 
     candidate_limits.clear()
     third_page = await search_service.search(
@@ -256,7 +256,7 @@ async def record_vector_query(
         limit=10,
         offset=20,
     )
-    assert candidate_limits == [180, 80]
+    assert candidate_limits == [180]
 
     candidate_limits.clear()
     fourth_page = await search_service.search(
@@ -269,7 +269,7 @@ async def record_vector_query(
         limit=10,
         offset=30,
     )
-    assert candidate_limits == [280, 80]
+    assert candidate_limits == [280]
 
     assert [row.permalink for row in third_page] == [row.permalink for row in stable_window[20:30]]
     assert [row.permalink for row in fourth_page] == [row.permalink for row in stable_window[30:40]]

test-int/test_multi_project_search.py

@@ -193,9 +193,11 @@ def record(_conn, _cursor, statement, _params, _context, _many):
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("mode", list(SearchRetrievalMode))
+@pytest.mark.parametrize("reranker_enabled", [False, True])
 async def test_global_pagination_is_independent_of_page_size(
-    corpus: Corpus, mode: SearchRetrievalMode
+    corpus: Corpus, mode: SearchRetrievalMode, reranker_enabled: bool
 ) -> None:
+    corpus.config.reranker_enabled = reranker_enabled
     repo = corpus.repository([project.id for project in corpus.projects[:2]])
     query = SearchService.prepare_query(SearchQuery(text="nebula", retrieval_mode=mode))
     assert query is not None
@@ -209,6 +211,8 @@ async def test_global_pagination_is_independent_of_page_size(
         (r.project_id, r.type, r.id, r.score) for r in complete
     ]
     assert await repo.search(query, offset=100) == []
+    # This database-scoped reader never invokes the single-project rerank pipeline.
+    assert corpus.provider.query_calls == (5 if mode != SearchRetrievalMode.FTS else 0)
 
 
 @pytest.mark.asyncio

test-int/test_stable_rerank_pagination.py

@@ -0,0 +1,78 @@
+"""Real API/MCP pagination over the same single-pass semantic repository."""
+
+import json
+
+import pytest
+from fastapi import FastAPI
+from fastmcp import Client, FastMCP
+from httpx import AsyncClient
+from sqlalchemy import func, update
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
+
+from basic_memory import db
+
+from basic_memory.deps.services import get_search_service_v2_external
+from basic_memory.models import Project
+from basic_memory.services.search_service import SearchService
+from semantic_search_helpers import (
+    BackendSearchRepository,
+    _FakeReranker,
+    rerank_search_repository as rerank_search_repository,
+    pagination_repository as pagination_repository,
+)
+
+
+@pytest.fixture
+def session_maker(
+    engine_factory: tuple[AsyncEngine, async_sessionmaker[AsyncSession]],
+) -> async_sessionmaker[AsyncSession]:
+    return engine_factory[1]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("mode", ["vector", "hybrid"])
+async def test_api_and_mcp_keep_probe_and_page_boundaries(
+    pagination_repository: BackendSearchRepository,
+    search_service: SearchService,
+    app: FastAPI,
+    client: AsyncClient,
+    test_project: Project,
+    mcp_server: FastMCP,
+    mode: str,
+) -> None:
+    v2_project_url = f"/v2/projects/{test_project.external_id}"
+    repo = pagination_repository
+    repo._rerank_provider = _FakeReranker({"Note 00": 0.8, "Note 01": 0.9})
+    async with db.scoped_session(repo.session_maker) as session:
+        await session.execute(
+            update(Project).where(Project.id == test_project.id).values(last_indexed_at=func.now())
+        )
+        await session.commit()
+    search_service.repository = repo
+    app.dependency_overrides[get_search_service_v2_external] = lambda: search_service
+    query = {"text": "auth session token", "retrieval_mode": mode, "min_similarity": 0.5}
+    response = await client.request(
+        "QUERY", f"{v2_project_url}/search/", json=query, params={"page_size": 100}
+    )
+    assert response.status_code == 200, response.text
+    expected = response.json()["results"]
+    pages = []
+    async with Client(mcp_server) as mcp:
+        for page in range(1, (len(expected) + 4) // 5 + 2):
+            result = await mcp.call_tool(
+                "search_notes",
+                {
+                    "project": test_project.name,
+                    "query": "auth session token",
+                    "search_type": mode,
+                    "min_similarity": 0.5,
+                    "page": page,
+                    "page_size": 5,
+                    "output_format": "json",
+                },
+            )
+            payload = json.loads(result.content[0].text)
+            assert payload["has_more"] == (page * 5 < len(expected))
+            assert payload["total_is_exact"] is False
+            pages.extend(payload["results"])
+    assert [row["permalink"] for row in pages] == [row["permalink"] for row in expected]

tests/repository/test_rerank_pipeline.py

@@ -29,66 +29,14 @@
 from basic_memory.repository.sqlite_search_repository import SQLiteSearchRepository
 from basic_memory.schemas.search import SearchItemType, SearchRetrievalMode
 from basic_memory.services.entity_service import EntityService
-
-type BackendSearchRepository = SQLiteSearchRepository | PostgresSearchRepository
-
-
-class _StubEmbeddingProvider:
-    """Deterministic embeddings that give the two auth notes DIFFERENT similarity.
-
-    An "auth" doc containing "deep" is tilted slightly off the query axis, so vector
-    retrieval ranks the plain-auth note strictly above it. That makes the pre-rerank
-    baseline a real ordering (not a tie), so a rerank that promotes the lower note is
-    a genuine "recover the below-cutoff doc" scenario (#950), not a coin flip.
-    """
-
-    model_name = "stub"
-    dimensions = 4
-
-    async def embed_query(self, text: str) -> list[float]:
-        return self._vectorize(text)
-
-    async def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        return [self._vectorize(t) for t in texts]
-
-    def runtime_log_attrs(self) -> dict[str, Any]:
-        return {}
-
-    @staticmethod
-    def _vectorize(text: str) -> list[float]:
-        lowered = text.lower()
-        if "auth" not in lowered:
-            return [0.0, 0.0, 0.0, 1.0]
-        # Unit vectors; cos with the query axis [1,0,0,0] is 1.0 vs 0.9.
-        if "deep" in lowered:
-            return [0.9, 0.4358898943540674, 0.0, 0.0]
-        return [1.0, 0.0, 0.0, 0.0]
-
-
-class _FakeReranker:
-    """Scores a document by the marker substring it contains; records call count."""
-
-    model_name = "fake-reranker"
-
-    def __init__(self, score_by_marker: dict[str, float]):
-        self.score_by_marker = score_by_marker
-        self.calls = 0
-        self.document_batches: list[list[str]] = []
-
-    async def rerank(self, query: str, documents: list[str]) -> list[float]:
-        self.calls += 1
-        self.document_batches.append(documents)
-        scores = []
-        for doc in documents:
-            score = 0.0
-            for marker, value in self.score_by_marker.items():
-                if marker in doc:
-                    score = value
-            scores.append(score)
-        return scores
-
-    def runtime_log_attrs(self) -> dict[str, Any]:
-        return {}
+from semantic_search_helpers import (
+    BackendSearchRepository,
+    _StubEmbeddingProvider,
+    _FakeReranker,
+    _entity_row,
+    _semantic_search_repository,
+    rerank_search_repository as rerank_search_repository,
+)
 
 
 class _BadReranker:
@@ -156,24 +104,6 @@ def test_validate_rerank_scores_rejects_non_numeric_value():
         validate_rerank_scores(["not-a-score"], expected_count=1)
 
 
-def _entity_row(*, project_id: int, row_id: int, title: str, permalink: str, content: str):
-    now = datetime.now(timezone.utc)
-    return SearchIndexRow(
-        project_id=project_id,
-        id=row_id,
-        type=SearchItemType.ENTITY.value,
-        title=title,
-        permalink=permalink,
-        file_path=f"{permalink}.md",
-        metadata={"note_type": "spec"},
-        entity_id=row_id,
-        content_stems=content,
-        content_snippet=content,
-        created_at=now,
-        updated_at=now,
-    )
-
-
 def _row(**overrides) -> SearchIndexRow:
     now = datetime.now(timezone.utc)
     base: dict[str, Any] = dict(
@@ -508,41 +438,6 @@ async def test_rerank_paginate_surfaces_permanent_faults(exc):
 # --- End-to-end through both repository backends ---
 
 
-def _semantic_search_repository(
-    session_maker: Any,
-    project_id: int,
-    app_config: BasicMemoryConfig,
-    **config_updates: object,
-) -> BackendSearchRepository:
-    config = app_config.model_copy(
-        update={
-            "semantic_search_enabled": True,
-            "semantic_min_similarity": 0.0,
-            **config_updates,
-        }
-    )
-    repository_type = (
-        PostgresSearchRepository
-        if config.database_backend == DatabaseBackend.POSTGRES
-        else SQLiteSearchRepository
-    )
-    return repository_type(
-        session_maker,
-        project_id=project_id,
-        app_config=config,
-        embedding_provider=_StubEmbeddingProvider(),
-    )
-
-
-@pytest.fixture
-def rerank_search_repository(
-    session_maker: Any,
-    test_project: Any,
-    app_config: BasicMemoryConfig,
-) -> BackendSearchRepository:
-    return _semantic_search_repository(session_maker, test_project.id, app_config)
-
-
 async def _index_two_auth_notes(repo: BackendSearchRepository) -> None:
     await repo.init_search_index()
     await repo.bulk_index_items(
@@ -620,7 +515,7 @@ async def record_vector_query(
     )
 
     assert [row.permalink for row in results] == ["specs/bravo", "specs/alpha"]
-    assert candidate_limits == [18, 8]
+    assert candidate_limits == [18]
 
 
 @pytest.mark.asyncio
@@ -795,7 +690,7 @@ async def record_vector_query(
 
     assert growing_prefix_results
     assert reranker.calls == 2
-    assert candidate_limits == [90, 80]
+    assert candidate_limits == [90]
 
 
 @pytest.mark.asyncio

tests/repository/test_search_trace.py

@@ -847,7 +847,7 @@ async def test_fts_vector_hybrid_and_rerank_trace_variants(
     )
     assert isinstance(reranked_trace, VectorQueryTrace)
     assert reranked_trace.rerank is not None
-    assert reranked_trace.rerank.stable_pool_refetched is True
+    assert reranked_trace.rerank.stable_pool_refetched is False
     assert reranked_trace.rerank.entries[0].key == ("entity", 3)
     alpha_rerank = next(
         entry for entry in reranked_trace.rerank.entries if entry.key == ("entity", 1)