pi-mega-compact 0.8.21 → 0.8.23

package/dist/extensions/mega-commands.js

@@ -12,6 +12,7 @@ import { decompressSmart } from "../src/store/compression.js";
 import { loadMetrics, fpRate, p95, defaultMetricsPath } from "../src/monitoring.js";
 import { C, recentUserQuery } from "./mega-runtime.js";
 import { runCompact, doRecall, doRecallAsync } from "./mega-pipeline.js";
+import { vectorStats, vectorRepoStats, vectorDataInvariant } from "../src/vectorStore.js";
 /** Resolve a checkpoint by id (or "recent"/"last") from this session's store. */
 export function findCheckpoint(runtime, sid, ref) {
     const all = listCheckpoints(sid, runtime.currentStateDir);
@@ -88,9 +89,9 @@ export function registerCommands(pi, runtime, config) {
                 const pct = usage?.percent != null ? `${usage.percent}%` : "n/a";
                 const tokens = usage?.tokens != null ? `${usage.tokens} tok` : "n/a";
                 const sid = normalizeSessionId(ctx.sessionManager.getSessionId());
-                const st = runtime.store.stats(sid);
-                const repo = runtime.store.repoStats();
-                const di = runtime.store.dataInvariant();
+                const st = vectorStats(runtime.store, sid);
+                const repo = vectorRepoStats(runtime.store);
+                const di = vectorDataInvariant(runtime.store);
                 const fmtB = (b) => b >= 1_048_576 ? `${(b / 1_048_576).toFixed(1)} MiB` :
                     b >= 1024 ? `${(b / 1024).toFixed(1)} KiB` : `${b} B`;
                 // Real cost: tokens saved × the captured model's input rate (USD/token),

package/dist/extensions/mega-events/agent-handlers.js

@@ -6,6 +6,7 @@ import { isMegaCache } from "../../src/game/scoring.js";
 import { resolveRepoRoot } from "../mega-config.js";
 import { classifyError } from "./error-classifier.js";
 import { safeSendUserMessage } from "./send-safe.js";
+import { vectorStats } from "../../src/vectorStore.js";
 /** Register agent/turn tracking event handlers. */
 export function registerAgentHandlers(pi, runtime, config) {
     // ---- Agent tracking for real-time widget + status-line updates ---------
@@ -215,7 +216,7 @@ export function registerAgentHandlers(pi, runtime, config) {
         try {
             if (runtime.getCachedGameState().game_mode_on) {
                 const repo = resolveRepoRoot(ctx.cwd) ?? runtime.currentStateDir;
-                const st = runtime.store.stats(runtime.rt.sessionId);
+                const st = vectorStats(runtime.store, runtime.rt.sessionId);
                 const cachePct = st.dedupHitRate * 100;
                 const modelId = runtime.currentModel?.modelId ?? "unknown";
                 recordScore(runtime.currentStateDir, {

package/dist/extensions/mega-events/compact-handlers.js

@@ -79,6 +79,32 @@ export function registerCompactHandlers(pi, runtime, config) {
         });
         if (!config.auto)
             return {}; // let pi run its own native compaction
+        // S38.5: COMPACT-DEDUP RACE GUARD — pi's _runAutoCompaction fires
+        // session_before_compact from within the agent loop (after agent_end,
+        // _handlePostAgentRun → _checkCompaction → _runAutoCompaction). If a
+        // compaction just completed (within the cooldown window), pi's
+        // prepareCompaction may still see the last entry as a compaction but
+        // the extension's session_compact handler hasn't been called yet, OR
+        // the setTimeout deferred ctx.compact() from the agent_end handler is
+        // about to race. Skip the durable trim and let pi run its own compact
+        // (or no-op if the branch is already compacted). Matches the cooldown
+        // in agent-handlers.ts so both call sites are consistent.
+        const cooldownMs = config.raceGuardStrict ? 30_000 : 10_000;
+        const sinceCompact = Date.now() - (runtime.rt.lastNativeCompactAt ?? 0);
+        if (sinceCompact < cooldownMs) {
+            runtime.logger.info("before-compact-skip-recent", {
+                sessionId: runtime.rt.sessionId,
+                reason: event.reason,
+                sinceCompactMs: sinceCompact,
+                cooldownMs,
+            });
+            runtime.diagBeforeCompactSupplied++;
+            const fb = fallbackCompaction(event);
+            if (fb) {
+                return { compaction: fb.compaction };
+            }
+            return {};
+        }
         try {
             const result = driveNativeCompaction(event, runtime, config);
             if (result && result.compaction.summary?.trim()) {

package/dist/extensions/mega-events/session-handlers.js

@@ -3,6 +3,7 @@ import { autoMaintain } from "../../src/store/sqlite.js";
 import { recentUserQuery, WIDGET_KEY, } from "../mega-runtime.js";
 import { doRecall, doRecallAsync, } from "../mega-pipeline.js";
 import { recallMemoriesAndInline } from "../../src/recall.js";
+import { vectorStats } from "../../src/vectorStore.js";
 /** Register session lifecycle event handlers. */
 export function registerSessionHandlers(pi, runtime, config) {
     // Capture model/provider whenever it changes (drives real cost estimation).
@@ -26,7 +27,7 @@ export function registerSessionHandlers(pi, runtime, config) {
         if (config.autoInline) {
             const sid = normalizeSessionId(ctx.sessionManager.getSessionId());
             const query = recentUserQuery(ctx);
-            if (query && runtime.store.stats(sid).checkpointCount > 0) {
+            if (query && vectorStats(runtime.store, sid).checkpointCount > 0) {
                 // S17: use the async variant on resume so cross-repo HNSW recall can
                 // augment when this repo's store is thin. session_start is an async-safe
                 // point (unlike the mid-turn context handler, which stays sync).

package/dist/extensions/mega-pipeline/compact.js

@@ -19,6 +19,7 @@ import { runRaptor } from "../../src/dedup/raptor/index.js";
 import { loadDedupConfig } from "../../src/config/dedup.js";
 import { upsertEmbedding as indexUpsertEmbedding } from "../../src/store/vectorIndex.js";
 import { runMemoryReview } from "./memory-review.js";
+import { vectorList } from "../../src/vectorStore.js";
 /** Run the full compaction pipeline and persist a checkpoint. Returns the result. */
 export function runCompact(pi, runtime, config, ctx, messages, opts = {}) {
     runtime.bindRepo(ctx.cwd);
@@ -171,7 +172,7 @@ function doCompact(view, keepFrom, opts, sid, config, pi, ctx, runtime) {
     if (config.raptorEnabled && !result.deduped) {
         try {
             const dd = loadDedupConfig();
-            const all = runtime.store.list(sid);
+            const all = vectorList(runtime.store, sid);
             const leaves = all.map((cp) => ({
                 id: cp.checkpointId,
                 messages: [],
@@ -204,7 +205,7 @@ function doCompact(view, keepFrom, opts, sid, config, pi, ctx, runtime) {
     // Non-fatal: a WASM init failure degrades to the sync scan silently.
     if (!result.deduped) {
         try {
-            const all = runtime.store.list(sid);
+            const all = vectorList(runtime.store, sid);
             const latest = all.find((cp) => cp.checkpointId === result.checkpointId);
             if (latest?.embedding) {
                 void indexUpsertEmbedding(runtime.currentStateDir, sid, latest.checkpointId, latest.embedding).catch(() => {

package/dist/extensions/mega-runtime/state.js

@@ -11,7 +11,7 @@
  */
 import { join } from "node:path";
 import { appendFileSync, mkdirSync } from "node:fs";
-import { VectorStore } from "../../src/vectorStore.js";
+import { VectorStore, vectorStats, vectorRepoStats, vectorDataInvariant } from "../../src/vectorStore.js";
 import { toEngineMessages } from "../../src/adapt.js";
 import { normalizeSessionId } from "../../src/store.js";
 import { Logger } from "../../src/log.js";
@@ -306,8 +306,8 @@ export class MegaRuntime {
         // never break the per-repo compaction path. Runs only on repo-switch
         // (this branch), so it's infrequent — not per-context-event.
         try {
-            const repo = this.store.repoStats();
-            const di = this.store.dataInvariant();
+            const repo = vectorRepoStats(this.store);
+            const di = vectorDataInvariant(this.store);
             const root = key !== dir ? key : (resolveRepoRoot(cwd ?? dir) ?? dir);
             upsertRepoRegistry({
                 repoRoot: root,
@@ -346,9 +346,9 @@ export class MegaRuntime {
             return;
         }
         const perfT0 = performance.now();
-        const st = this.store.stats(this.rt.sessionId);
-        const repo = this.store.repoStats();
-        const di = this.store.dataInvariant();
+        const st = vectorStats(this.store, this.rt.sessionId);
+        const repo = vectorRepoStats(this.store);
+        const di = vectorDataInvariant(this.store);
         // Live + store-wide cache-hit / compaction counters for the dashboard.
         const ds = getDedupStats(this.currentStateDir);
         const cacheHitsTotal = ds.deduped + getRecallInjected(this.currentStateDir);
@@ -492,7 +492,7 @@ export class MegaRuntime {
                 total: getCompactCount(this.currentStateDir),
             },
             timeSaved: {
-                compact: { sessionSec: sec(this.rt.tokensSaved), totalSec: sec(this.store.repoStats().tokensSaved) },
+                compact: { sessionSec: sec(this.rt.tokensSaved), totalSec: sec(vectorRepoStats(this.store).tokensSaved) },
                 cacheHit: { sessionSec: sec(this.rt.cacheHitTokens), totalSec: sec(cacheHitsTotalTokens) },
             },
             model,

package/dist/src/dedup/raptor/multilevel.js

@@ -0,0 +1,172 @@
+/**
+ * multilevel.ts — Multi-level RAPTOR retrieval engine (S42A).
+ *
+ * Upgrades the RAPTOR recall path from flat (leaf-only) to multi-level
+ * retrieval across the entire hierarchical tree. Searches ALL levels with
+ * configurable level weights, supports leaf expansion for cluster hits,
+ * and deduplicates overlapping results.
+ *
+ * PREVENT-PI-004: pure in-process math (cosine, BFS, extractive). No network.
+ * PREVENT-PI-001: produces SearchHit[] that feed into recallAndInline() —
+ * affects which checkpoints are recalled, not how messages are dropped.
+ */
+import { cosineSimilarity } from "../../embedder.js";
+import { mmrRerank } from "../mmr.js";
+import { leafDescendants } from "./retrieval.js";
+const DEFAULT_LEVEL_WEIGHTS = [1.0, 0.9, 0.8, 0.7, 0.5];
+// ── S42A-2: Level-weighted scoring ─────────────────────────────────────────
+/**
+ * Score all RAPTOR tree nodes by cosine similarity to the query, then apply
+ * level-specific weights. Returns hits sorted by weighted score descending.
+ *
+ * Level weights: leaves (level 0) get weight 1.0, level 1 gets 0.9, etc.
+ * This ensures detailed leaves score highest while still surfacing higher-level
+ * summaries when they're highly relevant.
+ */
+export function scoreTreeLevels(query, tree, opts) {
+    const { embedder } = opts;
+    const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
+    const qv = embedder.embed(query);
+    const hits = [];
+    // 1. Score all internal (summary) nodes.
+    for (const node of tree.nodes.values()) {
+        const rawScore = cosineSimilarity(qv, node.embedding);
+        const levelWeight = weights[Math.min(node.level, weights.length - 1)];
+        hits.push({
+            nodeId: node.id,
+            level: node.level,
+            score: rawScore * levelWeight,
+            rawScore,
+            isLeaf: false,
+            leafIds: node.children,
+            summary: node.summary,
+            embedding: node.embedding,
+        });
+    }
+    // 2. Score leaf nodes. Leaf ids are not in tree.nodes — they are children
+    //    referenced by internal nodes. Each leaf's embedding is the level-0
+    //    parent node that wraps it (same approach as stagedExpansion:95–102).
+    const seenLeaves = new Set();
+    for (const node of tree.nodes.values()) {
+        for (const leafId of node.children) {
+            if (seenLeaves.has(leafId) || tree.nodes.has(leafId))
+                continue;
+            seenLeaves.add(leafId);
+            const rawScore = cosineSimilarity(qv, node.embedding);
+            const leafWeight = weights[0];
+            hits.push({
+                nodeId: leafId,
+                level: 0,
+                score: rawScore * leafWeight,
+                rawScore,
+                isLeaf: true,
+                leafIds: [leafId],
+                summary: "", // leaves have no summary — they are raw checkpoint ids
+                embedding: node.embedding,
+            });
+        }
+    }
+    hits.sort((a, b) => b.score - a.score);
+    return hits;
+}
+// ── S42A-3: Leaf expansion ─────────────────────────────────────────────────
+/**
+ * Given a set of cluster-level hits, expand each one to include its leaf
+ * descendants. Deduplicates: if a leaf is already present as a direct hit,
+ * it is not duplicated. Returns the merged set (original hits + expanded leaves).
+ */
+export function expandLeafDescendants(hits, tree, maxPerCluster, _embedder, queryVector, levelWeights) {
+    const weights = levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
+    const existingIds = new Set(hits.map((h) => h.nodeId));
+    const expanded = [];
+    for (const hit of hits) {
+        if (hit.isLeaf) {
+            expanded.push(hit);
+            continue;
+        }
+        // Get all leaf descendants for this cluster node.
+        const node = tree.nodes.get(hit.nodeId);
+        if (!node) {
+            expanded.push(hit);
+            continue;
+        }
+        const rawLeafIds = leafDescendants(node, tree);
+        // Sort by cosine similarity to query, cap at maxPerCluster.
+        const leafHits = rawLeafIds
+            .map((lid) => {
+            // Leaf embedding = its nearest internal parent's embedding.
+            const parent = [...tree.nodes.values()].find((n) => n.children.includes(lid));
+            const sim = parent
+                ? cosineSimilarity(queryVector, parent.embedding)
+                : 0;
+            return { lid, sim, parent };
+        })
+            .sort((a, b) => b.sim - a.sim)
+            .slice(0, maxPerCluster)
+            .filter((l) => !existingIds.has(l.lid))
+            .map((l) => {
+            existingIds.add(l.lid);
+            const rawScore = l.sim;
+            return {
+                nodeId: l.lid,
+                level: 0,
+                score: rawScore * weights[0],
+                rawScore,
+                isLeaf: true,
+                leafIds: [l.lid],
+                summary: "",
+                embedding: l.parent?.embedding ?? hit.embedding,
+            };
+        });
+        expanded.push(hit, ...leafHits);
+    }
+    return expanded;
+}
+// ── S42A-4: Result dedup ───────────────────────────────────────────────────
+/**
+ * Deduplicate hits: if both a cluster node and its leaf children appear in
+ * results, remove the cluster hit (leaves provide more specific context).
+ * If no leaves are in the set, keep the cluster hit (it provides the abstract view).
+ */
+export function deduplicateMultilevelHits(hits) {
+    const leafIds = new Set(hits.filter((h) => h.isLeaf).map((h) => h.nodeId));
+    return hits.filter((h) => {
+        if (h.isLeaf)
+            return true;
+        // Cluster hit: keep only if none of its leaf children are present.
+        return !h.leafIds.some((lid) => leafIds.has(lid));
+    });
+}
+// ── S42A-5: Top-level pipeline ─────────────────────────────────────────────
+/**
+ * Full multi-level retrieval pipeline: score → expand → dedup → MMR → top-K.
+ * Drop-in replacement for `stagedExpansion()` in the RAPTOR recall path.
+ */
+export function multilevelRetrieval(query, tree, opts) {
+    if (!tree.rootId)
+        return [];
+    const { embedder } = opts;
+    const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
+    const leafExp = opts.leafExpansion !== false; // default true
+    const maxLeafExp = opts.maxLeafExpansion ?? 10;
+    const k = opts.k ?? 5;
+    const lambda = opts.mmrLambda ?? 0.5;
+    const qv = embedder.embed(query);
+    // 1. Score all nodes with level weights.
+    const scored = scoreTreeLevels(query, tree, { embedder, levelWeights: weights });
+    // 2. Top-N candidates for MMR diversity window.
+    const topN = scored.slice(0, k * 3);
+    // 3. Leaf expansion (optional).
+    const expanded = leafExp
+        ? expandLeafDescendants(topN, tree, maxLeafExp, embedder, qv, weights)
+        : topN;
+    // 4. Dedup: remove cluster hits when leaf children are present.
+    const deduped = deduplicateMultilevelHits(expanded);
+    // 5. MMR rerank to k.
+    const mmrItems = deduped.map((h) => ({
+        item: h,
+        vector: h.embedding,
+        relevance: h.score,
+    }));
+    return mmrRerank(mmrItems, k, lambda);
+}

package/dist/src/dedup/raptor/multilevel.test.js

@@ -0,0 +1,203 @@
+/**
+ * multilevel.test.ts — hermetic unit tests for S42A multi-level RAPTOR retrieval.
+ *
+ * Tests the scoreTreeLevels → expandLeafDescendants → deduplicateMultilevelHits
+ * → multilevelRetrieval pipeline. No network, no live store — uses TrigramEmbedder
+ * and synthetic RaptorTrees built from makeLeaves().
+ */
+import { test } from "node:test";
+import assert from "node:assert/strict";
+import { TrigramEmbedder } from "../../embedder.js";
+import { buildRaptorTree } from "./tree.js";
+import { scoreTreeLevels, expandLeafDescendants, deduplicateMultilevelHits, multilevelRetrieval, } from "./multilevel.js";
+function msg(text) {
+    return { role: "user", text };
+}
+/** Build N distinct leaves with deterministic content. */
+function makeLeaves(n, embedder = new TrigramEmbedder()) {
+    const leaves = [];
+    for (let i = 0; i < n; i++) {
+        const text = `topic ${i % 7}: the module ${i} validated the session token and refreshed the cache for region ${i}`;
+        leaves.push({
+            id: `leaf_${i}`,
+            messages: [msg(text)],
+            sourceText: text,
+            embedding: embedder.embed(text),
+        });
+    }
+    return leaves;
+}
+// ── test: scoreTreeLevels returns nodes at all levels ────────────────────────
+test("scoreTreeLevels returns results at multiple tree levels", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const hits = scoreTreeLevels("the auth module validates the session token", tree, {
+        embedder,
+    });
+    assert.ok(hits.length > 0, "should return hits");
+    // Should include both leaf (level 0) and cluster (level ≥ 1) hits.
+    const levels = new Set(hits.map((h) => h.level));
+    assert.ok(levels.has(0), "should have leaf-level hits");
+    if (tree.levels > 1) {
+        const hasCluster = [...levels].some((l) => l >= 1);
+        assert.ok(hasCluster, "should have cluster-level hits for multi-level tree");
+    }
+    // All hits should have valid scores.
+    for (const h of hits) {
+        assert.ok(h.score >= 0 && h.score <= 1, `score ${h.score} out of range`);
+        assert.ok(h.rawScore >= 0 && h.rawScore <= 1, `rawScore ${h.rawScore} out of range`);
+        assert.ok(h.score <= h.rawScore, "weighted score <= raw score (level weights ≤ 1)");
+    }
+    // Hits should be sorted by score descending.
+    for (let i = 1; i < hits.length; i++) {
+        assert.ok(hits[i - 1].score >= hits[i].score, `hits not sorted: hit[${i - 1}].score=${hits[i - 1].score} < hit[${i}].score=${hits[i].score}`);
+    }
+});
+// ── test: level weights affect scoring ───────────────────────────────────────
+test("level weights shift scores: higher weight for level → higher weighted score", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const query = "the auth module validates the session token";
+    // Uniform weights: all levels scored equally.
+    const uniform = scoreTreeLevels(query, tree, {
+        embedder,
+        levelWeights: [1.0, 1.0, 1.0, 1.0, 1.0],
+    });
+    // Penalized weights: higher levels penalized.
+    const penalized = scoreTreeLevels(query, tree, {
+        embedder,
+        levelWeights: [1.0, 0.1, 0.1, 0.1, 0.1],
+    });
+    // With penalized weights, cluster-level hits should score lower.
+    const clusterUniform = uniform.filter((h) => h.level >= 1);
+    const clusterPenalized = penalized.filter((h) => h.level >= 1);
+    if (clusterUniform.length > 0 && clusterPenalized.length > 0) {
+        const avgUniform = clusterUniform.reduce((s, h) => s + h.score, 0) / clusterUniform.length;
+        const avgPenalized = clusterPenalized.reduce((s, h) => s + h.score, 0) / clusterPenalized.length;
+        assert.ok(avgPenalized < avgUniform, `penalized cluster avg (${avgPenalized.toFixed(3)}) should be < uniform (${avgUniform.toFixed(3)})`);
+    }
+});
+// ── test: expandLeafDescendants adds leaf hits ──────────────────────────────
+test("expandLeafDescendants adds leaf descendants for cluster hits", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const query = "the auth module validates the session token";
+    const scored = scoreTreeLevels(query, tree, { embedder });
+    const qv = embedder.embed(query);
+    // Take only cluster hits (level ≥ 1).
+    const clusterHits = scored.filter((h) => !h.isLeaf).slice(0, 3);
+    assert.ok(clusterHits.length > 0, "should have cluster hits");
+    const expanded = expandLeafDescendants(clusterHits, tree, 5, // maxPerCluster
+    embedder, qv);
+    // Expanded set should include leaf hits.
+    const leafHits = expanded.filter((h) => h.isLeaf);
+    assert.ok(leafHits.length > 0, "should have expanded leaf hits");
+    assert.ok(expanded.length > clusterHits.length, `expanded (${expanded.length}) should be > cluster hits (${clusterHits.length})`);
+    // No duplicate ids.
+    const ids = new Set(expanded.map((h) => h.nodeId));
+    assert.equal(ids.size, expanded.length, "no duplicate node ids");
+});
+// ── test: deduplicateMultilevelHits removes cluster when leaves present ─────
+test("deduplicateMultilevelHits removes cluster hits when leaf children are present", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const query = "the auth module validates the session token";
+    const scored = scoreTreeLevels(query, tree, { embedder });
+    // Take a cluster hit and its leaf children.
+    const cluster = scored.find((h) => !h.isLeaf);
+    assert.ok(cluster, "should have a cluster hit");
+    const leafChildren = cluster.leafIds.slice(0, 2).map((lid) => ({
+        nodeId: lid,
+        level: 0,
+        score: 0.5,
+        rawScore: 0.5,
+        isLeaf: true,
+        leafIds: [lid],
+        summary: "",
+        embedding: cluster.embedding,
+    }));
+    const mixed = [cluster, ...leafChildren];
+    const deduped = deduplicateMultilevelHits(mixed);
+    // Cluster should be removed because its leaf children are present.
+    assert.ok(!deduped.find((h) => h.nodeId === cluster.nodeId), "cluster hit should be removed when leaf children are present");
+    assert.equal(deduped.length, leafChildren.length, "only leaf hits remain");
+});
+test("deduplicateMultilevelHits keeps cluster hits when no leaf children present", () => {
+    const clusterHit = {
+        nodeId: "cluster_1",
+        level: 1,
+        score: 0.8,
+        rawScore: 0.8,
+        isLeaf: false,
+        leafIds: ["leaf_1", "leaf_2"],
+        summary: "summarized content",
+        embedding: [1, 0, 0],
+    };
+    const deduped = deduplicateMultilevelHits([clusterHit]);
+    assert.equal(deduped.length, 1, "cluster should be kept when no leaf children present");
+    assert.equal(deduped[0].nodeId, "cluster_1");
+});
+// ── test: multilevelRetrieval returns cluster + leaf mix ────────────────────
+test("multilevelRetrieval returns a mix of cluster and leaf hits", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const hits = multilevelRetrieval("the auth module validates the session token", tree, {
+        embedder,
+        k: 5,
+        leafExpansion: true,
+        maxLeafExpansion: 3,
+    });
+    assert.ok(hits.length > 0, "should return hits");
+    assert.ok(hits.length <= 5, "should respect k=5");
+    // Should include leaf hits.
+    const leafHits = hits.filter((h) => h.isLeaf);
+    assert.ok(leafHits.length > 0, "should include leaf hits");
+    // All hits should have valid properties.
+    for (const h of hits) {
+        assert.ok(h.nodeId, "hit should have nodeId");
+        assert.ok(typeof h.level === "number", "hit should have numeric level");
+        assert.ok(h.score >= 0, "hit score should be non-negative");
+        assert.ok(Array.isArray(h.leafIds), "hit should have leafIds array");
+    }
+});
+test("multilevelRetrieval respects k parameter", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(100);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 8 });
+    for (const k of [1, 3, 5, 10]) {
+        const hits = multilevelRetrieval("the auth module validates the session token", tree, {
+            embedder,
+            k,
+        });
+        assert.ok(hits.length <= k, `k=${k}: got ${hits.length} hits, should be ≤ ${k}`);
+    }
+});
+test("multilevelRetrieval returns empty for empty tree", () => {
+    const embedder = new TrigramEmbedder();
+    const emptyTree = { nodes: new Map(), rootId: null, levels: 0, timedOut: false };
+    const hits = multilevelRetrieval("any query", emptyTree, { embedder });
+    assert.deepEqual(hits, []);
+});
+test("multilevelRetrieval with leafExpansion=false skips leaf expansion", () => {
+    const embedder = new TrigramEmbedder();
+    const leaves = makeLeaves(50);
+    const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
+    const hits = multilevelRetrieval("the auth module validates the session token", tree, {
+        embedder,
+        k: 5,
+        leafExpansion: false,
+    });
+    assert.ok(hits.length > 0, "should still return hits");
+    // With leaf expansion off and tree having multiple levels, we may get
+    // cluster hits that are NOT expanded. The mix depends on tree structure.
+    // Just verify we got valid results.
+    for (const h of hits) {
+        assert.ok(h.nodeId, "hit should have nodeId");
+        assert.ok(h.score >= 0, "hit score should be non-negative");
+    }
+});

package/dist/src/dedup/raptor/promote.test.js

@@ -11,7 +11,7 @@ import assert from "node:assert/strict";
 import { mkdtempSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { VectorStore } from "../../vectorStore.js";
+import { VectorStore, vectorList, vectorSearch } from "../../vectorStore.js";
 import { runRaptor } from "./index.js";
 import { compactSession } from "../../engine.js";
 import { Logger } from "../../log.js";
@@ -35,10 +35,10 @@ test("Fix D: vectorStore.search serves a persisted RAPTOR tree (broader recall)"
     }
     // No tree yet → flat search only, returns hits, no RAPTOR coverage.
     assert.equal(listRaptorNodes(SESS, stateDir).length, 0, "no tree initially");
-    const flat = s.search(SESS, "alpha wire bootstrap", 3);
+    const flat = vectorSearch(s, SESS, "alpha wire bootstrap", 3);
     assert.ok(flat.length > 0, "flat search returns hits");
     // Build + persist a RAPTOR tree for the session (mirrors runCompact refresh).
-    const all = s.list(SESS);
+    const all = vectorList(s, SESS);
     const leaves = all.map((cp) => ({
         id: cp.checkpointId,
         messages: [],
@@ -49,7 +49,7 @@ test("Fix D: vectorStore.search serves a persisted RAPTOR tree (broader recall)"
     assert.ok(tree && listRaptorNodes(SESS, stateDir).length > 0, "tree persisted");
     // With the tree live + RAPTOR_ENABLED, search still returns hits and now
     // exercises the RAPTOR-served path without regression.
-    const withTree = s.search(SESS, "alpha wire bootstrap", 3);
+    const withTree = vectorSearch(s, SESS, "alpha wire bootstrap", 3);
     assert.ok(withTree.length > 0, "search returns hits with RAPTOR promoted");
     // Every returned hit is a real checkpoint in the session.
     for (const h of withTree) {
@@ -60,7 +60,7 @@ test("Fix D: search still works for a session with <2 leaves (no tree)", () => {
     const stateDir = join(baseTmp, `run-${counter++}`);
     const s = new VectorStore({ dedupSim: 0.9, stateDir, config: raptorConfig() });
     compactSession({ sessionId: SESS, messages: [msg("only one topic here"), msg("ok", "Edit")], keepFrom: 2, timestamp: 1 }, s);
-    const r = s.search(SESS, "only one topic", 3);
+    const r = vectorSearch(s, SESS, "only one topic", 3);
     assert.ok(r.length > 0, "single-checkpoint search still works (no tree)");
     assert.equal(listRaptorNodes(SESS, stateDir).length, 0, "no tree built for <2 leaves");
 });

package/dist/src/dedup/raptor/retrieval.js

@@ -18,7 +18,7 @@ function isLeafId(id, tree) {
     return !tree.nodes.has(id);
 }
 /** All leaf (raw) ids reachable beneath a node via BFS. */
-function leafDescendants(node, tree) {
+export function leafDescendants(node, tree) {
     const out = [];
     const queue = [node];
     while (queue.length) {

package/dist/src/dedup/sprint12.test.js

@@ -3,7 +3,7 @@ import assert from "node:assert/strict";
 import { mkdtempSync, rmSync } from "node:fs";
 import { tmpdir } from "node:os";
 import { join } from "node:path";
-import { VectorStore, L2_ENABLED } from "../vectorStore.js";
+import { VectorStore, L2_ENABLED, vectorSemDedup, vectorList, vectorSearch } from "../vectorStore.js";
 import { mmrRerank } from "./mmr.js";
 import { topK } from "./topk.js";
 import { cosineSimilarity, defaultEmbedder } from "../embedder.js";
@@ -67,7 +67,7 @@ test("L2_ENABLED defaults true; search still returns hits", () => {
     assert.equal(L2_ENABLED, true);
     const s = store();
     s.add({ sessionId: "sess_l2", summary: "investigated the parser", regionText: "investigated src/parser.ts and added a tokenizer", timestamp: 1 });
-    const hits = s.search("sess_l2", "src/parser.ts tokenizer", 3);
+    const hits = vectorSearch(s, "sess_l2", "src/parser.ts tokenizer", 3);
     assert.ok(hits.length >= 1);
 });
 test("L2_ENABLED=false skips semantic tier but L0/L1 still work", () => {
@@ -106,14 +106,14 @@ test("semDedup marks redundant near-identical rows 'removed' and search excludes
         { id: "chkpt_001", text: "the cache stores parsed ast nodes for fast lookup", tok: 100 },
         { id: "chkpt_002", text: "the cache stores parsed ast nodes for fast lookup and reuse", tok: 900 },
     ]);
-    const removed = s.semDedup("sess_sd", 0.85);
+    const removed = vectorSemDedup(s, "sess_sd", 0.85);
     assert.equal(removed, 1);
-    const st = s.list("sess_sd");
+    const st = vectorList(s, "sess_sd");
     const dropped = st.find((c) => c.dedupStatus === "removed");
     assert.ok(dropped);
     assert.equal(dropped.checkpointId, "chkpt_001"); // lower tokenEstimate removed
     // Search excludes the removed row (only one active remains).
-    const hits = s.search("sess_sd", "cache parsed ast nodes", 5);
+    const hits = vectorSearch(s, "sess_sd", "cache parsed ast nodes", 5);
     assert.equal(hits.length, 1);
 });
 test("semDedup is idempotent (re-run removes nothing new)", () => {
@@ -123,8 +123,8 @@ test("semDedup is idempotent (re-run removes nothing new)", () => {
         { id: "chkpt_001", text: "identical region text for the dedup job now", tok: 100 },
         { id: "chkpt_002", text: "identical region text for the dedup job right now", tok: 200 },
     ]);
-    const first = s.semDedup("sess_sd2", 0.85);
-    const second = s.semDedup("sess_sd2", 0.85);
+    const first = vectorSemDedup(s, "sess_sd2", 0.85);
+    const second = vectorSemDedup(s, "sess_sd2", 0.85);
     assert.equal(first, 1);
     assert.equal(second, 0);
 });