pi-mega-compact 0.8.22 → 0.8.24

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

@@ -0,0 +1,291 @@
+/**
+ * openclaw-mega-compact — OpenClaw plugin adapter for the pi-mega-compact engine.
+ *
+ * Wires the pi-agnostic Trident engine (src/) into OpenClaw's plugin lifecycle:
+ *  - Registers a CompactionProvider that replaces the built-in summarizeInStages.
+ *  - Exposes `mega_status` and `mega_recall` tools for on-demand inspection.
+ *  - Hooks into `before_compaction` / `after_compaction` for diagnostics.
+ *
+ * Design constraints:
+ *  - NO imports from `@earendil-works/pi-coding-agent` or pi-agent-core.
+ *  - The engine core (src/) is pi-agnostic; this file is the sole OpenClaw boundary.
+ *  - No network at runtime — everything is local (stores + extractive summarizer).
+ */
+import { definePluginEntry } from "openclaw/plugin-sdk/plugin-entry";
+import { compactSession, setDefaultStore, } from "../src/engine.js";
+import { recallAndInline } from "../src/recall.js";
+import { VectorStore } from "../src/vectorStore.js";
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+const PLUGIN_ID = "mega-compact";
+const PLUGIN_LABEL = "Mega Compact (Trident)";
+/** Default state directory for vector store persistence. */
+const STATE_DIR = process.env.MEGA_COMPACT_STATE_DIR ?? undefined;
+/** Minimum messages before we bother compacting. */
+const MIN_MESSAGES_FOR_COMPACT = 6;
+// ---------------------------------------------------------------------------
+// Message conversion — OpenClaw unknown[] → EngineMessage[]
+// ---------------------------------------------------------------------------
+/**
+ * Best-effort conversion from OpenClaw's opaque message array to our
+ * EngineMessage shape. OpenClaw messages are typed as `unknown[]` so we
+ * handle whatever shape comes through gracefully.
+ */
+function toEngineMessages(messages) {
+    return messages.map((msg) => {
+        if (!msg || typeof msg !== "object") {
+            // Primitive fallback — treat as custom text.
+            return {
+                role: "custom",
+                text: String(msg ?? ""),
+            };
+        }
+        const m = msg;
+        const role = typeof m.role === "string" ? m.role : "custom";
+        // Normalize role to one of our four engine roles.
+        let engineRole;
+        switch (role) {
+            case "user":
+                engineRole = "user";
+                break;
+            case "assistant":
+                engineRole = "assistant";
+                break;
+            case "tool":
+            case "function":
+                engineRole = "tool";
+                break;
+            default:
+                engineRole = "custom";
+                break;
+        }
+        // Extract text content from common message shapes.
+        const text = typeof m.content === "string"
+            ? m.content
+            : typeof m.text === "string"
+                ? m.text
+                : Array.isArray(m.content)
+                    ? m.content
+                        .filter((part) => part.type === "text" && typeof part.text === "string")
+                        .map((part) => part.text)
+                        .join("\n")
+                    : "";
+        // Preserve tool metadata when present.
+        const toolName = typeof m.name === "string"
+            ? m.name
+            : typeof m.toolName === "string"
+                ? m.toolName
+                : undefined;
+        const input = typeof m.input === "string"
+            ? m.input
+            : typeof m.arguments === "string"
+                ? m.arguments
+                : m.arguments !== undefined
+                    ? JSON.stringify(m.arguments)
+                    : undefined;
+        const output = typeof m.output === "string"
+            ? m.output
+            : engineRole === "tool" && typeof m.content === "string"
+                ? m.content
+                : undefined;
+        return { role: engineRole, text, toolName, input, output };
+    });
+}
+// ---------------------------------------------------------------------------
+// Compaction provider
+// ---------------------------------------------------------------------------
+function createCompactionProvider(store) {
+    return {
+        id: PLUGIN_ID,
+        label: PLUGIN_LABEL,
+        async summarize({ messages, signal, compressionRatio, }) {
+            // Abort check — bail early if the caller cancelled.
+            if (signal?.aborted) {
+                throw new DOMException("Aborted", "AbortError");
+            }
+            const engineMessages = toEngineMessages(messages);
+            // Nothing meaningful to compact.
+            if (engineMessages.length < MIN_MESSAGES_FOR_COMPACT) {
+                return "";
+            }
+            // Map compression ratio → keepFrom boundary.
+            // compressionRatio=0.5 means "compact the oldest 50%".
+            // Default to compacting the oldest half if not specified.
+            const ratio = compressionRatio ?? 0.5;
+            const keepFrom = Math.max(MIN_MESSAGES_FOR_COMPACT, Math.floor(engineMessages.length * (1 - ratio)));
+            // Abort check after conversion (conversion is cheap but check anyway).
+            if (signal?.aborted) {
+                throw new DOMException("Aborted", "AbortError");
+            }
+            const sessionId = `openclaw-${Date.now()}`;
+            const input = {
+                sessionId,
+                messages: engineMessages,
+                keepFrom,
+            };
+            const result = compactSession(input, store);
+            if (result.skipped) {
+                return "";
+            }
+            return result.summary;
+        },
+    };
+}
+// ---------------------------------------------------------------------------
+// Plugin entry
+// ---------------------------------------------------------------------------
+export default definePluginEntry({
+    id: PLUGIN_ID,
+    name: "Mega Compact",
+    description: "Layered, local, vector-backed context compressor (Trident engine) for OpenClaw compaction.",
+    register(api) {
+        const logger = api.logger;
+        // Resolve state directory — prefer plugin config override.
+        const pluginCfg = (api.pluginConfig ?? {});
+        const stateDir = typeof pluginCfg.stateDir === "string" && pluginCfg.stateDir.length > 0
+            ? pluginCfg.stateDir
+            : STATE_DIR;
+        // Initialize vector store.
+        let store;
+        try {
+            store = new VectorStore({ stateDir });
+            setDefaultStore(store);
+            logger.info?.(`${PLUGIN_ID}: vector store initialized (stateDir=${stateDir ?? "default"})`);
+        }
+        catch (err) {
+            logger.error?.(`${PLUGIN_ID}: failed to init vector store:`, err);
+            return; // Hard bail — no point registering if store is broken.
+        }
+        // -----------------------------------------------------------------------
+        // Register compaction provider
+        // -----------------------------------------------------------------------
+        const provider = createCompactionProvider(store);
+        api.registerCompactionProvider(provider);
+        logger.info?.(`${PLUGIN_ID}: registered compaction provider "${provider.id}"`);
+        // -----------------------------------------------------------------------
+        // Hooks — before / after compaction diagnostics
+        // -----------------------------------------------------------------------
+        api.registerHook({
+            event: "before_compaction",
+            handler: async (ctx) => {
+                const msgCount = Array.isArray(ctx?.messages) ? ctx.messages.length : 0;
+                logger.info?.(`${PLUGIN_ID}: before_compaction — ${msgCount} messages in scope`);
+            },
+        });
+        api.registerHook({
+            event: "after_compaction",
+            handler: async (ctx) => {
+                const summaryLen = typeof ctx?.summary === "string" ? ctx.summary.length : 0;
+                logger.info?.(`${PLUGIN_ID}: after_compaction — summary ${summaryLen} chars`);
+            },
+        });
+        // -----------------------------------------------------------------------
+        // Tool: mega_status
+        // -----------------------------------------------------------------------
+        api.registerTool({
+            name: "mega_status",
+            description: "Show the current status of the mega-compact engine: vector store stats, checkpoint count, and recent compaction activity.",
+            parameters: {
+                type: "object",
+                properties: {
+                    sessionId: {
+                        type: "string",
+                        description: "Optional session ID to scope stats to.",
+                    },
+                },
+                additionalProperties: false,
+            },
+            handler: async (args) => {
+                const sessionId = args?.sessionId ?? "global";
+                try {
+                    const stats = store.stats(sessionId);
+                    const parts = [
+                        `**Mega Compact Status**`,
+                        `Session: ${sessionId}`,
+                        `Checkpoints: ${stats.checkpointCount}`,
+                        `Total tokens saved: ${stats.totalTokenEstimate}`,
+                        `Last checkpoint: ${stats.lastCheckpointId ?? "—"}`,
+                        `Injected count: ${stats.injectedCount}`,
+                        `Dedup hit rate: ${(stats.dedupHitRate * 100).toFixed(0)}%`,
+                    ];
+                    if (stats.lastSummary) {
+                        parts.push(`\nLast summary (truncated):\n  ${stats.lastSummary.slice(0, 120).replace(/\n/g, " ")}…`);
+                    }
+                    return { content: [{ type: "text", text: parts.join("\n") }] };
+                }
+                catch (err) {
+                    return {
+                        content: [{ type: "text", text: `Error reading mega-compact status: ${err}` }],
+                        isError: true,
+                    };
+                }
+            },
+        });
+        // -----------------------------------------------------------------------
+        // Tool: mega_recall
+        // -----------------------------------------------------------------------
+        api.registerTool({
+            name: "mega_recall",
+            description: "Recall and inline relevant context from the mega-compact vector store for the current session.",
+            parameters: {
+                type: "object",
+                properties: {
+                    sessionId: {
+                        type: "string",
+                        description: "Session ID to recall context for.",
+                    },
+                    query: {
+                        type: "string",
+                        description: "Natural language query for relevant context.",
+                    },
+                    limit: {
+                        type: "number",
+                        description: "Max checkpoints to recall (default 3).",
+                    },
+                },
+                required: ["sessionId", "query"],
+                additionalProperties: false,
+            },
+            handler: async (args) => {
+                const { sessionId, query, limit } = args;
+                if (!sessionId || !query) {
+                    return {
+                        content: [{ type: "text", text: "Both `sessionId` and `query` are required." }],
+                        isError: true,
+                    };
+                }
+                try {
+                    const result = recallAndInline({ sessionId, query, limit: limit ?? 3, source: "command", skipInjected: false }, store);
+                    if (result.toInject.length === 0) {
+                        return {
+                            content: [{ type: "text", text: "No relevant context found in the mega-compact store." }],
+                        };
+                    }
+                    const parts = [
+                        `**Recalled ${result.toInject.length} checkpoint(s):**`,
+                        ...result.report,
+                        "",
+                        "---",
+                        result.block,
+                    ];
+                    return { content: [{ type: "text", text: parts.join("\n") }] };
+                }
+                catch (err) {
+                    return {
+                        content: [{ type: "text", text: `Error during mega-recall: ${err}` }],
+                        isError: true,
+                    };
+                }
+            },
+        });
+        // -----------------------------------------------------------------------
+        // Cleanup on shutdown
+        // -----------------------------------------------------------------------
+        api.on("shutdown", () => {
+            logger.info?.(`${PLUGIN_ID}: shutting down — clearing default store`);
+            setDefaultStore(undefined);
+        });
+        logger.info?.(`${PLUGIN_ID}: plugin registered (tools: mega_status, mega_recall)`);
+    },
+});

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/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) {