fix: eliminate double IDB reads in hierarchical query routing, relax book membership assertion

Ranking functions (rankShelves/rankVolumes/rankBooks/rankPages/spillToWarm)
now return RankedResult { id, score, childIds } so Query.ts can use child
IDs from the ranking pass directly instead of re-fetching entities from IDB.

Test assertion for page-to-book membership relaxed from "exactly one" to
"at least one" per user feedback — pages can legitimately appear in
multiple books (e.g., legal corpora with cross-citations).

Co-authored-by: devlux76 <[email protected]>

Author: Copilot

lib/cortex/Query.ts

@@ -4,7 +4,7 @@ import type { EmbeddingRunner } from "../embeddings/EmbeddingRunner";
 import { runPromotionSweep } from "../core/SalienceEngine";
 import { computeSubgraphBounds } from "../core/HotpathPolicy";
 import type { QueryResult } from "./QueryResult";
-import { rankPages, rankBooks, rankVolumes, rankShelves, spillToWarm } from "./Ranking";
+import { rankPages, rankBooks, rankVolumes, rankShelves, spillToWarm, type RankedResult } from "./Ranking";
 import { buildMetroid } from "./MetroidBuilder";
 import { detectKnowledgeGap } from "./KnowledgeGapDetector";
 import { solveOpenTSP } from "./OpenTSPSolver";
@@ -64,10 +64,7 @@ export async function query(
       rankingOptions,
     );
     for (const s of topShelves) {
-      const shelf = await metadataStore.getShelf(s.id);
-      if (shelf) {
-        for (const vid of shelf.volumeIds) volumeIdsFromShelves.add(vid);
-      }
+      for (const vid of s.childIds) volumeIdsFromShelves.add(vid);
     }
   }
 
@@ -86,10 +83,7 @@ export async function query(
       rankingOptions,
     );
     for (const v of topVolumes) {
-      const volume = await metadataStore.getVolume(v.id);
-      if (volume) {
-        for (const bid of volume.bookIds) bookIdsFromVolumes.add(bid);
-      }
+      for (const bid of v.childIds) bookIdsFromVolumes.add(bid);
     }
   }
 
@@ -108,10 +102,7 @@ export async function query(
       rankingOptions,
     );
     for (const b of topBooks) {
-      const book = await metadataStore.getBook(b.id);
-      if (book) {
-        for (const pid of book.pageIds) pageIdsFromBooks.add(pid);
-      }
+      for (const pid of b.childIds) pageIdsFromBooks.add(pid);
     }
   }
 
@@ -123,7 +114,7 @@ export async function query(
   const seenIds = new Set(hotResults.map((r) => r.id));
 
   // --- Warm spill: fill up to topK if hot path is insufficient ---
-  let warmResults: Array<{ id: Hash; score: number }> = [];
+  let warmResults: RankedResult[] = [];
   if (hotResults.length < topK) {
     const allWarm = await spillToWarm("page", queryEmbedding, topK, rankingOptions);
     warmResults = allWarm.filter((r) => !seenIds.has(r.id));

lib/cortex/Ranking.ts

@@ -21,34 +21,42 @@ function cosineSimilarity(a: Float32Array, b: Float32Array): number {
   return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
 }
 
+export interface RankedResult {
+  id: Hash;
+  score: number;
+  /** Child IDs from the ranked entity (volumeIds / bookIds / pageIds). */
+  childIds: Hash[];
+}
+
 function pickTopK(
-  scored: Array<{ id: Hash; score: number }>,
+  scored: RankedResult[],
   k: number,
-): Array<{ id: Hash; score: number }> {
+): RankedResult[] {
   scored.sort((a, b) => b.score - a.score || a.id.localeCompare(b.id));
   return scored.slice(0, k);
 }
 
 /**
  * Ranks shelves by cosine similarity of their routing prototype to the query.
  * Uses routingPrototypeOffsets[0] as the representative vector.
+ * Returns child volumeIds alongside each scored shelf.
  */
 export async function rankShelves(
   queryEmbedding: Float32Array,
   residentShelfIds: Hash[],
   topK: number,
   options: RankingOptions,
-): Promise<Array<{ id: Hash; score: number }>> {
+): Promise<RankedResult[]> {
   if (residentShelfIds.length === 0) return [];
 
   const { vectorStore, metadataStore } = options;
-  const scored: Array<{ id: Hash; score: number }> = [];
+  const scored: RankedResult[] = [];
 
   for (const shelfId of residentShelfIds) {
     const shelf = await metadataStore.getShelf(shelfId);
     if (!shelf || shelf.routingPrototypeOffsets.length === 0) continue;
     const vec = await vectorStore.readVector(shelf.routingPrototypeOffsets[0], shelf.routingDim);
-    scored.push({ id: shelfId, score: cosineSimilarity(queryEmbedding, vec) });
+    scored.push({ id: shelfId, score: cosineSimilarity(queryEmbedding, vec), childIds: shelf.volumeIds });
   }
 
   return pickTopK(scored, topK);
@@ -57,49 +65,51 @@ export async function rankShelves(
 /**
  * Ranks volumes by cosine similarity of their first prototype to the query.
  * Uses prototypeOffsets[0] as the representative vector.
+ * Returns child bookIds alongside each scored volume.
  */
 export async function rankVolumes(
   queryEmbedding: Float32Array,
   residentVolumeIds: Hash[],
   topK: number,
   options: RankingOptions,
-): Promise<Array<{ id: Hash; score: number }>> {
+): Promise<RankedResult[]> {
   if (residentVolumeIds.length === 0) return [];
 
   const { vectorStore, metadataStore } = options;
-  const scored: Array<{ id: Hash; score: number }> = [];
+  const scored: RankedResult[] = [];
 
   for (const volumeId of residentVolumeIds) {
     const volume = await metadataStore.getVolume(volumeId);
     if (!volume || volume.prototypeOffsets.length === 0) continue;
     const vec = await vectorStore.readVector(volume.prototypeOffsets[0], volume.prototypeDim);
-    scored.push({ id: volumeId, score: cosineSimilarity(queryEmbedding, vec) });
+    scored.push({ id: volumeId, score: cosineSimilarity(queryEmbedding, vec), childIds: volume.bookIds });
   }
 
   return pickTopK(scored, topK);
 }
 
 /**
  * Ranks books by cosine similarity of their medoid page embedding to the query.
+ * Returns child pageIds alongside each scored book.
  */
 export async function rankBooks(
   queryEmbedding: Float32Array,
   residentBookIds: Hash[],
   topK: number,
   options: RankingOptions,
-): Promise<Array<{ id: Hash; score: number }>> {
+): Promise<RankedResult[]> {
   if (residentBookIds.length === 0) return [];
 
   const { vectorStore, metadataStore } = options;
-  const scored: Array<{ id: Hash; score: number }> = [];
+  const scored: RankedResult[] = [];
 
   for (const bookId of residentBookIds) {
     const book = await metadataStore.getBook(bookId);
     if (!book) continue;
     const medoidPage = await metadataStore.getPage(book.medoidPageId);
     if (!medoidPage) continue;
     const vec = await vectorStore.readVector(medoidPage.embeddingOffset, medoidPage.embeddingDim);
-    scored.push({ id: bookId, score: cosineSimilarity(queryEmbedding, vec) });
+    scored.push({ id: bookId, score: cosineSimilarity(queryEmbedding, vec), childIds: book.pageIds });
   }
 
   return pickTopK(scored, topK);
@@ -113,17 +123,17 @@ export async function rankPages(
   residentPageIds: Hash[],
   topK: number,
   options: RankingOptions,
-): Promise<Array<{ id: Hash; score: number }>> {
+): Promise<RankedResult[]> {
   if (residentPageIds.length === 0) return [];
 
   const { vectorStore, metadataStore } = options;
-  const scored: Array<{ id: Hash; score: number }> = [];
+  const scored: RankedResult[] = [];
 
   for (const pageId of residentPageIds) {
     const page = await metadataStore.getPage(pageId);
     if (!page) continue;
     const vec = await vectorStore.readVector(page.embeddingOffset, page.embeddingDim);
-    scored.push({ id: pageId, score: cosineSimilarity(queryEmbedding, vec) });
+    scored.push({ id: pageId, score: cosineSimilarity(queryEmbedding, vec), childIds: [] });
   }
 
   return pickTopK(scored, topK);
@@ -139,17 +149,17 @@ export async function spillToWarm(
   queryEmbedding: Float32Array,
   topK: number,
   options: RankingOptions,
-): Promise<Array<{ id: Hash; score: number }>> {
+): Promise<RankedResult[]> {
   if (tier !== "page") return [];
 
   const { vectorStore, metadataStore } = options;
   const allPages = await metadataStore.getAllPages();
   if (allPages.length === 0) return [];
 
-  const scored: Array<{ id: Hash; score: number }> = [];
+  const scored: RankedResult[] = [];
   for (const page of allPages) {
     const vec = await vectorStore.readVector(page.embeddingOffset, page.embeddingDim);
-    scored.push({ id: page.pageId, score: cosineSimilarity(queryEmbedding, vec) });
+    scored.push({ id: page.pageId, score: cosineSimilarity(queryEmbedding, vec), childIds: [] });
   }
 
   return pickTopK(scored, topK);

tests/integration/IngestQuery.test.ts

@@ -424,14 +424,11 @@ describe("integration (v0.5): hierarchical and dialectical ingest/query", () =>
     expect(result.books.length).toBeGreaterThanOrEqual(1);
     expect(result.book).toBeDefined();
 
-    // Every page must belong to exactly one book
+    // Every page must belong to at least one book
     const allBookPageIds = result.books.flatMap((b) => b.pageIds);
     for (const page of result.pages) {
       expect(allBookPageIds).toContain(page.pageId);
     }
-    // Enforce exactly-once membership (no page duplicated across books)
-    const uniqueBookPageIds = new Set(allBookPageIds);
-    expect(uniqueBookPageIds.size).toBe(allBookPageIds.length);
     // Every book's medoid must be one of its own pages
     for (const book of result.books) {
       const storedBook = await metadataStore.getBook(book.bookId);