fix(core): reuse retrieval for stable rerank pagination

Signed-off-by: phernandez <paul@basicmachines.co>

Author: phernandez

src/basic_memory/repository/search_repository_base.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/semantic_vector_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),

src/basic_memory/repository/sqlite_vec_index.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/semantic/test_semantic_coverage.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_multi_project_search.py

@@ -0,0 +1,80 @@
+"""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 tests.repository.test_rerank_pipeline import (
+    BackendSearchRepository,
+    _FakeReranker,
+    rerank_search_repository as rerank_search_repository,
+)
+from tests.repository.test_stable_rerank_pagination import (
+    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]

test-int/test_stable_rerank_pagination.py

@@ -620,7 +620,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 +795,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_rerank_pipeline.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)