From 3c87fb631b632b59328d14a22a7e694b84288d03 Mon Sep 17 00:00:00 2001 From: Vellum Assistant Date: Mon, 25 May 2026 01:54:11 -0500 Subject: [PATCH] feat(memory-v3): parallel-fan-out traversal with cycle/visited guards --- .../src/memory/v3/__tests__/traversal.test.ts | 395 ++++++++++++++++++ assistant/src/memory/v3/traversal.ts | 194 +++++++++ 2 files changed, 589 insertions(+) create mode 100644 assistant/src/memory/v3/__tests__/traversal.test.ts create mode 100644 assistant/src/memory/v3/traversal.ts diff --git a/assistant/src/memory/v3/__tests__/traversal.test.ts b/assistant/src/memory/v3/__tests__/traversal.test.ts new file mode 100644 index 00000000000..4a742815a94 --- /dev/null +++ b/assistant/src/memory/v3/__tests__/traversal.test.ts @@ -0,0 +1,395 @@ +/** + * Tests for `assistant/src/memory/v3/traversal.ts`. + * + * Provider-free: `descend` is always a deterministic stub. Coverage: + * - resolveChildren is a thin accessor (known node / leaf / unknown id). + * - linear descent collects the expected leaf pages and emits a TreeLevel per + * walked node in walk order. + * - a DAG (sub-node shared by two parents) is walked exactly once. + * - an injected cycle (A ↔ B) terminates. + * - breadthBudget caps the descents per level. + * - maxDepth halts the recursion at the right level. + * - seeds start the walk mid-tree (alongside / instead of the root). + * - reasoning from the descend result is threaded onto the level; defaults + * to "" when omitted. + * + * Fixtures are plain in-memory `TreeIndex` objects — no disk, no workspace. + */ + +import { describe, expect, test } from "bun:test"; + +import type { DescendResult } from "../traversal.js"; +import { resolveChildren, walkTree } from "../traversal.js"; +import type { ChildRef, TreeIndex } from "../tree-index.js"; + +// --------------------------------------------------------------------------- +// Fixture helpers +// --------------------------------------------------------------------------- + +function page(ref: string): ChildRef { + return { kind: "page", ref }; +} + +function node(ref: string): ChildRef { + return { kind: "node", ref }; +} + +/** + * Build a minimal in-memory `TreeIndex` from a forward-adjacency spec. Only + * `childrenByNode` and `root` are exercised by the traversal, so the reverse + * adjacency maps and `nodes` are left empty — the walk never reads them. + */ +function makeTree( + root: string, + childrenByNode: Record, +): TreeIndex { + return { + nodes: new Map(), + childrenByNode: new Map(Object.entries(childrenByNode)), + parentsByNode: new Map(), + pageParents: new Map(), + root, + }; +} + +/** Descend into every node child offered (mechanical "descend all" stub). */ +function descendAll( + _nodeId: string, + children: ReadonlyArray, +): DescendResult { + return { descend: children.filter((c) => c.kind === "node") }; +} + +// --------------------------------------------------------------------------- +// resolveChildren +// --------------------------------------------------------------------------- + +describe("resolveChildren", () => { + test("returns the ordered child refs for a known node", () => { + const tree = makeTree("_root", { + _root: [node("a"), page("p")], + }); + expect(resolveChildren(tree, "_root")).toEqual([node("a"), page("p")]); + }); + + test("returns [] for a leaf / unknown node id", () => { + const tree = makeTree("_root", { _root: [] }); + expect(resolveChildren(tree, "missing")).toEqual([]); + }); +}); + +// --------------------------------------------------------------------------- +// Linear descent +// --------------------------------------------------------------------------- + +describe("walkTree — linear descent", () => { + test("collects expected leaf pages and emits a level per walked node", async () => { + // _root → node:a → node:b → page:leaf (plus a page on each level) + const tree = makeTree("_root", { + _root: [page("p-root"), node("a")], + a: [page("p-a"), node("b")], + b: [page("leaf")], + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 8, + descend: descendAll, + }); + + expect([...pages].sort()).toEqual(["leaf", "p-a", "p-root"]); + expect(levels.map((l) => l.node)).toEqual(["_root", "a", "b"]); + + expect(levels[0]).toMatchObject({ + node: "_root", + considered: ["a"], + descended: ["a"], + skipped: [], + reasoning: "", + }); + expect(levels[2]).toMatchObject({ + node: "b", + considered: [], + descended: [], + skipped: [], + }); + }); + + test("defaults start to tree.root", async () => { + const tree = makeTree("home", { + home: [page("only")], + }); + const { pages, levels } = await walkTree(tree, { + breadthBudget: 4, + maxDepth: 4, + descend: descendAll, + }); + expect([...pages]).toEqual(["only"]); + expect(levels.map((l) => l.node)).toEqual(["home"]); + }); +}); + +// --------------------------------------------------------------------------- +// DAG dedup +// --------------------------------------------------------------------------- + +describe("walkTree — DAG dedup", () => { + test("a sub-node shared by two parents is walked exactly once", async () => { + // _root → {node:left, node:right}; both → node:shared → page:s + const tree = makeTree("_root", { + _root: [node("left"), node("right")], + left: [node("shared")], + right: [node("shared")], + shared: [page("s")], + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 8, + descend: descendAll, + }); + + expect([...pages]).toEqual(["s"]); + // `shared` appears once even though both left and right descend into it. + const walked = levels.map((l) => l.node); + expect(walked.filter((n) => n === "shared")).toHaveLength(1); + expect(walked.sort()).toEqual(["_root", "left", "right", "shared"]); + }); +}); + +// --------------------------------------------------------------------------- +// Cycle termination +// --------------------------------------------------------------------------- + +describe("walkTree — cycle termination", () => { + test("an injected A ↔ B cycle terminates and walks each once", async () => { + const tree = makeTree("a", { + a: [node("b"), page("pa")], + b: [node("a"), page("pb")], + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 100, + descend: descendAll, + }); + + expect([...pages].sort()).toEqual(["pa", "pb"]); + const walked = levels.map((l) => l.node).sort(); + expect(walked).toEqual(["a", "b"]); + }); + + test("a self-loop terminates", async () => { + const tree = makeTree("solo", { + solo: [node("solo"), page("p")], + }); + const { pages, levels } = await walkTree(tree, { + breadthBudget: 4, + maxDepth: 100, + descend: descendAll, + }); + expect([...pages]).toEqual(["p"]); + expect(levels.map((l) => l.node)).toEqual(["solo"]); + }); +}); + +// --------------------------------------------------------------------------- +// Breadth budget +// --------------------------------------------------------------------------- + +describe("walkTree — breadthBudget", () => { + test("caps the descents per node and records the rest as skipped", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b"), node("c"), node("d")], + a: [page("pa")], + b: [page("pb")], + c: [page("pc")], + d: [page("pd")], + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 2, + maxDepth: 8, + descend: descendAll, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.considered).toEqual(["a", "b", "c", "d"]); + expect(rootLevel.descended).toEqual(["a", "b"]); + expect(rootLevel.skipped).toEqual(["c", "d"]); + + // Only the first two children's pages are reached. + expect([...pages].sort()).toEqual(["pa", "pb"]); + expect(levels.map((l) => l.node).sort()).toEqual(["_root", "a", "b"]); + }); +}); + +// --------------------------------------------------------------------------- +// Depth budget +// --------------------------------------------------------------------------- + +describe("walkTree — maxDepth", () => { + test("halts recursion at the configured depth", async () => { + // _root(0) → a(1) → b(2) → c(3) + const tree = makeTree("_root", { + _root: [node("a")], + a: [node("b"), page("pa")], + b: [node("c"), page("pb")], + c: [page("pc")], + }); + + // maxDepth 1 walks depth 0 (_root) and depth 1 (a) only; b/c never walked. + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 1, + descend: descendAll, + }); + + expect(levels.map((l) => l.node)).toEqual(["_root", "a"]); + // `a`'s page is collected; b/c and their pages are not reached. + expect([...pages]).toEqual(["pa"]); + }); + + test("maxDepth 0 walks only the start level", async () => { + const tree = makeTree("_root", { + _root: [node("a"), page("pr")], + a: [page("pa")], + }); + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 0, + descend: descendAll, + }); + expect(levels.map((l) => l.node)).toEqual(["_root"]); + expect([...pages]).toEqual(["pr"]); + }); +}); + +// --------------------------------------------------------------------------- +// Seeds +// --------------------------------------------------------------------------- + +describe("walkTree — seeds", () => { + test("seeds start the walk mid-tree alongside start", async () => { + const tree = makeTree("_root", { + _root: [node("a"), page("pr")], + a: [page("pa")], + mid: [page("pm"), node("deep")], + deep: [page("pd")], + }); + + const { pages, levels } = await walkTree(tree, { + seeds: ["mid"], + breadthBudget: 8, + maxDepth: 8, + descend: descendAll, + }); + + // Both the root branch and the seeded `mid` subtree are explored. + expect([...pages].sort()).toEqual(["pa", "pd", "pm", "pr"]); + expect(levels.map((l) => l.node).sort()).toEqual([ + "_root", + "a", + "deep", + "mid", + ]); + }); + + test("a node that is both start and seed is walked once", async () => { + const tree = makeTree("dup", { + dup: [page("p")], + }); + const { levels } = await walkTree(tree, { + start: "dup", + seeds: ["dup"], + breadthBudget: 4, + maxDepth: 4, + descend: descendAll, + }); + expect(levels.map((l) => l.node)).toEqual(["dup"]); + }); +}); + +// --------------------------------------------------------------------------- +// Descend decision threading +// --------------------------------------------------------------------------- + +describe("walkTree — descend decision", () => { + test("threads the descend reasoning onto the level", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [page("pa")], + b: [page("pb")], + }); + + const descend = ( + _nodeId: string, + children: ReadonlyArray, + ): DescendResult => ({ + // Pick only "a". + descend: children.filter((c) => c.kind === "node" && c.ref === "a"), + reasoning: "a is more relevant", + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 8, + descend, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.reasoning).toBe("a is more relevant"); + expect(rootLevel.descended).toEqual(["a"]); + expect(rootLevel.skipped).toEqual(["b"]); + expect([...pages]).toEqual(["pa"]); + }); + + test("ignores descend picks that were not offered node children", async () => { + const tree = makeTree("_root", { + _root: [node("a"), page("pr")], + a: [page("pa")], + }); + + const descend = (): DescendResult => ({ + // "ghost" was never offered; "pr" is a page, not a node child. + descend: [node("ghost"), page("pr")], + }); + + const { pages, levels } = await walkTree(tree, { + breadthBudget: 8, + maxDepth: 8, + descend, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.considered).toEqual(["a"]); + expect(rootLevel.descended).toEqual([]); + expect(rootLevel.skipped).toEqual(["a"]); + // No node descent happened; only the root's own page is collected. + expect([...pages]).toEqual(["pr"]); + expect(levels.map((l) => l.node)).toEqual(["_root"]); + }); + + test("dedups repeated descend picks before applying breadthBudget", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [page("pa")], + b: [page("pb")], + }); + + const descend = (): DescendResult => ({ + // "a" repeated should count once; budget of 2 then still admits "b". + descend: [node("a"), node("a"), node("b")], + }); + + const { levels } = await walkTree(tree, { + breadthBudget: 2, + maxDepth: 8, + descend, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.descended).toEqual(["a", "b"]); + expect(rootLevel.skipped).toEqual([]); + }); +}); diff --git a/assistant/src/memory/v3/traversal.ts b/assistant/src/memory/v3/traversal.ts new file mode 100644 index 00000000000..2cd625f0a32 --- /dev/null +++ b/assistant/src/memory/v3/traversal.ts @@ -0,0 +1,194 @@ +/** + * Memory v3 — Tree traversal primitives. + * + * The *mechanical* half of the v3 read loop: a deterministic, provider-free + * walk over the {@link TreeIndex} DAG. The intelligence — *which* child nodes + * to recurse into at each level — is injected via the `descend` callback so + * this module stays pure and unit-testable without an LLM. The driver PR wires + * `descend` to the model's descend/skip decision; here `descend` is just a + * function `(nodeId, children) => chosen node-children`. + * + * `walkTree` fans out from a `start` node and any `seeds`, level by level: + * - At each node it resolves the ordered child refs, hands them to `descend`, + * and recurses into the chosen `node:` children (capped by `breadthBudget`). + * - Every `page:` child encountered anywhere in the walk is collected into the + * returned `pages` set — pages are leaves, never recursed into. + * - A `visited` set keyed by canonical id (`node:`) dedups shared + * sub-nodes (the DAG case) and terminates cycles (A ↔ B). A node is walked + * at most once regardless of how many parents reference it. + * - `maxDepth` bounds how deep the recursion goes; the start/seed level is + * depth 0. + * + * Each walked node emits one {@link TreeLevel} (the `harness/trace.ts` shape) + * recording what was considered, descended, and skipped. `reasoning` is + * supplied by the `descend` callback (the driver attaches the model's stated + * reason); the mechanical walk defaults it to `""`. + * + * Processing is strictly level-by-level so `visited` mutations are never raced: + * within a level the per-node `descend` calls run concurrently (`Promise.all`), + * but the chosen children for the *next* level are only dedup'd and enqueued + * after the whole level resolves. + */ + +import type { TreeLevel } from "../v2/harness/trace.js"; +import type { ChildRef, TreeIndex } from "./tree-index.js"; + +/** + * The descend decision injected into {@link walkTree}. Given a node id and its + * ordered child refs, return the subset of *node* children to recurse into. The + * driver PR wires this to the LLM; tests pass a deterministic stub. + * + * Returning a `reasoning` string is optional — when present it is threaded into + * the emitted {@link TreeLevel}; absent, the level's `reasoning` defaults to + * `""`. Returned refs that are not `node:` children of `nodeId`, or that repeat, + * are ignored by the walk (it only recurses into distinct node children it + * actually offered). + */ +export type DescendDecision = ( + nodeId: string, + children: ReadonlyArray, +) => Promise | DescendResult; + +/** + * The result of a {@link DescendDecision}. `descend` lists the `node:` children + * chosen for recursion; `reasoning` is the optional model rationale recorded on + * the level. + */ +export interface DescendResult { + descend: ChildRef[]; + reasoning?: string; +} + +/** Options controlling a {@link walkTree} run. */ +export interface WalkOptions { + /** Entry node id; defaults to `tree.root`. */ + start?: string; + /** Extra node ids to start from in parallel with `start`. */ + seeds?: string[]; + /** Max `node:` children to descend into per node (after the `descend` pick). */ + breadthBudget: number; + /** Max recursion depth; the start/seed level is depth 0. */ + maxDepth: number; + /** Injected descend decision (the LLM hook). */ + descend: DescendDecision; +} + +/** The result of a {@link walkTree} run. */ +export interface WalkResult { + /** Every `page:` slug encountered across the walk, dedup'd. */ + pages: Set; + /** One {@link TreeLevel} per walked node, in walk order. */ + levels: TreeLevel[]; +} + +/** + * Resolve the ordered child refs for `nodeId`. Thin accessor over + * `tree.childrenByNode`; returns an empty array for an unknown / leaf node id so + * callers never branch on `undefined`. + */ +export function resolveChildren( + tree: TreeIndex, + nodeId: string, +): ReadonlyArray { + return tree.childrenByNode.get(nodeId) ?? []; +} + +/** Canonical visited-set key for a node id. */ +function nodeKey(nodeId: string): string { + return `node:${nodeId}`; +} + +/** + * Walk the {@link TreeIndex} DAG from `start` (default `tree.root`) plus any + * `seeds`, driven by the injected `descend` decision. Deterministic and + * provider-free — see the module docstring for the full contract. + * + * Returns the collected leaf `pages` and the per-node `levels` trace. + */ +export async function walkTree( + tree: TreeIndex, + opts: WalkOptions, +): Promise { + const { breadthBudget, maxDepth, descend } = opts; + const start = opts.start ?? tree.root; + + const pages = new Set(); + const levels: TreeLevel[] = []; + const visited = new Set(); + + // Seed the frontier with `start` + `seeds`, dedup'd and marked visited up + // front so a node that is both the start and a seed is walked once. + let frontier: string[] = []; + for (const id of [start, ...(opts.seeds ?? [])]) { + const key = nodeKey(id); + if (visited.has(key)) continue; + visited.add(key); + frontier.push(id); + } + + // Depth 0 is the start/seed level; stop once we'd exceed `maxDepth`. + for (let depth = 0; depth <= maxDepth && frontier.length > 0; depth++) { + // Resolve every node on this level concurrently. `visited` is not mutated + // here — only after the whole level settles — so the concurrency is safe. + const levelResults = await Promise.all( + frontier.map(async (nodeId) => { + const children = resolveChildren(tree, nodeId); + const result = await descend(nodeId, children); + return { nodeId, children, result }; + }), + ); + + const nextFrontier: string[] = []; + + for (const { nodeId, children, result } of levelResults) { + // Collect every page child of this node as a leaf hit. + for (const child of children) { + if (child.kind === "page") pages.add(child.ref); + } + + // The set of node children this node legitimately offered, in order. The + // descend pick is intersected with this so a stub returning bogus or + // duplicate refs can't make the walk recurse into something not offered. + const offeredNodes = children.filter((c) => c.kind === "node"); + const offeredRefs = new Set(offeredNodes.map((c) => c.ref)); + + // Honor the descend pick in the order it was returned, dedup'd, filtered + // to genuinely-offered node children, and capped by `breadthBudget`. + const descended: string[] = []; + const descendedSet = new Set(); + for (const choice of result.descend) { + if (choice.kind !== "node") continue; + if (!offeredRefs.has(choice.ref)) continue; + if (descendedSet.has(choice.ref)) continue; + if (descended.length >= breadthBudget) break; + descendedSet.add(choice.ref); + descended.push(choice.ref); + } + + const considered = offeredNodes.map((c) => c.ref); + const skipped = considered.filter((ref) => !descendedSet.has(ref)); + + levels.push({ + node: nodeId, + considered, + descended, + skipped, + reasoning: result.reasoning ?? "", + }); + + // Enqueue chosen node children for the next level. Mark visited now (the + // level has fully resolved) so a shared sub-node or a cycle is enqueued at + // most once across the whole walk. + for (const ref of descended) { + const key = nodeKey(ref); + if (visited.has(key)) continue; + visited.add(key); + nextFrontier.push(ref); + } + } + + frontier = nextFrontier; + } + + return { pages, levels }; +}