diff --git a/assistant/src/memory/v3/__tests__/tree-walk.test.ts b/assistant/src/memory/v3/__tests__/tree-walk.test.ts new file mode 100644 index 00000000000..5b45021b1dd --- /dev/null +++ b/assistant/src/memory/v3/__tests__/tree-walk.test.ts @@ -0,0 +1,585 @@ +/** + * Tests for `assistant/src/memory/v3/tree-walk.ts`. + * + * The descent provider is always a scripted stub injected via the `provider` + * arg — no real LLM, no network, no `mock.module`, `~/.vellum/` untouched. The + * stub keys its scripted decision off the `` marker in the user + * message so one fixture provider can drive a whole multi-node walk with one + * call per visited node. + * + * Coverage: + * - scripted descent over a fixture tree collects the right leaf pages and + * records considered/descended/skipped + reasoning per node. + * - one descent call per *visited* node (not per offered child). + * - breadthBudget caps descents per node (skip the overflow). + * - maxDepth halts the walk. + * - scout page hits seed the start node set (deriveSeedNodes) so a subtree the + * root never reaches is still walked. + * - explicit seeds bias the start set. + * - scout hits are rendered into the descend prompt as pressure. + * - provider === null → fail-safe: descend nothing, walk still terminates and + * collects the pages it reached, reasoning records the failure. + * - leaf nodes (no node children) make no provider call. + * - request shape: forced tool_choice on `choose_branches`, abort signal + * forwarded. + */ + +import { describe, expect, test } from "bun:test"; + +import type { + Message, + Provider, + ProviderResponse, + SendMessageOptions, + ToolDefinition, +} from "../../../providers/types.js"; +import type { RetrievalInput } from "../../v2/harness/retriever.js"; +import type { ScoutResult } from "../../v2/harness/trace.js"; +import type { PageIndex } from "../../v2/page-index.js"; +import type { ChildRef, TreeIndex } from "../tree-index.js"; +import { createDescender, deriveSeedNodes, runTreeWalk } from "../tree-walk.js"; +import type { TreeNode } from "../types.js"; + +// --------------------------------------------------------------------------- +// Fixture helpers. +// --------------------------------------------------------------------------- + +function page(ref: string): ChildRef { + return { kind: "page", ref }; +} + +function node(ref: string): ChildRef { + return { kind: "node", ref }; +} + +interface ProviderCall { + messages: Message[]; + tools: ToolDefinition[] | undefined; + systemPrompt: string | undefined; + options: SendMessageOptions | undefined; +} + +/** + * Build a tree node with the given children refs. `summary` defaults to the id + * so `composeNodeIndex` produces deterministic, inspectable lines. + */ +function makeNode(id: string, children: ChildRef[]): TreeNode { + return { + id, + frontmatter: { + children: children.map((c) => `${c.kind}:${c.ref}`), + summary: `summary of ${id}`, + }, + body: "", + }; +} + +/** + * Build an in-memory `TreeIndex` from a forward-adjacency spec, materializing + * `nodes`, `childrenByNode`, and the `pageParents` reverse edges (the only maps + * `tree-walk.ts` reads). `parentsByNode` is left empty — the driver never reads + * it. + */ +function makeTree( + root: string, + childrenByNode: Record, +): TreeIndex { + const nodes = new Map(); + const children = new Map>(); + const pageParents = new Map>(); + for (const [id, refs] of Object.entries(childrenByNode)) { + nodes.set(id, makeNode(id, refs)); + children.set(id, refs); + for (const ref of refs) { + if (ref.kind !== "page") continue; + let parents = pageParents.get(ref.ref); + if (!parents) { + parents = new Set(); + pageParents.set(ref.ref, parents); + } + parents.add(id); + } + } + return { + nodes, + childrenByNode: children, + parentsByNode: new Map(), + pageParents, + root, + }; +} + +/** Empty page index — the driver only needs `bySlug` for page summaries. */ +function makePages(slugs: string[]): PageIndex { + const bySlug = new Map(); + const byId = new Map(); + let id = 1; + for (const slug of slugs) { + const entry = { + id, + slug, + summary: `summary of ${slug}`, + edges: [], + modifiedAt: 0, + }; + bySlug.set(slug, entry); + byId.set(id, entry); + id++; + } + return { entries: [...bySlug.values()], bySlug, byId, rendered: "" }; +} + +/** Minimal `RetrievalInput` carrying just the fields the driver reads. */ +function makeInput( + overrides?: Partial & { + breadthBudget?: number; + maxDepth?: number; + }, +): RetrievalInput { + const breadthBudget = overrides?.breadthBudget ?? 8; + const maxDepth = overrides?.maxDepth ?? 8; + const config = { + memory: { v3: { breadthBudget, maxDepth } }, + } as unknown as RetrievalInput["config"]; + const { breadthBudget: _b, maxDepth: _m, ...rest } = overrides ?? {}; + return { + workspaceDir: "/tmp/does-not-matter", + recentTurnPairs: [{ assistantMessage: "", userMessage: "tell me about a" }], + nowText: "2026-05-25 10:00 PT", + priorEverInjected: [], + config, + ...rest, + }; +} + +/** Pull the `` id out of a recorded descend prompt. */ +function nodeIdFromCall(call: ProviderCall): string | null { + for (const block of call.messages[0]?.content ?? []) { + if (block.type !== "text") continue; + const match = block.text.match(//); + if (match) return match[1]; + } + return null; +} + +/** + * A scripted descent provider. `script` maps a node id to the bare child-node + * ids to descend (and an optional reasoning string). Records every call and + * honors an already-aborted signal by throwing. + */ +function makeProvider( + script: Record, + calls: ProviderCall[], +): Provider { + return { + name: "stub", + sendMessage: async (messages, tools, systemPrompt, options) => { + calls.push({ messages, tools, systemPrompt, options }); + if (options?.signal?.aborted) { + const err = new Error("aborted"); + err.name = "AbortError"; + throw err; + } + const nodeId = + nodeIdFromCall({ messages, tools, systemPrompt, options }) ?? ""; + const decision = script[nodeId] ?? { descend: [] }; + const input: Record = { descend: decision.descend }; + if (decision.reasoning !== undefined) + input.reasoning = decision.reasoning; + const response: ProviderResponse = { + model: "stub-model", + stopReason: "tool_use", + usage: { inputTokens: 0, outputTokens: 0 }, + content: [ + { + type: "tool_use", + id: `tu-${nodeId}`, + name: "choose_branches", + input, + }, + ], + }; + return response; + }, + }; +} + +// --------------------------------------------------------------------------- +// deriveSeedNodes +// --------------------------------------------------------------------------- + +describe("deriveSeedNodes", () => { + test("maps scout page slugs to their parent nodes via pageParents", () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [page("pa")], + b: [page("pb")], + }); + const scouts: ScoutResult[] = [{ lane: "sparse", slugs: ["pb"] }]; + expect(deriveSeedNodes(tree, scouts, [])).toEqual(["b"]); + }); + + test("unions explicit seeds first, then scout-derived parents, dedup'd", () => { + const tree = makeTree("_root", { + _root: [node("a")], + a: [page("pa")], + }); + const scouts: ScoutResult[] = [{ lane: "hot", slugs: ["pa", "pa"] }]; + // "a" is both an explicit seed and the parent of pa — appears once, seeds first. + expect(deriveSeedNodes(tree, scouts, ["a", "x"])).toEqual(["a", "x"]); + }); + + test("ignores scout slugs with no parent node", () => { + const tree = makeTree("_root", { _root: [page("pr")] }); + const scouts: ScoutResult[] = [{ lane: "dense", slugs: ["orphan"] }]; + expect(deriveSeedNodes(tree, scouts, [])).toEqual([]); + }); +}); + +// --------------------------------------------------------------------------- +// runTreeWalk — scripted descent +// --------------------------------------------------------------------------- + +describe("runTreeWalk — scripted descent", () => { + test("collects the right leaf pages and records the descend/skip split", async () => { + // _root → {a, b}; a → leaf pa; b → leaf pb. Script descends only "a". + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [page("pa")], + b: [page("pb")], + }); + const pages = makePages(["pa", "pb"]); + const calls: ProviderCall[] = []; + const provider = makeProvider( + { _root: { descend: ["a"], reasoning: "a matches the turn" } }, + calls, + ); + + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts: [], + seeds: [], + provider, + }); + + // Only the descended branch's page is collected. + expect([...collected]).toEqual(["pa"]); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.considered).toEqual(["a", "b"]); + expect(rootLevel.descended).toEqual(["a"]); + expect(rootLevel.skipped).toEqual(["b"]); + expect(rootLevel.reasoning).toBe("a matches the turn"); + + // _root walked (has node children) + a walked (leaf, no call). b skipped. + expect(levels.map((l) => l.node).sort()).toEqual(["_root", "a"]); + }); + + test("makes exactly one descent call per visited node with node children", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [node("c"), page("pa")], + b: [page("pb")], + c: [page("pc")], + }); + const pages = makePages(["pa", "pb", "pc"]); + const calls: ProviderCall[] = []; + const provider = makeProvider( + { + _root: { descend: ["a", "b"] }, + a: { descend: ["c"] }, + // b and c are leaves of the descended set; c has no node children. + }, + calls, + ); + + await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts: [], + seeds: [], + provider, + }); + + // Calls happen for nodes that HAVE node children: _root, a. b (leaf) and + // c (leaf) are visited but short-circuit before the provider call. + const calledNodes = calls.map(nodeIdFromCall).sort(); + expect(calledNodes).toEqual(["_root", "a"]); + }); + + test("breadthBudget caps descents per node", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b"), node("c")], + a: [page("pa")], + b: [page("pb")], + c: [page("pc")], + }); + const pages = makePages(["pa", "pb", "pc"]); + const calls: ProviderCall[] = []; + // Model picks all three; budget 2 admits only the first two. + const provider = makeProvider( + { _root: { descend: ["a", "b", "c"] } }, + calls, + ); + + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput({ breadthBudget: 2 }), + tree, + pages, + scouts: [], + seeds: [], + provider, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.descended).toEqual(["a", "b"]); + expect(rootLevel.skipped).toEqual(["c"]); + expect([...collected].sort()).toEqual(["pa", "pb"]); + }); + + test("maxDepth halts the walk", async () => { + const tree = makeTree("_root", { + _root: [node("a")], + a: [node("b"), page("pa")], + b: [page("pb")], + }); + const pages = makePages(["pa", "pb"]); + const calls: ProviderCall[] = []; + const provider = makeProvider( + { _root: { descend: ["a"] }, a: { descend: ["b"] } }, + calls, + ); + + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput({ maxDepth: 1 }), + tree, + pages, + scouts: [], + seeds: [], + provider, + }); + + // Depth 0 (_root) and depth 1 (a) walked; b never reached. + expect(levels.map((l) => l.node)).toEqual(["_root", "a"]); + expect([...collected]).toEqual(["pa"]); + }); +}); + +// --------------------------------------------------------------------------- +// runTreeWalk — scout seeding +// --------------------------------------------------------------------------- + +describe("runTreeWalk — scout seeding", () => { + test("scout page hits seed a subtree the root never reaches", async () => { + // root only links to a; the "island" subtree is unreachable from root but a + // scout surfaced its leaf page, so deriveSeedNodes seeds `island`. + const tree = makeTree("_root", { + _root: [node("a")], + a: [page("pa")], + island: [page("treasure")], + }); + const pages = makePages(["pa", "treasure"]); + const calls: ProviderCall[] = []; + const provider = makeProvider({ _root: { descend: ["a"] } }, calls); + + const scouts: ScoutResult[] = [{ lane: "dense", slugs: ["treasure"] }]; + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts, + seeds: [], + provider, + }); + + // Both the root branch (pa) and the scout-seeded island (treasure) reached. + expect([...collected].sort()).toEqual(["pa", "treasure"]); + expect(levels.map((l) => l.node).sort()).toEqual(["_root", "a", "island"]); + }); + + test("explicit seeds bias the start set", async () => { + const tree = makeTree("_root", { + _root: [page("pr")], + mid: [page("pm")], + }); + const pages = makePages(["pr", "pm"]); + const calls: ProviderCall[] = []; + const provider = makeProvider({}, calls); + + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts: [], + seeds: ["mid"], + provider, + }); + + expect([...collected].sort()).toEqual(["pm", "pr"]); + expect(levels.map((l) => l.node).sort()).toEqual(["_root", "mid"]); + }); + + test("renders scout hits into the descend prompt as pressure", async () => { + const tree = makeTree("_root", { + _root: [node("a"), node("b")], + a: [page("pa")], + b: [page("pb")], + }); + const pages = makePages(["pa", "pb"]); + const calls: ProviderCall[] = []; + const provider = makeProvider({ _root: { descend: ["a"] } }, calls); + + const scouts: ScoutResult[] = [{ lane: "sparse", slugs: ["pb"] }]; + await runTreeWalk({ + input: makeInput(), + tree, + pages, + // Pass scouts but no parent-seed match so the start set stays root-only; + // we only assert the prompt rendering here. + scouts, + seeds: [], + provider, + }); + + const rootCall = calls.find((c) => nodeIdFromCall(c) === "_root")!; + const promptText = rootCall.messages[0].content + .filter((b) => b.type === "text") + .map((b) => (b as { text: string }).text) + .join("\n"); + expect(promptText).toContain(""); + expect(promptText).toContain("[sparse]: pb"); + }); +}); + +// --------------------------------------------------------------------------- +// runTreeWalk — fail-safe + request shape +// --------------------------------------------------------------------------- + +describe("runTreeWalk — fail-safe", () => { + test("provider null descends nothing but still terminates and collects reached pages", async () => { + const tree = makeTree("_root", { + _root: [node("a"), page("pr")], + a: [page("pa")], + }); + const pages = makePages(["pr", "pa"]); + + const { pages: collected, levels } = await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts: [], + seeds: [], + provider: null, + }); + + // Root's own page is collected; the undescended branch's page is not. + expect([...collected]).toEqual(["pr"]); + expect(levels.map((l) => l.node)).toEqual(["_root"]); + const rootLevel = levels[0]; + expect(rootLevel.descended).toEqual([]); + expect(rootLevel.skipped).toEqual(["a"]); + expect(rootLevel.reasoning).toContain("no provider"); + }); + + test("malformed tool input fails closed for that node", async () => { + const tree = makeTree("_root", { + _root: [node("a")], + a: [page("pa")], + }); + const pages = makePages(["pa"]); + const calls: ProviderCall[] = []; + // Provider returns a non-conforming tool input (descend is not an array). + const provider: Provider = { + name: "bad-schema", + sendMessage: async (messages, tools, systemPrompt, options) => { + calls.push({ messages, tools, systemPrompt, options }); + return { + model: "stub-model", + stopReason: "tool_use", + usage: { inputTokens: 0, outputTokens: 0 }, + content: [ + { + type: "tool_use", + id: "tu-1", + name: "choose_branches", + input: { descend: "not-an-array" }, + }, + ], + }; + }, + }; + + const { levels } = await runTreeWalk({ + input: makeInput(), + tree, + pages, + scouts: [], + seeds: [], + provider, + }); + + const rootLevel = levels.find((l) => l.node === "_root")!; + expect(rootLevel.descended).toEqual([]); + expect(rootLevel.reasoning).toContain("validation"); + }); +}); + +describe("createDescender — request shape", () => { + test("forces tool_choice on choose_branches and forwards the abort signal", async () => { + const tree = makeTree("_root", { + _root: [node("a")], + a: [page("pa")], + }); + const pages = makePages(["pa"]); + const calls: ProviderCall[] = []; + const provider = makeProvider({ _root: { descend: ["a"] } }, calls); + + const reasoningByNode = new Map(); + const descender = createDescender( + { + input: makeInput({ signal: AbortSignal.timeout(10_000) }), + tree, + pages, + scouts: [], + seeds: [], + provider, + }, + reasoningByNode, + ); + + await descender("_root", [...tree.childrenByNode.get("_root")!]); + + expect(calls).toHaveLength(1); + const call = calls[0]; + expect(call.tools?.[0]?.name).toBe("choose_branches"); + expect(call.options?.config?.tool_choice).toEqual({ + type: "tool", + name: "choose_branches", + }); + expect(call.options?.config?.callSite).toBe("memoryV3Descent"); + expect(call.options?.signal).toBeDefined(); + }); + + test("a node with no node children makes no provider call", async () => { + const tree = makeTree("leaf", { leaf: [page("p")] }); + const pages = makePages(["p"]); + const calls: ProviderCall[] = []; + const provider = makeProvider({}, calls); + + const reasoningByNode = new Map(); + const descender = createDescender( + { input: makeInput(), tree, pages, scouts: [], seeds: [], provider }, + reasoningByNode, + ); + + const chosen = await descender("leaf", [ + ...tree.childrenByNode.get("leaf")!, + ]); + expect(chosen).toEqual([]); + expect(calls).toHaveLength(0); + expect(reasoningByNode.get("leaf")).toBe(""); + }); +}); diff --git a/assistant/src/memory/v3/tree-walk.ts b/assistant/src/memory/v3/tree-walk.ts new file mode 100644 index 00000000000..c387a2a5bc4 --- /dev/null +++ b/assistant/src/memory/v3/tree-walk.ts @@ -0,0 +1,406 @@ +/** + * Memory v3 — tree-walk model driver. + * + * The *intelligence* half of the v3 tree descent. `traversal.ts` owns the + * mechanical, provider-free walk (`walkTree`); this module supplies the + * per-node `descend` decision that walk injects, and wires the whole thing into + * a single `runTreeWalk` entry point. + * + * Per visited node the driver makes one cheap LLM call (`memoryV3Descent`) over + * the node's *composed* index — `composeNodeIndex` renders one line per child + * (sub-node summary or leaf page summary) plus the node's routing hints — and + * asks which child *nodes* to descend into. The prompt also carries the + * conversation context (the just-arrived turn + NOW) and the surviving scout + * hits, so descent is **scout-seeded but not scout-bound**: the model sees where + * the cheap lanes already landed, yet still feels pressure to descend branches + * the scouts missed. A driver that only ratified the scouts would re-introduce + * the recall cliff the tree walk exists to avoid. + * + * Scout seeding works at two layers: + * 1. **Start set** — `runTreeWalk` derives seed *node* ids from scout-surfaced + * *page* slugs via the tree's `pageParents` reverse edges (a scout hit on + * `page:foo` seeds every node that lists `page:foo` as a child), unioned + * with any explicit `seeds`. `walkTree` fans out from `tree.root` + seeds. + * 2. **Descend pressure** — the surviving scout slugs are rendered into every + * descend prompt so the model can prefer (but is not forced onto) branches + * that contain them. + * + * Reasoning capture. The `createDescender` signature returns plain `ChildRef[]` + * (the chosen node children) to match the driver contract; the model's stated + * rationale is written into a side map keyed by node id. `runTreeWalk` adapts + * the descender into `walkTree`'s `DescendResult`-returning hook by pairing each + * node's chosen children with its recorded reasoning, so every emitted + * `TreeLevel` carries the model's reason for its descend/skip split — making a + * wrong high-level skip observable rather than silent. + * + * Fail-safe. When no provider is configured (or a per-node call errors / returns + * an unusable response) the descender descends *nothing* for that node and + * records the reason. The walk still terminates and still collects every leaf + * page it reached before the failure; it just stops exploring deeper from the + * affected node. Failing closed (descend nothing) rather than open (descend all) + * keeps a broken provider from blowing the breadth budget across the whole tree. + * + * This module is currently unwired — a later PR composes it into the loop. + */ + +import { z } from "zod"; + +import { + extractToolUse, + getConfiguredProvider, +} from "../../providers/provider-send-message.js"; +import type { + Message, + Provider, + ToolDefinition, +} from "../../providers/types.js"; +import { getLogger } from "../../util/logger.js"; +import type { RetrievalInput } from "../v2/harness/retriever.js"; +import type { ScoutResult } from "../v2/harness/trace.js"; +import type { PageIndex } from "../v2/page-index.js"; +import { composeNodeIndex } from "./index-composition.js"; +import type { WalkResult } from "./traversal.js"; +import { walkTree } from "./traversal.js"; +import type { ChildRef, TreeIndex } from "./tree-index.js"; + +const log = getLogger("memory-v3-tree-walk"); + +/** Tool name forced via `tool_choice`. Shared constant so tests can match it. */ +const DESCEND_TOOL_NAME = "choose_branches"; + +/** + * The descend decision the driver hands to `walkTree`. Returns the subset of + * `children` (node refs only) to recurse into. Matches the PR contract: a plain + * `ChildRef[]` promise. The model's reasoning is threaded out-of-band via the + * side map populated by {@link createDescender}, not the return value, so this + * signature stays small. + */ +export type Descender = ( + nodeId: string, + children: ChildRef[], +) => Promise; + +/** Arguments to {@link createDescender}. */ +export interface CreateDescenderArgs { + input: RetrievalInput; + tree: TreeIndex; + pages: PageIndex; + /** Surviving scout hits — rendered into the prompt as descend pressure. */ + scouts: ScoutResult[]; + /** Explicit seed node ids (folded into the prompt's seed context). */ + seeds: string[]; + /** + * Provider override seam for tests. Production omits it and the descender + * resolves `getConfiguredProvider("memoryV3Descent")` per call. Explicit + * `null` is distinct from `undefined`: it simulates "no provider configured" + * and exercises the fail-safe path without touching the real registry. + */ + provider?: Provider | null; +} + +/** Arguments to {@link runTreeWalk}. Identical to the descender's args. */ +export type RunTreeWalkArgs = CreateDescenderArgs; + +/** + * The forced-tool input schema. `descend` lists the bare node ids the model + * chose to recurse into; `reasoning` is its stated rationale for the + * descend/skip split. Mirrors v2's `select_pages_to_inject` forced-tool shape. + */ +const DescendToolResultSchema = z.object({ + descend: z.array(z.string()), + reasoning: z.string().optional(), +}); + +/** + * Build the forced tool definition for one node. `descend` is constrained to + * the node ids actually offered as `node:` children so the model can only pick + * from genuine branches (the walk filters anyway, but constraining the schema + * keeps the model honest and the trace clean). + */ +function buildDescendTool(offeredNodeIds: readonly string[]): ToolDefinition { + return { + name: DESCEND_TOOL_NAME, + description: + "Choose which child nodes of the current memory-tree node to descend " + + "into for the current turn. Prefer branches likely to contain pages " + + "that bear on the turn; you may favor branches the scout hits point at, " + + "but descend other promising branches too — missing a relevant subtree " + + "is worse than descending an extra one. Return an empty list only when " + + "no child node plausibly bears on the turn.", + input_schema: { + type: "object", + properties: { + descend: { + type: "array", + items: + offeredNodeIds.length > 0 + ? { type: "string", enum: [...offeredNodeIds] } + : { type: "string" }, + description: + "Bare ids of the child nodes to descend into. Choose only from " + + "the offered node children.", + }, + reasoning: { + type: "string", + description: + "One short sentence: why these branches were descended and the " + + "rest skipped.", + }, + }, + required: ["descend"], + }, + }; +} + +/** + * Render the recent-turn + NOW context the descend prompt needs. The just- + * arrived user turn is the last pair's `userMessage`; the prior assistant reply + * (when present) precedes it. NOW is passed verbatim. + */ +function renderConversationContext(input: RetrievalInput): string { + const lines: string[] = []; + const lastPair = input.recentTurnPairs[input.recentTurnPairs.length - 1]; + if (lastPair) { + if (lastPair.assistantMessage.trim().length > 0) { + lines.push(`[assistant]: ${lastPair.assistantMessage}`); + } + lines.push(`[user]: ${lastPair.userMessage}`); + } + return ( + `\n${input.nowText}\n\n\n` + + `\n${lines.join("\n")}\n` + ); +} + +/** + * Render the surviving scout hits as descend pressure — the page slugs each + * lane surfaced, grouped by lane. Empty string when there are no scout hits, so + * the prompt omits the block entirely. + */ +function renderScoutHits(scouts: readonly ScoutResult[]): string { + const lines: string[] = []; + for (const scout of scouts) { + if (scout.slugs.length === 0) continue; + lines.push(`[${scout.lane}]: ${scout.slugs.join(", ")}`); + } + if (lines.length === 0) return ""; + return `\n${lines.join("\n")}\n`; +} + +const DESCENT_SYSTEM_PROMPT = + "You are the descent driver for a hierarchical memory-retrieval walk. At each " + + "node you see its child index (one line per child sub-node or leaf page) and " + + "the current conversation turn. Choose which child *nodes* to descend into to " + + "find the pages that bear on the next reply. Leaf pages are collected " + + "automatically — you only decide which branches to explore deeper."; + +/** Fail-safe descend result: descend nothing, recording why on the side map. */ +function failClosed( + nodeId: string, + reasoning: string, + reasoningByNode: Map, +): ChildRef[] { + reasoningByNode.set(nodeId, reasoning); + return []; +} + +/** + * Create the per-node descend decision driving {@link walkTree}. + * + * The returned function makes one forced-tool `memoryV3Descent` call per node + * over its composed index, returning the chosen `node:` children. The model's + * reasoning for each node is written into `reasoningByNode` (keyed by node id) + * rather than the return value, so the small `Descender` signature is preserved + * and {@link runTreeWalk} can merge the reasoning into each `TreeLevel`. + * + * Provider resolution honors the `provider` arg (including explicit `null` for + * the fail-safe path) and otherwise resolves the configured call site once per + * call. Any failure — no provider, provider throw, missing/mismatched tool_use + * — fails closed (descend nothing) with the reason recorded. + */ +export function createDescender( + args: CreateDescenderArgs, + reasoningByNode: Map, +): Descender { + const { input, tree, pages, scouts } = args; + const conversationContext = renderConversationContext(input); + const scoutHits = renderScoutHits(scouts); + + return async (nodeId: string, children: ChildRef[]): Promise => { + const offeredNodes = children.filter((c) => c.kind === "node"); + // No node children to descend — nothing to ask the model. Record an empty + // reasoning so the level still reflects the (trivial) decision. + if (offeredNodes.length === 0) { + reasoningByNode.set(nodeId, ""); + return []; + } + + const provider = + args.provider !== undefined + ? args.provider + : await getConfiguredProvider("memoryV3Descent"); + if (!provider) { + log.warn( + { nodeId }, + "memoryV3Descent provider unavailable; descending nothing", + ); + return failClosed( + nodeId, + "no provider configured — descended nothing", + reasoningByNode, + ); + } + + const indexBlock = composeNodeIndex(nodeId, tree, pages); + const offeredNodeIds = offeredNodes.map((c) => c.ref); + + const userMsg: Message = { + role: "user", + content: [ + { type: "text", text: conversationContext }, + { + type: "text", + text: + (scoutHits ? `${scoutHits}\n\n` : "") + + `\n${indexBlock}\n`, + }, + ], + }; + + const descendTool = buildDescendTool(offeredNodeIds); + + let response; + try { + response = await provider.sendMessage( + [userMsg], + [descendTool], + DESCENT_SYSTEM_PROMPT, + { + config: { + callSite: "memoryV3Descent" as const, + tool_choice: { type: "tool" as const, name: DESCEND_TOOL_NAME }, + }, + ...(input.signal ? { signal: input.signal } : {}), + }, + ); + } catch (err) { + log.warn( + { err, nodeId }, + "Descent provider call threw; descending nothing", + ); + return failClosed( + nodeId, + "descent call failed — descended nothing", + reasoningByNode, + ); + } + + const toolBlock = extractToolUse(response); + if (!toolBlock || toolBlock.name !== DESCEND_TOOL_NAME) { + log.warn( + { stopReason: response.stopReason, nodeId }, + "Descent model returned no choose_branches tool_use; descending nothing", + ); + return failClosed( + nodeId, + "model returned no descend decision — descended nothing", + reasoningByNode, + ); + } + + const parsed = DescendToolResultSchema.safeParse(toolBlock.input); + if (!parsed.success) { + log.warn( + { error: parsed.error.message, nodeId }, + "Descent tool input did not match schema; descending nothing", + ); + return failClosed( + nodeId, + "descend decision failed validation — descended nothing", + reasoningByNode, + ); + } + + reasoningByNode.set(nodeId, parsed.data.reasoning ?? ""); + + // Map the chosen bare ids back to the offered ChildRefs. The walk filters + // bogus / unoffered refs anyway, but resolving against the offered set here + // keeps the returned ChildRefs canonical. + const offeredById = new Map(offeredNodes.map((c) => [c.ref, c])); + const chosen: ChildRef[] = []; + for (const id of parsed.data.descend) { + const ref = offeredById.get(id); + if (ref) chosen.push(ref); + } + return chosen; + }; +} + +/** + * Derive the seed *node* ids for the walk from the surviving scout *page* hits. + * + * Scouts surface concept-page slugs; the tree's `pageParents` reverse edges map + * each page slug to the node(s) that list it as a child. Seeding the walk at + * those parent nodes drops the model in near where the cheap lanes already + * landed (layer 1 of scout seeding), while the walk still fans out from the + * root and the descend pressure (layer 2) keeps it from collapsing onto the + * scouts. Explicit `seeds` are unioned in. Order is deterministic: explicit + * seeds first (in given order), then scout-derived parents in scout/slug order. + */ +export function deriveSeedNodes( + tree: TreeIndex, + scouts: readonly ScoutResult[], + seeds: readonly string[], +): string[] { + const out: string[] = []; + const seen = new Set(); + const push = (id: string): void => { + if (seen.has(id)) return; + seen.add(id); + out.push(id); + }; + for (const id of seeds) push(id); + for (const scout of scouts) { + for (const slug of scout.slugs) { + const parents = tree.pageParents.get(slug); + if (!parents) continue; + for (const parent of parents) push(parent); + } + } + return out; +} + +/** + * Drive a full scout-seeded tree walk for one retrieval pass. + * + * Wires {@link createDescender} into {@link walkTree} with `breadthBudget` / + * `maxDepth` drawn from `config.memory.v3` (on `input.config`) and the start set + * seeded by {@link deriveSeedNodes}. Returns the collected leaf pages and the + * per-node `TreeLevel[]`, each level carrying the model's recorded reasoning. + * + * The descender records reasoning into a node-keyed side map; this function + * adapts it into `walkTree`'s `DescendResult`-returning hook by pairing each + * node's chosen children with its recorded reason, so the walk threads the + * reasoning onto every emitted level. + */ +export async function runTreeWalk(args: RunTreeWalkArgs): Promise { + const { input, tree, scouts, seeds } = args; + const v3 = input.config.memory?.v3; + const breadthBudget = v3?.breadthBudget ?? 6; + const maxDepth = v3?.maxDepth ?? 6; + + const reasoningByNode = new Map(); + const descender = createDescender(args, reasoningByNode); + + const seedNodes = deriveSeedNodes(tree, scouts, seeds); + + return walkTree(tree, { + seeds: seedNodes, + breadthBudget, + maxDepth, + descend: async (nodeId, children) => { + const descend = await descender(nodeId, [...children]); + return { descend, reasoning: reasoningByNode.get(nodeId) ?? "" }; + }, + }); +}