fix(retrieval): every semantic search scopes first, then ranks - the note, milestone and system searches join the rule search on one shared shape, _rank_scoped (#4961)

Ordered straight off a *_embeddings table, the planner walks the HNSW index,
takes ~ef_search (40) nearest chunks across every owner and project, and only
then applies the scope: an in-scope record behind 40 nearer ones the caller
cannot see was silently dropped. #4958 fixed the rule search alone; the note,
milestone and system searches kept the fault.

_scoped_chunks builds the in-scope chunks with their distance; _rank_scoped
materializes them as a CTE and ranks exactly. All four searches go through it,
and the row shape is unchanged, so callers and mocks are untouched.

Tests: a structural guard that every semantic_search_* ranks through
_rank_scoped and orders nothing itself (with a replay of the old shape), a
compiled-SQL check of the MATERIALIZED CTE, and integration crowd tests - 80
nearer out-of-scope records - for notes (another user; the reader own other
project), milestones and systems.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

src/scribe/services/embeddings.py

@@ -856,6 +856,62 @@ def record_best_chunk(report: dict | None, chunks: dict[int, dict]) -> None:
     report["best_chunk"] = chunks
 
 
+# --- scope first, then rank (#4958, #4961) -----------------------------------
+#
+# Every semantic search here asks the same question of a shared table: the
+# nearest chunks AMONG THE ONES THIS CALLER MAY SEE. Ordered straight off a
+# `*_embeddings` table, the planner answers a different one. It walks the HNSW
+# index, takes about `hnsw.ef_search` (40) nearest chunks from every owner and
+# project, and only then applies the scope — so an in-scope record ranked past
+# the 40th chunk overall is silently gone, and on a shared install other
+# users' records are what fill those 40. The searches fail open, so nothing
+# reports it: the result is just shorter.
+#
+# The fix is one shape, shared by every search so it cannot be fixed in one
+# and missed in the next (which is how #4958 left three behind). The scope is
+# a MATERIALIZED CTE with the distance computed inside it; the index cannot
+# order a materialized CTE, so the ranking over it is exact.
+#
+# What exact costs is one distance per in-scope chunk. Rules, milestones and
+# Systems are hundreds of chunks. Notes are the big corpus, but the hook arms
+# that run on every prompt are project-scoped, so the pass is over one
+# project's chunks rather than the whole table; `retrieval_logs.duration_ms`
+# on those arms is where that claim is checked after the deploy. If it ever
+# stops holding, pgvector 0.8's `hnsw.iterative_scan` keeps the index and is
+# the next step — gated on the extension version, because an unknown `hnsw.*`
+# setting errors and these searches would read that error as "nothing found".
+
+
+def _scoped_chunks(embedding_model, record_key, distance):
+    """A select over an embedding table's chunks, each with its distance — the
+    caller joins its record and adds its scope, and `_rank_scoped` orders it.
+
+    The columns are labelled so every search's CTE carries the same four.
+    """
+    return select(
+        record_key.label("record_id"),
+        distance.label("distance"),
+        embedding_model.chunk_index.label("chunk_index"),
+        embedding_model.chunk_text.label("chunk_text"),
+    ).select_from(embedding_model)
+
+
+def _rank_scoped(record_model, scoped_chunks, *, name: str, limit: int):
+    """(record, distance, chunk_index, chunk_text) rows, nearest first, drawn
+    ONLY from the in-scope chunks — exactly, never through the index.
+
+    The row shape is what each search's collapse already unpacks, so a search
+    moves onto this without its callers or its mocks noticing.
+    """
+    scoped = scoped_chunks.cte(name).prefix_with("MATERIALIZED")
+    return (
+        select(record_model, scoped.c.distance, scoped.c.chunk_index, scoped.c.chunk_text)
+        .join(scoped, scoped.c.record_id == record_model.id)
+        .order_by(scoped.c.distance)
+        .limit(limit)
+    )
+
+
 async def semantic_search_notes(
     user_id: int,
     query: str,
@@ -924,11 +980,14 @@ async def semantic_search_notes(
     a caller that forgets is wrong in the safe direction.
 
     Ranking and the top-k cut happen in Postgres via pgvector's cosine-distance
-    operator (`<=>`, exposed as ``Vector.cosine_distance``) backed by the HNSW
-    index from migration 0067 — so this is an indexed ``ORDER BY ... LIMIT k``
-    rather than a full-table scan. Cosine distance is ``1 - cosine_similarity``,
-    so a similarity floor of *threshold* is a distance ceiling of
-    ``1 - threshold`` and similarity is recovered as ``1 - distance``.
+    operator (`<=>`, exposed as ``Vector.cosine_distance``), EXACTLY and over
+    the in-scope chunks only (#4961, see `_rank_scoped`). It used to be an
+    indexed ``ORDER BY ... LIMIT k`` straight off the HNSW index from
+    migration 0067, which ranks before it filters and so could drop an
+    in-scope record that sat behind ~40 nearer ones the caller cannot see.
+    Cosine distance is ``1 - cosine_similarity``, so a similarity floor of
+    *threshold* is a distance ceiling of ``1 - threshold`` and similarity is
+    recovered as ``1 - distance``.
 
     `demote_superseded` applies the supersession penalty (#278): a record a
     later note claims to have overtaken ranks below its equals. Callers asking
@@ -964,18 +1023,16 @@ async def semantic_search_notes(
             # Scope on Note, not NoteEmbedding.user_id: the embedding row belongs
             # to the note's owner, so filtering it would pin every scope to "own"
             # and leave shared records unreachable by meaning.
+            #
+            # Every filter below is the SCOPE, built into the chunks the
+            # ranking draws from rather than applied to what it returns
+            # (#4961) — see `_rank_scoped`.
             stmt = (
                 # chunk_index/chunk_text ride along so the collapse below can
                 # say WHICH passage matched. Without them the caller is left
                 # previewing the head of the body — a span this query has
                 # already determined is not why the record ranked (#4243).
-                select(
-                    Note,
-                    distance.label("distance"),
-                    NoteEmbedding.chunk_index,
-                    NoteEmbedding.chunk_text,
-                )
-                .select_from(NoteEmbedding)
+                _scoped_chunks(NoteEmbedding, NoteEmbedding.note_id, distance)
                 .join(Note, NoteEmbedding.note_id == Note.id)
                 .where(
                     notes_visibility_clause(user_id, scope),
@@ -1032,24 +1089,22 @@ async def semantic_search_notes(
             # the results and the live record that should have replaced it was
             # never fetched.
             #
-            # Ordering stays on RAW distance so pgvector's HNSW index still
-            # serves it (migration 0067). Ordering by `distance + penalty`
-            # instead would be exact, and would turn an indexed top-k into a
-            # scan-and-sort of every embedded note.
-            #
-            # The cost of that trade, stated plainly: a live record outside the
-            # over-fetch window cannot be promoted into the results. With a
-            # penalty far smaller than the window's score spread, that case
-            # needs the true answer to be more than _SUPERSESSION_OVERFETCH
-            # ranks down, which no observed query comes close to.
+            # Ordering is on RAW distance; the penalty is applied after the
+            # cut, over the over-fetched window. The cost of that, stated
+            # plainly: a live record outside the window cannot be promoted
+            # into the results. With a penalty far smaller than the window's
+            # score spread, that case needs the true answer to be more than
+            # _SUPERSESSION_OVERFETCH ranks down, which no observed query
+            # comes close to.
             fetch = limit * _CHUNK_OVERFETCH * (
                 _SUPERSESSION_OVERFETCH if demote_superseded else 1
             )
             # NO threshold predicate — see the note above this function. The
             # bar is applied after the collapse, where the rejected scores can
             # still be seen.
-            stmt = stmt.order_by(distance.asc()).limit(fetch)
-            rows = list((await session.execute(stmt)).all())
+            rows = list((await session.execute(
+                _rank_scoped(Note, stmt, name="scoped_note_chunks", limit=fetch)
+            )).all())
     except Exception:
         logger.warning("Failed to query note embeddings", exc_info=True)
         return []
@@ -1550,24 +1605,14 @@ async def semantic_search_rules(
         # function fails open.
         home = await rule_home(user_id, project_id, everywhere=everywhere)
 
-        # SCOPE FIRST, THEN RANK — exactly (#4958). Ordered straight off
-        # rule_embeddings, the planner walks the HNSW index, which hands back
-        # about `hnsw.ef_search` (40) nearest chunks from EVERY owner and
-        # project and only then applies the home filter: an in-scope rule
-        # ranked past the 40th chunk overall is silently gone, and on a
-        # shared install other users' rules fill those 40. A MATERIALIZED
-        # CTE cannot be ordered through the index, so the distance is
-        # computed for each in-scope chunk and the ranking is exact. A
-        # rulebook is hundreds of chunks, not millions — exact is cheap here.
+        # SCOPE FIRST, THEN RANK — exactly (#4958; the shape is
+        # `_rank_scoped`'s). The home filter is built into the chunks the
+        # ranking draws from: applied to what an HNSW walk returned, it ran
+        # after ~40 nearest chunks from every owner and project had been
+        # taken, and an in-scope rule ranked past them was silently gone.
         scoped = (
             joined_to_homes(
-                select(
-                    RuleEmbedding.rule_id.label("rule_id"),
-                    distance.label("distance"),
-                    RuleEmbedding.chunk_index.label("chunk_index"),
-                    RuleEmbedding.chunk_text.label("chunk_text"),
-                )
-                .select_from(RuleEmbedding)
+                _scoped_chunks(RuleEmbedding, RuleEmbedding.rule_id, distance)
                 .join(Rule, RuleEmbedding.rule_id == Rule.id)
             )
             .where(
@@ -1577,18 +1622,14 @@ async def semantic_search_rules(
                 home,
                 *( [Rule.kind == kind] if kind else [] ),
             )
-            .cte("scoped_rule_chunks")
-            .prefix_with("MATERIALIZED")
         )
 
         async with async_session() as session:
             rows = (await session.execute(
-                select(Rule, scoped.c.distance, scoped.c.chunk_index, scoped.c.chunk_text)
-                .join(scoped, scoped.c.rule_id == Rule.id)
                 # Overfetch so collapsing chunks to their best row still fills
                 # the page — the same reason the note search overfetches.
-                .order_by(scoped.c.distance)
-                .limit(limit * _CHUNK_OVERFETCH)
+                _rank_scoped(Rule, scoped, name="scoped_rule_chunks",
+                             limit=limit * _CHUNK_OVERFETCH)
             )).all()
     except Exception:
         logger.warning("Rule semantic search failed", exc_info=True)
@@ -1778,23 +1819,20 @@ async def semantic_search_milestones(
             scope = Milestone.project_id == project_id
         else:
             scope = Milestone.user_id == user_id
+        # Scope first, then rank (#4961): see `_rank_scoped`.
+        scoped = (
+            _scoped_chunks(MilestoneEmbedding, MilestoneEmbedding.milestone_id, distance)
+            .join(Milestone, MilestoneEmbedding.milestone_id == Milestone.id)
+            .where(
+                scope,
+                Milestone.deleted_at.is_(None),
+                *([Milestone.status == status] if status else []),
+            )
+        )
         async with async_session() as session:
             rows = (await session.execute(
-                select(
-                    Milestone,
-                    distance.label("distance"),
-                    MilestoneEmbedding.chunk_index,
-                    MilestoneEmbedding.chunk_text,
-                )
-                .select_from(MilestoneEmbedding)
-                .join(Milestone, MilestoneEmbedding.milestone_id == Milestone.id)
-                .where(
-                    scope,
-                    Milestone.deleted_at.is_(None),
-                    *([Milestone.status == status] if status else []),
-                )
-                .order_by(distance)
-                .limit(limit * _CHUNK_OVERFETCH)
+                _rank_scoped(Milestone, scoped, name="scoped_milestone_chunks",
+                             limit=limit * _CHUNK_OVERFETCH)
             )).all()
     except Exception:
         logger.warning("Milestone semantic search failed", exc_info=True)
@@ -1953,23 +1991,20 @@ async def semantic_search_systems(
             scope = System.project_id == project_id
         else:
             scope = System.user_id == user_id
+        # Scope first, then rank (#4961): see `_rank_scoped`.
+        scoped = (
+            _scoped_chunks(SystemEmbedding, SystemEmbedding.system_id, distance)
+            .join(System, SystemEmbedding.system_id == System.id)
+            .where(
+                scope,
+                System.deleted_at.is_(None),
+                System.status != "archived",
+            )
+        )
         async with async_session() as session:
             rows = (await session.execute(
-                select(
-                    System,
-                    distance.label("distance"),
-                    SystemEmbedding.chunk_index,
-                    SystemEmbedding.chunk_text,
-                )
-                .select_from(SystemEmbedding)
-                .join(System, SystemEmbedding.system_id == System.id)
-                .where(
-                    scope,
-                    System.deleted_at.is_(None),
-                    System.status != "archived",
-                )
-                .order_by(distance)
-                .limit(limit * _CHUNK_OVERFETCH)
+                _rank_scoped(System, scoped, name="scoped_system_chunks",
+                             limit=limit * _CHUNK_OVERFETCH)
             )).all()
     except Exception:
         logger.warning("System semantic search failed", exc_info=True)